upgradability functionality

Cargo.toml

@@ -13,7 +13,7 @@ version = "0.6.1"
 
 [dependencies]
 near-contract-tools-macros = {version = "=0.6.1", path = "./macros"}
-near-sdk = "4.0.0"
+near-sdk = {version = "4.0.0", features = ["unstable"]}
 serde = "1.0.144"
 serde_json = "1.0.85"
 thiserror = "1.0.35"

macros/src/migrate.rs

@@ -45,12 +45,10 @@ pub fn expand(meta: MigrateMeta) -> Result<TokenStream, darling::Error> {
         #[#near_sdk::near_bindgen]
         impl #imp #me::migrate::MigrateExternal for #ident #ty #wh {
             #[init(ignore_state)]
-            fn migrate(args: Option<String>) -> Self {
-                let old_state = <#ident as #me::migrate::MigrateController>::deserialize_old_schema();
-
+            fn migrate() -> Self {
+                let old_state = <#ident as ::near_contract_tools::migrate::MigrateController>::deserialize_old_schema();
                 <#ident as #me::migrate::MigrateHook>::on_migrate(
                     old_state,
-                    args,
                 )
             }
         }

src/lib.rs

@@ -8,6 +8,7 @@ pub mod owner;
 pub mod pause;
 pub mod rbac;
 pub mod slot;
+pub mod upgrade;
 pub mod utils;
 
 pub use near_contract_tools_macros::*;

src/migrate.rs

@@ -33,13 +33,12 @@ pub trait MigrateHook: MigrateController {
     /// arguments from caller, and replaces it with the new schema.
     fn on_migrate(
         old_schema: <Self as MigrateController>::OldSchema,
-        args: Option<String>,
     ) -> <Self as MigrateController>::NewSchema;
 }
 
 /// Migrate-able contracts expose this trait publically
 #[ext_contract(ext_migrate)]
 pub trait MigrateExternal {
     /// Perform the migration with optional arguments
-    fn migrate(args: Option<String>) -> Self;
+    fn migrate() -> Self;
 }

src/owner.rs

@@ -104,6 +104,18 @@ pub trait Owner {
         proposed_owner.set(new.as_ref());
     }
 
+    /// Same as require_owner but as a method
+    fn assert_owner(&self) {
+        require!(
+            &env::predecessor_account_id()
+                == Self::slot_owner()
+                    .read()
+                    .as_ref()
+                    .unwrap_or_else(|| env::panic_str(NO_OWNER_FAIL_MESSAGE)),
+            ONLY_OWNER_FAIL_MESSAGE,
+        );
+    }
+
     /// Initializes the contract owner. Can only be called once.
     ///
     /// Emits an `OwnerEvent::Transfer` event.

src/upgrade.rs

@@ -0,0 +1,54 @@
+//! Upgrade Mod
+//!
+//! Makes it easier to upgrade your contract by providing a simple interface for upgrading the code and the state of your contract.
+
+use near_sdk::{env, sys, Gas};
+
+use near_sdk::borsh::{BorshDeserialize, BorshSerialize};
+
+/// Upgrade Trait
+pub trait Upgrade {
+    /// upgrade_all - Upgrades the code and the state of the contract
+    fn upgrade_all();
+}
+
+/// Contracts may implement this trait to inject code into Upgrade functions.
+pub trait UpgradeHook {
+    /// Executed before a upgrade call is conducted
+    fn on_upgrade();
+}
+
+/// Naked upgrade function which calls migrate method on the contract
+pub fn upgrade<T>()
+where
+    T: BorshDeserialize + BorshSerialize,
+{
+    env::setup_panic_hook();
+
+    const MIGRATE_METHOD_NAME: &[u8; 7] = b"migrate";
+    const UPDATE_GAS_LEFTOVER: Gas = Gas(5_000_000_000_000);
+
+    unsafe {
+        // Load code into register 0 result from the input argument if factory call or from promise if callback.
+        sys::input(0);
+        // Create a promise batch to update current contract with code from register 0.
+        let promise_id = sys::promise_batch_create(
+            env::current_account_id().as_bytes().len() as u64,
+            env::current_account_id().as_bytes().as_ptr() as u64,
+        );
+        // Deploy the contract code from register 0.
+        sys::promise_batch_action_deploy_contract(promise_id, u64::MAX, 0);
+        // Call promise to migrate the state.
+        // Batched together to fail upgrade if migration fails.
+        sys::promise_batch_action_function_call(
+            promise_id,
+            MIGRATE_METHOD_NAME.len() as u64,
+            MIGRATE_METHOD_NAME.as_ptr() as u64,
+            0,
+            0,
+            0,
+            (env::prepaid_gas() - env::used_gas() - UPDATE_GAS_LEFTOVER).0,
+        );
+        sys::promise_return(promise_id);
+    }
+}

tests/macros/migrate.rs

@@ -29,7 +29,7 @@ struct MyContract {
 }
 
 impl MigrateHook for MyContract {
-    fn on_migrate(old: Old, _args: Option<String>) -> Self {
+    fn on_migrate(old: Old) -> Self {
         Self { bar: old.foo }
     }
 }
@@ -43,7 +43,7 @@ fn default_from() {
 
     assert_eq!(old.foo, 99);
 
-    let migrated = <MyContract as MigrateExternal>::migrate(None);
+    let migrated = <MyContract as MigrateExternal>::migrate();
 
     assert_eq!(migrated.bar, 99);
 }

tests/macros/mod.rs

@@ -123,7 +123,7 @@ struct MigrateIntegration {
 }
 
 impl MigrateHook for MigrateIntegration {
-    fn on_migrate(old: Integration, _args: Option<String>) -> Self {
+    fn on_migrate(old: Integration) -> Self {
         Self::require_owner();
         Self::require_unpaused();
 
@@ -217,7 +217,7 @@ fn integration() {
     // Perform migration
     env::state_write(&c);
 
-    let mut migrated = <MigrateIntegration as MigrateExternal>::migrate(None);
+    let mut migrated = <MigrateIntegration as MigrateExternal>::migrate();
 
     assert_eq!(migrated.moved_value, 25);
     assert_eq!(migrated.get_value(), 25);
@@ -330,7 +330,7 @@ fn integration_fail_migrate_allow() {
 
     testing_env!(context);
 
-    <MigrateIntegration as MigrateExternal>::migrate(None);
+    <MigrateIntegration as MigrateExternal>::migrate();
 }
 
 #[test]
@@ -348,7 +348,7 @@ fn integration_fail_migrate_paused() {
 
     env::state_write(&c);
 
-    <MigrateIntegration as MigrateExternal>::migrate(None);
+    <MigrateIntegration as MigrateExternal>::migrate();
 }
 
 #[cfg(test)]

workspaces-tests/Cargo.toml

@@ -5,6 +5,15 @@ name = "workspaces-tests"
 publish = false
 version = "0.1.0"
 
+[[bin]]
+name = "upgrade_bad"
+
+[[bin]]
+name = "upgrade_old"
+
+[[bin]]
+name = "upgrade_new"
+
 [[bin]]
 name = "counter_multisig"
 
@@ -22,7 +31,7 @@ name = "simple_multisig"
 
 [dependencies]
 near-contract-tools = {path = "../"}
-near-sdk = "4.0.0"
+near-sdk = {version = "4.0.0", features = ["unstable"]}
 strum = "0.24.1"
 strum_macros = "0.24.3"
 thiserror = "1.0.34"

workspaces-tests/src/bin/upgrade_bad.rs

@@ -0,0 +1,30 @@
+#![allow(missing_docs)]
+
+use near_sdk::{
+    borsh::{self, BorshDeserialize, BorshSerialize},
+    near_bindgen, PanicOnDefault,
+};
+
+pub fn main() {} // Ignore
+
+#[derive(BorshSerialize, BorshDeserialize, PanicOnDefault)]
+#[near_bindgen]
+pub struct ContractBad {
+    pub foo: u32,
+}
+
+#[near_bindgen]
+impl ContractBad {
+    #[init]
+    pub fn new() -> Self {
+        Self { foo: 0 }
+    }
+
+    pub fn increment_foo(&mut self) {
+        self.foo += 1;
+    }
+
+    pub fn get_foo(&self) -> u32 {
+        self.foo
+    }
+}

workspaces-tests/src/bin/upgrade_new.rs

@@ -0,0 +1,45 @@
+#![allow(missing_docs)]
+
+use near_contract_tools::{migrate::MigrateExternal, migrate::MigrateHook, Migrate};
+use near_sdk::{
+    borsh::{self, BorshDeserialize, BorshSerialize},
+    near_bindgen, PanicOnDefault,
+};
+
+pub fn main() {}
+
+#[derive(BorshSerialize, BorshDeserialize, PanicOnDefault)]
+pub struct ContractOld {
+    pub foo: u32,
+}
+
+#[derive(BorshSerialize, BorshDeserialize, PanicOnDefault, Migrate)]
+#[migrate(from = "ContractOld")]
+#[near_bindgen]
+pub struct ContractNew {
+    pub bar: u64,
+}
+
+impl MigrateHook for ContractNew {
+    fn on_migrate(old_schema: ContractOld) -> Self {
+        Self {
+            bar: old_schema.foo as u64,
+        }
+    }
+}
+
+#[near_bindgen]
+impl ContractNew {
+    #[init]
+    pub fn new() -> Self {
+        Self { bar: 0 }
+    }
+
+    pub fn decrement_bar(&mut self) {
+        self.bar -= 1;
+    }
+
+    pub fn get_bar(&self) -> u64 {
+        self.bar
+    }
+}

workspaces-tests/src/bin/upgrade_old.rs

@@ -0,0 +1,50 @@
+#![allow(missing_docs)]
+
+use near_contract_tools::upgrade::{upgrade as other_upgrade, Upgrade, UpgradeHook};
+
+use near_contract_tools::{owner::Owner, owner::OwnerExternal, Owner};
+
+use near_sdk::{
+    borsh::{self, BorshDeserialize, BorshSerialize},
+    env, near_bindgen, PanicOnDefault,
+};
+pub fn main() {}
+
+#[derive(BorshSerialize, BorshDeserialize, PanicOnDefault, Owner)]
+#[near_bindgen]
+pub struct ContractOld {
+    pub foo: u32,
+}
+
+#[near_bindgen]
+impl ContractOld {
+    #[init]
+    pub fn new() -> Self {
+        let mut contract = Self { foo: 0 };
+
+        Owner::init(&mut contract, &env::predecessor_account_id());
+        contract
+    }
+
+    pub fn increment_foo(&mut self) {
+        self.foo += 1;
+    }
+
+    pub fn get_foo(&self) -> u32 {
+        self.foo
+    }
+}
+
+impl UpgradeHook for ContractOld {
+    fn on_upgrade() {
+        Self::require_owner();
+    }
+}
+
+impl Upgrade for ContractOld {
+    #[no_mangle]
+    fn upgrade_all() {
+        Self::require_owner();
+        other_upgrade::<ContractOld>();
+    }
+}