feat(zk_toolbox): resume functionality (matter-labs#2376)

## What ❔

<!-- What are the changes this PR brings about? -->
<!-- Example: This PR adds a PR template to the repo. -->
<!-- (For bigger PRs adding more context is appreciated) -->

## Why ❔

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.

---------

Signed-off-by: Danil <[email protected]>

zk_toolbox/crates/common/src/cmd.rs

@@ -1,5 +1,6 @@
 use std::{
     ffi::OsStr,
+    fmt::{Display, Formatter},
     io,
     process::{Command, Output, Stdio},
     string::FromUtf8Error,
@@ -21,10 +22,19 @@ pub struct Cmd<'a> {
 }
 
 #[derive(thiserror::Error, Debug)]
-#[error("Cmd error: {source} {stderr:?}")]
 pub struct CmdError {
-    stderr: Option<String>,
-    source: anyhow::Error,
+    pub stderr: Option<String>,
+    pub source: anyhow::Error,
+}
+
+impl Display for CmdError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let mut data = format!("{}", &self.source);
+        if let Some(stderr) = &self.stderr {
+            data = format!("{data}\n{stderr}");
+        }
+        write!(f, "{}", data)
+    }
 }
 
 impl From<xshell::Error> for CmdError {

zk_toolbox/crates/common/src/forge.rs

@@ -5,16 +5,20 @@ use std::{
 
 use clap::{Parser, ValueEnum};
 use ethers::{
+    core::types::Bytes,
     middleware::Middleware,
     prelude::{LocalWallet, Signer},
     types::{Address, H256, U256},
-    utils::hex::ToHex,
+    utils::{hex, hex::ToHex},
 };
 use serde::{Deserialize, Serialize};
 use strum_macros::Display;
 use xshell::{cmd, Shell};
 
-use crate::{cmd::Cmd, ethereum::create_ethers_client};
+use crate::{
+    cmd::{Cmd, CmdResult},
+    ethereum::create_ethers_client,
+};
 
 /// Forge is a wrapper around the forge binary.
 pub struct Forge {
@@ -54,8 +58,24 @@ impl ForgeScript {
     pub fn run(mut self, shell: &Shell) -> anyhow::Result<()> {
         let _dir_guard = shell.push_dir(&self.base_path);
         let script_path = self.script_path.as_os_str();
-        let args = self.args.build();
-        Ok(Cmd::new(cmd!(shell, "forge script {script_path} --legacy {args...}")).run()?)
+        let args_no_resume = self.args.build();
+        if self.args.resume {
+            let mut args = args_no_resume.clone();
+            args.push(ForgeScriptArg::Resume.to_string());
+            let res = Cmd::new(cmd!(shell, "forge script {script_path} --legacy {args...}")).run();
+            if !res.resume_not_successful_because_has_not_began() {
+                return Ok(res?);
+            }
+        }
+        let res = Cmd::new(cmd!(
+            shell,
+            "forge script {script_path} --legacy {args_no_resume...}"
+        ))
+        .run();
+        if res.proposal_error() {
+            return Ok(());
+        }
+        Ok(res?)
     }
 
     pub fn wallet_args_passed(&self) -> bool {
@@ -87,6 +107,13 @@ impl ForgeScript {
         self
     }
 
+    pub fn with_calldata(mut self, calldata: &Bytes) -> Self {
+        self.args.add_arg(ForgeScriptArg::Sig {
+            sig: hex::encode(calldata),
+        });
+        self
+    }
+
     /// Makes sure a transaction is sent, only after its previous one has been confirmed and succeeded.
     pub fn with_slow(mut self) -> Self {
         self.args.add_arg(ForgeScriptArg::Slow);
@@ -208,6 +235,7 @@ pub enum ForgeScriptArg {
         url: String,
     },
     Verify,
+    Resume,
 }
 
 /// ForgeScriptArgs is a set of arguments that can be passed to the forge script command.
@@ -229,6 +257,8 @@ pub struct ForgeScriptArgs {
     /// Verifier API key
     #[clap(long)]
     pub verifier_api_key: Option<String>,
+    #[clap(long)]
+    pub resume: bool,
     /// List of additional arguments that can be passed through the CLI.
     ///
     /// e.g.: `zk_inception init -a --private-key=<PRIVATE_KEY>`
@@ -348,3 +378,35 @@ pub enum ForgeVerifier {
     Blockscout,
     Oklink,
 }
+
+// Trait for handling forge errors. Required for implementing method for CmdResult
+trait ForgeErrorHandler {
+    // Resume doesn't work if the forge script has never been started on this chain before.
+    // So we want to catch it and try again without resume arg if it's the case
+    fn resume_not_successful_because_has_not_began(&self) -> bool;
+    // Catch the error if upgrade tx has already been processed. We do execute much of
+    // txs using upgrade mechanism and if this particular upgrade has already been processed we could assume
+    // it as a success
+    fn proposal_error(&self) -> bool;
+}
+
+impl ForgeErrorHandler for CmdResult<()> {
+    fn resume_not_successful_because_has_not_began(&self) -> bool {
+        let text = "Deployment not found for chain";
+        check_error(self, text)
+    }
+
+    fn proposal_error(&self) -> bool {
+        let text = "revert: Operation with this proposal id already exists";
+        check_error(self, text)
+    }
+}
+
+fn check_error(cmd_result: &CmdResult<()>, error_text: &str) -> bool {
+    if let Err(cmd_error) = &cmd_result {
+        if let Some(stderr) = &cmd_error.stderr {
+            return stderr.contains(error_text);
+        }
+    }
+    false
+}

zk_toolbox/crates/zk_inception/src/accept_ownership.rs

@@ -2,21 +2,30 @@ use common::{
     forge::{Forge, ForgeScript, ForgeScriptArgs},
     spinner::Spinner,
 };
-use config::{
-    forge_interface::{
-        accept_ownership::AcceptOwnershipInput, script_params::ACCEPT_GOVERNANCE_SCRIPT_PARAMS,
-    },
-    traits::SaveConfig,
-    EcosystemConfig,
+use config::{forge_interface::script_params::ACCEPT_GOVERNANCE_SCRIPT_PARAMS, EcosystemConfig};
+use ethers::{
+    abi::parse_abi,
+    contract::BaseContract,
+    types::{Address, H256},
 };
-use ethers::types::{Address, H256};
+use lazy_static::lazy_static;
 use xshell::Shell;
 
 use crate::{
     messages::MSG_ACCEPTING_GOVERNANCE_SPINNER,
     utils::forge::{check_the_balance, fill_forge_private_key},
 };
 
+lazy_static! {
+    static ref ACCEPT_ADMIN: BaseContract = BaseContract::from(
+        parse_abi(&[
+            "function acceptOwner(address governor, address target) public",
+            "function acceptAdmin(address governor, address target) public"
+        ])
+        .unwrap(),
+    );
+}
+
 pub async fn accept_admin(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
@@ -26,6 +35,15 @@ pub async fn accept_admin(
     forge_args: &ForgeScriptArgs,
     l1_rpc_url: String,
 ) -> anyhow::Result<()> {
+    // Resume for accept admin doesn't work properly. Foundry assumes that if signature of the function is the same,
+    // than it's the same call, but because we are calling this function multiple times during the init process,
+    // code assumes that doing only once is enough, but actually we need to accept admin multiple times
+    let mut forge_args = forge_args.clone();
+    forge_args.resume = false;
+
+    let calldata = ACCEPT_ADMIN
+        .encode("acceptAdmin", (governor_contract, target_address))
+        .unwrap();
     let foundry_contracts_path = ecosystem_config.path_to_foundry();
     let forge = Forge::new(&foundry_contracts_path)
         .script(
@@ -35,16 +53,8 @@ pub async fn accept_admin(
         .with_ffi()
         .with_rpc_url(l1_rpc_url)
         .with_broadcast()
-        .with_signature("acceptAdmin()");
-    accept_ownership(
-        shell,
-        ecosystem_config,
-        governor_contract,
-        governor,
-        target_address,
-        forge,
-    )
-    .await
+        .with_calldata(&calldata);
+    accept_ownership(shell, governor, forge).await
 }
 
 pub async fn accept_owner(
@@ -56,6 +66,13 @@ pub async fn accept_owner(
     forge_args: &ForgeScriptArgs,
     l1_rpc_url: String,
 ) -> anyhow::Result<()> {
+    // resume doesn't properly work here.
+    let mut forge_args = forge_args.clone();
+    forge_args.resume = false;
+
+    let calldata = ACCEPT_ADMIN
+        .encode("acceptOwner", (governor_contract, target_address))
+        .unwrap();
     let foundry_contracts_path = ecosystem_config.path_to_foundry();
     let forge = Forge::new(&foundry_contracts_path)
         .script(
@@ -65,37 +82,16 @@ pub async fn accept_owner(
         .with_ffi()
         .with_rpc_url(l1_rpc_url)
         .with_broadcast()
-        .with_signature("acceptOwner()");
-    accept_ownership(
-        shell,
-        ecosystem_config,
-        governor_contract,
-        governor,
-        target_address,
-        forge,
-    )
-    .await
+        .with_calldata(&calldata);
+    accept_ownership(shell, governor, forge).await
 }
 
 async fn accept_ownership(
     shell: &Shell,
-    ecosystem_config: &EcosystemConfig,
-    governor_contract: Address,
     governor: Option<H256>,
-    target_address: Address,
     mut forge: ForgeScript,
 ) -> anyhow::Result<()> {
-    let input = AcceptOwnershipInput {
-        target_addr: target_address,
-        governor: governor_contract,
-    };
-    input.save(
-        shell,
-        ACCEPT_GOVERNANCE_SCRIPT_PARAMS.input(&ecosystem_config.link_to_code),
-    )?;
-
     forge = fill_forge_private_key(forge, governor)?;
-
     check_the_balance(&forge).await?;
     let spinner = Spinner::new(MSG_ACCEPTING_GOVERNANCE_SPINNER);
     forge.run(shell)?;