1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

use near_sdk::json_types::U64;
use near_sdk::{env, near, require, AccountId, Duration, Promise, Timestamp};

type WrappedDuration = U64;

pub trait Ownable {
    fn assert_owner(&self) {
        require!(env::predecessor_account_id() == self.get_owner(), "Owner must be predecessor");
    }
    fn get_owner(&self) -> AccountId;
    fn set_owner(&mut self, owner: AccountId);
}

pub trait Upgradable {
    fn get_staging_duration(&self) -> WrappedDuration;
    fn stage_code(&mut self, code: Vec<u8>, timestamp: Timestamp);
    fn deploy_code(&mut self) -> Promise;

    /// Implement migration for the next version.
    /// Should be `unimplemented` for a new contract.
    /// TODO: consider adding version of the contract stored in the storage?
    fn migrate(&mut self) {
        unimplemented!();
    }
}

#[near]
pub struct Upgrade {
    pub owner: AccountId,
    pub staging_duration: Duration,
    pub staging_timestamp: Timestamp,
}

impl Upgrade {
    pub fn new(owner: AccountId, staging_duration: Duration) -> Self {
        Self { owner, staging_duration, staging_timestamp: 0 }
    }
}

impl Ownable for Upgrade {
    fn get_owner(&self) -> AccountId {
        self.owner.clone()
    }

    fn set_owner(&mut self, owner: AccountId) {
        self.assert_owner();
        self.owner = owner;
    }
}

impl Upgradable for Upgrade {
    fn get_staging_duration(&self) -> WrappedDuration {
        self.staging_duration.into()
    }

    fn stage_code(&mut self, code: Vec<u8>, timestamp: Timestamp) {
        self.assert_owner();
        require!(
            env::block_timestamp() + self.staging_duration < timestamp,
            "Timestamp must be later than staging duration"
        );
        // Writes directly into storage to avoid serialization penalty by using default struct.
        env::storage_write(b"upgrade", &code);
        self.staging_timestamp = timestamp;
    }

    fn deploy_code(&mut self) -> Promise {
        if self.staging_timestamp < env::block_timestamp() {
            env::panic_str(
                format!(
                    "Deploy code too early: staging ends on {}",
                    self.staging_timestamp + self.staging_duration
                )
                .as_str(),
            );
        }
        let code = env::storage_read(b"upgrade")
            .unwrap_or_else(|| env::panic_str("No upgrade code available"));
        env::storage_remove(b"upgrade");
        Promise::new(env::current_account_id()).deploy_contract(code)
    }
}