feat: upgrade contract with one transaction

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

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<u8>) -> CallUpgrade {
+        CallUpgrade::call(&self.contract).args(bytes)
+    }
+
     pub fn stage_upgrade(&self, bytes: Vec<u8>) -> CallStageUpgrade {
         CallStageUpgrade::call(&self.contract).args(bytes)
     }

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<str> 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",

engine/src/contract_methods/admin.rs

@@ -30,23 +30,25 @@ 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::{
         engine::{
             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<I: IO + Copy, E: Env>(mut io: I, env: &E) -> Result<(), ContractError> {
@@ -178,6 +180,38 @@ pub fn stage_upgrade<I: IO + Copy, E: Env>(io: I, env: &E) -> Result<(), Contrac
     })
 }
 
+pub fn upgrade<I: IO + Copy, E: Env, H: PromiseHandler>(
+    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<I: IO + Copy, E: Env>(io: I, env: &E) -> Result<(), ContractError> {
     with_hashchain(io, env, function_name!(), |io| {

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() {