Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(upgrader): clean up the UpgraderInterface and simplify the contract #84

Merged
merged 5 commits into from
Nov 22, 2024
Merged

refactor(upgrader): clean up the UpgraderInterface and simplify the contract #84

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c963d4a
refactor(upgrader): clean up the UpgraderInterface and simplify the c…
cgorenflo Nov 21, 2024
a85ca84
Merge branch 'main' into upgrader-refactor
cgorenflo Nov 21, 2024
6b2e0d1
Merge branch 'main' into upgrader-refactor
cgorenflo Nov 22, 2024
60b43a8
update Cargo.lock
cgorenflo Nov 22, 2024
d917e2c
introduce test fixture
cgorenflo Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 7 additions & 13 deletions contracts/axelar-gateway/src/contract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,15 @@ 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> {
fn upgrade(env: Env, new_wasm_hash: BytesN<32>) {
Self::owner(&env).require_auth();

env.deployer().update_current_contract_wasm(new_wasm_hash);
Self::start_migration(&env);

Ok(())
}
}

Expand Down Expand Up @@ -315,25 +311,23 @@ impl AxelarGateway {
}

fn ensure_is_migrating(env: &Env) -> Result<(), ContractError> {
let is_migrating = env
.storage()
.instance()
.get::<DataKey, bool>(&DataKey::Migrating)
.unwrap_or(false);
ensure!(
env.storage().instance().has(&DataKey::Migrating),
ContractError::MigrationNotAllowed
);

ensure!(is_migrating, ContractError::MigrationNotAllowed);
Ok(())
}

fn start_migration(env: &Env) {
env.storage().instance().set(&DataKey::Migrating, &true);
env.storage().instance().set(&DataKey::Migrating, &());
}

// 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);
env.storage().instance().remove(&DataKey::Migrating);
}
}
306 changes: 18 additions & 288 deletions contracts/upgrader/src/contract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,10 @@
#![allow(dead_code)]

use crate::error::ContractError;
use axelar_soroban_std::traits::ThenOk;
use soroban_sdk::{contract, contractimpl, Address, BytesN, Env, String, Symbol, Val};
use axelar_soroban_std::{ensure, UpgradeableClient};
use soroban_sdk::{
contract, contractimpl, symbol_short, Address, BytesN, Env, String, Symbol, Val,
};

const MIGRATE: &str = "migrate";
const UPGRADE: &str = "upgrade";
const VERSION: &str = "version";
const MIGRATE: Symbol = symbol_short!("migrate");

#[contract]
pub struct Upgrader;
Expand All @@ -22,291 +20,23 @@ impl Upgrader {
new_wasm_hash: BytesN<32>,
migration_data: soroban_sdk::Vec<Val>,
) -> Result<(), ContractError> {
ensure_new_version_is_different(&env, &contract_address, &new_version)?;

// it's safe to map the true return value of the upgrade and migrate functions to (),
// because we don't care about it, and in case of failure the contract will panic anyway
env.invoke_contract::<()>(
&contract_address,
&Symbol::new(&env, UPGRADE),
soroban_sdk::vec![&env, new_wasm_hash.into()],
);

env.invoke_contract::<()>(
&contract_address,
&Symbol::new(&env, MIGRATE),
migration_data,
);

ensure_new_version_matches_expected(&env, &contract_address, &new_version)
}
}

fn ensure_new_version_is_different(
env: &Env,
contract_address: &Address,
new_version: &String,
) -> Result<(), ContractError> {
let no_match = current_version(env, contract_address) != *new_version;
no_match.then_ok((), ContractError::SameVersion)
}

fn ensure_new_version_matches_expected(
env: &Env,
contract_address: &Address,
new_version: &String,
) -> Result<(), ContractError> {
let versions_match = current_version(env, contract_address) == *new_version;
versions_match.then_ok((), ContractError::UnexpectedNewVersion)
}

fn current_version(env: &Env, contract_address: &Address) -> String {
env.invoke_contract(
contract_address,
&Symbol::new(env, VERSION),
soroban_sdk::vec![env],
)
}

#[cfg(test)]
mod tests {
use crate::contract::{Upgrader, UpgraderClient, MIGRATE, UPGRADE, VERSION};
use axelar_soroban_std::UpgradeableInterface;
use soroban_sdk::testutils::{Address as _, MockAuth, MockAuthInvoke};
use soroban_sdk::{contract, contracterror, contractimpl, contracttype, Address, Symbol};
use soroban_sdk::{BytesN, Env, String};

/// A simple contract to test the upgrader
#[contract]
struct DummyContract;

/// Dummy contract logic before upgrade
#[contractimpl]
impl UpgradeableInterface for DummyContract {
type Error = ContractError;

fn version(env: Env) -> String {
String::from_str(&env, "0.1.0")
}

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);
Ok(())
}
}

#[contractimpl]
impl DummyContract {
pub fn __constructor(env: Env, owner: Address) {
env.storage().instance().set(&DataKey::Owner, &owner)
}

fn owner(env: &Env) -> Address {
env.storage().instance().get(&DataKey::Owner).unwrap()
}
}

#[contracttype]
enum DataKey {
Data,
Owner,
}

#[contracterror]
enum ContractError {
SomeFailure = 1,
}
let contract_client = UpgradeableClient::new(&env, &contract_address);

/// Dummy contract logic after upgrade is loaded into WASM_AFTER_UPGRADE
///
/// #[contractimpl]
/// impl UpgradeableInterface for DummyContract {
/// type Error = ContractError;
///
/// fn version(env: Env) -> String {
/// String::from_str(&env, "0.2.0")
/// }
///
/// 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);
/// Ok(())
/// }
/// }
///
/// #[contractimpl]
/// impl DummyContract {
/// pub fn migrate(env: Env, migration_data: String) {
/// Self::owner(&env).require_auth();
///
/// env.storage()
/// .instance()
/// .set(&DataKey::Data, &migration_data);
/// }
///
/// fn owner(env: &Env) -> Address {
/// env.storage().instance().get(&DataKey::Owner).unwrap()
/// }
/// }
const WASM_AFTER_UPGRADE: &[u8] = include_bytes!("testdata/dummy.wasm");

#[test]
fn upgrade_and_migrate_are_atomic() {
let env = Env::default();

let owner = Address::generate(&env);
let contract_address = env.register(DummyContract, (&owner,));
let upgrader_address = env.register(Upgrader, ());

let original_version: String = query_version(&env, &contract_address);
assert_eq!(original_version, String::from_str(&env, "0.1.0"));

let hash_after_upgrade = env.deployer().upload_contract_wasm(WASM_AFTER_UPGRADE);
let expected_data = String::from_str(&env, "migration successful");
let expected_version = String::from_str(&env, "0.2.0");

let (upgrade_auth, migrate_auth) =
build_invocation_auths(&env, &contract_address, &hash_after_upgrade, &expected_data);

// add the owner to the set of authenticated addresses
env.mock_auths(&[
MockAuth {
address: &owner,
invoke: &upgrade_auth,
},
MockAuth {
address: &owner,
invoke: &migrate_auth,
},
]);

UpgraderClient::new(&env, &upgrader_address).upgrade(
&contract_address,
&expected_version,
&hash_after_upgrade,
&soroban_sdk::vec![&env, expected_data.to_val()],
);

// ensure new version is set correctly
let upgraded_version: String = env.invoke_contract(
&contract_address,
&Symbol::new(&env, VERSION),
soroban_sdk::vec![&env],
);
assert_eq!(upgraded_version, expected_version);

// ensure migration was successful
env.as_contract(&contract_address, || {
let data: String = env.storage().instance().get(&DataKey::Data).unwrap();
assert_eq!(data, expected_data);
});
}

#[test]
#[should_panic(expected = "HostError: Error(Auth, InvalidAction)")]
fn upgrade_fails_if_caller_is_authenticated_but_not_owner() {
let env = Env::default();

let owner = Address::generate(&env);
let contract_address = env.register(DummyContract, (&owner,));
let upgrader_address = env.register(Upgrader, ());

let hash_after_upgrade = env.deployer().upload_contract_wasm(WASM_AFTER_UPGRADE);
let expected_data = String::from_str(&env, "migration successful");
let expected_version = String::from_str(&env, "0.2.0");

let (upgrade_auth, migrate_auth) =
build_invocation_auths(&env, &contract_address, &hash_after_upgrade, &expected_data);

// add the caller to the set of authenticated addresses
let caller = Address::generate(&env);
env.mock_auths(&[
MockAuth {
address: &caller,
invoke: &upgrade_auth,
},
MockAuth {
address: &caller,
invoke: &migrate_auth,
},
]);

// should panic: caller is authenticated, but not the owner
UpgraderClient::new(&env, &upgrader_address).upgrade(
&contract_address,
&expected_version,
&hash_after_upgrade,
&soroban_sdk::vec![&env, expected_data.to_val()],
ensure!(
contract_client.version() != new_version,
ContractError::SameVersion
);
}

#[test]
#[should_panic(expected = "HostError: Error(Auth, InvalidAction)")]
fn upgrade_fails_if_correct_owner_is_not_authenticated_for_full_invocation_tree() {
let env = Env::default();

let owner = Address::generate(&env);
let contract_address = env.register(DummyContract, (&owner,));

let upgrader_address = env.register(Upgrader, ());
contract_client.upgrade(&new_wasm_hash);
// The types of the arguments to the migrate function are unknown to this contract, so we need to call it with invoke_contract.
// The migrate function's return value can be safely cast to () no matter what it really is,
// because it will panic on failure anyway
env.invoke_contract::<()>(&contract_address, &MIGRATE, migration_data);

let hash_after_upgrade = env.deployer().upload_contract_wasm(WASM_AFTER_UPGRADE);
let expected_data = String::from_str(&env, "migration successful");
let expected_version = String::from_str(&env, "0.2.0");

let (upgrade_auth, migrate_auth) =
build_invocation_auths(&env, &contract_address, &hash_after_upgrade, &expected_data);

// only add the owner to the set of authenticated addresses for the upgrade function, and the caller for the migrate function
let caller = Address::generate(&env);
env.mock_auths(&[
MockAuth {
address: &owner,
invoke: &upgrade_auth,
},
MockAuth {
address: &caller,
invoke: &migrate_auth,
},
]);

UpgraderClient::new(&env, &upgrader_address).upgrade(
&contract_address,
&expected_version,
&hash_after_upgrade,
&soroban_sdk::vec![&env, expected_data.to_val()],
ensure!(
contract_client.version() == new_version,
ContractError::UnexpectedNewVersion
);
}

fn build_invocation_auths<'a>(
env: &Env,
contract_address: &'a Address,
hash_after_upgrade: &'a BytesN<32>,
expected_data: &'a String,
) -> (MockAuthInvoke<'a>, MockAuthInvoke<'a>) {
let upgrade = MockAuthInvoke {
contract: contract_address,
fn_name: UPGRADE,
args: soroban_sdk::vec![&env, hash_after_upgrade.to_val()],
sub_invokes: &[],
};
let migrate = MockAuthInvoke {
contract: contract_address,
fn_name: MIGRATE,
args: soroban_sdk::vec![&env, expected_data.to_val()],
sub_invokes: &[],
};
(upgrade, migrate)
}

fn query_version(env: &Env, contract_address: &Address) -> String {
env.invoke_contract(
contract_address,
&Symbol::new(env, VERSION),
soroban_sdk::vec![&env],
)
Ok(())
}
}
Binary file removed contracts/upgrader/src/testdata/dummy.wasm
View file
Open in desktop
Binary file not shown.
Loading
Loading