feat(upgrader): add generalized atomic upgrade and migration capabili…

…ties

contracts/axelar-gateway/src/contract.rs

@@ -101,15 +101,15 @@ impl AxelarGateway {
 
         message_approval
             == Self::message_approval_hash(
-                &env,
-                Message {
-                    source_chain,
-                    message_id,
-                    source_address,
-                    contract_address,
-                    payload_hash,
-                },
-            )
+            &env,
+            Message {
+                source_chain,
+                message_id,
+                source_address,
+                contract_address,
+                payload_hash,
+            },
+        )
     }
 
     /// Checks if a message is executed.
@@ -258,10 +258,33 @@ impl AxelarGateway {
         Self::owner(&env)?.require_auth();
 
         env.deployer().update_current_contract_wasm(new_wasm_hash);
+        Self::start_migrating(&env);
 
         Ok(())
     }
 
+    pub fn migrate(env: Env, _migration_data: ()) -> Result<(), ContractError> {
+        // DO NOT REMOVE THIS LINE, IT PREVENTS THE CONTRACT FROM BEING MIGRATED MULTIPLE TIMES
+        Self::assert_is_migrating(&env)?;
+
+        // Add migration logic here as needed
+
+
+        // DO NOT REMOVE THIS LINE, IT PREVENTS THE CONTRACT FROM BEING MIGRATED MULTIPLE TIMES
+        Self::complete_migration(&env);
+        Ok(())
+    }
+
+    fn assert_is_migrating(env: &Env) -> Result<(), ContractError> {
+        let is_migrating = env.storage().instance().get::<DataKey, bool>(&DataKey::Migrating).unwrap_or(false);
+
+        if !is_migrating {
+            Err(ContractError::MigrationAlreadyCompleted)
+        } else {
+            Ok(())
+        }
+    }
+
     pub fn transfer_ownership(env: Env, new_owner: Address) -> Result<(), ContractError> {
         let owner: Address = Self::owner(&env)?;
         owner.require_auth();
@@ -290,9 +313,7 @@ impl AxelarGateway {
     pub fn signers_hash_by_epoch(env: &Env, epoch: u64) -> Result<BytesN<32>, ContractError> {
         auth::signers_hash_by_epoch(env, epoch)
     }
-}
 
-impl AxelarGateway {
     /// Get the message approval value by `source_chain` and `message_id`, defaulting to `MessageNotApproved`
     fn message_approval(
         env: &Env,
@@ -324,4 +345,12 @@ impl AxelarGateway {
             .instance()
             .extend_ttl(INSTANCE_TTL_THRESHOLD, INSTANCE_TTL_EXTEND_TO);
     }
+
+    fn start_migrating(env: &Env) {
+        env.storage().instance().set(&DataKey::Migrating, &true);
+    }
+
+    fn complete_migration(env: &Env) {
+        env.storage().instance().set(&DataKey::Migrating, &false);
+    }
 }

contracts/axelar-gateway/src/error.rs

@@ -23,4 +23,5 @@ pub enum ContractError {
     EmptyMessages = 14,
     // Executable
     NotApproved = 15,
+    MigrationAlreadyCompleted = 16,
 }

contracts/axelar-gateway/src/storage_types.rs

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

contracts/axelar-gateway/src/types.rs

@@ -93,6 +93,7 @@ impl Proof {
     }
 }
 
+
 #[cfg(test)]
 mod test {
     use crate::types::{CommandType, Message, WeightedSigner, WeightedSigners};
@@ -166,7 +167,7 @@ mod test {
             hex!("90e3761c0794fbbd8b563a0d05d83395e7f88f64f30eebb7c5533329f6653e84"),
             hex!("60e146cb9c548ba6e614a87910d8172c9d21279a3f8f4da256ff36e15b80ea30"),
         ]
-        .map(|hash| BytesN::<32>::from_array(&env, &hash));
+            .map(|hash| BytesN::<32>::from_array(&env, &hash));
 
         let mut messages = soroban_sdk::Vec::new(&env);
 

contracts/upgrader/Cargo.toml

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

contracts/upgrader/src/contract.rs

@@ -0,0 +1,68 @@
+use crate::error::ContractError;
+use soroban_sdk::{contract, contractimpl, Address, BytesN, Env, String, Symbol, Val};
+
+#[contract]
+pub struct Upgrader;
+
+#[contractimpl]
+impl Upgrader {
+    pub fn __constructor(_env: Env) {}
+    pub fn upgrade(env: Env, contract_address: Address, new_version: String, new_wasm_hash: BytesN<32>, migration_data: soroban_sdk::Vec<Val>) -> Result<(), ContractError> {
+        assert_new_version_is_different(&env, &contract_address, &new_version)?;
+
+        env.invoke_contract::<()>(&contract_address, &upgrade_fn(&env), soroban_sdk::vec![&env, new_wasm_hash.into()]);
+        env.invoke_contract::<()>(&contract_address, &migrate_fn(&env), migration_data);
+
+        assert_new_version_matches_expected(&env, &contract_address, &new_version)
+    }
+}
+
+fn migrate_fn(env: &Env) -> Symbol {
+    Symbol::new(&env, "migrate")
+}
+
+fn upgrade_fn(env: &Env) -> Symbol {
+    Symbol::new(&env, "upgrade")
+}
+
+fn version_fn(env: &Env) -> Symbol {
+    Symbol::new(&env, "version")
+}
+
+fn assert_new_version_is_different(env: &Env, contract_address: &Address, new_version: &String) -> Result<(), ContractError> {
+    let cur_version: String = env.invoke_contract(&contract_address, &version_fn(&env), soroban_sdk::vec![&env]);
+
+    if cur_version != *new_version {
+        Ok(())
+    } else {
+        Err(ContractError::SameVersion)
+    }
+}
+
+fn assert_new_version_matches_expected(env: &Env, contract_address: &Address, new_version: &String) -> Result<(), ContractError> {
+    let cur_version: String = env.invoke_contract(&contract_address, &version_fn(&env), soroban_sdk::vec![&env]);
+
+    if cur_version == *new_version {
+        Ok(())
+    } else {
+        Err(ContractError::VersionMismatch)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::contract::{Upgrader, UpgraderClient};
+    use axelar_gateway::AxelarGateway;
+    use soroban_sdk::{BytesN, Env, String};
+
+    const _WASM: &[u8] = include_bytes!("testdata/axelar_gateway.wasm");
+    #[test]
+    fn test() {
+        let env = Env::default();
+
+        let contract_id = env.register(Upgrader, ());
+        let gateway_id = env.register(AxelarGateway, ());
+        let client = UpgraderClient::new(&env, &contract_id);
+        client.upgrade(&gateway_id, &String::from_str(&env, "1.0"), &BytesN::from_array(&env, &[0; 32]), &soroban_sdk::vec![&env]);
+    }
+}

contracts/upgrader/src/error.rs

@@ -0,0 +1,8 @@
+use soroban_sdk::contracterror;
+
+#[contracterror]
+#[repr(u32)]
+pub enum ContractError {
+    SameVersion = 1,
+    VersionMismatch = 2,
+}

contracts/upgrader/src/lib.rs

@@ -0,0 +1,4 @@
+#![no_std]
+
+mod contract;
+mod error;