Add tests for verifying governance vaa

aptos/contracts/sources/error.move

@@ -78,4 +78,8 @@ module pyth::error {
     public fun positive_value(): u64 {
         error::invalid_state(19)
     }
+    
+   public fun invalid_governance_magic_value(): u64 {
+    error::invalid_argument(20)
+   }
 }

aptos/contracts/sources/governance/governance.move

@@ -3,6 +3,7 @@ module pyth::governance {
     use pyth::data_source;
     use wormhole::u16;
     use pyth::governance_instruction;
+    use pyth::pyth;
     use pyth::governance_action;
     use pyth::contract_upgrade;
     use pyth::set_governance_data_source;
@@ -11,6 +12,7 @@ module pyth::governance {
     use pyth::error;
     use pyth::set_update_fee;
     use pyth::state;
+    use std::account;
 
     public entry fun execute_governance_instruction(vaa_bytes: vector<u8>) {
         let parsed_vaa = parse_and_verify_governance_vaa(vaa_bytes);
@@ -52,4 +54,70 @@ module pyth::governance {
 
         parsed_vaa
     }
+
+    #[test_only]
+    fun setup_test(
+        chain_id: u64,
+        stale_price_threshold: u64,
+        governance_emitter_chain_id: u64,
+        governance_emitter_address: vector<u8>,
+        update_fee: u64,
+    ) {
+        // Initialize wormhole with a large message collection fee
+        wormhole::wormhole_test::setup(100000);
+
+        // Deploy and initialize a test instance of the Pyth contract
+        let deployer = account::create_signer_with_capability(&
+            account::create_test_signer_cap(@0x277fa055b6a73c42c0662d5236c65c864ccbf2d4abd21f174a30c8b786eab84b));
+        let (_pyth, signer_capability) = account::create_resource_account(&deployer, b"pyth");
+        pyth::init_test(signer_capability, chain_id, stale_price_threshold, governance_emitter_chain_id, governance_emitter_address, update_fee);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 6)]
+    fun execute_governance_instruction_invalid_vaa() {
+        setup_test(20, 50, 24, x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf", 100);
+        let vaa_bytes = x"6c436741b108";
+        execute_governance_instruction(vaa_bytes);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 65550)]
+    fun execute_governance_instruction_invalid_data_source() {
+        setup_test(20, 100, 50, x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf", 100);
+
+        // A VAA with:
+        // - Emitter chain ID of 20
+        // - Emitter address of x"ed67fcc21620d1bf9f69db61ea65ea36ae2df4f86c8e1b9503f0da287c24ed41"
+        let vaa_bytes = x"0100000000010066359039306c20c8e6d0047ca82aef1b3d1059a3196ab9b21ee9eb8d8438c4e06c3f181d86687cf52f8c4a167ce8af6a5dbadad22253a4016dc28a25f181a37301527e4f9b000000010014ed67fcc21620d1bf9f69db61ea65ea36ae2df4f86c8e1b9503f0da287c24ed410000000000000000005054474eb01087a85361f738f19454e66664d3c9";
+        execute_governance_instruction(vaa_bytes);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 65551)]
+    fun execute_governance_instruction_invalid_sequence_number_0() {
+        setup_test(20, 100, 50, x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf", 100);
+
+        // A VAA with:
+        // - Emitter chain ID 50
+        // - Emitter address x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf"
+        // - Sequence number 0
+        let vaa_bytes = x"010000000001004d7facf7151ada96a35a3f099843c5f13bd0e0a6cbf50722d4e456d370bbce8641ecc16450979d4c403888f9f08d5975503d810732dc95575880d2a4c64d40aa01527e4f9b000000010032f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf0000000000000000005054474eb01087a85361f738f19454e66664d3c9";
+        execute_governance_instruction(vaa_bytes);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 65556)]
+    fun execute_governance_instruction_invalid_instruction_magic() {
+        setup_test(20, 100, 50, x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf", 100);
+
+        // A VAA with:
+        // - Emitter chain ID 50
+        // - Emitter address x"f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf"
+        // - Sequence number 1
+        // - A payload with the value x"5054474eb01087a85361f738f19454e66664d3c9", so the magic number will be 5054474e
+        let vaa_bytes = x"01000000000100583334c65aff30780bf7f2ac783398a2a985e3e4873264e46c3cddfdfb2eaa484365e9f4a3ecc14d059ac1cf0a7b6a58075749ad17a3bfd4153d8f45b9084a3500527e4f9b000000010032f06413c0148c78916554f134dcd17a7c8029a3a2bda475a4a1182305c53078bf0000000000000001005054474eb01087a85361f738f19454e66664d3c9";
+        execute_governance_instruction(vaa_bytes);
+    }
+
 }

aptos/contracts/sources/governance/governance_instruction.move

@@ -3,6 +3,7 @@ module pyth::governance_instruction {
     use pyth::deserialize;
     use pyth::error;
 
+    const MAGIC: vector<u8> = x"5054474d"; // "PTGM": Pyth Governance Message
     const MODULE: u8 = 2;
     const TARGET_CHAIN_ID: u64 = 3;
 
@@ -21,6 +22,8 @@ module pyth::governance_instruction {
  
     public fun from_byte_vec(bytes: vector<u8>): GovernanceInstruction {
         let cursor = cursor::init(bytes);
+        let magic = deserialize::deserialize_vector(&mut cursor, 4);
+        assert!(magic == MAGIC, error::invalid_governance_magic_value());
         let module_ = deserialize::deserialize_u8(&mut cursor);
         let action = deserialize::deserialize_u8(&mut cursor);
         let target_chain_id = deserialize::deserialize_u16(&mut cursor);
@@ -58,4 +61,25 @@ module pyth::governance_instruction {
         } = instruction;
         payload
     }
+
+    #[test]
+    #[expected_failure(abort_code = 65556)]
+    fun test_from_byte_vec_invalid_magic() {
+        let bytes = x"5054474eb01087a85361f738f19454e66664d3c9";
+        destroy(from_byte_vec(bytes));
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 65548)]
+    fun test_from_byte_vec_invalid_module() {
+        let bytes = x"5054474db05087a85361f738f19454e66664d3c9";
+        destroy(from_byte_vec(bytes));
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 65548)]
+    fun test_from_byte_vec_invalid_target_chain_id() {
+        let bytes = x"5054474db05087a85361f738f19454e66664d3c9";
+        destroy(from_byte_vec(bytes));
+    }
 }