diff --git a/engine-tests/src/tests/upgrade.rs b/engine-tests/src/tests/upgrade.rs index 3930b363f..c8aa2fa84 100644 --- a/engine-tests/src/tests/upgrade.rs +++ b/engine-tests/src/tests/upgrade.rs @@ -4,6 +4,29 @@ use crate::utils::workspace::deploy_engine; #[tokio::test] async fn test_code_upgrade() { + let aurora = deploy_engine().await; + // do upgrade + let result = aurora + .upgrade(contract_bytes()) + .max_gas() + .transact() + .await + .unwrap(); + assert!(result.is_success()); + + // call a new method + let result = aurora + .as_raw_contract() + .view("some_new_fancy_function") + .await + .unwrap(); + + let output: [u32; 7] = result.borsh().unwrap(); + assert_eq!(output, [3, 1, 4, 1, 5, 9, 2]); +} + +#[tokio::test] +async fn test_code_upgrade_with_stage() { let aurora = deploy_engine().await; // do upgrade let result = aurora diff --git a/engine-workspace/src/contract.rs b/engine-workspace/src/contract.rs index ccef318b4..8c9ba30ca 100644 --- a/engine-workspace/src/contract.rs +++ b/engine-workspace/src/contract.rs @@ -11,14 +11,14 @@ use crate::operation::{ CallSetEthConnectorContractAccount, CallSetEthConnectorContractData, CallSetFixedGas, CallSetKeyManager, CallSetOwner, CallSetPausedFlags, CallSetSiloParams, CallSetWhitelistStatus, CallStageUpgrade, CallStateMigration, CallStorageDeposit, CallStorageUnregister, - CallStorageWithdraw, CallSubmit, CallWithdraw, ViewAccountsCounter, ViewBalance, ViewBlockHash, - ViewBridgeProver, ViewChainId, ViewCode, ViewErc20FromNep141, ViewFactoryWnearAddress, - ViewFtBalanceOf, ViewFtBalanceOfEth, ViewFtMetadata, ViewFtTotalEthSupplyOnAurora, - ViewFtTotalEthSupplyOnNear, ViewFtTotalSupply, ViewGetErc20Metadata, - ViewGetEthConnectorContractAccount, ViewGetFixedGas, ViewGetSiloParams, ViewGetWhitelistStatus, - ViewIsUsedProof, ViewNep141FromErc20, ViewNonce, ViewOwner, ViewPausedFlags, - ViewPausedPrecompiles, ViewStorageAt, ViewStorageBalanceOf, ViewUpgradeIndex, ViewVersion, - ViewView, + CallStorageWithdraw, CallSubmit, CallUpgrade, CallWithdraw, ViewAccountsCounter, ViewBalance, + ViewBlockHash, ViewBridgeProver, ViewChainId, ViewCode, ViewErc20FromNep141, + ViewFactoryWnearAddress, ViewFtBalanceOf, ViewFtBalanceOfEth, ViewFtMetadata, + ViewFtTotalEthSupplyOnAurora, ViewFtTotalEthSupplyOnNear, ViewFtTotalSupply, + ViewGetErc20Metadata, ViewGetEthConnectorContractAccount, ViewGetFixedGas, ViewGetSiloParams, + ViewGetWhitelistStatus, ViewIsUsedProof, ViewNep141FromErc20, ViewNonce, ViewOwner, + ViewPausedFlags, ViewPausedPrecompiles, ViewStorageAt, ViewStorageBalanceOf, ViewUpgradeIndex, + ViewVersion, ViewView, }; use crate::transaction::{CallTransaction, ViewTransaction}; use aurora_engine_types::account_id::AccountId; @@ -267,6 +267,10 @@ impl EngineContract { CallFactorySetWNearAddress::call(&self.contract).args_borsh(address) } + pub fn upgrade(&self, bytes: Vec) -> CallUpgrade { + CallUpgrade::call(&self.contract).args(bytes) + } + pub fn stage_upgrade(&self, bytes: Vec) -> CallStageUpgrade { CallStageUpgrade::call(&self.contract).args(bytes) } diff --git a/engine-workspace/src/operation.rs b/engine-workspace/src/operation.rs index c0ed09867..f25e7924f 100644 --- a/engine-workspace/src/operation.rs +++ b/engine-workspace/src/operation.rs @@ -36,6 +36,7 @@ impl_call_return![ (CallDeployUpgrade, Call::DeployUpgrade), (CallResumePrecompiles, Call::ResumePrecompiles), (CallPausePrecompiles, Call::PausePrecompiles), + (CallUpgrade, Call::Upgrade), (CallStageUpgrade, Call::StageUpgrade), (CallStateMigration, Call::StateMigration), (CallMintAccount, Call::MintAccount), @@ -127,6 +128,7 @@ pub(crate) enum Call { StorageUnregister, StorageWithdraw, PausePrecompiles, + Upgrade, StageUpgrade, DeployUpgrade, StateMigration, @@ -176,6 +178,7 @@ impl AsRef for Call { Call::StorageUnregister => "storage_unregister", Call::StorageWithdraw => "storage_withdraw", Call::PausePrecompiles => "pause_precompiles", + Call::Upgrade => "upgrade", Call::StageUpgrade => "stage_upgrade", Call::DeployUpgrade => "deploy_upgrade", Call::StateMigration => "state_migration", diff --git a/engine/src/contract_methods/admin.rs b/engine/src/contract_methods/admin.rs index a0db2abe1..50eaf4694 100644 --- a/engine/src/contract_methods/admin.rs +++ b/engine/src/contract_methods/admin.rs @@ -30,6 +30,7 @@ use aurora_engine_sdk::{ promise::PromiseHandler, }; use aurora_engine_types::parameters::engine::FullAccessKeyArgs; +use aurora_engine_types::types::{NearGas, ZERO_YOCTO}; use aurora_engine_types::{ borsh::BorshDeserialize, parameters::{ @@ -37,16 +38,17 @@ use aurora_engine_types::{ NewCallArgs, PausePrecompilesCallArgs, RelayerKeyArgs, RelayerKeyManagerArgs, SetOwnerArgs, SetUpgradeDelayBlocksArgs, StartHashchainArgs, }, - promise::{PromiseAction, PromiseBatchAction}, + promise::{PromiseAction, PromiseBatchAction, PromiseCreateArgs}, }, storage::{self, KeyPrefix}, types::{Address, Yocto}, - vec, + vec, ToString, }; use function_name::named; const CODE_KEY: &[u8; 4] = b"CODE"; const CODE_STAGE_KEY: &[u8; 10] = b"CODE_STAGE"; +const GAS_FOR_STATE_MIGRATION: NearGas = NearGas::new(100_000_000_000_000); #[named] pub fn new(mut io: I, env: &E) -> Result<(), ContractError> { @@ -178,6 +180,38 @@ pub fn stage_upgrade(io: I, env: &E) -> Result<(), Contrac }) } +pub fn upgrade( + io: I, + env: &E, + handler: &mut H, +) -> Result<(), ContractError> { + let state = state::get_state(&io)?; + require_running(&state)?; + require_owner_only(&state, &env.predecessor_account_id())?; + + let code = io.read_input().to_vec(); + let current_account_id = env.current_account_id(); + let batch = PromiseBatchAction { + target_account_id: current_account_id.clone(), + actions: vec![PromiseAction::DeployContract { code }], + }; + let state_migration_callback = PromiseCreateArgs { + target_account_id: current_account_id, + method: "state_migration".to_string(), + args: vec![], + attached_balance: ZERO_YOCTO, + attached_gas: GAS_FOR_STATE_MIGRATION, + }; + let promise_id = unsafe { + let base_id = handler.promise_create_batch(&batch); + handler.promise_attach_callback(base_id, &state_migration_callback) + }; + + handler.promise_return(promise_id); + + Ok(()) +} + #[named] pub fn resume_precompiles(io: I, env: &E) -> Result<(), ContractError> { with_hashchain(io, env, function_name!(), |io| { diff --git a/engine/src/lib.rs b/engine/src/lib.rs index fdc8346e8..b3864f51b 100644 --- a/engine/src/lib.rs +++ b/engine/src/lib.rs @@ -182,6 +182,18 @@ mod contract { .sdk_unwrap(); } + /// Upgrade the contract with the provided code bytes. + #[no_mangle] + pub extern "C" fn upgrade() { + let io = Runtime; + let env = Runtime; + let mut handler = Runtime; + + contract_methods::admin::upgrade(io, &env, &mut handler) + .map_err(ContractError::msg) + .sdk_unwrap(); + } + /// Stage new code for deployment. #[no_mangle] pub extern "C" fn stage_upgrade() {