WIP: wait for multisig upgrades

contracts/rust/deployer/src/builder.rs

@@ -75,6 +75,8 @@ pub struct DeployerArgs<P: Provider + WalletProvider> {
     #[builder(default)]
     dry_run: bool,
     #[builder(default)]
+    wait_for_multisig_upgrade: bool,
+    #[builder(default)]
     genesis_lc_state: Option<LightClientStateSol>,
     #[builder(default)]
     genesis_st_state: Option<StakeTableStateSol>,
@@ -201,6 +203,7 @@ impl<P: Provider + WalletProvider> DeployerArgs<P> {
                         contracts,
                         self.rpc_url.to_string(),
                         Some(self.dry_run),
+                        self.wait_for_multisig_upgrade,
                     )
                     .await?;
                 } else {
@@ -294,6 +297,7 @@ impl<P: Provider + WalletProvider> DeployerArgs<P> {
                         use_mock,
                         rpc_url.to_string(),
                         Some(dry_run),
+                        self.wait_for_multisig_upgrade,
                     )
                     .await?;
                 } else {
@@ -322,6 +326,7 @@ impl<P: Provider + WalletProvider> DeployerArgs<P> {
                         use_mock,
                         rpc_url.to_string(),
                         Some(dry_run),
+                        self.wait_for_multisig_upgrade,
                     )
                     .await?;
                 } else {
@@ -400,6 +405,7 @@ impl<P: Provider + WalletProvider> DeployerArgs<P> {
                         )?,
                         multisig_pauser,
                         Some(dry_run),
+                        self.wait_for_multisig_upgrade,
                     )
                     .await?;
                 } else {
@@ -731,6 +737,7 @@ impl<P: Provider + WalletProvider> DeployerArgs<P> {
                 safe_addr: multisig,
                 use_hardware_wallet,
                 dry_run,
+                wait_for_execution: self.wait_for_multisig_upgrade,
             },
         )
         .await?;

contracts/rust/deployer/src/lib.rs

@@ -2276,6 +2276,7 @@ mod tests {
             options.is_mock,
             sepolia_rpc_url.clone(),
             Some(dry_run),
+            false,
         )
         .await?;
         tracing::info!(
@@ -2608,6 +2609,7 @@ mod tests {
             multisig_admin,
             pauser,
             Some(dry_run),
+            false,
         )
         .await?;
 
@@ -2900,6 +2902,7 @@ mod tests {
             &mut contracts,
             sepolia_rpc_url.clone(),
             Some(dry_run),
+            false,
         )
         .await?;
         tracing::info!(
@@ -3184,6 +3187,7 @@ mod tests {
                 safe_addr: multisig_admin,
                 use_hardware_wallet: false,
                 dry_run,
+                wait_for_execution: false,
             },
         )
         .await?;
@@ -3454,4 +3458,121 @@ mod tests {
 
         Ok(())
     }
+
+    #[test_log::test(tokio::test)]
+    async fn test_wait_for_upgrade_execution_with_background_upgrade() -> Result<()> {
+        use tokio::time::sleep;
+
+        use crate::proposals::multisig::wait_for_upgrade_execution;
+
+        let provider = ProviderBuilder::new().connect_anvil_with_wallet();
+        let mut contracts = Contracts::new();
+        let admin = provider.get_accounts().await?[0];
+        let prover = Address::random();
+
+        let genesis_state = LightClientStateSol::dummy_genesis();
+        let genesis_stake = StakeTableStateSol::dummy_genesis();
+
+        let proxy_addr = deploy_light_client_proxy(
+            &provider,
+            &mut contracts,
+            false,
+            genesis_state.clone(),
+            genesis_stake.clone(),
+            admin,
+            Some(prover),
+        )
+        .await?;
+
+        let initial_impl = read_proxy_impl(&provider, proxy_addr).await?;
+        tracing::info!("Initial implementation: {initial_impl:#x}");
+
+        let blocks_per_epoch = 10;
+        let epoch_start_block = 22;
+        let pv2_addr = contracts
+            .deploy(
+                Contract::PlonkVerifierV2,
+                PlonkVerifierV2::deploy_builder(&provider),
+            )
+            .await?;
+
+        let target_lcv2_bytecode = LightClientV2::BYTECODE.encode_hex();
+        let lcv2_linked_bytecode = {
+            match target_lcv2_bytecode
+                .matches(LIBRARY_PLACEHOLDER_ADDRESS)
+                .count()
+            {
+                0 => return Err(anyhow!("lib placeholder not found")),
+                1 => Bytes::from_hex(target_lcv2_bytecode.replacen(
+                    LIBRARY_PLACEHOLDER_ADDRESS,
+                    &pv2_addr.encode_hex(),
+                    1,
+                ))?,
+                _ => {
+                    return Err(anyhow!(
+                        "more than one lib placeholder found, consider using a different value"
+                    ))
+                },
+            }
+        };
+
+        let lcv2_addr = contracts
+            .deploy(
+                Contract::LightClientV2,
+                LightClientV2::deploy_builder(&provider)
+                    .map(|req| req.with_deploy_code(lcv2_linked_bytecode)),
+            )
+            .await?;
+
+        let provider_clone = provider.clone();
+        let proxy_addr_clone = proxy_addr;
+        let lcv2_addr_clone = lcv2_addr;
+
+        let upgrade_task = tokio::spawn(async move {
+            sleep(Duration::from_secs(2)).await;
+            tracing::info!("Background task: performing upgrade now");
+
+            let proxy = LightClient::new(proxy_addr_clone, &provider_clone);
+            let init_data = LightClientV2::new(lcv2_addr_clone, &provider_clone)
+                .initializeV2(blocks_per_epoch, epoch_start_block)
+                .calldata()
+                .to_owned();
+
+            let receipt = proxy
+                .upgradeToAndCall(lcv2_addr_clone, init_data)
+                .send()
+                .await
+                .expect("upgrade failed")
+                .get_receipt()
+                .await
+                .expect("receipt failed");
+
+            tracing::info!(
+                "Background task: upgrade completed with tx {}",
+                receipt.transaction_hash
+            );
+        });
+
+        let poll_interval_secs = Some(1);
+        tracing::info!(
+            "Main task: starting wait_for_upgrade_execution with 1-second poll interval"
+        );
+        wait_for_upgrade_execution(&provider, proxy_addr, lcv2_addr, poll_interval_secs).await?;
+
+        upgrade_task.await?;
+
+        let final_impl = read_proxy_impl(&provider, proxy_addr).await?;
+        tracing::info!("Final implementation: {final_impl:#x}");
+        assert_eq!(
+            final_impl, lcv2_addr,
+            "Implementation should be updated to LCV2"
+        );
+
+        let lc_v2 = LightClientV2::new(proxy_addr, &provider);
+        assert_eq!(lc_v2.getVersion().call().await?.majorVersion, 2);
+        assert_eq!(lc_v2.blocksPerEpoch().call().await?, blocks_per_epoch);
+        assert_eq!(lc_v2.epochStartBlock().call().await?, epoch_start_block);
+
+        Ok(())
+    }
 }

contracts/rust/deployer/src/proposals/multisig.rs

@@ -27,6 +27,7 @@ pub struct TransferOwnershipParams {
     pub safe_addr: Address,
     pub use_hardware_wallet: bool,
     pub dry_run: bool,
+    pub wait_for_execution: bool,
 }
 
 /// Call the transfer ownership script to transfer ownership of a contract to a new owner
@@ -98,7 +99,7 @@ pub async fn transfer_ownership_from_multisig_to_timelock(
 
     let owner_addr = proxy_instance.owner().call().await?;
 
-    if !params.dry_run && !crate::is_contract(provider, owner_addr).await? {
+    if !params.dry_run && !crate::is_contract(&provider, owner_addr).await? {
         tracing::error!("Proxy owner is not a contract. Expected: {owner_addr:#x}");
         anyhow::bail!(
             "Proxy owner is not a contract. Expected: {owner_addr:#x}. Use --dry-run to skip this \
@@ -118,7 +119,10 @@ pub async fn transfer_ownership_from_multisig_to_timelock(
     if !params.dry_run {
         tracing::info!("Transfer Ownership proposal sent to {}", contract);
         tracing::info!("Send this link to the signers to sign the proposal: https://app.safe.global/transactions/queue?safe={}", params.safe_addr);
-        // IDEA: add a function to wait for the proposal to be executed
+
+        if params.wait_for_execution {
+            wait_for_ownership_transfer(&provider, proxy_addr, params.new_owner, None).await?;
+        }
     } else {
         tracing::info!("Dry run, skipping proposal");
     }
@@ -222,6 +226,69 @@ pub async fn verify_node_js_files() -> Result<()> {
     Ok(())
 }
 
+/// Wait for a multisig upgrade to be executed by polling the implementation address
+///
+/// Polls every `poll_interval_secs` seconds (default 12) without timeout until the implementation address changes
+pub async fn wait_for_upgrade_execution(
+    provider: impl Provider,
+    proxy_addr: Address,
+    expected_impl_addr: Address,
+    poll_interval_secs: Option<u64>,
+) -> Result<()> {
+    use std::{thread::sleep, time::Duration};
+
+    let poll_interval = poll_interval_secs.unwrap_or(12);
+    tracing::info!(
+        "Waiting for multisig upgrade to be executed. Proxy: {proxy_addr:#x}, Expected \
+         implementation: {expected_impl_addr:#x}"
+    );
+    tracing::info!("Polling every {} seconds (no timeout)...", poll_interval);
+
+    loop {
+        let current_impl = crate::read_proxy_impl(&provider, proxy_addr).await?;
+
+        if current_impl == expected_impl_addr {
+            tracing::info!("Upgrade executed! Implementation is now {current_impl:#x}");
+            return Ok(());
+        }
+
+        tracing::info!("Current implementation: {current_impl:#x}, waiting...");
+        sleep(Duration::from_secs(poll_interval));
+    }
+}
+
+/// Wait for ownership transfer to be executed by polling the owner address
+///
+/// Polls every `poll_interval_secs` seconds (default 12) without timeout until the owner address changes
+pub async fn wait_for_ownership_transfer(
+    provider: impl Provider,
+    proxy_addr: Address,
+    expected_owner: Address,
+    poll_interval_secs: Option<u64>,
+) -> Result<()> {
+    use std::{thread::sleep, time::Duration};
+
+    let poll_interval = poll_interval_secs.unwrap_or(12);
+    tracing::info!(
+        "Waiting for ownership transfer to be executed. Proxy: {proxy_addr:#x}, Expected owner: \
+         {expected_owner:#x}"
+    );
+    tracing::info!("Polling every {} seconds (no timeout)...", poll_interval);
+
+    loop {
+        let proxy_instance = OwnableUpgradeable::new(proxy_addr, &provider);
+        let current_owner = proxy_instance.owner().call().await?;
+
+        if current_owner == expected_owner {
+            tracing::info!("Ownership transfer executed! Owner is now {current_owner:#x}");
+            return Ok(());
+        }
+
+        tracing::info!("Current owner: {current_owner:#x}, waiting...");
+        sleep(Duration::from_secs(poll_interval));
+    }
+}
+
 /// Parameters for upgrading LightClient to V2
 pub struct LightClientV2UpgradeParams {
     pub blocks_per_epoch: u64,
@@ -245,6 +312,7 @@ pub async fn upgrade_light_client_v2_multisig_owner(
     is_mock: bool,
     rpc_url: String,
     dry_run: Option<bool>,
+    wait_for_execution: bool,
 ) -> Result<String> {
     let expected_major_version: u8 = 2;
     let dry_run = dry_run.unwrap_or_else(|| {
@@ -369,8 +437,11 @@ pub async fn upgrade_light_client_v2_multisig_owner(
                 "LightClientProxy upgrade proposal sent. Send this link to the signers to sign the proposal: https://app.safe.global/transactions/queue?safe={}",
                 owner_addr
             );
+
+        if wait_for_execution {
+            wait_for_upgrade_execution(&provider, proxy_addr, lcv2_addr, None).await?;
+        }
     }
-    // IDEA: add a function to wait for the proposal to be executed
 
     Ok(result)
 }
@@ -391,6 +462,7 @@ pub async fn upgrade_light_client_v3_multisig_owner(
     is_mock: bool,
     rpc_url: String,
     dry_run: Option<bool>,
+    wait_for_execution: bool,
 ) -> Result<String> {
     let expected_major_version: u8 = 3;
     let dry_run = dry_run.unwrap_or_else(|| {
@@ -510,8 +582,11 @@ pub async fn upgrade_light_client_v3_multisig_owner(
                 "LightClientProxy upgrade proposal sent. Send this link to the signers to sign the proposal: https://app.safe.global/transactions/queue?safe={}",
                 owner_addr
             );
+
+        if wait_for_execution {
+            wait_for_upgrade_execution(&provider, proxy_addr, lcv3_addr, None).await?;
+        }
     }
-    // IDEA: add a function to wait for the proposal to be executed
 
     Ok(result)
 }
@@ -529,6 +604,7 @@ pub async fn upgrade_esp_token_v2_multisig_owner(
     contracts: &mut Contracts,
     rpc_url: String,
     dry_run: Option<bool>,
+    wait_for_execution: bool,
 ) -> Result<String> {
     let dry_run = dry_run.unwrap_or_else(|| {
         tracing::warn!("Dry run not specified, defaulting to false");
@@ -591,6 +667,10 @@ pub async fn upgrade_esp_token_v2_multisig_owner(
                  sign the proposal: https://app.safe.global/transactions/queue?safe={}",
                 owner_addr
             );
+
+        if wait_for_execution {
+            wait_for_upgrade_execution(&provider, proxy_addr, esp_token_v2_addr, None).await?;
+        }
     }
 
     Ok(result)
@@ -612,6 +692,7 @@ pub async fn upgrade_stake_table_v2_multisig_owner(
     multisig_address: Address,
     pauser: Address,
     dry_run: Option<bool>,
+    wait_for_execution: bool,
 ) -> Result<()> {
     tracing::info!("Upgrading StakeTableProxy to StakeTableV2 using multisig owner");
     let dry_run = dry_run.unwrap_or(false);
@@ -658,5 +739,9 @@ pub async fn upgrade_stake_table_v2_multisig_owner(
 
     tracing::info!("StakeTableProxy upgrade proposal sent");
 
+    if !dry_run && wait_for_execution {
+        wait_for_upgrade_execution(&provider, proxy_addr, stake_table_v2_addr, None).await?;
+    }
+
     Ok(())
 }

sequencer/src/bin/deploy.rs

@@ -204,6 +204,10 @@ struct Options {
     #[clap(long, default_value = "false")]
     pub verify_node_js_files: bool,
 
+    /// Wait for multisig upgrade to be executed (polls every 12 seconds with no timeout)
+    #[clap(long, env = "ESPRESSO_WAIT_FOR_MULTISIG_UPGRADE", default_value = "false")]
+    pub wait_for_multisig_upgrade: bool,
+
     /// Stake table capacity for the prover circuit
     #[clap(short, long, env = "ESPRESSO_SEQUENCER_STAKE_TABLE_CAPACITY", default_value_t = DEFAULT_STAKE_TABLE_CAPACITY)]
     pub stake_table_capacity: usize,
@@ -435,6 +439,7 @@ async fn main() -> anyhow::Result<()> {
         .mock_light_client(opt.use_mock)
         .use_multisig(opt.use_multisig)
         .dry_run(opt.dry_run)
+        .wait_for_multisig_upgrade(opt.wait_for_multisig_upgrade)
         .rpc_url(opt.rpc_url.clone());
     if let Some(multisig) = opt.multisig_address {
         args_builder.multisig(multisig);