feat(upgrader): add generalized atomic upgrade and migration capabilities

contracts/axelar-gateway/src/contract.rs

@@ -1,13 +1,13 @@
 use crate::error::ContractError;
+use crate::interface::AxelarGatewayInterface;
+use crate::storage_types::{DataKey, MessageApprovalKey, MessageApprovalValue};
 use crate::types::{CommandType, Message, Proof, WeightedSigners};
+use crate::{auth, event};
 use axelar_soroban_std::ensure;
+pub use axelar_soroban_std::UpgradeableInterface;
 use soroban_sdk::xdr::ToXdr;
 use soroban_sdk::{contract, contractimpl, Address, Bytes, BytesN, Env, String, Vec};
 
-use crate::interface::AxelarGatewayInterface;
-use crate::storage_types::{DataKey, MessageApprovalKey, MessageApprovalValue};
-use crate::{auth, event};
-
 const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 /// Parameters for extending the contract instance and its instance storage.
@@ -22,6 +22,24 @@ const INSTANCE_TTL_EXTEND_TO: u32 = 60 * LEDGERS_PER_DAY;
 #[contract]
 pub struct AxelarGateway;
 
+#[contractimpl]
+impl UpgradeableInterface for AxelarGateway {
+    type Error = ContractError;
+
+    fn version(env: Env) -> String {
+        String::from_str(&env, CONTRACT_VERSION)
+    }
+
+    fn upgrade(env: Env, new_wasm_hash: BytesN<32>) -> Result<(), ContractError> {
+        Self::owner(&env).require_auth();
+
+        env.deployer().update_current_contract_wasm(new_wasm_hash);
+        Self::start_migration(&env);
+
+        Ok(())
+    }
+}
+
 #[contractimpl]
 impl AxelarGateway {
     /// Initialize the gateway
@@ -47,6 +65,18 @@ impl AxelarGateway {
 
         Ok(())
     }
+
+    /// Migrate the contract state after upgrading the contract code. the migration_data type can be adjusted as needed.
+    pub fn migrate(env: Env, migration_data: ()) -> Result<(), ContractError> {
+        // This function should not get modified.
+        // Custom migration logic that changes from version to version should be added in the run_migration function
+        Self::ensure_is_migrating(&env)?;
+
+        Self::run_migration(&env, migration_data);
+        Self::complete_migration(&env);
+
+        Ok(())
+    }
 }
 
 #[contractimpl]
@@ -223,16 +253,6 @@ impl AxelarGatewayInterface for AxelarGateway {
         auth::epoch(env)
     }
 
-    fn version(env: &Env) -> String {
-        String::from_str(env, CONTRACT_VERSION)
-    }
-
-    fn upgrade(env: Env, new_wasm_hash: BytesN<32>) {
-        Self::owner(&env).require_auth();
-
-        env.deployer().update_current_contract_wasm(new_wasm_hash);
-    }
-
     fn transfer_ownership(env: Env, new_owner: Address) {
         let owner: Address = Self::owner(&env);
         owner.require_auth();
@@ -290,4 +310,27 @@ impl AxelarGateway {
             .instance()
             .extend_ttl(INSTANCE_TTL_THRESHOLD, INSTANCE_TTL_EXTEND_TO);
     }
+
+    fn ensure_is_migrating(env: &Env) -> Result<(), ContractError> {
+        let is_migrating = env
+            .storage()
+            .instance()
+            .get::<DataKey, bool>(&DataKey::Migrating)
+            .unwrap_or(false);
+
+        ensure!(is_migrating, ContractError::MigrationNotAllowed);
+        Ok(())
+    }
+
+    fn start_migration(env: &Env) {
+        env.storage().instance().set(&DataKey::Migrating, &true);
+    }
+
+    // Modify this function to add migration logic
+    #[allow(clippy::missing_const_for_fn)] // exclude no-op implementations from this lint
+    fn run_migration(_env: &Env, _migration_data: ()) {}
+
+    fn complete_migration(env: &Env) {
+        env.storage().instance().set(&DataKey::Migrating, &false);
+    }
 }

contracts/axelar-gateway/src/error.rs

@@ -20,4 +20,5 @@ pub enum ContractError {
     EmptyMessages = 12,
     // Executable
     NotApproved = 13,
+    MigrationNotAllowed = 14,
 }

contracts/axelar-gateway/src/interface.rs

@@ -1,12 +1,12 @@
-use soroban_sdk::{contractclient, Address, Bytes, BytesN, Env, String, Vec};
-
 use crate::{
     error::ContractError,
     types::{Message, Proof, WeightedSigners},
 };
+use axelar_soroban_std::UpgradeableInterface;
+use soroban_sdk::{contractclient, Address, Bytes, BytesN, Env, String, Vec};
 
 #[contractclient(name = "AxelarGatewayClient")]
-pub trait AxelarGatewayInterface {
+pub trait AxelarGatewayInterface: UpgradeableInterface {
     /// Sends a message to the specified destination chain and contarct address with a given payload.
     ///
     /// This function is the entry point for general message passing between chains.
@@ -76,12 +76,6 @@ pub trait AxelarGatewayInterface {
     /// Returns the epoch of the gateway.
     fn epoch(env: &Env) -> u64;
 
-    /// Returns the version of the gateway.
-    fn version(env: &Env) -> String;
-
-    /// Upgrades the gateway to a new wasm hash.
-    fn upgrade(env: Env, new_wasm_hash: BytesN<32>);
-
     /// Transfers ownership of the gateway to a new address.
     fn transfer_ownership(env: Env, new_owner: Address);
 

contracts/axelar-gateway/src/storage_types.rs

@@ -22,6 +22,7 @@ pub enum DataKey {
     Owner,
     Operator,
     MessageApproval(MessageApprovalKey),
+    Migrating,
     /// Auth Module
     PreviousSignerRetention,
     DomainSeparator,

contracts/upgrader/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "upgrader"
+version = "0.1.0"
+edition = { workspace = true }
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+axelar-soroban-std = { workspace = true }
+soroban-sdk = { workspace = true }
+
+[dev-dependencies]
+soroban-sdk = { workspace = true, features = ["testutils"] }
+
+[features]
+
+[lints]
+workspace = true