near: internal feedback

target_chains/near/receiver/Cargo.lock

@@ -2070,7 +2070,7 @@ dependencies = [
 
 [[package]]
 name = "pyth-wormhole-attester-sdk"
-version = "0.1.1"
+version = "0.1.2"
 dependencies = [
  "hex 0.4.3",
  "pyth-sdk 0.5.0",

target_chains/near/receiver/src/governance.rs

@@ -65,7 +65,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.
@@ -268,11 +269,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
@@ -340,6 +336,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)?,
@@ -378,19 +384,21 @@ impl Pyth {
                             .refund_vaa(env::predecessor_account_id(), env::attached_deposit()),
                     );
 
-                // 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(());
             }
         }
 
-        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,
         )
     }
 
@@ -413,6 +421,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, rest): (wormhole::Vaa<()>, _) =
             serde_wormhole::from_slice_with_payload(&claim_vaa).expect("Failed to deserialize VAA");
 
@@ -455,11 +466,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,
         )
     }
 
@@ -475,7 +487,9 @@ impl Pyth {
     #[handle_result]
     pub(crate) fn upgrade(&mut self, new_code: Vec<u8>) -> Result<Promise, Error> {
         let signature = TryInto::<[u8; 32]>::try_into(env::sha256(&new_code)).ok();
+        ensure!(self.codehash.is_some(), UnauthorizedUpgrade);
         ensure!(signature == self.codehash, UnauthorizedUpgrade);
+        self.codehash = None;
 
         Ok(Promise::new(env::current_account_id())
             .deploy_contract(new_code)
@@ -497,7 +511,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

@@ -123,7 +123,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
@@ -145,13 +145,15 @@ 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,
         //        ...
         //     }
         // }
         // ```
-        let mut state: Self = env::state_read().expect("Failed to read state");
-        state.codehash = None;
+        let state: Self = env::state_read().expect("Failed to read state");
         state
     }
 
@@ -230,14 +232,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
@@ -277,11 +272,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),
         )
     }
 
@@ -402,18 +398,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 {
@@ -429,7 +429,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);
         }