near: internal feedback

Author: Reisen

target_chains/near/receiver/src/governance.rs

@@ -66,7 +66,8 @@ pub enum GovernanceModule {
 }
 
 /// A `GovernanceAction` represents the different actions that can be voted on and executed by the
-/// governance system.
+/// governance system. Note that this implementation is NEAR specific, for example within the
+/// UpgradeContract variant we use a codehash unlike a code_id in Cosmwasm, or a Pubkey in Solana.
 ///
 /// [ref:chain_structure] This type uses a [u8; 32] for contract upgrades which differs from other
 /// chains, see the reference for more details.
@@ -269,11 +270,6 @@ impl Pyth {
             // Convert to local VAA type to catch API changes.
             let vaa = Vaa::from(vaa);
 
-            ensure!(
-                self.executed_governance_vaa < vaa.sequence,
-                VaaVerificationFailed
-            );
-
             // Confirm the VAA is coming from a trusted source chain.
             ensure!(
                 self.gov_source
@@ -339,6 +335,16 @@ impl Pyth {
             InvalidPayload
         );
 
+        // Ensure the VAA is ahead in sequence, this check is here instead of during
+        // `execute_governance_instruction` as otherwise someone would be able to slip
+        // competing actions into the execution stream before the sequence is updated.
+        ensure!(
+            self.executed_governance_vaa < vaa.sequence,
+            VaaVerificationFailed
+        );
+
+        self.executed_governance_vaa = vaa.sequence;
+
         match GovernanceInstruction::deserialize(rest)?.action {
             SetDataSources { data_sources } => self.set_sources(data_sources),
             SetFee { base, expo } => self.set_update_fee(base, expo)?,
@@ -360,10 +366,13 @@ impl Pyth {
             AuthorizeGovernanceDataSourceTransfer { claim_vaa } => {
                 let claim_vaa = hex::encode(claim_vaa);
 
-                // Return early, the callback has to perform the rest of the processing.
+                // Return early, the callback has to perform the rest of the processing. Normally
+                // VAA processing will complete and the code below this match statement will
+                // execute. But because the VAA verification is async, we must return here instead
+                // and the logic below is duplicated within the authorize_gov_source_transfer function.
                 return Ok(PromiseOrValue::Promise(
                     ext_wormhole::ext(self.wormhole.clone())
-                        .with_static_gas(Gas(30_000_000_000_000))
+                        .with_static_gas(Gas(10_000_000_000_000))
                         .verify_vaa(claim_vaa.clone())
                         .then(
                             Self::ext(env::current_account_id())
@@ -384,14 +393,13 @@ impl Pyth {
             }
         }
 
-        self.executed_governance_vaa = vaa.sequence;
-
         // Refund storage difference to `account_id` after storage execution.
-        self.refund_storage_usage(
+        Self::refund_storage_usage(
             account_id,
             storage,
             env::storage_usage(),
             env::attached_deposit(),
+            None,
         )
         .map(|v| PromiseOrValue::Value(v))
     }
@@ -416,6 +424,9 @@ impl Pyth {
         storage: u64,
         #[callback_result] _result: Result<u32, near_sdk::PromiseError>,
     ) -> Result<(), Error> {
+        // If VAA verification failed we should bail.
+        ensure!(is_promise_success(), VaaVerificationFailed);
+
         let vaa = hex::decode(claim_vaa).map_err(|_| InvalidPayload)?;
         let (vaa, rest): (wormhole::Vaa<()>, _) =
             serde_wormhole::from_slice_with_payload(&vaa).expect("Failed to deserialize VAA");
@@ -459,11 +470,12 @@ impl Pyth {
         }
 
         // Refund storage difference to `account_id` after storage execution.
-        self.refund_storage_usage(
+        Self::refund_storage_usage(
             account_id,
             storage,
             env::storage_usage(),
             env::attached_deposit(),
+            None,
         )
     }
 
@@ -503,7 +515,7 @@ impl Pyth {
         amount: u128,
         storage: u64,
     ) -> Result<(), Error> {
-        self.refund_storage_usage(account_id, storage, env::storage_usage(), amount)
+        Self::refund_storage_usage(account_id, storage, env::storage_usage(), amount, None)
     }
 }
 

target_chains/near/receiver/src/lib.rs

@@ -148,7 +148,7 @@ impl Pyth {
     #[init(ignore_state)]
     pub fn migrate() -> Self {
         // This currently deserializes and produces the same state, I.E migration is a no-op to the
-        // current state. We only update the codehash to prevent re-upgrading.
+        // current state.
         //
         // In the case where we want to actually migrate to a new state, we can do this by defining
         // the old State struct here and then deserializing into that, then migrating into the new
@@ -170,7 +170,10 @@ impl Pyth {
         //
         //     // Construct new Pyth State from old, perform any migrations needed.
         //     let old: OldPyth = env::state_read().expect("Failed to read state");
+        //
         //     Self {
+        //        sources:    old.sources,
+        //        gov_source: old.gov_source,
         //        ...
         //     }
         // }
@@ -264,14 +267,7 @@ impl Pyth {
         // forces the caller to add the required fee to the deposit. The protocol defines the fee
         // as a u128, but storage is a u64, so we need to check that the fee does not overflow the
         // storage cost as well.
-        let storage = (env::storage_usage() as u128)
-            .checked_sub(
-                self.update_fee
-                    .checked_div(env::storage_byte_cost())
-                    .ok_or(Error::ArithmeticOverflow)?,
-            )
-            .ok_or(Error::InsufficientDeposit)
-            .and_then(|s| u64::try_from(s).map_err(|_| Error::ArithmeticOverflow))?;
+        let storage = env::storage_usage();
 
         // Deserialize VAA, note that we already deserialized and verified the VAA in `process_vaa`
         // at this point so we only care about the `rest` component which contains bytes we can
@@ -318,11 +314,12 @@ impl Pyth {
         );
 
         // Refund storage difference to `account_id` after storage execution.
-        self.refund_storage_usage(
+        Self::refund_storage_usage(
             account_id,
             storage,
             env::storage_usage(),
             env::attached_deposit(),
+            Some(self.update_fee),
         )
     }
 
@@ -408,11 +405,12 @@ impl Pyth {
         );
 
         // Refund storage difference to `account_id` after storage execution.
-        self.refund_storage_usage(
+        Self::refund_storage_usage(
             account_id,
             storage,
             env::storage_usage(),
             env::attached_deposit(),
+            Some(self.update_fee),
         )
     }
 
@@ -586,18 +584,22 @@ impl Pyth {
         }
     }
 
-    /// Checks storage usage invariants and additionally refunds the caller if they overpay.
+    /// Checks storage usage invariants and additionally refunds the caller if they overpay. This
+    /// method can optionally charge a fee to the caller which is removed from their deposit during
+    /// refund.
     fn refund_storage_usage(
-        &self,
         recipient: AccountId,
         before: StorageUsage,
         after: StorageUsage,
         deposit: Balance,
+        additional_fee: Option<Balance>,
     ) -> Result<(), Error> {
+        let fee = additional_fee.unwrap_or(0);
+
         if let Some(diff) = after.checked_sub(before) {
             // Handle storage increases if checked_sub succeeds.
             let cost = Balance::from(diff);
-            let cost = cost * env::storage_byte_cost();
+            let cost = (cost * env::storage_byte_cost()) + fee;
 
             // If the cost is higher than the deposit we bail.
             if cost > deposit {
@@ -613,7 +615,7 @@ impl Pyth {
             // the amount reduced, as well the original deposit they sent.
             let refund = Balance::from(before - after);
             let refund = refund * env::storage_byte_cost();
-            let refund = refund + deposit;
+            let refund = refund + deposit - fee;
             Promise::new(recipient).transfer(refund);
         }
 

target_chains/near/receiver/tests/workspaces.rs

@@ -239,10 +239,9 @@ async fn test_set_governance_source() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact_async()
+        .transact()
         .await
         .expect("Failed to submit VAA")
-        .await
         .unwrap()
         .failures()
         .is_empty());