Add support for the delayed upgrade

.gitignore

@@ -1,6 +1,7 @@
 /target
 Cargo.lock
 near-plugins/tests/contracts/*/target
+examples/target
 
 # Ignore IDE data
 .vscode/

examples/upgradable-examples/README.md

@@ -105,7 +105,7 @@ But it doesn't work in that way because we can't provide in Bash so long args...
 For running `up_satge_code` take a look on `up_stage_code/src/main.rs` script.
 ```shell
 $ cd up_stage_code
-$ cargo run -- "<PATH_TO_KEY_FOR_CONTRACT_ACCOUNT>"
+$ cargo run -- -p '<PATH_TO_KEY_FOR_CONTRACT_ACCOUNT>'
 $ cd ..
 ```
 Where `<PATH_TO_KEY_FOR_CONTRACT_ACCOUNT>` is `$HOME/.near-credentials/testnet/<CONTRACT_ACCOUNT>.json`

examples/upgradable-examples/up_stage_code/src/main.rs

@@ -33,9 +33,12 @@ async fn main() {
 
     let contract: Account = match &*args.network {
         "testnet" => get_contract!(testnet, args.path_to_key),
-        "mainnet" =>  get_contract!(mainnet, args.path_to_key),
+        "mainnet" => get_contract!(mainnet, args.path_to_key),
         "betanet" => get_contract!(betanet, args.path_to_key),
-        network => panic!("Unknown network {}. Possible networks: testnet, mainnet, betanet", network)
+        network => panic!(
+            "Unknown network {}. Possible networks: testnet, mainnet, betanet",
+            network
+        ),
     };
 
     let wasm = std::fs::read(&args.wasm).unwrap();

examples/upgradable-examples/upgradable_base/src/lib.rs

@@ -15,6 +15,7 @@ impl Counter {
     pub fn new() -> Self {
         let mut contract = Self { counter: 0 };
         contract.owner_set(Some(near_sdk::env::predecessor_account_id()));
+        contract.up_init_staging_duration(std::time::Duration::from_secs(60).as_nanos().try_into().unwrap()); // 1 minute
         contract
     }
 

near-plugins-derive/src/upgradable.rs

@@ -8,6 +8,10 @@ use syn::{parse_macro_input, DeriveInput};
 #[darling(default, attributes(upgradable), forward_attrs(allow, doc, cfg))]
 struct Opts {
     code_storage_key: Option<String>,
+    staging_timestamp_storage_key: Option<String>,
+    staging_duration_storage_key: Option<String>,
+    update_staging_duration_storage_key: Option<String>,
+    update_staging_duration_timestamp_storage_key: Option<String>,
 }
 
 pub fn derive_upgradable(input: TokenStream) -> TokenStream {
@@ -21,20 +25,56 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream {
         .code_storage_key
         .unwrap_or_else(|| "__CODE__".to_string());
 
+    let staging_timestamp_storage_key = opts
+        .staging_timestamp_storage_key
+        .unwrap_or_else(|| "__TIMESTAMP__".to_string());
+
+    let staging_duration_storage_key = opts
+        .staging_duration_storage_key
+        .unwrap_or_else(|| "__DURATION__".to_string());
+
+    let update_staging_duration_storage_key = opts
+        .update_staging_duration_storage_key
+        .unwrap_or_else(|| "__UPDATE_DURATION__".to_string());
+
+    let update_staging_duration_timestamp_storage_key = opts
+        .update_staging_duration_timestamp_storage_key
+        .unwrap_or_else(|| "__UPDATE_DURATION_TIMESTAMP__".to_string());
+
     let output = quote! {
         #[near_bindgen]
         impl Upgradable for #ident {
-            fn up_storage_key(&self) -> Vec<u8>{
+            fn up_storage_key(&self) -> Vec<u8> {
                 (#code_storage_key).as_bytes().to_vec()
             }
 
+            fn up_staging_timestamp_storage_key(&self) -> Vec<u8> {
+                (#staging_timestamp_storage_key).as_bytes().to_vec()
+            }
+
+            fn up_staging_duration_storage_key(&self) -> Vec<u8> {
+                (#staging_duration_storage_key).as_bytes().to_vec()
+            }
+
+            fn up_update_staging_duration_storage_key(&self) -> Vec<u8> {
+                (#update_staging_duration_storage_key).as_bytes().to_vec()
+            }
+
+            fn up_update_staging_duration_timestamp_storage_key(&self) -> Vec<u8> {
+                (#update_staging_duration_timestamp_storage_key).as_bytes().to_vec()
+            }
+
             #[#cratename::only(owner)]
             fn up_stage_code(&mut self, #[serializer(borsh)] code: Vec<u8>) {
+                let timestamp = near_sdk::env::block_timestamp() + self.up_get_staging_duration().unwrap_or(0);
+
                 if code.is_empty() {
                     near_sdk::env::storage_remove(self.up_storage_key().as_ref());
                 } else {
                     near_sdk::env::storage_write(self.up_storage_key().as_ref(), code.as_ref());
                 }
+
+                near_sdk::env::storage_write(self.up_staging_timestamp_storage_key().as_ref(), &timestamp.to_be_bytes());
             }
 
             #[result_serializer(borsh)]
@@ -49,9 +89,84 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream {
 
             #[#cratename::only(owner)]
             fn up_deploy_code(&mut self) -> near_sdk::Promise {
+                let staging_timestamp = self.up_get_staging_timestamp().unwrap_or(0);
+                if staging_timestamp < near_sdk::env::block_timestamp() {
+                    near_sdk::env::panic_str(
+                        format!(
+                            "Upgradable: Deploy code too early: staging ends on {}",
+                            staging_timestamp
+                        )
+                        .as_str(),
+                    );
+                }
+
                 near_sdk::Promise::new(near_sdk::env::current_account_id())
                     .deploy_contract(self.up_staged_code().unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: No staged code")))
             }
+
+            fn up_get_staging_timestamp(&self) -> Option<near_sdk::Timestamp> {
+                near_sdk::env::storage_read(self.up_staging_timestamp_storage_key().as_ref()).map(|staging_timestamp_bytes| {
+                    u64::from_be_bytes(staging_timestamp_bytes.try_into().unwrap_or_else(|_|
+                        near_sdk::env::panic_str("Upgradable: Invalid u64 timestamp format"))
+                    )
+                })
+            }
+
+            fn up_get_staging_duration(&self) -> Option<near_sdk::Duration> {
+                near_sdk::env::storage_read(self.up_staging_duration_storage_key().as_ref()).map(|staging_duration_bytes| {
+                    u64::from_be_bytes(staging_duration_bytes.try_into().unwrap_or_else(|_|
+                        near_sdk::env::panic_str("Upgradable: Invalid u64 Duration format"))
+                    )
+                })
+            }
+
+            #[#cratename::only(owner)]
+            fn up_init_staging_duration(&self, staging_duration: near_sdk::Duration) {
+                near_sdk::require!(self.up_get_staging_duration().is_none(), "Upgradable: staging duration was already initialized");
+                near_sdk::env::storage_write(self.up_staging_duration_storage_key().as_ref(), &staging_duration.to_be_bytes());
+            }
+
+            #[#cratename::only(owner)]
+            fn up_stage_update_staging_duration(&self, staging_duration: near_sdk::Duration) {
+                let staging_duration_timestamp = near_sdk::env::block_timestamp() + self.up_get_staging_duration().unwrap_or(0);
+                near_sdk::env::storage_write(self.up_update_staging_duration_storage_key().as_ref(), &staging_duration.to_be_bytes());
+                near_sdk::env::storage_write(self.up_update_staging_duration_timestamp_storage_key().as_ref(), &staging_duration_timestamp.to_be_bytes());
+            }
+
+            #[#cratename::only(owner)]
+            fn up_apply_update_staging_duration(&self) {
+                let staging_timestamp = self.up_get_update_staging_duration_timestamp()
+                    .unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: No staged update"));
+
+                if staging_timestamp < near_sdk::env::block_timestamp() {
+                    near_sdk::env::panic_str(
+                        format!(
+                            "Upgradable: Update duration too early: staging ends on {}",
+                            staging_timestamp
+                        )
+                        .as_str(),
+                    );
+                }
+
+                near_sdk::env::storage_write(self.up_staging_duration_storage_key().as_ref(), &staging_timestamp.to_be_bytes());
+                near_sdk::env::storage_remove(self.up_update_staging_duration_storage_key().as_ref());
+            }
+
+            fn up_get_update_staging_duration_timestamp(&self) -> Option<near_sdk::Timestamp> {
+                near_sdk::env::storage_read(self.up_update_staging_duration_timestamp_storage_key().as_ref()).map(|timestamp_bytes| {
+                    u64::from_be_bytes(timestamp_bytes.try_into().unwrap_or_else(|_|
+                        near_sdk::env::panic_str("Upgradable: Invalid u64 timestamp format"))
+                    )
+                })
+            }
+
+            fn up_get_update_staging_duration(&self) -> Option<near_sdk::Duration> {
+                near_sdk::env::storage_read(self.up_update_staging_duration_storage_key().as_ref()).map(|duration_bytes| {
+                    u64::from_be_bytes(duration_bytes.try_into().unwrap_or_else(|_|
+                        near_sdk::env::panic_str("Upgradable: Invalid u64 Duration format"))
+                    )
+                })
+            }
         }
     };
 

near-plugins/src/upgradable.rs

@@ -20,14 +20,30 @@
 //! After the code is deployed, it should be removed from staging. This will prevent an old code
 //! with a security vulnerability to be deployed, in case it was upgraded using other mechanism.
 use crate::events::{AsEvent, EventMetadata};
-use near_sdk::{AccountId, CryptoHash, Promise};
+use near_sdk::{AccountId, CryptoHash, Duration, Promise};
 use serde::Serialize;
 
 pub trait Upgradable {
     /// Key of storage slot to save the staged code.
     /// By default b"__CODE__" is used.
     fn up_storage_key(&self) -> Vec<u8>;
 
+    /// Key of storage slot to save the allowed timestamp to deploy the staged code.
+    /// By default b"__TIMESTAMP__" is used.
+    fn up_staging_timestamp_storage_key(&self) -> Vec<u8>;
+
+    /// Key of storage slot to save the delay duration of deploying staged code.
+    /// By default b"__DURATION__" is used.
+    fn up_staging_duration_storage_key(&self) -> Vec<u8>;
+
+    /// Key of storage slot to save the staged delay duration update.
+    /// By default b"__UPDATE_DURATION__" is used.
+    fn up_update_staging_duration_storage_key(&self) -> Vec<u8>;
+
+    /// Key of storage slot to save the allowed timestamp to apply the staged duration update.
+    /// By default b"__UPDATE_DURATION_TIMESTAMP__" is used.
+    fn up_update_staging_duration_timestamp_storage_key(&self) -> Vec<u8>;
+
     /// Allows authorized account to stage some code to be potentially deployed later.
     /// If a previous code was staged but not deployed, it is discarded.
     fn up_stage_code(&mut self, code: Vec<u8>);
@@ -40,6 +56,27 @@ pub trait Upgradable {
 
     /// Allows authorized account to deploy staged code. If no code is staged the method fails.
     fn up_deploy_code(&mut self) -> Promise;
+
+    /// Initialize the duration of the delay for deploying the staged code.
+    fn up_init_staging_duration(&self, staging_duration: near_sdk::Duration);
+
+    /// Returns the staging delay duration.
+    fn up_get_staging_duration(&self) -> Option<Duration>;
+
+    /// Returns the timestamp until which deploying the last staged code is not allowed.
+    fn up_get_staging_timestamp(&self) -> Option<Duration>;
+
+    /// Allows authorized account to stage update of the staging duration.
+    fn up_stage_update_staging_duration(&self, staging_duration: near_sdk::Duration);
+
+    /// Allows authorized account to apply the staging duration update.
+    fn up_apply_update_staging_duration(&self);
+
+    /// Returns the timestamp until which applying the last staged duration is not allowed.
+    fn up_get_update_staging_duration_timestamp(&self) -> Option<near_sdk::Timestamp>;
+
+    /// Returns the staged duration update.
+    fn up_get_update_staging_duration(&self) -> Option<near_sdk::Duration>;
 }
 
 /// Event emitted when the code is staged
@@ -139,4 +176,75 @@ mod tests {
 
         counter.up_deploy_code();
     }
+
+    #[test]
+    fn test_stage_code_with_delay() {
+        let (mut counter, mut ctx) = setup_basic();
+
+        ctx.predecessor_account_id = "eli.test".to_string().try_into().unwrap();
+        testing_env!(ctx.clone());
+
+        assert_eq!(counter.up_staged_code(), None);
+
+        let staging_duration: u64 = std::time::Duration::from_secs(60)
+            .as_nanos()
+            .try_into()
+            .unwrap();
+        counter.up_init_staging_duration(staging_duration);
+
+        let staging_timestamp = ctx.block_timestamp + staging_duration;
+        counter.up_stage_code(vec![1]);
+        assert_eq!(
+            counter.up_get_staging_timestamp().unwrap(),
+            staging_timestamp
+        );
+
+        assert_eq!(counter.up_staged_code(), Some(vec![1]));
+
+        ctx.block_timestamp = ctx.block_timestamp + staging_duration;
+        testing_env!(ctx);
+
+        assert_eq!(
+            counter.up_staged_code_hash(),
+            Some(sha256(vec![1].as_slice()).try_into().unwrap())
+        );
+
+        counter.up_deploy_code();
+    }
+
+    #[test]
+    #[should_panic(expected = "Upgradable: Deploy code too early: staging ends on")]
+    fn test_panic_stage_code_with_delay() {
+        let (mut counter, mut ctx) = setup_basic();
+
+        ctx.predecessor_account_id = "eli.test".to_string().try_into().unwrap();
+        testing_env!(ctx.clone());
+
+        assert_eq!(counter.up_staged_code(), None);
+
+        let staging_duration: u64 = std::time::Duration::from_secs(60)
+            .as_nanos()
+            .try_into()
+            .unwrap();
+        counter.up_init_staging_duration(staging_duration);
+
+        let staging_timestamp = ctx.block_timestamp + staging_duration;
+        counter.up_stage_code(vec![1]);
+        assert_eq!(
+            counter.up_get_staging_timestamp().unwrap(),
+            staging_timestamp
+        );
+
+        assert_eq!(counter.up_staged_code(), Some(vec![1]));
+
+        assert_eq!(
+            counter.up_staged_code_hash(),
+            Some(sha256(vec![1].as_slice()).try_into().unwrap())
+        );
+
+        ctx.block_timestamp = staging_timestamp + 1;
+        testing_env!(ctx);
+
+        counter.up_deploy_code();
+    }
 }