[near] Fix governance inconsistencies. (#453)

* near: update governance inconsistencies
* near: add small script for automated testing
* near: address feedback and update tests
* near: cannot trust outcome(), use failure len

target-chains/near/receiver/.gitignore

@@ -0,0 +1,2 @@
+# Ignore *wasm build artifacts which may enter the directory for workspaces tests.
+*.wasm

target-chains/near/receiver/Cargo.toml

@@ -13,9 +13,13 @@ crate-type      = ["cdylib", "lib"]
 byteorder       = { version = "1.4.3" }
 hex             = { version = "0.4.3" }
 near-sdk        = { version = "4.1.1" }
-p2w-sdk         = { path = "../../third_party/pyth/p2w-sdk/rust" }
+nom             = { version = "7.1.2" }
+num-traits      = { version = "0.2.15" }
+num-derive      = { version = "0.3.3" }
+p2w-sdk         = { path = "../../../third_party/pyth/p2w-sdk/rust" }
 pyth-sdk        = { version = "0.7.0" }
 serde_wormhole  = { git = "https://github.com/wormhole-foundation/wormhole" }
+strum           = { version = "0.24.1", features = ["derive"] }
 thiserror       = { version = "1.0.38" }
 wormhole-core   = { git = "https://github.com/wormhole-foundation/wormhole" }
 

target-chains/near/receiver/src/error.rs

@@ -6,6 +6,19 @@ use {
     thiserror::Error,
 };
 
+/// Small macro for throwing errors in the contract when a boolean condition is not met.
+///
+/// It would be nice to have anyhow::ensure!() here, but the contract acts as a library and a
+/// concrete error type is a better API for this case.
+#[macro_export]
+macro_rules! ensure {
+    ($cond:expr, $err:expr) => {
+        if !$cond {
+            return Err($err);
+        }
+    };
+}
+
 #[derive(Error, Debug, Serialize, FunctionError)]
 #[serde(crate = "near_sdk::serde")]
 pub enum Error {
@@ -18,10 +31,16 @@ pub enum Error {
     #[error("A VAA payload could not be deserialized.")]
     InvalidPayload,
 
+    #[error("Governance Module ID not valid.")]
+    InvalidGovernanceModule,
+
+    #[error("Governance Module Action not valid.")]
+    InvalidGovernanceAction,
+
     #[error("Source for attestation is not allowed.")]
     UnknownSource,
 
-    #[error("Unauthorized Upgrade")]
+    #[error("Unauthorized Upgrade.")]
     UnauthorizedUpgrade,
 
     #[error("Insufficient tokens deposited to cover storage.")]
@@ -32,6 +51,12 @@ pub enum Error {
 
     #[error("Fee is too large.")]
     FeeTooLarge,
+
+    #[error("Arithmetic overflow.")]
+    ArithmeticOverflow,
+
+    #[error("Unknown error.")]
+    Unknown,
 }
 
 /// Convert IO errors into Payload errors, the only I/O we do is parsing with `Cursor` so this is a
@@ -41,3 +66,10 @@ impl From<std::io::Error> for Error {
         Error::InvalidPayload
     }
 }
+
+/// Convert `nom` errors into local crate `InvalidPayload` errors.
+impl From<nom::Err<nom::error::Error<&[u8]>>> for Error {
+    fn from(_: nom::Err<nom::error::Error<&[u8]>>) -> Self {
+        Error::InvalidPayload
+    }
+}

target-chains/near/receiver/src/governance.rs

@@ -2,7 +2,11 @@
 
 use {
     crate::{
-        error::Error,
+        ensure,
+        error::Error::{
+            self,
+            *,
+        },
         ext::ext_wormhole,
         state::{
             Chain,
@@ -12,11 +16,6 @@ use {
         Pyth,
         PythExt,
     },
-    byteorder::{
-        BigEndian,
-        ReadBytesExt,
-        WriteBytesExt,
-    },
     near_sdk::{
         borsh::{
             self,
@@ -34,162 +33,218 @@ use {
         Gas,
         Promise,
     },
-    std::io::Read,
+    num_traits::FromPrimitive,
+    strum::EnumDiscriminants,
     wormhole::Chain as WormholeChain,
 };
 
 /// Magic Header for identifying Governance VAAs.
-const GOVERNANCE_MAGIC: [u8; 4] = [0x50, 0x54, 0x47, 0x4d];
+const GOVERNANCE_MAGIC: [u8; 4] = *b"PTGM";
 
-/// ID for the module this contract identifies as: Pyth Receiver (0x1).
+/// The type of contract that can accept a governance instruction.
-const GOVERNANCE_MODULE: u8 = 0x01;
+#[derive(
-
+    BorshDeserialize,
-/// Enumeration of IDs for different governance actions.
+    BorshSerialize,
+    Clone,
+    Copy,
+    Debug,
+    Deserialize,
+    Eq,
+    PartialEq,
+    Serialize,
+    num_derive::FromPrimitive,
+    num_derive::ToPrimitive,
+)]
+#[serde(crate = "near_sdk::serde")]
 #[repr(u8)]
-pub enum ActionId {
+pub enum GovernanceModule {
-    ContractUpgrade        = 0,
+    /// The PythNet executor contract
-    SetDataSources         = 1,
+    Executor = 0,
-    SetGovernanceSource    = 2,
+    /// A target chain contract (like this one!)
-    SetStalePriceThreshold = 3,
+    Target   = 1,
-    SetUpdateFee           = 4,
-}
-
-impl TryInto<ActionId> for u8 {
-    type Error = Error;
-    fn try_into(self) -> Result<ActionId, Error> {
-        match self {
-            0 => Ok(ActionId::ContractUpgrade),
-            1 => Ok(ActionId::SetDataSources),
-            2 => Ok(ActionId::SetGovernanceSource),
-            3 => Ok(ActionId::SetStalePriceThreshold),
-            4 => Ok(ActionId::SetUpdateFee),
-            _ => Err(Error::InvalidPayload),
-        }
-    }
-}
-
-impl From<ActionId> for u8 {
-    fn from(val: ActionId) -> Self {
-        val as u8
-    }
 }
 
 /// A `GovernanceAction` represents the different actions that can be voted on and executed by the
 /// governance system.
-#[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize)]
+///
+/// [ref:chain_structure] This type uses a [u8; 32] for contract upgrades which differs from other
+/// chains, see the reference for more details.
+///
+/// [ref:action_discriminants] The discriminants for this enum are duplicated into a separate enum
+/// containing only the discriminants with no fields called `GovernanceActionId`. This allow for
+/// type-safe matching IDs during deserialization. When new actions are added, this will force the
+/// developer to update the parser.
+#[derive(
+    BorshDeserialize,
+    BorshSerialize,
+    Debug,
+    Deserialize,
+    EnumDiscriminants,
+    Eq,
+    PartialEq,
+    Serialize,
+)]
+#[strum_discriminants(derive(num_derive::ToPrimitive, num_derive::FromPrimitive))]
+#[strum_discriminants(name(GovernanceActionId))]
 #[serde(crate = "near_sdk::serde")]
 pub enum GovernanceAction {
-    ContractUpgrade([u8; 32]),
+    UpgradeContract { codehash: [u8; 32] },
-    SetDataSources(Vec<Source>),
+    AuthorizeGovernanceDataSourceTransfer { claim_vaa: Vec<u8> },
-    SetGovernanceSource(Source),
+    SetDataSources { data_sources: Vec<Source> },
-    SetStalePriceThreshold(u64),
+    SetFee { base: u64, expo: u64 },
-    SetUpdateFee(u64),
+    SetValidPeriod { valid_seconds: u64 },
+    RequestGovernanceDataSourceTransfer { governance_data_source_index: u32 },
 }
 
-impl GovernanceAction {
+#[derive(BorshDeserialize, BorshSerialize, Debug, Deserialize, Eq, PartialEq, Serialize)]
-    pub fn id(&self) -> ActionId {
+#[serde(crate = "near_sdk::serde")]
-        match self {
+pub struct GovernanceInstruction {
-            GovernanceAction::ContractUpgrade(_) => ActionId::ContractUpgrade,
+    pub module: GovernanceModule,
-            GovernanceAction::SetDataSources(_) => ActionId::SetDataSources,
+    pub action: GovernanceAction,
-            GovernanceAction::SetGovernanceSource(_) => ActionId::SetGovernanceSource,
+    pub target: Chain,
-            GovernanceAction::SetStalePriceThreshold(_) => ActionId::SetStalePriceThreshold,
+}
-            GovernanceAction::SetUpdateFee(_) => ActionId::SetUpdateFee,
+
+impl GovernanceInstruction {
+    /// Implements a `deserialize` method for the `GovernanceAction` enum using `nom` to
+    /// deserialize the payload. The use of `nom` gives us parser safety, error handling, full
+    /// buffer consumption, and a more readable implementation while staying efficient.
+    pub fn deserialize(input: impl AsRef<[u8]>) -> Result<Self, Error> {
+        use nom::{
+            bytes::complete::take,
+            combinator::all_consuming,
+            multi::length_count,
+            number::complete::{
+                be_u16,
+                be_u32,
+                be_u64,
+                be_u8,
+            },
+        };
+
+        let input = input.as_ref();
+
+        // Verify Governance header is as expected so we can bail to avoid more parsing.
+        let (input, magic) = take(4usize)(input)?;
+        let (input, module) = be_u8(input)?;
+        let (input, action) = be_u8(input)?;
+        let (input, chain) = be_u16(input)?;
+        let module = GovernanceModule::from_u8(module).ok_or(InvalidGovernanceModule)?;
+        let chain = Chain::from(WormholeChain::from(chain));
+
+        // Safely parse the action ID. [ref:action_discriminants]
+        let action = GovernanceActionId::from_u8(action).ok_or(InvalidGovernanceAction)?;
+
+        ensure!(magic == GOVERNANCE_MAGIC, InvalidGovernanceModule);
+        ensure!(module == GovernanceModule::Target, InvalidGovernanceModule);
+
+        Ok(GovernanceInstruction {
+            module,
+            target: chain,
+            action: match action {
+                GovernanceActionId::UpgradeContract => {
+                    let (_input, bytes) = all_consuming(take(32usize))(input)?;
+                    let mut codehash = [0u8; 32];
+                    codehash.copy_from_slice(bytes);
+                    GovernanceAction::UpgradeContract { codehash }
+                }
+
+                GovernanceActionId::AuthorizeGovernanceDataSourceTransfer => {
+                    let (_input, claim_vaa) = all_consuming(take(input.len()))(input)?;
+                    GovernanceAction::AuthorizeGovernanceDataSourceTransfer {
+                        claim_vaa: claim_vaa.to_vec(),
                     }
                 }
 
-    pub fn deserialize(data: &[u8]) -> Result<Self, Error> {
+                GovernanceActionId::SetDataSources => {
-        let mut cursor = std::io::Cursor::new(data);
+                    let (_input, data_sources) = all_consuming(length_count(be_u8, |input| {
-        let magic = cursor.read_u32::<BigEndian>()?;
+                        let (input, chain) = be_u16(input)?;
-        let module = cursor.read_u8()?;
+                        let (input, bytes) = take(32usize)(input)?;
-        let action = cursor.read_u8()?.try_into()?;
+                        let chain = Chain::from(WormholeChain::from(chain));
-        let target = cursor.read_u16::<BigEndian>()?;
-
-        assert!(module == GOVERNANCE_MODULE);
-        assert!(target == 0 || target == u16::from(WormholeChain::Near));
-        assert!(magic == u32::from_le_bytes(GOVERNANCE_MAGIC));
-
-        Ok(match action {
-            ActionId::ContractUpgrade => {
-                let mut hash = [0u8; 32];
-                cursor.read_exact(&mut hash)?;
-                Self::ContractUpgrade(hash)
-            }
-
-            ActionId::SetDataSources => {
-                let mut sources = Vec::new();
-                let count = cursor.read_u8()?;
-
-                for _ in 0..count {
                         let mut emitter = [0u8; 32];
-                    cursor.read_exact(&mut emitter)?;
+                        emitter.copy_from_slice(bytes);
-                    sources.push(Source {
+                        Ok((input, Source { chain, emitter }))
-                        emitter,
+                    }))(input)?;
-                        pyth_emitter_chain: Chain::from(WormholeChain::from(
+                    GovernanceAction::SetDataSources { data_sources }
-                            cursor.read_u16::<BigEndian>()?,
-                        )),
-                    });
                 }
 
-                Self::SetDataSources(sources)
+                GovernanceActionId::SetFee => {
+                    let (_input, (val, expo)) = all_consuming(|input| {
+                        let (input, val) = be_u64(input)?;
+                        let (input, expo) = be_u64(input)?;
+                        Ok((input, (val, expo)))
+                    })(input)?;
+                    GovernanceAction::SetFee { base: val, expo }
                 }
 
-            ActionId::SetGovernanceSource => {
+                GovernanceActionId::SetValidPeriod => {
-                let mut emitter = [0u8; 32];
+                    let (_input, valid_seconds) = all_consuming(be_u64)(input)?;
-                cursor.read_exact(&mut emitter)?;
+                    GovernanceAction::SetValidPeriod { valid_seconds }
-                Self::SetGovernanceSource(Source {
+                }
-                    emitter,
+
-                    pyth_emitter_chain: Chain(cursor.read_u16::<BigEndian>()?),
+                GovernanceActionId::RequestGovernanceDataSourceTransfer => {
+                    let (_input, governance_data_source_index) = all_consuming(be_u32)(input)?;
+                    GovernanceAction::RequestGovernanceDataSourceTransfer {
+                        governance_data_source_index,
+                    }
+                }
+            },
         })
     }
 
-            ActionId::SetStalePriceThreshold => {
+    /// Implements a `serialize` method for the `GovernanceAction` enum. The `nom` library doesn't
-                let stale_price_threshold = cursor.read_u64::<BigEndian>()?;
+    /// provide serialization but serialization is a safer operation, so we can just use a simple
-                Self::SetStalePriceThreshold(stale_price_threshold)
+    /// push buffer to serialize.
+    pub fn serialize(&self) -> Result<Vec<u8>, Error> {
+        let mut buf = Vec::new();
+        buf.extend_from_slice(&GOVERNANCE_MAGIC);
+        buf.push(self.module as u8);
+
+        match &self.action {
+            GovernanceAction::UpgradeContract { codehash } => {
+                buf.push(GovernanceActionId::UpgradeContract as u8);
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
+                buf.extend_from_slice(codehash);
             }
 
-            ActionId::SetUpdateFee => {
+            GovernanceAction::AuthorizeGovernanceDataSourceTransfer { claim_vaa } => {
-                let update_fee = cursor.read_u64::<BigEndian>()?;
+                buf.push(GovernanceActionId::AuthorizeGovernanceDataSourceTransfer as u8);
-                Self::SetUpdateFee(update_fee)
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
-            }
+                buf.extend_from_slice(claim_vaa);
-        })
             }
 
-    pub fn serialize(&self) -> Vec<u8> {
+            GovernanceAction::SetDataSources { data_sources } => {
-        let mut data = Vec::new();
+                buf.push(GovernanceActionId::SetDataSources as u8);
-        let magic = u32::from_le_bytes(GOVERNANCE_MAGIC);
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
-        data.write_u32::<BigEndian>(magic).unwrap();
+                buf.push(u8::try_from(data_sources.len()).map_err(|_| InvalidPayload)?);
-        data.push(GOVERNANCE_MODULE);
+                for source in data_sources {
-        data.push(self.id() as u8);
+                    buf.extend_from_slice(&(u16::from(source.chain).to_be_bytes()));
-        data.extend_from_slice(&0u16.to_le_bytes());
+                    buf.extend_from_slice(&source.emitter);
-
-        match self {
-            Self::ContractUpgrade(hash) => {
-                data.extend_from_slice(hash);
-            }
-
-            Self::SetDataSources(sources) => {
-                data.push(sources.len() as u8);
-                for source in sources {
-                    data.extend_from_slice(&source.emitter);
-                    data.extend_from_slice(&source.pyth_emitter_chain.0.to_le_bytes());
                 }
             }
 
-            Self::SetGovernanceSource(source) => {
+            GovernanceAction::SetFee { base: val, expo } => {
-                data.extend_from_slice(&source.emitter);
+                buf.push(GovernanceActionId::SetFee as u8);
-                data.extend_from_slice(&source.pyth_emitter_chain.0.to_le_bytes());
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
+                buf.extend_from_slice(&val.to_be_bytes());
+                buf.extend_from_slice(&expo.to_be_bytes());
             }
 
-            Self::SetStalePriceThreshold(stale_price_threshold) => {
+            GovernanceAction::SetValidPeriod { valid_seconds } => {
-                data.extend_from_slice(&stale_price_threshold.to_le_bytes());
+                buf.push(GovernanceActionId::SetValidPeriod as u8);
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
+                buf.extend_from_slice(&valid_seconds.to_be_bytes());
             }
 
-            Self::SetUpdateFee(update_fee) => {
+            GovernanceAction::RequestGovernanceDataSourceTransfer {
-                data.extend_from_slice(&update_fee.to_le_bytes());
+                governance_data_source_index,
+            } => {
+                buf.push(GovernanceActionId::RequestGovernanceDataSourceTransfer as u8);
+                buf.extend_from_slice(&u16::from(self.target).to_be_bytes());
+                buf.extend_from_slice(&governance_data_source_index.to_be_bytes());
             }
         }
 
-        data
+        Ok(buf)
     }
 }
 
@@ -205,33 +260,28 @@ impl Pyth {
         // Verify the VAA is coming from a trusted source chain before attempting to verify VAA
         // signatures. Avoids a cross-contract call early.
         {
-            let vaa = hex::decode(&vaa).map_err(|_| Error::InvalidHex)?;
+            let vaa = hex::decode(&vaa).map_err(|_| InvalidHex)?;
             let vaa = serde_wormhole::from_slice_with_payload::<wormhole::Vaa<()>>(&vaa);
-            let vaa = vaa.map_err(|_| Error::InvalidVaa)?;
+            let vaa = vaa.map_err(|_| InvalidVaa)?;
             let (vaa, _rest) = vaa;
 
-            // Convert to local VAA type to catch APi changes.
+            // Convert to local VAA type to catch API changes.
             let vaa = Vaa::from(vaa);
 
-            // Prevent VAA re-execution.
+            ensure!(
-            if self.executed_gov_sequences.contains(&vaa.sequence) {
+                self.executed_governance_vaa < vaa.sequence,
-                return Err(Error::VaaVerificationFailed);
+                VaaVerificationFailed
-            }
+            );
 
             // Confirm the VAA is coming from a trusted source chain.
-            if self.gov_source
+            ensure!(
-                != (Source {
+                self.gov_source
+                    == (Source {
                         emitter: vaa.emitter_address,
-                    pyth_emitter_chain: vaa.emitter_chain,
+                        chain:   vaa.emitter_chain,
-                })
+                    }),
-            {
+                UnknownSource
-                return Err(Error::UnknownSource);
+            );
-            }
-
-            // Insert before calling Wormhole to prevent re-execution. If we wait until after the
-            // Wormhole call we could end up with multiple VAA's with the same sequence being
-            // executed in parallel.
-            self.executed_gov_sequences.insert(&vaa.sequence);
         }
 
         // Verify VAA and refund the caller in case of failure.
@@ -254,6 +304,10 @@ impl Pyth {
     }
 
     /// Invoke handler upon successful verification of a VAA action.
+    ///
+    /// IMPORTANT: These functions should be idempotent otherwise NEAR's async model would allow
+    /// for two VAA's to run in parallel before the sequence is updated. Another fix for this would
+    /// be to pass the previous index and update during failure.
     #[payable]
     #[private]
     #[handle_result]
@@ -265,9 +319,7 @@ impl Pyth {
     ) -> Result<(), Error> {
         use GovernanceAction::*;
 
-        if !is_promise_success() {
+        ensure!(is_promise_success(), VaaVerificationFailed);
-            return Err(Error::VaaVerificationFailed);
-        }
 
         // Get Storage Usage before execution.
         let storage = env::storage_usage();
@@ -275,17 +327,39 @@ impl Pyth {
         // 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
         // deserialize into an Action.
-        let vaa = hex::decode(&vaa).unwrap();
+        let vaa = hex::decode(vaa).map_err(|_| InvalidPayload)?;
-        let (_, rest): (wormhole::Vaa<()>, _) =
+        let (vaa, rest): (wormhole::Vaa<()>, _) =
-            serde_wormhole::from_slice_with_payload(&vaa).map_err(|_| Error::InvalidPayload)?;
+            serde_wormhole::from_slice_with_payload(&vaa).map_err(|_| InvalidPayload)?;
 
-        match GovernanceAction::deserialize(rest)? {
+        // Deserialize and verify the action is destined for this chain.
-            ContractUpgrade(codehash) => self.set_upgrade_hash(codehash),
+        let instruction = GovernanceInstruction::deserialize(rest)?;
-            SetDataSources(sources) => self.set_sources(sources),
+
-            SetGovernanceSource(source) => self.set_gov_source(source),
+        ensure!(
-            SetStalePriceThreshold(threshold) => self.set_stale_price_threshold(threshold),
+            instruction.target == Chain::from(WormholeChain::Near)
-            SetUpdateFee(fee) => self.set_update_fee(fee),
+                || instruction.target == Chain::from(WormholeChain::Any),
+            InvalidPayload
+        );
+
+        match GovernanceInstruction::deserialize(rest)?.action {
+            SetDataSources { data_sources } => self.set_sources(data_sources),
+            SetFee { base, expo } => self.set_update_fee(base, expo)?,
+            SetValidPeriod { valid_seconds } => self.set_valid_period(valid_seconds),
+            RequestGovernanceDataSourceTransfer { .. } => Err(InvalidPayload)?,
+            AuthorizeGovernanceDataSourceTransfer { claim_vaa } => {
+                self.authorize_gov_source_transfer(claim_vaa)?
             }
+            UpgradeContract { codehash } => {
+                // Additionally restrict to only Near for upgrades. This is a safety measure to
+                // prevent accidental upgrades to the wrong contract.
+                ensure!(
+                    instruction.target == Chain::from(WormholeChain::Near),
+                    InvalidPayload
+                );
+                self.set_upgrade_hash(codehash)
+            }
+        }
+
+        self.executed_governance_vaa = vaa.sequence;
 
         // Refund storage difference to `account_id` after storage execution.
         self.refund_storage_usage(
@@ -312,18 +386,67 @@ impl Pyth {
     }
 
     #[private]
-    pub fn set_gov_source(&mut self, source: Source) {
+    #[handle_result]
-        self.gov_source = source;
+    pub fn authorize_gov_source_transfer(&mut self, claim_vaa: Vec<u8>) -> Result<(), Error> {
+        let (vaa, rest): (wormhole::Vaa<()>, _) =
+            serde_wormhole::from_slice_with_payload(&claim_vaa).expect("Failed to deserialize VAA");
+
+        // Convert to local VAA type to catch API changes.
+        let vaa = Vaa::from(vaa);
+
+        // Parse GovernanceInstruction from Payload.
+        let instruction =
+            GovernanceInstruction::deserialize(rest).expect("Failed to deserialize action");
+
+        // Execute the embedded VAA action.
+        match instruction.action {
+            GovernanceAction::RequestGovernanceDataSourceTransfer {
+                governance_data_source_index,
+            } => {
+                ensure!(
+                    self.executed_governance_change_vaa < governance_data_source_index as u64,
+                    Unknown
+                );
+
+                // Additionally restrict to only Near for Authorizations.
+                ensure!(
+                    instruction.target == Chain::from(WormholeChain::Near)
+                        || instruction.target == Chain::from(WormholeChain::Any),
+                    InvalidPayload
+                );
+
+                self.executed_governance_change_vaa = governance_data_source_index as u64;
+
+                // Update Governance Source
+                self.gov_source = Source {
+                    emitter: vaa.emitter_address,
+                    chain:   vaa.emitter_chain,
+                };
+            }
+
+            _ => Err(Unknown)?,
+        }
+
+        Ok(())
     }
 
     #[private]
-    pub fn set_stale_price_threshold(&mut self, threshold: u64) {
+    pub fn set_valid_period(&mut self, threshold: u64) {
         self.stale_threshold = threshold;
     }
 
     #[private]
-    pub fn set_update_fee(&mut self, fee: u64) {
+    #[handle_result]
-        self.update_fee = fee;
+    pub fn set_update_fee(&mut self, fee: u64, expo: u64) -> Result<(), Error> {
+        self.update_fee = (fee as u128)
+            .checked_mul(
+                10_u128
+                    .checked_pow(u32::try_from(expo).map_err(|_| ArithmeticOverflow)?)
+                    .ok_or(ArithmeticOverflow)?,
+            )
+            .ok_or(ArithmeticOverflow)?;
+
+        Ok(())
     }
 
     #[private]
@@ -347,10 +470,7 @@ impl Pyth {
     #[handle_result]
     pub(crate) fn upgrade(&mut self, new_code: Vec<u8>) -> Result<Promise, Error> {
         let signature = env::sha256(&new_code);
-
+        ensure!(signature == self.codehash, UnauthorizedUpgrade);
-        if signature != self.codehash {
-            return Err(Error::UnauthorizedUpgrade);
-        }
 
         Ok(Promise::new(env::current_account_id())
             .deploy_contract(new_code)
@@ -380,6 +500,223 @@ impl Pyth {
     fn is_valid_governance_source(&self, source: &Source) -> Result<(), Error> {
         (self.gov_source == *source)
             .then_some(())
-            .ok_or(Error::UnknownSource)
+            .ok_or(UnknownSource)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use {
+        super::*,
+        crate::governance::GovernanceActionId,
+        near_sdk::{
+            test_utils::{
+                accounts,
+                VMContextBuilder,
+            },
+            testing_env,
+        },
+        std::io::{
+            Cursor,
+            Write,
+        },
+    };
+
+    fn get_context() -> VMContextBuilder {
+        let mut context = VMContextBuilder::new();
+        context
+            .current_account_id(accounts(0))
+            .signer_account_id(accounts(0))
+            .predecessor_account_id(accounts(0))
+            .attached_deposit(0)
+            .is_view(false);
+        context
+    }
+
+    #[test]
+    fn test_upgrade() {
+        let mut context = get_context();
+        context.is_view(false);
+        testing_env!(context.build());
+
+        let mut contract = Pyth::new(
+            near_sdk::AccountId::new_unchecked("pyth.near".to_owned()),
+            [0; 32],
+            Source::default(),
+            Source::default(),
+            0.into(),
+            32,
+        );
+
+        contract.codehash = env::sha256(&[1, 2, 3]).try_into().unwrap();
+        contract.upgrade(vec![1, 2, 3]).expect("Upgrade failed");
+    }
+
+    #[test]
+    #[should_panic(expected = "UnauthorizedUpgrade")]
+    fn test_upgrade_fail() {
+        let mut context = get_context();
+        context.is_view(false);
+        testing_env!(context.build());
+
+        let mut contract = Pyth::new(
+            near_sdk::AccountId::new_unchecked("pyth.near".to_owned()),
+            [0; 32],
+            Source::default(),
+            Source::default(),
+            0.into(),
+            32,
+        );
+
+        contract.codehash = env::sha256(&[1, 2, 3]).try_into().unwrap();
+        contract.upgrade(vec![1, 2, 3, 4]).expect("Upgrade failed");
+    }
+
+    #[test]
+    fn test_set_valid_period() {
+        let mut context = get_context();
+        context.is_view(false);
+        testing_env!(context.build());
+
+        let mut contract = Pyth::new(
+            near_sdk::AccountId::new_unchecked("pyth.near".to_owned()),
+            [0; 32],
+            Source::default(),
+            Source::default(),
+            0.into(),
+            32,
+        );
+
+        contract.set_valid_period(100);
+        assert_eq!(contract.stale_threshold, 100);
+    }
+
+    #[test]
+    fn test_set_update_fee() {
+        let mut context = get_context();
+        context.is_view(false);
+        testing_env!(context.build());
+
+        let mut contract = Pyth::new(
+            near_sdk::AccountId::new_unchecked("pyth.near".to_owned()),
+            [0; 32],
+            Source::default(),
+            Source::default(),
+            0.into(),
+            32,
+        );
+
+        contract.set_update_fee(100, 2).expect("Failed to set fee");
+        assert_eq!(contract.update_fee, 10000);
+    }
+
+    #[test]
+    fn test_governance_serialize_matches_deserialize() {
+        // We match on the GovernanceActionId so that when new variants are added the test is
+        // forced to be updated. There's nothing special about SetFee we just need a concrete value
+        // to match on.
+        match GovernanceActionId::SetFee {
+            GovernanceActionId::SetValidPeriod => {
+                let instruction = GovernanceInstruction {
+                    module: GovernanceModule::Target,
+                    target: Chain::from(WormholeChain::Near),
+                    action: GovernanceAction::SetValidPeriod { valid_seconds: 100 },
+                };
+
+                assert_eq!(
+                    instruction,
+                    GovernanceInstruction::deserialize(instruction.serialize().unwrap()).unwrap()
+                );
+            }
+
+            GovernanceActionId::SetDataSources => {
+                let instruction = GovernanceInstruction {
+                    module: GovernanceModule::Target,
+                    target: Chain::from(WormholeChain::Near),
+                    action: GovernanceAction::SetDataSources {
+                        data_sources: vec![Source::default()],
+                    },
+                };
+
+                assert_eq!(
+                    instruction,
+                    GovernanceInstruction::deserialize(instruction.serialize().unwrap()).unwrap()
+                );
+            }
+
+            GovernanceActionId::SetFee => {
+                let instruction = GovernanceInstruction {
+                    module: GovernanceModule::Target,
+                    target: Chain::from(WormholeChain::Near),
+                    action: GovernanceAction::SetFee {
+                        base: 100,
+                        expo: 100,
+                    },
+                };
+
+                assert_eq!(
+                    instruction,
+                    GovernanceInstruction::deserialize(instruction.serialize().unwrap()).unwrap()
+                );
+            }
+
+            GovernanceActionId::UpgradeContract => {
+                let instruction = GovernanceInstruction {
+                    module: GovernanceModule::Target,
+                    target: Chain::from(WormholeChain::Near),
+                    action: GovernanceAction::UpgradeContract { codehash: [1; 32] },
+                };
+
+                assert_eq!(
+                    instruction,
+                    GovernanceInstruction::deserialize(instruction.serialize().unwrap()).unwrap()
+                );
+            }
+
+            GovernanceActionId::AuthorizeGovernanceDataSourceTransfer => {
+                let vaa = {
+                    let vaa = wormhole::Vaa {
+                        emitter_chain: wormhole::Chain::Any,
+                        emitter_address: wormhole::Address([0; 32]),
+                        sequence: 1,
+                        payload: (),
+                        ..Default::default()
+                    };
+
+                    let mut cur = Cursor::new(Vec::new());
+                    serde_wormhole::to_writer(&mut cur, &vaa).expect("Failed to serialize VAA");
+                    cur.write_all(
+                        &GovernanceInstruction {
+                            target: Chain::from(WormholeChain::Near),
+                            module: GovernanceModule::Target,
+                            action: GovernanceAction::RequestGovernanceDataSourceTransfer {
+                                governance_data_source_index: 1,
+                            },
+                        }
+                        .serialize()
+                        .unwrap(),
+                    )
+                    .expect("Failed to write Payload");
+                    cur.into_inner()
+                };
+
+                let instruction = GovernanceInstruction {
+                    module: GovernanceModule::Target,
+                    target: Chain::from(WormholeChain::Near),
+                    action: GovernanceAction::AuthorizeGovernanceDataSourceTransfer {
+                        claim_vaa: vaa,
+                    },
+                };
+
+                assert_eq!(
+                    instruction,
+                    GovernanceInstruction::deserialize(instruction.serialize().unwrap()).unwrap()
+                );
+            }
+
+            GovernanceActionId::RequestGovernanceDataSourceTransfer => {
+                unimplemented!()
+            }
+        }
     }
 }

target-chains/near/receiver/src/lib.rs

@@ -15,6 +15,7 @@ use {
         },
         env,
         is_promise_success,
+        json_types::U128,
         log,
         near_bindgen,
         AccountId,
@@ -47,7 +48,6 @@ pub mod tests;
 enum StorageKeys {
     Source,
     Prices,
-    Governance,
 }
 
 /// The `State` contains all persisted state for the contract. This includes runtime configuration.
@@ -63,8 +63,11 @@ pub struct Pyth {
     /// The Governance Source.
     gov_source: Source,
 
-    /// The last executed sequence number for governance actions.
+    /// The last executed sequence number for all governance actions.
-    executed_gov_sequences: UnorderedSet<u64>,
+    executed_governance_vaa: u64,
+
+    /// The last executed sequence number only for governance change actions.
+    executed_governance_change_vaa: u64,
 
     /// A Mapping from PriceFeed ID to Price Info.
     prices: UnorderedMap<PriceIdentifier, PriceFeed>,
@@ -83,7 +86,7 @@ pub struct Pyth {
     stale_threshold: Duration,
 
     /// Fee for updating price.
-    update_fee: u64,
+    update_fee: u128,
 }
 
 #[near_bindgen]
@@ -95,7 +98,7 @@ impl Pyth {
         codehash: [u8; 32],
         initial_source: Source,
         gov_source: Source,
-        update_fee: u64,
+        update_fee: U128,
         stale_threshold: u64,
     ) -> Self {
         // Add an initial Source so that the contract can be used.
@@ -103,13 +106,14 @@ impl Pyth {
         sources.insert(&initial_source);
         Self {
             prices: UnorderedMap::new(StorageKeys::Prices),
-            executed_gov_sequences: UnorderedSet::new(StorageKeys::Governance),
+            executed_governance_vaa: 0,
+            executed_governance_change_vaa: 0,
             stale_threshold,
             gov_source,
             sources,
             wormhole,
             codehash,
-            update_fee,
+            update_fee: update_fee.into(),
         }
     }
 
@@ -138,7 +142,7 @@ impl Pyth {
 
         if !self.sources.contains(&Source {
             emitter: vaa.emitter_address,
-            pyth_emitter_chain: vaa.emitter_chain,
+            chain:   vaa.emitter_chain,
         }) {
             return Err(Error::UnknownSource);
         }
@@ -162,11 +166,19 @@ impl Pyth {
         Ok(())
     }
 
-    /// Return the deposit required to update a price feed.
+    /// Return the deposit required to update a price feed. This is the upper limit for an update
-    pub fn get_update_fee_estimate(&self) -> u64 {
+    /// call and any remaining deposit not consumed for storage will be refunded.
+    #[allow(unused_variables)]
+    pub fn get_update_fee_estimate(&self, vaa: String) -> U128 {
         let byte_cost = env::storage_byte_cost();
         let data_cost = byte_cost * std::mem::size_of::<PriceFeed>() as u128;
-        4u64 * u64::try_from(data_cost).unwrap() + self.update_fee
+
+        // The const multiplications here are to provide additional headway for any unexpected data
+        // costs in NEAR's storage calculations.
+        //
+        // 5 is the upper limit for PriceFeed amount in a single update.
+        // 4 is the value obtained through testing for headway.
+        (5u128 * 4u128 * data_cost + self.update_fee).into()
     }
 
     #[payable]
@@ -183,10 +195,17 @@ impl Pyth {
         }
 
         // Get Storage Usage before execution, subtracting the fee from the deposit has the effect
-        // forces the caller to add the required fee to the deposit.
+        // forces the caller to add the required fee to the deposit. The protocol defines the fee
-        let storage = env::storage_usage()
+        // as a u128, but storage is a u64, so we need to check that the fee does not overflow the
-            .checked_sub(self.update_fee)
+        // storage cost as well.
-            .ok_or(Error::InsufficientDeposit)?;
+        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))?;
 
         // 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

target-chains/near/receiver/src/state.rs

@@ -60,7 +60,6 @@ pub struct PriceFeed {
     pub ema_price: Price,
 }
 
-// TODO: Source the Timestamp
 impl From<&PriceAttestation> for PriceFeed {
     fn from(price_attestation: &PriceAttestation) -> Self {
         Self {
@@ -88,6 +87,7 @@ impl From<&PriceAttestation> for PriceFeed {
     BorshDeserialize,
     BorshSerialize,
     Clone,
+    Copy,
     Debug,
     Default,
     Deserialize,
@@ -99,14 +99,22 @@ impl From<&PriceAttestation> for PriceFeed {
 )]
 #[serde(crate = "near_sdk::serde")]
 #[repr(transparent)]
-pub struct Chain(pub u16);
+pub struct Chain(u16);
 
+/// Converts from a WormholeChain, rather than a u16. This lets us rely on Wormhole's SDK to
+/// validate the chain identifier.
 impl From<WormholeChain> for Chain {
     fn from(chain: WormholeChain) -> Self {
         Self(u16::from(chain))
     }
 }
 
+impl From<Chain> for u16 {
+    fn from(chain: Chain) -> Self {
+        chain.0
+    }
+}
+
 /// A `Source` describes an origin chain from which Pyth attestations are allowed.
 ///
 /// This allows for example Pyth prices to be sent from either Pythnet or Solana, but can be used
@@ -127,7 +135,7 @@ impl From<WormholeChain> for Chain {
 #[serde(crate = "near_sdk::serde")]
 pub struct Source {
     pub emitter: WormholeAddress,
-    pub pyth_emitter_chain: Chain,
+    pub chain:   Chain,
 }
 
 /// A local `Vaa` type converted to from the Wormhole definition, this helps catch any upstream

target-chains/near/receiver/src/tests.rs

@@ -19,7 +19,7 @@ mod tests {
             [0; 32],
             Source::default(),
             Source::default(),
-            1,
+            1.into(),
             32,
         )
     }

target-chains/near/receiver/tests/workspaces.rs

@@ -7,8 +7,13 @@ use {
         PriceStatus,
     },
     pyth::{
-        governance::GovernanceAction,
+        governance::{
+            GovernanceAction,
+            GovernanceInstruction,
+            GovernanceModule,
+        },
         state::{
+            Chain,
             Price,
             PriceIdentifier,
             Source,
@@ -19,6 +24,7 @@ use {
         Cursor,
         Write,
     },
+    wormhole::Chain as WormholeChain,
 };
 
 async fn initialize_chain() -> (
@@ -44,19 +50,20 @@ async fn initialize_chain() -> (
         .expect("Failed to deploy wormhole_stub.wasm");
 
     // Initialize Wormhole.
-    wormhole
+    let _ = wormhole
         .call("new")
         .args_json(&json!({}))
         .gas(300_000_000_000_000)
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to initialize Wormhole")
+        .await
         .unwrap();
 
     // Initialize Pyth, one time operation that sets the Wormhole contract address.
     let codehash = [0u8; 32];
 
-    contract
+    let _ = contract
         .call("new")
         .args_json(&json!({
             "wormhole":        wormhole.id(),
@@ -67,22 +74,24 @@ async fn initialize_chain() -> (
             "stale_threshold": 32,
         }))
         .gas(300_000_000_000_000)
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to initialize Pyth")
+        .await
         .unwrap();
 
     (worker, contract, wormhole)
 }
 
 #[tokio::test]
-async fn test_source_add() {
+async fn test_set_sources() {
     let (_, contract, _) = initialize_chain().await;
 
     // Submit a new Source to the contract, this will trigger a cross-contract call to wormhole
     let vaa = wormhole::Vaa {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
+        sequence: 1,
         payload: (),
         ..Default::default()
     };
@@ -91,14 +100,21 @@ async fn test_source_add() {
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).expect("Failed to serialize VAA");
         cur.write_all(
-            &GovernanceAction::SetDataSources(vec![
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Any),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetDataSources {
+                    data_sources: vec![
                         Source::default(),
                         Source {
                             emitter: [1; 32],
-                    pyth_emitter_chain: pyth::state::Chain(1),
+                            chain:   Chain::from(WormholeChain::Solana),
                         },
-            ])
+                    ],
-            .serialize(),
+                },
+            }
+            .serialize()
+            .unwrap(),
         )
         .expect("Failed to write Payload");
         hex::encode(cur.into_inner())
@@ -111,11 +127,13 @@ async fn test_source_add() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // There should now be a two sources in the contract state.
     assert_eq!(
@@ -125,7 +143,7 @@ async fn test_source_add() {
             Source::default(),
             Source {
                 emitter: [1; 32],
-                pyth_emitter_chain: pyth::state::Chain(1),
+                chain:   Chain::from(WormholeChain::Solana),
             },
         ]
     );
@@ -140,18 +158,51 @@ async fn test_set_governance_source() {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
         payload: (),
+        sequence: 2,
         ..Default::default()
     };
 
     let vaa = {
+        let request_vaa = wormhole::Vaa {
+            emitter_chain: wormhole::Chain::Solana,
+            emitter_address: wormhole::Address([1; 32]),
+            payload: (),
+            sequence: 1,
+            ..Default::default()
+        };
+
+        // Data Source Upgrades are submitted with an embedded VAA, generate that one here first
+        // before we embed it.
+        let request_vaa = {
+            let mut cur = Cursor::new(Vec::new());
+            serde_wormhole::to_writer(&mut cur, &request_vaa).expect("Failed to serialize VAA");
+            cur.write_all(
+                &GovernanceInstruction {
+                    target: Chain::from(WormholeChain::Near),
+                    module: GovernanceModule::Target,
+                    action: GovernanceAction::RequestGovernanceDataSourceTransfer {
+                        governance_data_source_index: 1,
+                    },
+                }
+                .serialize()
+                .unwrap(),
+            )
+            .expect("Failed to write Payload");
+            cur.into_inner()
+        };
+
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).expect("Failed to serialize VAA");
         cur.write_all(
-            &GovernanceAction::SetGovernanceSource(Source {
+            &GovernanceInstruction {
-                emitter:            [1; 32],
+                target: Chain::from(WormholeChain::Near),
-                pyth_emitter_chain: pyth::state::Chain(1),
+                module: GovernanceModule::Target,
-            })
+                action: GovernanceAction::AuthorizeGovernanceDataSourceTransfer {
-            .serialize(),
+                    claim_vaa: request_vaa,
+                },
+            }
+            .serialize()
+            .unwrap(),
         )
         .expect("Failed to write Payload");
         hex::encode(cur.into_inner())
@@ -164,15 +215,17 @@ async fn test_set_governance_source() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // An action from the new source should now be accepted.
     let vaa = wormhole::Vaa {
-        sequence: 1, // NOTE: Incremented Governance Sequence
+        sequence: 3, // NOTE: Incremented Governance Sequence
         emitter_chain: wormhole::Chain::Solana,
         emitter_address: wormhole::Address([1; 32]),
         payload: (),
@@ -183,14 +236,21 @@ async fn test_set_governance_source() {
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).expect("Failed to serialize VAA");
         cur.write_all(
-            &GovernanceAction::SetDataSources(vec![
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetDataSources {
+                    data_sources: vec![
                         Source::default(),
                         Source {
                             emitter: [2; 32],
-                    pyth_emitter_chain: pyth::state::Chain(2),
+                            chain:   Chain::from(WormholeChain::Solana),
                         },
-            ])
+                    ],
-            .serialize(),
+                },
+            }
+            .serialize()
+            .unwrap(),
         )
         .expect("Failed to write Payload");
         hex::encode(cur.into_inner())
@@ -203,15 +263,17 @@ async fn test_set_governance_source() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // But not from the old source.
     let vaa = wormhole::Vaa {
-        sequence: 1, // NOTE: Incremented Governance Sequence
+        sequence: 4, // NOTE: Incremented Governance Sequence
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
         payload: (),
@@ -222,31 +284,40 @@ async fn test_set_governance_source() {
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).expect("Failed to serialize VAA");
         cur.write_all(
-            &GovernanceAction::SetDataSources(vec![
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetDataSources {
+                    data_sources: vec![
                         Source::default(),
                         Source {
                             emitter: [2; 32],
-                    pyth_emitter_chain: pyth::state::Chain(2),
+                            chain:   Chain::from(WormholeChain::Solana),
                         },
-            ])
+                    ],
-            .serialize(),
+                },
+            }
+            .serialize()
+            .unwrap(),
         )
         .expect("Failed to write Payload");
         hex::encode(cur.into_inner())
     };
 
-    assert!(!contract
+    assert!(contract
         .call("execute_governance_instruction")
         .gas(300_000_000_000_000)
         .deposit(300_000_000_000_000_000_000_000)
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
+        .await
+        .unwrap()
         .outcome()
-        .is_failure());
+        .is_success());
 }
 
 #[tokio::test]
@@ -258,6 +329,7 @@ async fn test_stale_threshold() {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
         payload: (),
+        sequence: 1,
         ..Default::default()
     };
 
@@ -302,13 +374,16 @@ async fn test_stale_threshold() {
     let update_fee = serde_json::from_slice::<U128>(
         &contract
             .view("get_update_fee_estimate")
-            .args(vec![])
+            .args_json(&json!({
+                "vaa": vaa,
+            }))
             .await
             .unwrap()
             .result,
     )
     .unwrap();
 
+    // Submit price. As there are no prices this should succeed despite being old.
     assert!(contract
         .call("update_price_feed")
         .gas(300_000_000_000_000)
@@ -316,14 +391,16 @@ async fn test_stale_threshold() {
         .args_json(&json!({
             "vaa_hex": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
-    // Assert Price cannot be requested, 60 seconds in the past should be considered stale.
+    // Despite succeeding, assert Price cannot be requested, 60 seconds in the past should be
-    // [tag:failed_price_check]
+    // considered stale. [tag:failed_price_check]
     assert_eq!(
         None,
         serde_json::from_slice::<Option<Price>>(
@@ -337,11 +414,12 @@ async fn test_stale_threshold() {
         .unwrap(),
     );
 
-    // Submit another Price Attestation to the contract with an even older timestamp.
+    // Submit another Price Attestation to the contract with an even older timestamp. Which
+    // should now fail due to the existing newer price.
     let vaa = wormhole::Vaa {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
-        sequence: 1,
+        sequence: 2,
         payload: (),
         ..Default::default()
     };
@@ -376,7 +454,7 @@ async fn test_stale_threshold() {
         hex::encode(cur.into_inner())
     };
 
-    // The update handler should succeed even if price is old, but simply not update the price.
+    // The update handler should now succeed even if price is old, but simply not update the price.
     assert!(contract
         .call("update_price_feed")
         .gas(300_000_000_000_000)
@@ -384,11 +462,13 @@ async fn test_stale_threshold() {
         .args_json(&json!({
             "vaa_hex": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // The price however should _not_ have updated and if we check the unsafe stored price the
     // timestamp and price should be unchanged.
@@ -414,7 +494,7 @@ async fn test_stale_threshold() {
     let vaa = wormhole::Vaa {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
-        sequence: 2,
+        sequence: 3,
         payload: (),
         ..Default::default()
     };
@@ -422,7 +502,15 @@ async fn test_stale_threshold() {
     let vaa = {
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
-        cur.write_all(&GovernanceAction::SetStalePriceThreshold(256).serialize())
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetValidPeriod { valid_seconds: 256 },
+            }
+            .serialize()
+            .unwrap(),
+        )
         .unwrap();
         hex::encode(cur.into_inner())
     };
@@ -434,11 +522,13 @@ async fn test_stale_threshold() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // It should now be possible to request the price that previously returned None.
     // [ref:failed_price_check]
@@ -470,34 +560,45 @@ async fn test_contract_fees() {
         .expect("Failed to get UNIX timestamp")
         .as_secs();
 
-    // Fetch Update fee before changing it.
-    let update_fee = serde_json::from_slice::<U128>(
-        &contract
-            .view("get_update_fee_estimate")
-            .args(vec![])
-            .await
-            .unwrap()
-            .result,
-    )
-    .unwrap();
-
     // Set a high fee for the contract needed to submit a price.
     let vaa = wormhole::Vaa {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
         payload: (),
+        sequence: 1,
         ..Default::default()
     };
 
     let vaa = {
         let mut cur = Cursor::new(Vec::new());
         serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
-        cur.write_all(&GovernanceAction::SetUpdateFee(u64::MAX).serialize())
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
         .unwrap();
         hex::encode(cur.into_inner())
     };
 
-    // Now set the update_fee too high for the deposit to cover.
+    // Fetch Update fee before changing it.
+    let update_fee = serde_json::from_slice::<U128>(
+        &contract
+            .view("get_update_fee_estimate")
+            .args_json(&json!({
+                "vaa": vaa,
+            }))
+            .await
+            .unwrap()
+            .result,
+    )
+    .unwrap();
+
+    // Now set the update_fee so that it is too high for the deposit to cover.
     assert!(contract
         .call("execute_governance_instruction")
         .gas(300_000_000_000_000)
@@ -505,30 +606,37 @@ async fn test_contract_fees() {
         .args_json(&json!({
             "vaa": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
+    // Check the state has actually changed before we try and execute another VAA.
     assert_ne!(
-        update_fee,
+        u128::from(update_fee),
+        u128::from(
             serde_json::from_slice::<U128>(
                 &contract
                     .view("get_update_fee_estimate")
-                .args(vec![])
+                    .args_json(&json!({
+                        "vaa": vaa,
+                    }))
                     .await
                     .unwrap()
                     .result,
             )
             .unwrap()
+        )
     );
 
     // Attempt to update the price feed with a now too low deposit.
     let vaa = wormhole::Vaa {
         emitter_chain: wormhole::Chain::Any,
         emitter_address: wormhole::Address([0; 32]),
-        sequence: 1,
+        sequence: 2,
         payload: (),
         ..Default::default()
     };
@@ -570,11 +678,13 @@ async fn test_contract_fees() {
         .args_json(&json!({
             "vaa_hex": vaa,
         }))
-        .transact()
+        .transact_async()
         .await
         .expect("Failed to submit VAA")
-        .outcome()
+        .await
-        .is_success());
+        .unwrap()
+        .failures()
+        .is_empty());
 
     // Submitting a Price should have failed because the fee was not enough.
     assert_eq!(
@@ -590,3 +700,243 @@ async fn test_contract_fees() {
         .unwrap(),
     );
 }
+
+// A test that attempts to SetFee twice with the same governance action, the first should succeed,
+// the second should fail.
+#[tokio::test]
+async fn test_same_governance_sequence_fails() {
+    let (_, contract, _) = initialize_chain().await;
+
+    // Set a high fee for the contract needed to submit a price.
+    let vaa = wormhole::Vaa {
+        emitter_chain: wormhole::Chain::Any,
+        emitter_address: wormhole::Address([0; 32]),
+        payload: (),
+        sequence: 1,
+        ..Default::default()
+    };
+
+    let vaa = {
+        let mut cur = Cursor::new(Vec::new());
+        serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
+        .unwrap();
+        hex::encode(cur.into_inner())
+    };
+
+    // Attempt our first SetFee.
+    assert!(contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+
+    // Attempt to run the same VAA again.
+    assert!(!contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+}
+
+// A test that attempts to SetFee twice with the same governance action, the first should succeed,
+// the second should fail.
+#[tokio::test]
+async fn test_out_of_order_sequences_fail() {
+    let (_, contract, _) = initialize_chain().await;
+
+    // Set a high fee for the contract needed to submit a price.
+    let vaa = wormhole::Vaa {
+        emitter_chain: wormhole::Chain::Any,
+        emitter_address: wormhole::Address([0; 32]),
+        payload: (),
+        sequence: 1,
+        ..Default::default()
+    };
+
+    let vaa = {
+        let mut cur = Cursor::new(Vec::new());
+        serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
+        .unwrap();
+        hex::encode(cur.into_inner())
+    };
+
+    // Attempt our first SetFee.
+    assert!(contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+
+    // Generate another VAA with sequence 3.
+    let vaa = wormhole::Vaa {
+        emitter_chain: wormhole::Chain::Any,
+        emitter_address: wormhole::Address([0; 32]),
+        payload: (),
+        sequence: 3,
+        ..Default::default()
+    };
+
+    let vaa = {
+        let mut cur = Cursor::new(Vec::new());
+        serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
+        .unwrap();
+        hex::encode(cur.into_inner())
+    };
+
+    // This should succeed.
+    assert!(contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+
+    // Generate another VAA with sequence 2.
+    let vaa = wormhole::Vaa {
+        emitter_chain: wormhole::Chain::Any,
+        emitter_address: wormhole::Address([0; 32]),
+        payload: (),
+        sequence: 2,
+        ..Default::default()
+    };
+
+    let vaa = {
+        let mut cur = Cursor::new(Vec::new());
+        serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Near),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
+        .unwrap();
+        hex::encode(cur.into_inner())
+    };
+
+    // This should fail due to being out of order.
+    assert!(!contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+}
+
+// A test that fails if the governance action payload target is not NEAR.
+#[tokio::test]
+async fn test_governance_target_fails_if_not_near() {
+    let (_, contract, _) = initialize_chain().await;
+
+    let vaa = wormhole::Vaa {
+        emitter_chain: wormhole::Chain::Any,
+        emitter_address: wormhole::Address([0; 32]),
+        payload: (),
+        sequence: 1,
+        ..Default::default()
+    };
+
+    let vaa = {
+        let mut cur = Cursor::new(Vec::new());
+        serde_wormhole::to_writer(&mut cur, &vaa).unwrap();
+        cur.write_all(
+            &GovernanceInstruction {
+                target: Chain::from(WormholeChain::Solana),
+                module: GovernanceModule::Target,
+                action: GovernanceAction::SetFee { base: 128, expo: 8 },
+            }
+            .serialize()
+            .unwrap(),
+        )
+        .unwrap();
+        hex::encode(cur.into_inner())
+    };
+
+    // This should fail as the target is Solana, when Near is expected.
+    assert!(!contract
+        .call("execute_governance_instruction")
+        .gas(300_000_000_000_000)
+        .deposit(300_000_000_000_000_000_000_000)
+        .args_json(&json!({
+            "vaa": vaa,
+        }))
+        .transact_async()
+        .await
+        .expect("Failed to submit VAA")
+        .await
+        .unwrap()
+        .failures()
+        .is_empty());
+}

target-chains/near/receiver/workspace-test.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+#
+# This script is used to prepare the environment in order to run the NEAR
+# workspaces based tests. It relies on the relative position of the wormhole-
+# stub contract to this directory.
+set -euo pipefail
+
+# Setup rust to build wasm.
+rustup target add wasm32-unknown-unknown
+
+cargo build --release --target wasm32-unknown-unknown
+cp target/wasm32-unknown-unknown/release/pyth.wasm .
+
+(
+    cd ../wormhole-stub
+    cargo build --release --target wasm32-unknown-unknown
+    cp target/wasm32-unknown-unknown/release/wormhole_stub.wasm ../receiver
+)
+
+RUST_LOG=info cargo nextest run