Skip to main content

unc_contract_standards/upgrade/
mod.rs

1use unc_sdk::json_types::U64;
2use unc_sdk::{env, require, unc, AccountId, Duration, Promise, Timestamp};
3
4type WrappedDuration = U64;
5
6pub trait Ownable {
7    fn assert_owner(&self) {
8        require!(env::predecessor_account_id() == self.get_owner(), "Owner must be predecessor");
9    }
10    fn get_owner(&self) -> AccountId;
11    fn set_owner(&mut self, owner: AccountId);
12}
13
14pub trait Upgradable {
15    fn get_staging_duration(&self) -> WrappedDuration;
16    fn stage_code(&mut self, code: Vec<u8>, timestamp: Timestamp);
17    fn deploy_code(&mut self) -> Promise;
18
19    /// Implement migration for the next version.
20    /// Should be `unimplemented` for a new contract.
21    /// TODO: consider adding version of the contract stored in the storage?
22    fn migrate(&mut self) {
23        unimplemented!();
24    }
25}
26
27#[unc]
28pub struct Upgrade {
29    pub owner: AccountId,
30    pub staging_duration: Duration,
31    pub staging_timestamp: Timestamp,
32}
33
34impl Upgrade {
35    pub fn new(owner: AccountId, staging_duration: Duration) -> Self {
36        Self { owner, staging_duration, staging_timestamp: 0 }
37    }
38}
39
40impl Ownable for Upgrade {
41    fn get_owner(&self) -> AccountId {
42        self.owner.clone()
43    }
44
45    fn set_owner(&mut self, owner: AccountId) {
46        self.assert_owner();
47        self.owner = owner;
48    }
49}
50
51impl Upgradable for Upgrade {
52    fn get_staging_duration(&self) -> WrappedDuration {
53        self.staging_duration.into()
54    }
55
56    fn stage_code(&mut self, code: Vec<u8>, timestamp: Timestamp) {
57        self.assert_owner();
58        require!(
59            env::block_timestamp() + self.staging_duration < timestamp,
60            "Timestamp must be later than staging duration"
61        );
62        // Writes directly into storage to avoid serialization penalty by using default struct.
63        env::storage_write(b"upgrade", &code);
64        self.staging_timestamp = timestamp;
65    }
66
67    fn deploy_code(&mut self) -> Promise {
68        if self.staging_timestamp < env::block_timestamp() {
69            env::panic_str(
70                format!(
71                    "Deploy code too early: staging ends on {}",
72                    self.staging_timestamp + self.staging_duration
73                )
74                .as_str(),
75            );
76        }
77        let code = env::storage_read(b"upgrade")
78            .unwrap_or_else(|| env::panic_str("No upgrade code available"));
79        env::storage_remove(b"upgrade");
80        Promise::new(env::current_account_id()).deploy_contract(code)
81    }
82}