templar_common/
upgrade.rs

1use near_sdk::{
2    env,
3    json_types::{Base58CryptoHash, Base64VecU8},
4    near, Gas, NearToken, Promise,
5};
6
7/// The post-deploy state-migration method every contract exposes.
8pub const MIGRATE_METHOD: &str = "migrate";
9
10/// Where an upgrade's new code comes from — a raw blob, or a global contract pinned by code hash.
11/// Both are immutable: exactly this code is deployed, then `migrate` runs. (Global-*by-account-id* is
12/// deliberately excluded: it is a live delegation to the publisher, so a later republish would swap
13/// the code with no proposal, timelock, or `migrate` — bypassing governance and bricking versioned
14/// state.)
15///
16/// JSON: `Code` is untagged (a bare base64 string, matching the pre-`UpgradeSource` wire), so serde
17/// requires it declared last; `GlobalHash` is externally tagged. Borsh tags are pinned via explicit
18/// discriminants (`use_discriminant`), decoupling the persisted tag from declaration order —
19/// `UpgradeSource` lives in pending governance proposals, so a future variant just takes its own
20/// discriminant without disturbing the stored ones.
21#[derive(Debug, Clone, PartialEq, Eq)]
22#[near(serializers = [json, borsh(use_discriminant = true)])]
23#[repr(u8)]
24pub enum UpgradeSource {
25    /// A NEAR global contract referenced by its (immutable) code hash.
26    GlobalHash(Base58CryptoHash) = 1,
27    /// A raw WASM blob deployed onto the account.
28    #[serde(untagged)]
29    Code(Base64VecU8) = 0,
30}
31
32/// A compact, loggable stand-in for an [`UpgradeSource`] — a hash, never the (potentially large) blob.
33#[derive(Debug, Clone, PartialEq, Eq)]
34#[near(serializers = [json])]
35pub enum UpgradeSummary {
36    /// sha256 of the deployed wasm blob.
37    CodeHash(Base58CryptoHash),
38    /// A global contract referenced by its code hash.
39    GlobalHash(Base58CryptoHash),
40}
41
42impl UpgradeSource {
43    /// A `Code` variant carrying an empty blob is never a valid deploy.
44    pub fn is_empty_code(&self) -> bool {
45        matches!(self, UpgradeSource::Code(code) if code.0.is_empty())
46    }
47
48    /// A compact [`UpgradeSummary`] for event logging (`Code` → its sha256, via the on-chain host).
49    pub fn summary(&self) -> UpgradeSummary {
50        match self {
51            UpgradeSource::Code(blob) => {
52                UpgradeSummary::CodeHash(env::sha256_array(&blob.0).into())
53            }
54            UpgradeSource::GlobalHash(hash) => UpgradeSummary::GlobalHash(*hash),
55        }
56    }
57
58    /// Atomically deploy the new code and run `migrate_method` on it in a single receipt: a failed
59    /// migration reverts the code change too. The deploy always targets the current account, so this
60    /// is a self-upgrade primitive; the caller owns access control.
61    pub fn deploy_and_migrate(
62        self,
63        migrate_method: impl Into<String>,
64        migrate_args: Base64VecU8,
65        gas: Gas,
66    ) -> Promise {
67        let promise = Promise::new(env::current_account_id());
68        let deployed = match self {
69            UpgradeSource::Code(code) => promise.deploy_contract(code.0),
70            UpgradeSource::GlobalHash(hash) => promise.use_global_contract(hash),
71        };
72        deployed.function_call(
73            migrate_method.into(),
74            migrate_args.0,
75            NearToken::from_yoctonear(0),
76            gas,
77        )
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use near_sdk::serde_json::{self, json};
85
86    #[test]
87    fn code_is_untagged_bare_base64_in_json() {
88        let code = UpgradeSource::Code(Base64VecU8(vec![0xde, 0xad, 0xbe, 0xef]));
89        let value = serde_json::to_value(&code).unwrap();
90        // Bare base64 string, exactly as the pre-`UpgradeSource` `code` field serialized.
91        assert_eq!(value, json!("3q2+7w=="));
92        assert_eq!(
93            serde_json::from_value::<UpgradeSource>(value).unwrap(),
94            code
95        );
96    }
97
98    #[test]
99    fn global_hash_stays_externally_tagged_in_json() {
100        let hash = UpgradeSource::GlobalHash(Base58CryptoHash::from([0u8; 32]));
101        let value = serde_json::to_value(&hash).unwrap();
102        assert_eq!(
103            value,
104            json!({ "GlobalHash": "11111111111111111111111111111111" })
105        );
106        assert_eq!(
107            serde_json::from_value::<UpgradeSource>(value).unwrap(),
108            hash
109        );
110    }
111
112    #[test]
113    fn both_variants_borsh_roundtrip() {
114        for source in [
115            UpgradeSource::Code(Base64VecU8(vec![1, 2, 3])),
116            UpgradeSource::GlobalHash(Base58CryptoHash::from([7u8; 32])),
117        ] {
118            let bytes = near_sdk::borsh::to_vec(&source).unwrap();
119            assert_eq!(
120                near_sdk::borsh::from_slice::<UpgradeSource>(&bytes).unwrap(),
121                source
122            );
123        }
124    }
125
126    /// Golden bytes for the persisted borsh format — `UpgradeSource` lives inside pending governance
127    /// proposals, so these explicit discriminants must not shift.
128    #[test]
129    fn borsh_discriminants_are_stable() {
130        // Code = tag 0, then a Vec<u8> (u32 LE length + bytes).
131        assert_eq!(
132            near_sdk::borsh::to_vec(&UpgradeSource::Code(Base64VecU8(vec![0xaa, 0xbb]))).unwrap(),
133            vec![0, 2, 0, 0, 0, 0xaa, 0xbb],
134        );
135        // GlobalHash = tag 1, then the 32-byte hash.
136        assert_eq!(
137            near_sdk::borsh::to_vec(&UpgradeSource::GlobalHash(Base58CryptoHash::from(
138                [0u8; 32]
139            )))
140            .unwrap(),
141            [&[1u8][..], &[0u8; 32][..]].concat(),
142        );
143    }
144}