[sui 7/x] - contract upgrades, version control (#762)

* state getters and setters, change Move.toml dependency to sui/integration_v2

* finish state.move

* add new line to pyth

* use deployer cap pattern for state module

* sui pyth

* update price feeds, dynamic object fields, Sui object PriceInfoObject

* register price info object with pyth state after creation

* sui governance

* some newlines

* error codes

* update and comment

* unit tests for pyth.move, add UpgradeCap to Pyth State (will be used for contract upgrades)

* updates

* test_get_update_fee test passes

* fix test_get_update_fee and test_update_price_feeds_corrupt_vaa

* test_update_price_feeds_invalid_data_source

* test_create_and_update_price_feeds

* test_create_and_update_price_feeds_success and test_create_and_update_price_feeds_price_info_object_not_found_failure

* test_update_cache

* update

* test_update_cache_old_update

* update_price_feeds_if_fresh

* comment

* contract upgrades start

* contract upgradeability

* update clock stuff

* edits

* use clone of sui/integration_v2 for stability

* make contract_upgrade::execute a public(friend) fun, remove clock arg

* E_INCORRECT_IDENTIFIER_LENGTH

* comment and edit

* add a single comment

target_chains/sui/contracts/Move.toml

@@ -5,12 +5,12 @@ version = "0.0.1"
 [dependencies.Sui]
 git = "https://github.com/MystenLabs/sui.git"
 subdir = "crates/sui-framework/packages/sui-framework"
-rev = "ddfc3fa0768a38286787319603a5458a9ff91cc1"
+rev = "a63f425b9999c7fdfe483598720a9effc0acdc9e"
 
 [dependencies.Wormhole]
 git = "https://github.com/wormhole-foundation/wormhole.git"
 subdir = "sui/wormhole"
-rev = "sui/integration_v2"
+rev = "sui/integration_v2_stable"
 
 [addresses]
 pyth = "0x250"

target_chains/sui/contracts/sources/batch_price_attestation.move

@@ -165,14 +165,11 @@ module pyth::batch_price_attestation {
     fun test_deserialize_batch_price_attestation_invalid_magic() {
         use sui::test_scenario::{Self, take_shared, return_shared, ctx};
         let test = test_scenario::begin(@0x1234);
-        clock::create_for_testing(ctx(&mut test));
+        let test_clock = clock::create_for_testing(ctx(&mut test));
-        test_scenario::next_tx(&mut test, @0x1234);
-        let test_clock = take_shared<Clock>(&test);
-
         // A batch price attestation with a magic number of 0x50325749
         let bytes = x"5032574900030000000102000400951436e0be37536be96f0896366089506a59763d036728332d3e3038047851aea7c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1000000000000049a0000000000000008fffffffb00000000000005dc0000000000000003000000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000006150000000000000007215258d81468614f6b7e194c5d145609394f67b041e93e6695dcc616faadd0603b9551a68d01d954d6387aff4df1529027ffb2fee413082e509feb29cc4904fe000000000000041a0000000000000003fffffffb00000000000005cb0000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e4000000000000048600000000000000078ac9cf3ab299af710d735163726fdae0db8465280502eb9f801f74b3c1bd190333832fad6e36eb05a8972fe5f219b27b5b2bb2230a79ce79beb4c5c5e7ecc76d00000000000003f20000000000000002fffffffb00000000000005e70000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e40000000000000685000000000000000861db714e9ff987b6fedf00d01f9fea6db7c30632d6fc83b7bc9459d7192bc44a21a28b4c6619968bd8c20e95b0aaed7df2187fd310275347e0376a2cd7427db800000000000006cb0000000000000001fffffffb00000000000005e40000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000007970000000000000001";
         let _ = destroy(deserialize(bytes, &test_clock));
-        return_shared(test_clock);
+        clock::destroy_for_testing(test_clock);
         test_scenario::end(test);
     }
 
@@ -181,9 +178,8 @@ module pyth::batch_price_attestation {
         use sui::test_scenario::{Self, take_shared, return_shared, ctx};
         // Set the arrival time
         let test = test_scenario::begin(@0x1234);
-        clock::create_for_testing(ctx(&mut test));
+        let test_clock = clock::create_for_testing(ctx(&mut test));
         test_scenario::next_tx(&mut test, @0x1234);
-        let test_clock = take_shared<Clock>(&test);
         let arrival_time_in_seconds = clock::timestamp_ms(&test_clock) / 1000;
 
         // let arrival_time = tx_context::epoch(ctx(&mut test));
@@ -244,7 +240,7 @@ module pyth::batch_price_attestation {
         assert!(&expected == &deserialized, 1);
         destroy(expected);
         destroy(deserialized);
-        return_shared(test_clock);
+        clock::destroy_for_testing(test_clock);
         test_scenario::end(test);
     }
 }

target_chains/sui/contracts/sources/governance/contract_upgrade.move

@@ -1,12 +1,101 @@
-module pyth::contract_upgrade {
+// SPDX-License-Identifier: Apache 2
-    use pyth::state::{State};
 
-    use wormhole::state::{State as WormState};
+/// Note: This module is based on the upgrade_contract module
+/// from the Sui Wormhole package:
+/// https://github.com/wormhole-foundation/wormhole/blob/sui/integration_v2/sui/wormhole/sources/governance/upgrade_contract.move
+
+/// This module implements handling a governance VAA to enact upgrading the
+/// Pyth contract to a new build. The procedure to upgrade this contract
+/// requires a Programmable Transaction, which includes the following procedure:
+/// 1.  Load new build.
+/// 2.  Authorize upgrade.
+/// 3.  Upgrade.
+/// 4.  Commit upgrade.
+module pyth::contract_upgrade {
+    use sui::event::{Self};
+    use sui::object::{Self, ID};
+    use sui::package::{UpgradeReceipt, UpgradeTicket};
+
+    use pyth::state::{Self, State};
+
+    use wormhole::bytes32::{Self, Bytes32};
+    use wormhole::cursor::{Self};
 
     friend pyth::governance;
 
-    /// Payload should be the bytes digest of the new contract.
+    /// Digest is all zeros.
-    public(friend) fun execute(_worm_state: &WormState, _pyth_state: &State, _payload: vector<u8>){
+    const E_DIGEST_ZERO_BYTES: u64 = 0;
-        // TODO
+    /// Specific governance payload ID (action) to complete upgrading the
+    /// contract.
+    const ACTION_UPGRADE_CONTRACT: u8 = 1;
+
+    // Event reflecting package upgrade.
+    struct ContractUpgraded has drop, copy {
+        old_contract: ID,
+        new_contract: ID
+    }
+
+    struct UpgradeContract {
+        digest: Bytes32
+    }
+
+    /// Redeem governance VAA to issue an `UpgradeTicket` for the upgrade given
+    /// a contract upgrade VAA. This governance message is only relevant for Sui
+    /// because a contract upgrade is only relevant to one particular network
+    /// (in this case Sui), whose build digest is encoded in this message.
+    ///
+    /// NOTE: This method is guarded by a minimum build version check. This
+    /// method could break backward compatibility on an upgrade.
+    public(friend) fun execute(
+        pyth_state: &mut State,
+        payload: vector<u8>,
+    ): UpgradeTicket {
+        // Proceed with processing new implementation version.
+        handle_upgrade_contract(pyth_state, payload)
+    }
+
+    fun handle_upgrade_contract(
+        pyth_state: &mut State,
+        payload: vector<u8>
+    ): UpgradeTicket {
+
+        let UpgradeContract { digest } = deserialize(payload);
+
+        state::authorize_upgrade(pyth_state, digest)
+    }
+
+    /// Finalize the upgrade that ran to produce the given `receipt`. This
+    /// method invokes `state::commit_upgrade` which interacts with
+    /// `sui::package`.
+    public fun commit_upgrade(
+        self: &mut State,
+        receipt: UpgradeReceipt,
+    ) {
+        let latest_package_id = state::commit_upgrade(self, receipt);
+
+        // Emit an event reflecting package ID change.
+        event::emit(
+            ContractUpgraded {
+                old_contract: object::id_from_address(@pyth),
+                new_contract: latest_package_id
+            }
+        );
+    }
+
+    fun deserialize(payload: vector<u8>): UpgradeContract {
+        let cur = cursor::new(payload);
+
+        // This amount cannot be greater than max u64.
+        let digest = bytes32::take_bytes(&mut cur);
+        assert!(bytes32::is_nonzero(&digest), E_DIGEST_ZERO_BYTES);
+
+        cursor::destroy_empty(cur);
+
+        UpgradeContract { digest }
+    }
+
+    #[test_only]
+    public fun action(): u8 {
+        ACTION_UPGRADE_CONTRACT
     }
 }

target_chains/sui/contracts/sources/governance/governance.move

@@ -1,5 +1,7 @@
 module pyth::governance {
     use sui::clock::{Clock};
+    use sui::package::{UpgradeTicket};
+    use sui::tx_context::{TxContext};
 
     use pyth::data_source::{Self};
     use pyth::governance_instruction;
@@ -8,6 +10,7 @@ module pyth::governance {
     use pyth::set_governance_data_source;
     use pyth::set_data_sources;
     use pyth::set_stale_price_threshold;
+    use pyth::transfer_fee;
     use pyth::state::{State};
     use pyth::set_update_fee;
     use pyth::state;
@@ -18,23 +21,47 @@ module pyth::governance {
     const E_GOVERNANCE_CONTRACT_UPGRADE_CHAIN_ID_ZERO: u64 = 0;
     const E_INVALID_GOVERNANCE_ACTION: u64 = 1;
     const E_INVALID_GOVERNANCE_DATA_SOURCE: u64 = 2;
-    const E_INVALID_GOVERNANCE_SEQUENCE_NUMBER: u64 = 3;
+    const E_MUST_USE_EXECUTE_CONTRACT_UPGRADE_GOVERNANCE_INSTRUCTION_CALLSITE: u64 = 3;
+    const E_GOVERNANCE_ACTION_MUST_BE_CONTRACT_UPGRADE: u64 = 4;
 
-    public entry fun execute_governance_instruction(
+    /// Rather than having execute_governance_instruction handle the contract
+    /// upgrade governance instruction, we have a separate function that processes
+    /// contract upgrade instructions, because doing contract upgrades is a
+    /// multi-step process, and the first step of doing a contract upgrade
+    /// yields a return value, namely the upgrade ticket, which is non-droppable.
+    public fun execute_contract_upgrade_governance_instruction(
         pyth_state : &mut State,
         worm_state: &WormState,
         vaa_bytes: vector<u8>,
         clock: &Clock
-    ) {
+    ): UpgradeTicket {
-        let parsed_vaa = parse_and_verify_governance_vaa(pyth_state, worm_state, vaa_bytes, clock);
+        let parsed_vaa = parse_and_verify_and_replay_protect_governance_vaa(pyth_state, worm_state, vaa_bytes, clock);
         let instruction = governance_instruction::from_byte_vec(vaa::take_payload(parsed_vaa));
-
-        // Dispatch the instruction to the appropiate handler
         let action = governance_instruction::get_action(&instruction);
-        if (action == governance_action::new_contract_upgrade()) {
+        assert!(action == governance_action::new_contract_upgrade(),
+            E_GOVERNANCE_ACTION_MUST_BE_CONTRACT_UPGRADE);
         assert!(governance_instruction::get_target_chain_id(&instruction) != 0,
             E_GOVERNANCE_CONTRACT_UPGRADE_CHAIN_ID_ZERO);
-            contract_upgrade::execute(worm_state, pyth_state, governance_instruction::destroy(instruction));
+        contract_upgrade::execute(pyth_state, governance_instruction::destroy(instruction))
+    }
+
+    /// Execute a governance instruction.
+    public entry fun execute_governance_instruction(
+        pyth_state : &mut State,
+        worm_state: &WormState,
+        vaa_bytes: vector<u8>,
+        clock: &Clock,
+        ctx: &mut TxContext
+    ) {
+        let parsed_vaa = parse_and_verify_and_replay_protect_governance_vaa(pyth_state, worm_state, vaa_bytes, clock);
+        let instruction = governance_instruction::from_byte_vec(vaa::take_payload(parsed_vaa));
+
+        // Get the governance action.
+        let action = governance_instruction::get_action(&instruction);
+
+        // Dispatch the instruction to the appropiate handler.
+        if (action == governance_action::new_contract_upgrade()) {
+            abort(E_MUST_USE_EXECUTE_CONTRACT_UPGRADE_GOVERNANCE_INSTRUCTION_CALLSITE)
         } else if (action == governance_action::new_set_governance_data_source()) {
             set_governance_data_source::execute(pyth_state, governance_instruction::destroy(instruction));
         } else if (action == governance_action::new_set_data_sources()) {
@@ -43,13 +70,15 @@ module pyth::governance {
             set_update_fee::execute(pyth_state, governance_instruction::destroy(instruction));
         } else if (action == governance_action::new_set_stale_price_threshold()) {
             set_stale_price_threshold::execute(pyth_state, governance_instruction::destroy(instruction));
+        } else if (action == governance_action::new_transfer_fee()) {
+            transfer_fee::execute(pyth_state, governance_instruction::destroy(instruction), ctx);
         } else {
             governance_instruction::destroy(instruction);
             assert!(false, E_INVALID_GOVERNANCE_ACTION);
         }
     }
 
-    fun parse_and_verify_governance_vaa(
+    fun parse_and_verify_and_replay_protect_governance_vaa(
         pyth_state: &mut State,
         worm_state: &WormState,
         bytes: vector<u8>,
@@ -66,11 +95,8 @@ module pyth::governance {
                     vaa::emitter_address(&parsed_vaa))),
             E_INVALID_GOVERNANCE_DATA_SOURCE);
 
-        // Check that the sequence number is greater than the last executed governance VAA
+        // Prevent replay attacks by consuming the VAA digest (adding it to a set)
-        let sequence = vaa::sequence(&parsed_vaa);
+        state::consume_vaa(pyth_state, vaa::digest(&parsed_vaa));
-        assert!(sequence > state::get_last_executed_governance_sequence(pyth_state), E_INVALID_GOVERNANCE_SEQUENCE_NUMBER);
-        state::set_last_executed_governance_sequence(pyth_state, sequence);
-
         parsed_vaa
     }
 }

target_chains/sui/contracts/sources/governance/governance_action.move

@@ -6,6 +6,7 @@ module pyth::governance_action {
     const SET_DATA_SOURCES: u8 = 2;
     const SET_UPDATE_FEE: u8 = 3;
     const SET_STALE_PRICE_THRESHOLD: u8 = 4;
+    const TRANSFER_FEE: u8 = 5;
 
     const E_INVALID_GOVERNANCE_ACTION: u64 = 5;
 
@@ -37,4 +38,8 @@ module pyth::governance_action {
     public fun new_set_stale_price_threshold(): GovernanceAction {
         GovernanceAction { value: SET_STALE_PRICE_THRESHOLD }
     }
+
+    public fun new_transfer_fee(): GovernanceAction {
+        GovernanceAction { value: TRANSFER_FEE }
+    }
 }

target_chains/sui/contracts/sources/governance/set_data_sources.move

@@ -8,19 +8,22 @@ module pyth::set_data_sources {
     use pyth::deserialize;
     use pyth::data_source::{Self, DataSource};
     use pyth::state::{Self, State};
+    use pyth::version_control::{SetDataSources};
 
     friend pyth::governance;
 
-    struct SetDataSources {
+    struct DataSources {
         sources: vector<DataSource>,
     }
 
     public(friend) fun execute(state: &mut State, payload: vector<u8>) {
-        let SetDataSources { sources } = from_byte_vec(payload);
+        state::check_minimum_requirement<SetDataSources>(state);
+
+        let DataSources { sources } = from_byte_vec(payload);
         state::set_data_sources(state, sources);
     }
 
-    fun from_byte_vec(bytes: vector<u8>): SetDataSources {
+    fun from_byte_vec(bytes: vector<u8>): DataSources {
         let cursor = cursor::new(bytes);
         let data_sources_count = deserialize::deserialize_u8(&mut cursor);
 
@@ -37,7 +40,7 @@ module pyth::set_data_sources {
 
         cursor::destroy_empty(cursor);
 
-        SetDataSources {
+        DataSources {
             sources
         }
     }

target_chains/sui/contracts/sources/governance/set_governance_data_source.move

@@ -2,33 +2,35 @@ module pyth::set_governance_data_source {
     use pyth::deserialize;
     use pyth::data_source;
     use pyth::state::{Self, State};
+    use pyth::version_control::SetGovernanceDataSource;
 
     use wormhole::cursor;
     use wormhole::external_address::{Self, ExternalAddress};
     use wormhole::bytes32::{Self};
-    //use wormhole::state::{Self}
 
     friend pyth::governance;
 
-    struct SetGovernanceDataSource {
+    struct GovernanceDataSource {
         emitter_chain_id: u64,
         emitter_address: ExternalAddress,
         initial_sequence: u64,
     }
 
     public(friend) fun execute(pyth_state: &mut State, payload: vector<u8>) {
-        let SetGovernanceDataSource { emitter_chain_id, emitter_address, initial_sequence } = from_byte_vec(payload);
+        state::check_minimum_requirement<SetGovernanceDataSource>(pyth_state);
+
+        // TODO - What is GovernanceDataSource initial_sequence used for?
+        let GovernanceDataSource { emitter_chain_id, emitter_address, initial_sequence: _initial_sequence } = from_byte_vec(payload);
         state::set_governance_data_source(pyth_state, data_source::new(emitter_chain_id, emitter_address));
-        state::set_last_executed_governance_sequence(pyth_state, initial_sequence);
     }
 
-    fun from_byte_vec(bytes: vector<u8>): SetGovernanceDataSource {
+    fun from_byte_vec(bytes: vector<u8>): GovernanceDataSource {
         let cursor = cursor::new(bytes);
         let emitter_chain_id = deserialize::deserialize_u16(&mut cursor);
         let emitter_address = external_address::new(bytes32::from_bytes(deserialize::deserialize_vector(&mut cursor, 32)));
         let initial_sequence = deserialize::deserialize_u64(&mut cursor);
         cursor::destroy_empty(cursor);
-        SetGovernanceDataSource {
+        GovernanceDataSource {
             emitter_chain_id: (emitter_chain_id as u64),
             emitter_address,
             initial_sequence

target_chains/sui/contracts/sources/governance/set_stale_price_threshold.move

@@ -2,23 +2,26 @@ module pyth::set_stale_price_threshold {
     use wormhole::cursor;
     use pyth::deserialize;
     use pyth::state::{Self, State};
+    use pyth::version_control::SetStalePriceThreshold;
 
     friend pyth::governance;
 
-    struct SetStalePriceThreshold {
+    struct StalePriceThreshold {
         threshold: u64,
     }
 
     public(friend) fun execute(state: &mut State, payload: vector<u8>) {
-        let SetStalePriceThreshold { threshold } = from_byte_vec(payload);
+        state::check_minimum_requirement<SetStalePriceThreshold>(state);
+
+        let StalePriceThreshold { threshold } = from_byte_vec(payload);
         state::set_stale_price_threshold_secs(state, threshold);
     }
 
-    fun from_byte_vec(bytes: vector<u8>): SetStalePriceThreshold {
+    fun from_byte_vec(bytes: vector<u8>): StalePriceThreshold {
         let cursor = cursor::new(bytes);
         let threshold = deserialize::deserialize_u64(&mut cursor);
         cursor::destroy_empty(cursor);
-        SetStalePriceThreshold {
+        StalePriceThreshold {
             threshold
         }
     }

target_chains/sui/contracts/sources/governance/set_update_fee.move

@@ -3,6 +3,7 @@ module pyth::set_update_fee {
 
     use pyth::deserialize;
     use pyth::state::{Self, State};
+    use pyth::version_control::SetUpdateFee;
 
     use wormhole::cursor;
 
@@ -12,24 +13,26 @@ module pyth::set_update_fee {
     const MAX_U64: u128 = (1 << 64) - 1;
     const E_EXPONENT_DOES_NOT_FIT_IN_U8: u64 = 0;
 
-    struct SetUpdateFee {
+    struct UpdateFee {
         mantissa: u64,
         exponent: u64,
     }
 
     public(friend) fun execute(pyth_state: &mut State, payload: vector<u8>) {
-        let SetUpdateFee { mantissa, exponent } = from_byte_vec(payload);
+        state::check_minimum_requirement<SetUpdateFee>(pyth_state);
+
+        let UpdateFee { mantissa, exponent } = from_byte_vec(payload);
         assert!(exponent <= 255, E_EXPONENT_DOES_NOT_FIT_IN_U8);
         let fee = apply_exponent(mantissa, (exponent as u8));
         state::set_base_update_fee(pyth_state, fee);
     }
 
-    fun from_byte_vec(bytes: vector<u8>): SetUpdateFee {
+    fun from_byte_vec(bytes: vector<u8>): UpdateFee {
         let cursor = cursor::new(bytes);
         let mantissa = deserialize::deserialize_u64(&mut cursor);
         let exponent = deserialize::deserialize_u64(&mut cursor);
         cursor::destroy_empty(cursor);
-        SetUpdateFee {
+        UpdateFee {
             mantissa,
             exponent,
         }

target_chains/sui/contracts/sources/governance/transfer_fee.move

@@ -0,0 +1,51 @@
+module pyth::transfer_fee {
+
+    use sui::transfer::Self;
+    use sui::coin::Self;
+    use sui::tx_context::TxContext;
+
+    use wormhole::cursor;
+    use wormhole::external_address::{Self};
+    use wormhole::bytes32::{Self};
+
+    use pyth::state::{Self, State};
+    use pyth::version_control::{TransferFee};
+
+    friend pyth::governance;
+
+    struct PythFee {
+        amount: u64,
+        recipient: address
+    }
+
+    public(friend) fun execute(state: &mut State, payload: vector<u8>, ctx: &mut TxContext) {
+        state::check_minimum_requirement<TransferFee>(state);
+
+        let PythFee { amount, recipient } = from_byte_vec(payload);
+
+        transfer::public_transfer(
+            coin::from_balance(
+                state::withdraw_fee(state, amount),
+                ctx
+            ),
+            recipient
+        );
+    }
+
+    fun from_byte_vec(payload: vector<u8>): PythFee {
+        let cur = cursor::new(payload);
+
+        // This amount cannot be greater than max u64.
+        let amount = bytes32::to_u64_be(bytes32::take_bytes(&mut cur));
+
+        // Recipient must be non-zero address.
+        let recipient = external_address::take_nonzero(&mut cur);
+
+        cursor::destroy_empty(cur);
+
+        PythFee {
+            amount: (amount as u64),
+            recipient: external_address::to_address(recipient)
+        }
+    }
+}

target_chains/sui/contracts/sources/migrate.move

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache 2
+
+/// Note: This module is largely taken from the Sui Wormhole package:
+/// https://github.com/wormhole-foundation/wormhole/blob/sui/integration_v2/sui/wormhole/sources/migrate.move
+
+/// This module implements an entry method intended to be called after an
+/// upgrade has been commited. The purpose is to add one-off migration logic
+/// that would alter pyth `State`.
+///
+/// Included in migration is the ability to ensure that breaking changes for
+/// any of pyth's methods by enforcing the current build version as their
+/// required minimum version.
+module pyth::migrate {
+    use pyth::state::{Self, State};
+
+    // This import is only used when `state::require_current_version` is used.
+    // use pyth::version_control::{Self as control};
+
+    /// Upgrade procedure is not complete (most likely due to an upgrade not
+    /// being initialized since upgrades can only be performed via programmable
+    /// transaction).
+    const E_CANNOT_MIGRATE: u64 = 0;
+
+    /// Execute migration logic. See `pyth::migrate` description for more
+    /// info.
+    public entry fun migrate(pyth_state: &mut State) {
+        // pyth `State` only allows one to call `migrate` after the upgrade
+        // procedure completed.
+        assert!(state::can_migrate(pyth_state), E_CANNOT_MIGRATE);
+
+        ////////////////////////////////////////////////////////////////////////
+        //
+        // If there are any methods that require the current build, we need to
+        // explicity require them here.
+        //
+        // Calls to `require_current_version` are commented out for convenience.
+        //
+        ////////////////////////////////////////////////////////////////////////
+
+
+        // state::require_current_version<control::SetDataSources>(pyth_state);
+        // state::require_current_version<control::SetGovernanceDataSource>(pyth_state);
+        // state::require_current_version<control::SetStalePriceThreshold>(pyth_state);
+        // state::require_current_version<control::SetUpdateFee>(pyth_state);
+        // state::require_current_version<control::TransferFee>(pyth_state);
+        // state::require_current_version<control::UpdatePriceFeeds>(pyth_state);
+        // state::require_current_version<control::CreatePriceFeeds>(pyth_state);
+
+        ////////////////////////////////////////////////////////////////////////
+        //
+        // NOTE: Put any one-off migration logic here.
+        //
+        // Most upgrades likely won't need to do anything, in which case the
+        // rest of this function's body may be empty. Make sure to delete it
+        // after the migration has gone through successfully.
+        //
+        // WARNING: The migration does *not* proceed atomically with the
+        // upgrade (as they are done in separate transactions).
+        // If the nature of your migration absolutely requires the migration to
+        // happen before certain other functionality is available, then guard
+        // that functionality with the `assert!` from above.
+        //
+        ////////////////////////////////////////////////////////////////////////
+
+        ////////////////////////////////////////////////////////////////////////
+        // Ensure that `migrate` cannot be called again.
+        state::disable_migration(pyth_state);
+    }
+}

target_chains/sui/contracts/sources/price_identifier.move

@@ -1,16 +1,15 @@
 module pyth::price_identifier {
     use std::vector;
-    //use pyth::error;
 
     const IDENTIFIER_BYTES_LENGTH: u64 = 32;
+    const E_INCORRECT_IDENTIFIER_LENGTH: u64 = 0;
 
     struct PriceIdentifier has copy, drop, store {
         bytes: vector<u8>,
     }
 
     public fun from_byte_vec(bytes: vector<u8>): PriceIdentifier {
-        assert!(vector::length(&bytes) == IDENTIFIER_BYTES_LENGTH, 0); //error::incorrect_identifier_length()
+        assert!(vector::length(&bytes) == IDENTIFIER_BYTES_LENGTH, E_INCORRECT_IDENTIFIER_LENGTH);
-
         PriceIdentifier {
             bytes: bytes
         }

target_chains/sui/contracts/sources/pyth.move

@@ -15,6 +15,7 @@ module pyth::pyth {
     use pyth::price_feed::{Self};
     use pyth::price::{Self, Price};
     use pyth::price_identifier::{PriceIdentifier};
+    use pyth::version_control::{UpdatePriceFeeds, CreatePriceFeeds};
 
     use wormhole::external_address::{Self};
     use wormhole::vaa::{Self};
@@ -93,6 +94,8 @@ module pyth::pyth {
         clock: &Clock,
         ctx: &mut TxContext
     ){
+        // Version control.
+        state::check_minimum_requirement<CreatePriceFeeds>(pyth_state);
         while (!vector::is_empty(&vaas)) {
             let vaa = vector::pop_back(&mut vaas);
 
@@ -137,8 +140,8 @@ module pyth::pyth {
         };
     }
 
-    /// Update PriceInfo objects and corresponding price feeds with the
+    /// Update Pyth Price Info objects (containing price feeds) with the
-    /// data in the given VAAs.
+    /// price data in the given VAAs.
     ///
     /// The vaas argument is a vector of VAAs encoded as bytes.
     ///
@@ -158,6 +161,8 @@ module pyth::pyth {
         fee: Coin<SUI>,
         clock: &Clock
     ){
+        // Version control.
+        state::check_minimum_requirement<UpdatePriceFeeds>(pyth_state);
         // Charge the message update fee
         assert!(get_total_update_fee(pyth_state, &vaas) <= coin::value(&fee), E_INSUFFICIENT_FEE);
 
@@ -176,9 +181,11 @@ module pyth::pyth {
         };
     }
 
-    /// Precondition: A Sui object of type PriceInfoObject must exist for each update
+    /// Make sure that a Sui object of type PriceInfoObject exists for each update
     /// encoded in the worm_vaa (batch_attestation_vaa). These should be passed in
-    /// via the price_info_objects argument.
+    /// via the price_info_objects argument. If for any price feed update, a
+    /// a PriceInfoObject with a matching price identifier is not found, the update_cache
+    /// function will revert, causing this function to revert.
     fun update_price_feed_from_single_vaa(
         worm_state: &WormState,
         pyth_state: &PythState,
@@ -216,10 +223,9 @@ module pyth::pyth {
             let update = vector::pop_back(&mut updates);
             let i = 0;
             let found = false;
-            // Find PriceInfoObjects corresponding to the current update (PriceInfo).
+            // Note - Would it be worth it to construct an in-memory hash-map to make look-ups faster?
-            // TODO - Construct an in-memory table to make look-ups faster?
+            //        This loop might be expensive if there are a large number of price_info_objects
-            //        This loop might be expensive if there are a large number
+            //        passed in.
-            //        of updates and/or price_info_objects we are updating.
             while (i < vector::length<PriceInfoObject>(price_info_objects) && found == false){
                 // Check if the current price info object corresponds to the price feed that
                 // the update is meant for.
@@ -386,10 +392,10 @@ module pyth::pyth_tests{
 
     use sui::sui::SUI;
     use sui::coin::{Self, Coin};
-    use sui::clock::{Self, Clock};
     use sui::test_scenario::{Self, Scenario, ctx, take_shared, return_shared};
     use sui::package::Self;
     use sui::object::{Self, ID};
+    use sui::clock::{Self, Clock};
 
     use pyth::state::{Self, State as PythState};
     use pyth::price_identifier::{Self};
@@ -426,7 +432,7 @@ module pyth::pyth_tests{
         data_sources: vector<DataSource>,
         base_update_fee: u64,
         to_mint: u64
-    ): (Scenario, Coin<SUI>) {
+    ): (Scenario, Coin<SUI>, Clock) {
 
         let scenario = test_scenario::begin(DEPLOYER);
 
@@ -470,9 +476,6 @@ module pyth::pyth_tests{
             test_scenario::ctx(&mut scenario)
         );
 
-        // Create and share a global clock object for timekeeping.
-        clock::create_for_testing(ctx(&mut scenario));
-
         // Initialize Pyth state.
         let pyth_upgrade_cap=
             package::test_publish(
@@ -498,7 +501,8 @@ module pyth::pyth_tests{
         );
 
         let coins = coin::mint_for_testing<SUI>(to_mint, ctx(&mut scenario));
-        (scenario, coins)
+        let clock = clock::create_for_testing(ctx(&mut scenario));
+        (scenario, coins, clock)
     }
 
     #[test_only]
@@ -571,7 +575,7 @@ module pyth::pyth_tests{
     #[test]
     fun test_get_update_fee() {
         let single_update_fee = 50;
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], single_update_fee, 0);
+        let (scenario, test_coins, _clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], single_update_fee, 0);
         test_scenario::next_tx(&mut scenario, DEPLOYER, );
         let pyth_state = take_shared<PythState>(&scenario);
         // Pass in a single VAA
@@ -590,17 +594,17 @@ module pyth::pyth_tests{
 
         return_shared(pyth_state);
         coin::burn_for_testing<SUI>(test_coins);
+        clock::destroy_for_testing(_clock);
         test_scenario::end(scenario);
     }
 
     #[test]
     #[expected_failure(abort_code = wormhole::vaa::E_WRONG_VERSION)]
     fun test_create_price_feeds_corrupt_vaa() {
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], 50, 0);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], 50, 0);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         // Pass in a corrupt VAA, which should fail deseriaizing
         let corrupt_vaa = x"90F8bf6A479f320ead074411a4B0e7944Ea8c9C1";
@@ -616,7 +620,7 @@ module pyth::pyth_tests{
 
         return_shared(pyth_state);
         return_shared(worm_state);
-        return_shared(clock);
+        clock::destroy_for_testing(clock);
         coin::burn_for_testing<SUI>(test_coins);
         test_scenario::end(scenario);
     }
@@ -634,12 +638,11 @@ module pyth::pyth_tests{
             )
         ];
 
-        let (scenario, test_coins) = setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, 50, 0);
+        let (scenario, test_coins, clock) = setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, 50, 0);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         pyth::create_price_feeds(
             &mut worm_state,
@@ -651,7 +654,7 @@ module pyth::pyth_tests{
 
         return_shared(pyth_state);
         return_shared(worm_state);
-        return_shared(clock);
+        clock::destroy_for_testing(clock);
         coin::burn_for_testing<SUI>(test_coins);
         test_scenario::end(scenario);
     }
@@ -675,12 +678,11 @@ module pyth::pyth_tests{
         let base_update_fee = 50;
         let coins_to_mint = 5000;
 
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         pyth::create_price_feeds(
             &mut worm_state,
@@ -730,7 +732,7 @@ module pyth::pyth_tests{
         return_shared(price_info_object_3);
         return_shared(price_info_object_4);
 
-        return_shared(clock);
+        clock::destroy_for_testing(clock);
         test_scenario::end(scenario);
     }
 
@@ -741,12 +743,11 @@ module pyth::pyth_tests{
         let base_update_fee = 50;
         let coins_to_mint = 5000;
 
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         pyth::create_price_feeds(
             &mut worm_state,
@@ -805,12 +806,11 @@ module pyth::pyth_tests{
         let base_update_fee = 50;
         let coins_to_mint = 5;
 
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         pyth::create_price_feeds(
             &mut worm_state,
@@ -849,13 +849,13 @@ module pyth::pyth_tests{
         let base_update_fee = 50;
         let coins_to_mint = 5000;
 
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
+        // Update cache is called by create_price_feeds.
         pyth::create_price_feeds(
             &mut worm_state,
             &mut pyth_state,
@@ -898,7 +898,7 @@ module pyth::pyth_tests{
         return_shared(price_info_object_4);
         coin::burn_for_testing<SUI>(test_coins);
 
-        return_shared(clock);
+        clock::destroy_for_testing(clock);
         test_scenario::end(scenario);
     }
 
@@ -908,12 +908,11 @@ module pyth::pyth_tests{
         let base_update_fee = 50;
         let coins_to_mint = 5000;
 
-        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        let (scenario, test_coins, clock) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
         test_scenario::next_tx(&mut scenario, DEPLOYER);
 
         let pyth_state = take_shared<PythState>(&scenario);
         let worm_state = take_shared<WormState>(&scenario);
-        let clock = take_shared<Clock>(&scenario);
 
         pyth::create_price_feeds(
             &mut worm_state,
@@ -965,7 +964,6 @@ module pyth::pyth_tests{
 
         vector::destroy_empty(price_info_object_vec);
         
-        // Confirm that the current price and ema price didn't change
         let current_price_info = price_info::get_price_info_from_price_info_object(&price_info_object_1);
         let current_price_feed = price_info::get_price_feed(&current_price_info);
         let current_price = price_feed::get_price(current_price_feed);
@@ -1007,7 +1005,7 @@ module pyth::pyth_tests{
 
         vector::destroy_empty(price_info_object_vec);
 
-        // Confirm that the Pyth cached price got updated to fresh_price
+        // Confirm that the Pyth cached price got updated to fresh_price.
         let current_price_info = price_info::get_price_info_from_price_info_object(&price_info_object_1);
         let current_price_feed = price_info::get_price_feed(&current_price_info);
         let current_price = price_feed::get_price(current_price_feed);
@@ -1024,7 +1022,7 @@ module pyth::pyth_tests{
         return_shared(price_info_object_4);
         coin::burn_for_testing<SUI>(test_coins);
 
-        return_shared(clock);
+        clock::destroy_for_testing(clock);
         test_scenario::end(scenario);
     }
 }

target_chains/sui/contracts/sources/required_version.move

@@ -0,0 +1,241 @@
+// SPDX-License-Identifier: Apache 2
+
+/// Note: This module is based on the required_version module
+/// from the Sui Wormhole package:
+/// https://github.com/wormhole-foundation/wormhole/blob/sui/integration_v2/sui/wormhole/sources/resources/required_version.move
+
+/// This module implements a mechanism for version control. While keeping track
+/// of the latest version of a package build, `RequiredVersion` manages the
+/// minimum required version number for any method in that package. For any
+/// upgrade where a particular method can have backward compatibility, the
+/// minimum version would not have to change (because the method should work the
+/// same way with the previous version or current version).
+///
+/// If there happens to be a breaking change for a particular method, this
+/// module can force that the method's minimum requirement be the latest build.
+/// If a previous build were used, the method would abort if a check is in place
+/// with `RequiredVersion`.
+///
+/// There is no magic behind the way ths module works. `RequiredVersion` is
+/// intended to live in a package's shared object that gets passed into its
+/// methods (e.g. pyth's `State` object).
+module pyth::required_version {
+    use sui::dynamic_field::{Self as field};
+    use sui::object::{Self, UID};
+    use sui::package::{Self, UpgradeCap};
+    use sui::tx_context::{TxContext};
+
+    /// Build version passed does not meet method's minimum required version.
+    const E_OUTDATED_VERSION: u64 = 0;
+
+    /// Container to keep track of latest build version. Dynamic fields are
+    /// associated with its `id`.
+    struct RequiredVersion has store {
+        id: UID,
+        latest_version: u64
+    }
+
+    struct Key<phantom MethodType> has store, drop, copy {}
+
+    /// Create new `RequiredVersion` with a configured starting version.
+    public fun new(version: u64, ctx: &mut TxContext): RequiredVersion {
+        RequiredVersion {
+            id: object::new(ctx),
+            latest_version: version
+        }
+    }
+
+    /// Retrieve latest build version.
+    public fun current(self: &RequiredVersion): u64 {
+        self.latest_version
+    }
+
+    /// Add specific method handling via custom `MethodType`. At the time a
+    /// method is added, the minimum build version associated with this method
+    /// by default is the latest version.
+    public fun add<MethodType>(self: &mut RequiredVersion) {
+        field::add(&mut self.id, Key<MethodType> {}, self.latest_version)
+    }
+
+    /// This method will abort if the version for a particular `MethodType` is
+    /// not up-to-date with the version of the current build.
+    ///
+    /// For example, if the minimum requirement for `foobar` module (with an
+    /// appropriately named `MethodType` like `FooBar`) is `1` and the current
+    /// implementation version is `2`, this method will succeed because the
+    /// build meets the minimum required version of `1` in order for `foobar` to
+    /// work. So if someone were to use an older build like version `1`, this
+    /// method will succeed.
+    ///
+    /// But if `check_minimum_requirement` were invoked for `foobar` when the
+    /// minimum requirement is `2` and the current build is only version `1`,
+    /// then this method will abort because the build does not meet the minimum
+    /// version requirement for `foobar`.
+    ///
+    /// This method also assumes that the `MethodType` being checked for is
+    /// already a dynamic field (using `add`) during initialization.
+    public fun check_minimum_requirement<MethodType>(
+        self: &RequiredVersion,
+        build_version: u64
+    ) {
+        assert!(
+            build_version >= minimum_for<MethodType>(self),
+            E_OUTDATED_VERSION
+        );
+    }
+
+    /// At `commit_upgrade`, use this method to update the tracker's knowledge
+    /// of the latest upgrade (build) version, which is obtained from the
+    /// `UpgradeCap` in `sui::package`.
+    public fun update_latest(
+        self: &mut RequiredVersion,
+        upgrade_cap: &UpgradeCap
+    ) {
+        self.latest_version = package::version(upgrade_cap);
+    }
+
+    /// Once the global version is updated via `commit_upgrade` and there is a
+    /// particular method that has a breaking change, use this method to uptick
+    /// that method's minimum required version to the latest.
+    public fun require_current_version<MethodType>(self: &mut RequiredVersion) {
+        let min_version = field::borrow_mut(&mut self.id, Key<MethodType> {});
+        *min_version = self.latest_version;
+    }
+
+    /// Retrieve the minimum required version for a particular method (via
+    /// `MethodType`).
+    public fun minimum_for<MethodType>(self: &RequiredVersion): u64 {
+        *field::borrow(&self.id, Key<MethodType> {})
+    }
+
+    #[test_only]
+    public fun set_required_version<MethodType>(
+        self: &mut RequiredVersion,
+        version: u64
+    ) {
+        *field::borrow_mut(&mut self.id, Key<MethodType> {}) = version;
+    }
+
+    #[test_only]
+    public fun destroy(req: RequiredVersion) {
+        let RequiredVersion { id, latest_version: _} = req;
+        object::delete(id);
+    }
+}
+
+#[test_only]
+module pyth::required_version_test {
+    use sui::hash::{keccak256};
+    use sui::object::{Self};
+    use sui::package::{Self};
+    use sui::tx_context::{Self};
+
+    use pyth::required_version::{Self};
+
+    struct SomeMethod {}
+    struct AnotherMethod {}
+
+    #[test]
+    public fun test_check_minimum_requirement() {
+        let ctx = &mut tx_context::dummy();
+
+        let version = 1;
+        let req = required_version::new(version, ctx);
+        assert!(required_version::current(&req) == version, 0);
+
+        required_version::add<SomeMethod>(&mut req);
+        assert!(required_version::minimum_for<SomeMethod>(&req) == version, 0);
+
+        // Should not abort here.
+        required_version::check_minimum_requirement<SomeMethod>(&req, version);
+
+        // And should not abort if the version is anything greater than the
+        // current.
+        let new_version = version + 1;
+        required_version::check_minimum_requirement<SomeMethod>(
+            &req,
+            new_version
+        );
+
+        // Uptick based on new upgrade.
+        let upgrade_cap = package::test_publish(
+            object::id_from_address(@pyth),
+            ctx
+        );
+        let digest = keccak256(&x"DEADBEEF");
+        let policy = package::upgrade_policy(&upgrade_cap);
+        let upgrade_ticket =
+            package::authorize_upgrade(&mut upgrade_cap, policy, digest);
+        let upgrade_receipt = package::test_upgrade(upgrade_ticket);
+        package::commit_upgrade(&mut upgrade_cap, upgrade_receipt);
+        assert!(package::version(&upgrade_cap) == new_version, 0);
+
+        // Update to the latest version.
+        required_version::update_latest(&mut req, &upgrade_cap);
+        assert!(required_version::current(&req) == new_version, 0);
+
+        // Should still not abort here.
+        required_version::check_minimum_requirement<SomeMethod>(
+            &req,
+            new_version
+        );
+
+        // Require new version for `SomeMethod` and show that
+        // `check_minimum_requirement` still succeeds.
+        required_version::require_current_version<SomeMethod>(&mut req);
+        assert!(
+            required_version::minimum_for<SomeMethod>(&req) == new_version,
+            0
+        );
+        required_version::check_minimum_requirement<SomeMethod>(
+            &req,
+            new_version
+        );
+
+        // If another method gets added to the mix, it should automatically meet
+        // the minimum requirement because its version will be the latest.
+        required_version::add<AnotherMethod>(&mut req);
+        assert!(
+            required_version::minimum_for<AnotherMethod>(&req) == new_version,
+            0
+        );
+        required_version::check_minimum_requirement<SomeMethod>(
+            &req,
+            new_version
+        );
+
+        // Clean up.
+        package::make_immutable(upgrade_cap);
+        required_version::destroy(req);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = required_version::E_OUTDATED_VERSION)]
+    public fun test_cannot_check_minimum_requirement_with_outdated_version() {
+        let ctx = &mut tx_context::dummy();
+
+        let version = 1;
+        let req = required_version::new(version, ctx);
+        assert!(required_version::current(&req) == version, 0);
+
+        required_version::add<SomeMethod>(&mut req);
+
+        // Should not abort here.
+        required_version::check_minimum_requirement<SomeMethod>(&req, version);
+
+        // Uptick minimum requirement and fail at `check_minimum_requirement`.
+        let new_version = 10;
+        required_version::set_required_version<SomeMethod>(
+            &mut req,
+            new_version
+        );
+        let old_version = new_version - 1;
+        required_version::check_minimum_requirement<SomeMethod>(
+            &req,
+            old_version
+        );
+
+        // Clean up.
+        required_version::destroy(req);
+    }
+}

target_chains/sui/contracts/sources/state.move

@@ -3,11 +3,21 @@ module pyth::state {
     use sui::object::{Self, UID, ID};
     use sui::transfer::{Self};
     use sui::tx_context::{Self, TxContext};
-    use sui::package::{Self, UpgradeCap};
+    use sui::dynamic_field::{Self as field};
+    use sui::package::{Self, UpgradeCap, UpgradeReceipt, UpgradeTicket};
+    use sui::balance::{Balance};
+    use sui::sui::SUI;
 
     use pyth::data_source::{Self, DataSource};
     use pyth::price_info::{Self};
     use pyth::price_identifier::{PriceIdentifier};
+    use pyth::required_version::{Self, RequiredVersion};
+    use pyth::version_control::{Self as control};
+
+    use wormhole::setup::{assert_package_upgrade_cap};
+    use wormhole::consumed_vaas::{Self, ConsumedVAAs};
+    use wormhole::bytes32::{Self, Bytes32};
+    use wormhole::fee_collector::{Self, FeeCollector};
 
     use wormhole::setup::{assert_package_upgrade_cap};
 
@@ -19,6 +29,13 @@ module pyth::state {
     friend pyth::set_data_sources;
     friend pyth::governance;
     friend pyth::set_governance_data_source;
+    friend pyth::migrate;
+    friend pyth::contract_upgrade;
+    friend pyth::transfer_fee;
+
+    const E_BUILD_VERSION_MISMATCH: u64 = 0;
+    const E_INVALID_BUILD_VERSION: u64 = 1;
+    const E_VAA_ALREADY_CONSUMED: u64 = 2;
 
     /// Capability for creating a bridge state object, granted to sender when this
     /// module is deployed
@@ -26,13 +43,27 @@ module pyth::state {
         id: UID
     }
 
+    /// Used as key for dynamic field reflecting whether `migrate` can be
+    /// called.
+    ///
+    /// See `migrate` module for more info.
+    struct MigrationControl has store, drop, copy {}
+
     struct State has key {
         id: UID,
         governance_data_source: DataSource,
-        last_executed_governance_sequence: u64,
         stale_price_threshold: u64,
         base_update_fee: u64,
-        upgrade_cap: UpgradeCap
+        consumed_vaas: ConsumedVAAs,
+
+        // Upgrade capability.
+        upgrade_cap: UpgradeCap,
+
+        // Fee collector.
+        fee_collector: FeeCollector,
+
+        /// Contract build version tracker.
+        required_version: RequiredVersion
     }
 
     fun init(ctx: &mut TxContext) {
@@ -56,7 +87,7 @@ module pyth::state {
         );
     }
 
-    // Initialization
+    /// Initialization
     public(friend) fun init_and_share_state(
         deployer: DeployerCap,
         upgrade_cap: UpgradeCap,
@@ -66,9 +97,9 @@ module pyth::state {
         sources: vector<DataSource>,
         ctx: &mut TxContext
     ) {
-        // TODO - version control
+        // Only init and share state once (in the initial deployment).
-        // let version = wormhole::version_control::version();
+        let version = wormhole::version_control::version();
-        //assert!(version == 1, E_INVALID_BUILD_VERSION);
+        assert!(version == 1, E_INVALID_BUILD_VERSION);
 
         let DeployerCap { id } = deployer;
         object::delete(id);
@@ -81,6 +112,8 @@ module pyth::state {
 
         let uid = object::new(ctx);
 
+        field::add(&mut uid, MigrationControl {}, false);
+
         // Create a set that contains all registered data sources and
         // attach it to uid as a dynamic field to minimize the
         // size of State.
@@ -97,19 +130,131 @@ module pyth::state {
             data_source::add(&mut uid, vector::pop_back(&mut sources));
         };
 
+        let consumed_vaas = consumed_vaas::new(ctx);
+
+        let required_version = required_version::new(control::version(), ctx);
+        required_version::add<control::SetDataSources>(&mut required_version);
+        required_version::add<control::SetGovernanceDataSource>(&mut required_version);
+        required_version::add<control::SetStalePriceThreshold>(&mut required_version);
+        required_version::add<control::SetUpdateFee>(&mut required_version);
+        required_version::add<control::TransferFee>(&mut required_version);
+        required_version::add<control::UpdatePriceFeeds>(&mut required_version);
+        required_version::add<control::CreatePriceFeeds>(&mut required_version);
+
         // Share state so that is a shared Sui object.
         transfer::share_object(
             State {
                 id: uid,
                 upgrade_cap,
                 governance_data_source,
-                last_executed_governance_sequence: 0,
                 stale_price_threshold,
                 base_update_fee,
+                consumed_vaas,
+                fee_collector: fee_collector::new(base_update_fee),
+                required_version
             }
         );
     }
 
+    /// Retrieve current build version of latest upgrade.
+    public fun current_version(self: &State): u64 {
+        required_version::current(&self.required_version)
+    }
+
+    /// Issue an `UpgradeTicket` for the upgrade.
+    public(friend) fun authorize_upgrade(
+        self: &mut State,
+        implementation_digest: Bytes32
+    ): UpgradeTicket {
+        let policy = package::upgrade_policy(&self.upgrade_cap);
+
+        // TODO: grab package ID from `UpgradeCap` and store it
+        // in a dynamic field. This will be the old ID after the upgrade.
+        // Both IDs will be emitted in a `ContractUpgraded` event.
+        package::authorize_upgrade(
+            &mut self.upgrade_cap,
+            policy,
+            bytes32::to_bytes(implementation_digest),
+        )
+    }
+
+    /// Finalize the upgrade that ran to produce the given `receipt`.
+    public(friend) fun commit_upgrade(
+        self: &mut State,
+        receipt: UpgradeReceipt
+    ): ID {
+        // Uptick the upgrade cap version number using this receipt.
+        package::commit_upgrade(&mut self.upgrade_cap, receipt);
+
+        // Check that the upticked hard-coded version version agrees with the
+        // upticked version number.
+        assert!(
+            package::version(&self.upgrade_cap) == control::version() + 1,
+            E_BUILD_VERSION_MISMATCH
+        );
+
+        // Update global version.
+        required_version::update_latest(
+            &mut self.required_version,
+            &self.upgrade_cap
+        );
+
+        // Enable `migrate` to be called after commiting the upgrade.
+        //
+        // A separate method is required because `state` is a dependency of
+        // `migrate`. This method warehouses state modifications required
+        // for the new implementation plus enabling any methods required to be
+        // gated by the current implementation version. In most cases `migrate`
+        // is a no-op. But it still must be called in order to reset the
+        // migration control to `false`.
+        //
+        // See `migrate` module for more info.
+        enable_migration(self);
+
+        package::upgrade_package(&self.upgrade_cap)
+    }
+
+    /// Enforce a particular method to use the current build version as its
+    /// minimum required version. This method ensures that a method is not
+    /// backwards compatible with older builds.
+    public(friend) fun require_current_version<ControlType>(self: &mut State) {
+        required_version::require_current_version<ControlType>(
+            &mut self.required_version,
+        )
+    }
+
+    /// Check whether a particular method meets the minimum build version for
+    /// the latest Wormhole implementation.
+    public(friend) fun check_minimum_requirement<ControlType>(self: &State) {
+        required_version::check_minimum_requirement<ControlType>(
+            &self.required_version,
+            control::version()
+        )
+    }
+
+    // Upgrade and migration-related functionality
+
+    /// Check whether `migrate` can be called.
+    ///
+    /// See `wormhole::migrate` module for more info.
+    public fun can_migrate(self: &State): bool {
+        *field::borrow(&self.id, MigrationControl {})
+    }
+
+    /// Allow `migrate` to be called after upgrade.
+    ///
+    /// See `wormhole::migrate` module for more info.
+    public(friend) fun enable_migration(self: &mut State) {
+        *field::borrow_mut(&mut self.id, MigrationControl {}) = true;
+    }
+
+    /// Disallow `migrate` to be called.
+    ///
+    /// See `wormhole::migrate` module for more info.
+    public(friend) fun disable_migration(self: &mut State) {
+        *field::borrow_mut(&mut self.id, MigrationControl {}) = false;
+    }
+
     // Accessors
     public fun get_stale_price_threshold_secs(s: &State): u64 {
         s.stale_price_threshold
@@ -127,15 +272,35 @@ module pyth::state {
         s.governance_data_source == source
     }
 
-    public fun get_last_executed_governance_sequence(s: &State): u64 {
-        s.last_executed_governance_sequence
-    }
-
     public fun price_feed_object_exists(s: &State, p: PriceIdentifier): bool {
         price_info::contains(&s.id, p)
     }
 
-    // Setters
+    // Mutators and Setters
+
+    /// Withdraw collected fees when governance action to transfer fees to a
+    /// particular recipient.
+    ///
+    /// See `pyth::transfer_fee` for more info.
+    public(friend) fun withdraw_fee(
+        self: &mut State,
+        amount: u64
+    ): Balance<SUI> {
+        fee_collector::withdraw_balance(&mut self.fee_collector, amount)
+    }
+
+    public(friend) fun deposit_fee(self: &mut State, fee: Balance<SUI>) {
+        fee_collector::deposit_balance(&mut self.fee_collector, fee);
+    }
+
+    public(friend) fun set_fee_collector_fee(self: &mut State, amount: u64) {
+        fee_collector::change_fee(&mut self.fee_collector, amount);
+    }
+
+    public(friend) fun consume_vaa(state: &mut State, vaa_digest: Bytes32){
+        consumed_vaas::consume(&mut state.consumed_vaas, vaa_digest);
+    }
+
     public(friend) fun set_data_sources(s: &mut State, new_sources: vector<DataSource>) {
         // Empty the existing table of data sources registered in state.
         data_source::empty(&mut s.id);
@@ -149,10 +314,6 @@ module pyth::state {
         price_info::add(&mut s.id, price_identifier, id);
     }
 
-    public(friend) fun set_last_executed_governance_sequence(s: &mut State, sequence: u64) {
-        s.last_executed_governance_sequence = sequence;
-    }
-
     public(friend) fun set_governance_data_source(s: &mut State, source: DataSource) {
         s. governance_data_source = source;
     }

target_chains/sui/contracts/sources/version_control.move

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache 2
+
+/// Note: This module is based on the version_control module in
+/// the Sui Wormhole package:
+/// https://github.com/wormhole-foundation/wormhole/blob/sui/integration_v2/sui/wormhole/sources/version_control.move
+
+/// This module implements dynamic field keys as empty structs. These keys with
+/// `RequiredVersion` are used to determine minimum build requirements for
+/// particular Pyth methods and breaking backward compatibility for these
+/// methods if an upgrade requires the latest upgrade version for its
+/// functionality.
+///
+/// See `pyth::required_version` and `pyth::state` for more info.
+module pyth::version_control {
+    /// This value tracks the current Pyth contract version. We are
+    /// placing this constant value at the top, which goes against Move style
+    /// guides so that we bring special attention to changing this value when
+    /// a new implementation is built for a contract upgrade.
+    const CURRENT_BUILD_VERSION: u64 = 1;
+
+    /// Key used to check minimum version requirement for `set_data_sources`
+    struct SetDataSources {}
+
+    /// Key used to check minimum version requirement for `set_governance_data_source`
+    struct SetGovernanceDataSource {}
+
+    /// Key used to check minimum version requirement for `set_stale_price_threshold`
+    struct SetStalePriceThreshold {}
+
+    /// Key used to check minimum version requirement for `set_update_fee`
+    struct SetUpdateFee {}
+
+    /// Key used to check minimum version requirement for `transfer_fee`
+    struct TransferFee {}
+
+    /// Key used to check minimum version requirement for `update_price_feeds`
+    struct UpdatePriceFeeds {}
+
+    /// Key used to check minimum version requirement for `create_price_feeds`
+    struct CreatePriceFeeds {}
+
+    //=======================================================================
+
+    /// Return const value `CURRENT_BUILD_VERSION` for this particular build.
+    /// This value is used to determine whether this implementation meets
+    /// minimum requirements for various Pyth methods required by `State`.
+    public fun version(): u64 {
+        CURRENT_BUILD_VERSION
+    }
+}