diff --git a/near-plugins-derive/src/upgradable.rs b/near-plugins-derive/src/upgradable.rs index f42cb5b..c1397a6 100644 --- a/near-plugins-derive/src/upgradable.rs +++ b/near-plugins-derive/src/upgradable.rs @@ -140,6 +140,13 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream { fn up_set_staging_duration_unchecked(&self, staging_duration: near_sdk::Duration) { self.up_storage_write(__UpgradableStorageKey::StagingDuration, &::near_sdk::borsh::to_vec(&staging_duration).unwrap()); } + + /// Computes the `sha256` hash of `code` and panics if the conversion to `CryptoHash` fails. + fn up_hash_code(code: &[u8]) -> ::near_sdk::CryptoHash { + let hash = near_sdk::env::sha256(code); + std::convert::TryInto::try_into(hash) + .expect("sha256 should convert to CryptoHash") + } } #[near] @@ -176,11 +183,11 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream { fn up_staged_code_hash(&self) -> Option<::near_sdk::CryptoHash> { self.up_staged_code() - .map(|code| std::convert::TryInto::try_into(::near_sdk::env::sha256(code.as_ref())).unwrap()) + .map(|code| Self::up_hash_code(code.as_ref())) } #[#cratename::access_control_any(roles(#(#acl_roles_code_deployers),*))] - fn up_deploy_code(&mut self, function_call_args: Option<#cratename::upgradable::FunctionCallArgs>) -> near_sdk::Promise { + fn up_deploy_code(&mut self, hash: String, function_call_args: Option<#cratename::upgradable::FunctionCallArgs>) -> near_sdk::Promise { let staging_timestamp = self.up_get_timestamp(__UpgradableStorageKey::StagingTimestamp) .unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: staging timestamp isn't set")); @@ -195,6 +202,17 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream { } let code = self.up_staged_code().unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: No staged code")); + let expected_hash = ::near_sdk::base64::encode(Self::up_hash_code(code.as_ref())); + if hash != expected_hash { + near_sdk::env::panic_str( + format!( + "Upgradable: Cannot deploy due to wrong hash: expected hash: {}", + expected_hash, + ) + .as_str(), + ) + } + let promise = ::near_sdk::Promise::new(::near_sdk::env::current_account_id()) .deploy_contract(code); match function_call_args { diff --git a/near-plugins-derive/tests/common/upgradable_contract.rs b/near-plugins-derive/tests/common/upgradable_contract.rs index 999e10d..6d1e239 100644 --- a/near-plugins-derive/tests/common/upgradable_contract.rs +++ b/near-plugins-derive/tests/common/upgradable_contract.rs @@ -71,11 +71,13 @@ impl UpgradableContract { pub async fn up_deploy_code( &self, caller: &Account, + hash: String, function_call_args: Option, ) -> near_workspaces::Result { caller .call(self.contract.id(), "up_deploy_code") .args_json(json!({ + "hash": hash, "function_call_args": function_call_args, })) .max_gas() diff --git a/near-plugins-derive/tests/upgradable.rs b/near-plugins-derive/tests/upgradable.rs index 101ec48..6c2fde3 100644 --- a/near-plugins-derive/tests/upgradable.rs +++ b/near-plugins-derive/tests/upgradable.rs @@ -109,13 +109,13 @@ impl Setup { } /// Asserts staged code equals `expected_code`. - async fn assert_staged_code(&self, expected_code: Option>) { + async fn assert_staged_code(&self, expected_code: Option<&Vec>) { let staged = self .upgradable_contract .up_staged_code(&self.unauth_account) .await .expect("Call to up_staged_code should succeed"); - assert_eq!(staged, expected_code); + assert_eq!(staged.as_ref(), expected_code); } /// Asserts the staging duration of the `Upgradable` contract equals the `expected_duration`. @@ -210,6 +210,14 @@ fn convert_code_to_crypto_hash(code: &[u8]) -> CryptoHash { .expect("Code should be converted to CryptoHash") } +/// Computes the hash `code` according the to requirements of the `hash` parameter of +/// `Upgradable::up_deploy_code`. +fn convert_code_to_deploy_hash(code: &[u8]) -> String { + use near_sdk::base64::Engine; + let hash = near_sdk::env::sha256(code); + near_sdk::base64::prelude::BASE64_STANDARD.encode(hash) +} + /// Smoke test of contract setup. #[tokio::test] async fn test_setup() -> anyhow::Result<()> { @@ -332,7 +340,7 @@ async fn test_staging_empty_code_clears_storage() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Verify staging empty code removes it. let res = setup @@ -436,11 +444,72 @@ async fn test_deploy_code_without_delay() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; + + // Deploy staged code. + let res = setup + .upgradable_contract + .up_deploy_code(&dao, convert_code_to_deploy_hash(&code), None) + .await?; + assert_success_with_unit_return(res); + + Ok(()) +} + +#[tokio::test] +async fn test_deploy_code_with_hash_success() -> anyhow::Result<()> { + let worker = near_workspaces::sandbox().await?; + let dao = worker.dev_create_account().await?; + let setup = Setup::new(worker.clone(), Some(dao.id().clone()), None).await?; + + // Stage some code. + let code = vec![1, 2, 3]; + let res = setup + .upgradable_contract + .up_stage_code(&dao, code.clone()) + .await?; + assert_success_with_unit_return(res); + setup.assert_staged_code(Some(&code)).await; // Deploy staged code. - let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?; + let hash = convert_code_to_deploy_hash(&code); + let res = setup + .upgradable_contract + .up_deploy_code(&dao, hash, None) + .await?; + assert_success_with_unit_return(res); + + Ok(()) +} + +/// Verifies failure of `up_deploy_code(hash, ...)` when `hash` does not correspond to the +/// hash of staged code. +#[tokio::test] +async fn test_deploy_code_with_hash_invalid_hash() -> anyhow::Result<()> { + let worker = near_workspaces::sandbox().await?; + let dao = worker.dev_create_account().await?; + let setup = Setup::new(worker.clone(), Some(dao.id().clone()), None).await?; + + // Stage some code. + let code = vec![1, 2, 3]; + let res = setup + .upgradable_contract + .up_stage_code(&dao, code.clone()) + .await?; assert_success_with_unit_return(res); + setup.assert_staged_code(Some(&code)).await; + + // Deployment is aborted if an invalid hash is provided. + let res = setup + .upgradable_contract + .up_deploy_code(&dao, "invalid_hash".to_owned(), None) + .await?; + let actual_hash = convert_code_to_deploy_hash(&code); + let expected_err = format!( + "Upgradable: Cannot deploy due to wrong hash: expected hash: {}", + actual_hash + ); + assert_failure_with(res, &expected_err); Ok(()) } @@ -465,10 +534,13 @@ async fn test_deploy_code_and_call_method() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Deploy staged code. - let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?; + let res = setup + .upgradable_contract + .up_deploy_code(&dao, convert_code_to_deploy_hash(&code), None) + .await?; assert_success_with_unit_return(res); // The newly deployed contract defines the function `is_upgraded`. Calling it successfully @@ -502,7 +574,7 @@ async fn test_deploy_code_with_migration() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Deploy staged code and call the new contract's `migrate` method. let function_call_args = FunctionCallArgs { @@ -513,7 +585,11 @@ async fn test_deploy_code_with_migration() -> anyhow::Result<()> { }; let res = setup .upgradable_contract - .up_deploy_code(&dao, Some(function_call_args)) + .up_deploy_code( + &dao, + convert_code_to_deploy_hash(&code), + Some(function_call_args), + ) .await?; assert_success_with_unit_return(res); @@ -545,7 +621,7 @@ async fn test_deploy_code_with_migration_failure_rollback() -> anyhow::Result<() .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Deploy staged code and call the new contract's `migrate_with_failure` method. let function_call_args = FunctionCallArgs { @@ -556,7 +632,11 @@ async fn test_deploy_code_with_migration_failure_rollback() -> anyhow::Result<() }; let res = setup .upgradable_contract - .up_deploy_code(&dao, Some(function_call_args)) + .up_deploy_code( + &dao, + convert_code_to_deploy_hash(&code), + Some(function_call_args), + ) .await?; assert_failure_with(res, "Failing migration on purpose"); @@ -591,21 +671,23 @@ async fn test_deploy_code_in_batch_transaction_pitfall() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Construct the function call actions to be executed in a batch transaction. // Note that we are attaching a call to `migrate_with_failure`, which will fail. let fn_call_deploy = near_workspaces::operations::Function::new("up_deploy_code") - .args_json(json!({ "function_call_args": FunctionCallArgs { + .args_json(json!({ + "hash": convert_code_to_deploy_hash(&code), + "function_call_args": FunctionCallArgs { function_name: "migrate_with_failure".to_string(), arguments: Vec::new(), amount: NearToken::from_yoctonear(0), gas: Gas::from_tgas(2), } })) - .gas(Gas::from_tgas(201)); + .gas(Gas::from_tgas(220)); let fn_call_remove_code = near_workspaces::operations::Function::new("up_stage_code") .args_borsh(Vec::::new()) - .gas(Gas::from_tgas(90)); + .gas(Gas::from_tgas(80)); let res = dao .batch(setup.contract.id()) @@ -658,13 +740,16 @@ async fn test_deploy_code_with_delay() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Let the staging duration pass. fast_forward_beyond(&worker, staging_duration).await; // Deploy staged code. - let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?; + let res = setup + .upgradable_contract + .up_deploy_code(&dao, convert_code_to_deploy_hash(&code), None) + .await?; assert_success_with_unit_return(res); Ok(()) @@ -688,13 +773,16 @@ async fn test_deploy_code_with_delay_failure_too_early() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Let some time pass but not enough. fast_forward_beyond(&worker, sdk_duration_from_secs(1)).await; // Verify trying to deploy staged code fails. - let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?; + let res = setup + .upgradable_contract + .up_deploy_code(&dao, convert_code_to_deploy_hash(&code), None) + .await?; assert_failure_with(res, ERR_MSG_DEPLOY_CODE_TOO_EARLY); // Verify `code` wasn't deployed by calling a function that is defined only in the initial @@ -717,13 +805,17 @@ async fn test_deploy_code_permission_failure() -> anyhow::Result<()> { .up_stage_code(&dao, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Only the roles passed as `code_deployers` to the `Upgradable` derive macro may successfully // call this method. let res = setup .upgradable_contract - .up_deploy_code(&setup.unauth_account, None) + .up_deploy_code( + &setup.unauth_account, + convert_code_to_deploy_hash(&code), + None, + ) .await?; assert_insufficient_acl_permissions( res, @@ -762,7 +854,10 @@ async fn test_deploy_code_empty_failure() -> anyhow::Result<()> { // The staging timestamp is set when staging code and removed when unstaging code. So when there // is no code staged, there is no staging timestamp. Hence the error message regarding a missing // staging timestamp is expected. - let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?; + let res = setup + .upgradable_contract + .up_deploy_code(&dao, "".to_owned(), None) + .await?; assert_failure_with(res, ERR_MSG_NO_STAGING_TS); Ok(()) @@ -1003,14 +1098,18 @@ async fn test_acl_permission_scope() -> anyhow::Result<()> { .up_stage_code(&code_stager, code.clone()) .await?; assert_success_with_unit_return(res); - setup.assert_staged_code(Some(code)).await; + setup.assert_staged_code(Some(&code)).await; // Verify `code_stager` is not authorized to deploy staged code. Only grantees of at least one // of the roles passed as `code_deployers` to the `Upgradable` derive macro are authorized to // deploy code. let res = setup .upgradable_contract - .up_deploy_code(&setup.unauth_account, None) + .up_deploy_code( + &setup.unauth_account, + convert_code_to_deploy_hash(&code), + None, + ) .await?; assert_insufficient_acl_permissions( res, diff --git a/near-plugins/src/upgradable.rs b/near-plugins/src/upgradable.rs index 1df00ad..a32b020 100644 --- a/near-plugins/src/upgradable.rs +++ b/near-plugins/src/upgradable.rs @@ -102,6 +102,17 @@ pub trait Upgradable { /// Allows an authorized account to deploy the staged code. It panics if no code is staged. /// + /// # Verifying the hash of staged code + /// + /// Some workflows (e.g. when a DAO interacts with an `Upgradable` contract) are facilitated if + /// deployment succeeds only in case the hash of staged code corresponds to a given hash. This + /// behavior can be enabled with the `hash` parameter. In case it is `h`, the deployment + /// succeeds only if `h` equals the base64 encoded string of the staged code's `sha256` hash. In + /// particular, the encoding according to [`near_sdk::base64::encode`] is expected. Note that + /// `near_sdk` uses a rather dated version of the `base64` crate whose API differs from current + /// versions. + /// + /// /// # Attaching a function call /// /// If `function_call_args` are provided, code is deployed in a batch promise that contains the @@ -143,7 +154,11 @@ pub trait Upgradable { /// [asynchronous design]: https://docs.near.org/concepts/basics/transactions/overview /// [state migration]: https://docs.near.org/develop/upgrade#migrating-the-state /// [storage staked]: https://docs.near.org/concepts/storage/storage-staking#btw-you-can-remove-data-to-unstake-some-tokens - fn up_deploy_code(&mut self, function_call_args: Option) -> Promise; + fn up_deploy_code( + &mut self, + hash: String, + function_call_args: Option, + ) -> Promise; /// Initializes the duration of the delay for deploying the staged code. It defaults to zero if /// code is staged before the staging duration is initialized. Once the staging duration has