switch terra contracts over to using governance packet structure

Change-Id: I10b53cb0cfdb86ca2aa57cab20ec375a26fd252e

terra/contracts/token-bridge/src/contract.rs

@@ -8,7 +8,7 @@ use crate::msg::{HandleMsg, InitMsg, QueryMsg};
 use crate::state::{
     bridge_contracts, bridge_contracts_read, config, config_read, wrapped_asset,
     wrapped_asset_address, wrapped_asset_address_read, wrapped_asset_read, Action, AssetMeta,
-    ConfigInfo, TokenBridgeMessage, TransferInfo,
+    ConfigInfo, RegisterChain, TokenBridgeMessage, TransferInfo,
 };
 use wormhole::byte_utils::ByteUtils;
 use wormhole::byte_utils::{extend_address_to_32, extend_string_to_32};
@@ -20,7 +20,7 @@ use cw20_base::msg::QueryMsg as TokenQuery;
 use wormhole::msg::HandleMsg as WormholeHandleMsg;
 use wormhole::msg::QueryMsg as WormholeQueryMsg;
 
-use wormhole::state::ParsedVAA;
+use wormhole::state::{GovernancePacket, ParsedVAA};
 
 use cw20::TokenInfoResponse;
 
@@ -43,7 +43,8 @@ pub fn init<S: Storage, A: Api, Q: Querier>(
 ) -> StdResult<InitResponse> {
     // Save general wormhole info
     let state = ConfigInfo {
-        owner: msg.owner,
+        gov_chain: msg.gov_chain,
+        gov_address: msg.gov_address.as_slice().to_vec(),
         wormhole_contract: msg.wormhole_contract,
         wrapped_asset_code_id: msg.wrapped_asset_code_id,
     };
@@ -93,10 +94,6 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
             nonce,
         ),
         HandleMsg::SubmitVaa { data } => submit_vaa(deps, env, &data),
-        HandleMsg::RegisterChain {
-            chain_id,
-            chain_address,
-        } => handle_register_chain(deps, env, chain_id, chain_address.as_slice().to_vec()),
         HandleMsg::CreateAssetMeta {
             asset_address,
             nonce,
@@ -104,38 +101,6 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
     }
 }
 
-fn handle_register_chain<S: Storage, A: Api, Q: Querier>(
-    deps: &mut Extern<S, A, Q>,
-    env: Env,
-    chain_id: u16,
-    chain_address: Vec<u8>,
-) -> StdResult<HandleResponse> {
-    let cfg = config_read(&deps.storage).load()?;
-
-    if env.message.sender != cfg.owner {
-        return Err(StdError::unauthorized());
-    }
-
-    let existing = bridge_contracts_read(&deps.storage).load(&chain_id.to_be_bytes());
-    if existing.is_ok() {
-        return Err(StdError::generic_err(
-            "bridge contract already exists for this chain",
-        ));
-    }
-
-    let mut bucket = bridge_contracts(&mut deps.storage);
-    bucket.save(&chain_id.to_be_bytes(), &chain_address)?;
-
-    Ok(HandleResponse {
-        messages: vec![],
-        log: vec![
-            log("chain_id", chain_id),
-            log("chain_address", hex::encode(chain_address)),
-        ],
-        data: None,
-    })
-}
-
 /// Handle wrapped asset registration messages
 fn handle_register_asset<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
@@ -269,9 +234,16 @@ fn submit_vaa<S: Storage, A: Api, Q: Querier>(
     env: Env,
     data: &Binary,
 ) -> StdResult<HandleResponse> {
+    let state = config_read(&deps.storage).load()?;
+
     let vaa = parse_vaa(deps, env.block.time, data)?;
     let data = vaa.payload;
 
+    // check if vaa is from governance
+    if state.gov_chain == vaa.emitter_chain && state.gov_address == vaa.emitter_address {
+        return handle_governance_payload(deps, env, &data);
+    }
+
     let message = TokenBridgeMessage::deserialize(&data)?;
 
     let result = match message.action {
@@ -288,6 +260,56 @@ fn submit_vaa<S: Storage, A: Api, Q: Querier>(
     return result;
 }
 
+fn handle_governance_payload<S: Storage, A: Api, Q: Querier>(
+    deps: &mut Extern<S, A, Q>,
+    env: Env,
+    data: &Vec<u8>,
+) -> StdResult<HandleResponse> {
+    let gov_packet = GovernancePacket::deserialize(&data)?;
+
+    let module = String::from_utf8(gov_packet.module).unwrap();
+    let module: String = module.chars().filter(|c| !c.is_whitespace()).collect();
+
+    if module != "token_bridge" {
+        return Err(StdError::generic_err("this is not a valid module"))
+    }
+
+    match gov_packet.action {
+        0u8 => handle_register_chain(deps, env, &gov_packet.payload),
+        _ => ContractError::InvalidVAAAction.std_err(),
+    }
+}
+
+fn handle_register_chain<S: Storage, A: Api, Q: Querier>(
+    deps: &mut Extern<S, A, Q>,
+    env: Env,
+    data: &Vec<u8>,
+) -> StdResult<HandleResponse> {
+    let RegisterChain {
+        chain_id,
+        chain_address,
+    } = RegisterChain::deserialize(&data)?;
+
+    let existing = bridge_contracts_read(&deps.storage).load(&chain_id.to_be_bytes());
+    if existing.is_ok() {
+        return Err(StdError::generic_err(
+            "bridge contract already exists for this chain",
+        ));
+    }
+
+    let mut bucket = bridge_contracts(&mut deps.storage);
+    bucket.save(&chain_id.to_be_bytes(), &chain_address)?;
+
+    Ok(HandleResponse {
+        messages: vec![],
+        log: vec![
+            log("chain_id", chain_id),
+            log("chain_address", hex::encode(chain_address)),
+        ],
+        data: None,
+    })
+}
+
 fn handle_complete_transfer<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
     env: Env,

terra/contracts/token-bridge/src/msg.rs

@@ -4,7 +4,11 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct InitMsg {
-    pub owner: HumanAddr,
+
+    // governance contract details
+    pub gov_chain: u16,
+    pub gov_address: Binary,
+
     pub wormhole_contract: HumanAddr,
     pub wrapped_asset_code_id: u64,
 }
@@ -30,11 +34,6 @@ pub enum HandleMsg {
         data: Binary,
     },
 
-    RegisterChain {
-        chain_id: u16,
-        chain_address: Binary,
-    },
-
     CreateAssetMeta {
         asset_address: HumanAddr,
         nonce: u32,

terra/contracts/token-bridge/src/state.rs

@@ -1,7 +1,7 @@
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-use cosmwasm_std::{HumanAddr, StdResult, Storage};
+use cosmwasm_std::{HumanAddr, StdResult, Storage, Binary};
 use cosmwasm_storage::{
     bucket, bucket_read, singleton, singleton_read, Bucket, ReadonlyBucket, ReadonlySingleton,
     Singleton,
@@ -18,8 +18,10 @@ pub static BRIDGE_CONTRACTS: &[u8] = b"bridge_contracts";
 // Guardian set information
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct ConfigInfo {
-    // Current active guardian set
+    // governance contract details
-    pub owner: HumanAddr,
+    pub gov_chain: u16,
+    pub gov_address: Vec<u8>,
+
     pub wormhole_contract: HumanAddr,
     pub wrapped_asset_code_id: u64,
 }
@@ -184,3 +186,23 @@ impl AssetMeta {
         .concat()
     }
 }
+
+pub struct RegisterChain {
+    pub chain_id: u16,
+    pub chain_address: Vec<u8>,
+}
+
+impl RegisterChain {
+
+    pub fn deserialize(data: &Vec<u8>) -> StdResult<Self> {
+        let data = data.as_slice();
+        let chain_id = data.get_u16(0);
+        let chain_address = data[2..].to_vec();
+
+        Ok(RegisterChain {
+            chain_id,
+            chain_address
+        })
+    }
+
+}

terra/contracts/wormhole/src/contract.rs

@@ -1,17 +1,17 @@
 use cosmwasm_std::{
     has_coins, log, to_binary, Api, BankMsg, Binary, Coin, CosmosMsg, Env, Extern, HandleResponse,
     HumanAddr, InitResponse, Querier, StdError, StdResult, Storage,
+    WasmMsg
 };
 
 use crate::byte_utils::extend_address_to_32;
 use crate::byte_utils::ByteUtils;
 use crate::error::ContractError;
-use crate::msg::{
+use crate::msg::{GetAddressHexResponse, GetStateResponse, GuardianSetInfoResponse, HandleMsg, InitMsg, QueryMsg};
-    GetAddressHexResponse, GetStateResponse, GuardianSetInfoResponse, HandleMsg, InitMsg, QueryMsg,
-};
 use crate::state::{
     config, config_read, guardian_set_get, guardian_set_set, vaa_archive_add, vaa_archive_check,
-    ConfigInfo, GuardianAddress, GuardianSetInfo, ParsedVAA, WormholeGovernance,
+    ConfigInfo, GovernancePacket, GuardianAddress, GuardianSetInfo, GuardianSetUpgrade, ParsedVAA,
+    TransferFee,
 };
 
 use k256::ecdsa::recoverable::Id as RecoverableId;
@@ -22,7 +22,6 @@ use k256::EncodedPoint;
 use sha3::{Digest, Keccak256};
 
 use generic_array::GenericArray;
-
 use std::convert::TryFrom;
 
 // Chain ID of Terra
@@ -39,9 +38,10 @@ pub fn init<S: Storage, A: Api, Q: Querier>(
 ) -> StdResult<InitResponse> {
     // Save general wormhole info
     let state = ConfigInfo {
+        gov_chain: msg.gov_chain,
+        gov_address: msg.gov_address.as_slice().to_vec(),
         guardian_set_index: 0,
         guardian_set_expirity: msg.guardian_set_expirity,
-        owner: deps.api.canonical_address(&env.message.sender)?,
         fee: Coin::new(FEE_AMOUNT, FEE_DENOMINATION), // 0.01 Luna (or 10000 uluna) fee by default
     };
     config(&mut deps.storage).save(&state)?;
@@ -65,10 +65,6 @@ pub fn handle<S: Storage, A: Api, Q: Querier>(
         HandleMsg::PostMessage { message, nonce } => {
             handle_post_message(deps, env, &message.as_slice(), nonce)
         }
-        // HandleMsg::SubmitVAA { vaa } => handle_submit_vaa(deps, env, &vaa.as_slice()),
-        HandleMsg::TransferFee { amount, recipient } => {
-            handle_transfer_fee(deps, env, amount, recipient)
-        }
         HandleMsg::SubmitVAA { vaa } => handle_submit_vaa(deps, env, vaa.as_slice()),
     }
 }
@@ -82,30 +78,33 @@ fn handle_submit_vaa<S: Storage, A: Api, Q: Querier>(
     let state = config_read(&deps.storage).load()?;
 
     let vaa = parse_and_verify_vaa(&deps.storage, data, env.block.time)?;
-    if vaa.emitter_chain != 0u16 {
+    if state.gov_chain == vaa.emitter_chain && state.gov_address == vaa.emitter_address {
-        // chain 0 is the wormhole chain ?
+        return handle_governance_payload(deps, env, &vaa.payload);
-        return Err(StdError::generic_err(
-            "governance actions may only come from chain 0",
-        ));
     }
 
-    let gov = WormholeGovernance::deserialize(&vaa.payload)?;
+    ContractError::InvalidVAAAction.std_err()
+}
 
-    let result = match gov.action {
+fn handle_governance_payload<S: Storage, A: Api, Q: Querier>(
-        0u8 => {
+    deps: &mut Extern<S, A, Q>,
-            if vaa.guardian_set_index != state.guardian_set_index {
+    env: Env,
-                return ContractError::NotCurrentGuardianSet.std_err();
+    data: &Vec<u8>,
-            }
+) -> StdResult<HandleResponse> {
-            vaa_update_guardian_set(deps, env, gov.payload.as_slice())
+    let gov_packet = GovernancePacket::deserialize(&data)?;
-        }
+
+    let module = String::from_utf8(gov_packet.module).unwrap();
+    let module: String = module.chars().filter(|c| !c.is_whitespace()).collect();
+
+    if module != "core" {
+        return Err(StdError::generic_err("this is not a valid module"))
+    }
+
+    match gov_packet.action {
+        // 0 is reserved for upgrade / migration
+        1u8 => vaa_update_guardian_set(deps, env, &gov_packet.payload),
+        2u8 => handle_transfer_fee(deps, env, &gov_packet.payload),
         _ => ContractError::InvalidVAAAction.std_err(),
-    };
-
-    if result.is_ok() {
-        vaa_archive_add(&mut deps.storage, vaa.hash.as_slice())?;
     }
-
-    result
 }
 
 /// Parses raw VAA data into a struct and verifies whether it contains sufficient signatures of an
@@ -182,7 +181,7 @@ fn parse_and_verify_vaa<S: Storage>(
 fn vaa_update_guardian_set<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
     env: Env,
-    data: &[u8],
+    data: &Vec<u8>,
 ) -> StdResult<HandleResponse> {
     /* Payload format
     0   uint32 new_index
@@ -190,41 +189,19 @@ fn vaa_update_guardian_set<S: Storage, A: Api, Q: Querier>(
     5   [][20]uint8 guardian addresses
     */
 
-    const GUARDIAN_INDEX_POS: usize = 0;
-    const LENGTH_POS: usize = 4;
-    const ADDRESS_POS: usize = 5;
-    const ADDRESS_LEN: usize = 20;
-
-    if ADDRESS_POS >= data.len() {
-        return ContractError::InvalidVAA.std_err();
-    }
-
     let mut state = config_read(&deps.storage).load()?;
 
-    let new_guardian_set_index = data.get_u32(GUARDIAN_INDEX_POS);
+    let GuardianSetUpgrade {
+        new_guardian_set_index,
+        new_guardian_set,
+    } = GuardianSetUpgrade::deserialize(&data)?;
 
     if new_guardian_set_index != state.guardian_set_index + 1 {
         return ContractError::GuardianSetIndexIncreaseError.std_err();
     }
-    let len = data.get_u8(LENGTH_POS);
-
-    let mut new_guardian_set = GuardianSetInfo {
-        addresses: vec![],
-        expiration_time: 0,
-    };
-    let mut pos = ADDRESS_POS;
-    for _ in 0..len {
-        if pos + ADDRESS_LEN > data.len() {
-            return ContractError::InvalidVAA.std_err();
-        }
-
-        new_guardian_set.addresses.push(GuardianAddress {
-            bytes: data[pos..pos + ADDRESS_LEN].to_vec().into(),
-        });
-        pos += ADDRESS_LEN;
-    }
 
     let old_guardian_set_index = state.guardian_set_index;
+
     state.guardian_set_index = new_guardian_set_index;
 
     guardian_set_set(
@@ -232,6 +209,7 @@ fn vaa_update_guardian_set<S: Storage, A: Api, Q: Querier>(
         state.guardian_set_index,
         &new_guardian_set,
     )?;
+
     config(&mut deps.storage).save(&state)?;
 
     let mut old_guardian_set = guardian_set_get(&deps.storage, old_guardian_set_index)?;
@@ -249,6 +227,24 @@ fn vaa_update_guardian_set<S: Storage, A: Api, Q: Querier>(
     })
 }
 
+pub fn handle_transfer_fee<S: Storage, A: Api, Q: Querier>(
+    deps: &mut Extern<S, A, Q>,
+    env: Env,
+    data: &Vec<u8>,
+) -> StdResult<HandleResponse> {
+    let transfer_msg = TransferFee::deserialize(&data)?;
+
+    Ok(HandleResponse {
+        messages: vec![CosmosMsg::Bank(BankMsg::Send {
+            from_address: env.contract.address,
+            to_address: deps.api.human_address(&transfer_msg.recipient)?,
+            amount: vec![transfer_msg.amount],
+        })],
+        log: vec![],
+        data: None,
+    })
+}
+
 fn handle_post_message<S: Storage, A: Api, Q: Querier>(
     deps: &mut Extern<S, A, Q>,
     env: Env,
@@ -280,29 +276,6 @@ fn handle_post_message<S: Storage, A: Api, Q: Querier>(
     })
 }
 
-pub fn handle_transfer_fee<S: Storage, A: Api, Q: Querier>(
-    deps: &mut Extern<S, A, Q>,
-    env: Env,
-    amount: Coin,
-    recipient: HumanAddr,
-) -> StdResult<HandleResponse> {
-    let state = config_read(&deps.storage).load()?;
-
-    if deps.api.canonical_address(&env.message.sender)? != state.owner {
-        return ContractError::PermissionDenied.std_err();
-    }
-
-    Ok(HandleResponse {
-        messages: vec![CosmosMsg::Bank(BankMsg::Send {
-            from_address: env.contract.address,
-            to_address: recipient,
-            amount: vec![amount],
-        })],
-        log: vec![],
-        data: None,
-    })
-}
-
 pub fn query<S: Storage, A: Api, Q: Querier>(
     deps: &Extern<S, A, Q>,
     msg: QueryMsg,

terra/contracts/wormhole/src/msg.rs

@@ -6,6 +6,9 @@ use crate::state::{GuardianAddress, GuardianSetInfo};
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct InitMsg {
+    pub gov_chain: u16,
+    pub gov_address: Binary,
+
     pub initial_guardian_set: GuardianSetInfo,
     pub guardian_set_expirity: u64,
 }
@@ -20,10 +23,6 @@ pub enum HandleMsg {
         message: Binary,
         nonce: u32
     },
-    TransferFee {
-        amount: Coin,
-        recipient: HumanAddr,
-    },
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]

terra/contracts/wormhole/src/state.rs

@@ -1,7 +1,7 @@
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-use cosmwasm_std::{Binary, CanonicalAddr, HumanAddr, StdResult, Storage, Coin};
+use cosmwasm_std::{Binary, CanonicalAddr, HumanAddr, StdResult, Storage, Coin, Uint128};
 use cosmwasm_storage::{
     bucket, bucket_read, singleton, singleton_read, Bucket, ReadonlyBucket, ReadonlySingleton,
     Singleton,
@@ -26,8 +26,9 @@ pub struct ConfigInfo {
     // Period for which a guardian set stays active after it has been replaced
     pub guardian_set_expirity: u64,
 
-    // Contract owner address, it can make contract active/inactive
+    // governance contract details
-    pub owner: CanonicalAddr,
+    pub gov_chain: u16,
+    pub gov_address: Vec<u8>,
 
     // Asset locking fee
     pub fee: Coin,
@@ -219,20 +220,96 @@ pub fn wrapped_asset_address_read<S: Storage>(storage: &S) -> ReadonlyBucket<S,
 }
 
 
-pub struct WormholeGovernance {
+pub struct GovernancePacket {
+    pub module: Vec<u8>,
+    pub chain: u16,
     pub action: u8,
     pub payload: Vec<u8>,
 }
 
-impl WormholeGovernance {
+impl GovernancePacket {
     pub fn deserialize(data: &Vec<u8>) -> StdResult<Self> {
         let data = data.as_slice();
-        let action = data.get_u8(0);
+        let module = data.get_bytes32(0).to_vec();
-        let payload = &data[1..];
+        let chain = data.get_u16(32);
+        let action = data.get_u8(34);
+        let payload = data[35..].to_vec();
 
-        Ok(WormholeGovernance {
+        Ok(GovernancePacket {
+            module,
+            chain,
             action,
-            payload: payload.to_vec(),
+            payload
+        })
+    }
+}
+
+// action 1
+pub struct GuardianSetUpgrade {
+    pub new_guardian_set_index: u32,
+    pub new_guardian_set: GuardianSetInfo,
+}
+
+impl GuardianSetUpgrade {
+    pub fn deserialize(data: &Vec<u8>) -> StdResult<Self> {
+
+        const ADDRESS_LEN: usize = 20;
+
+        let data = data.as_slice();
+        let new_guardian_set_index = data.get_u32(0);
+
+        let n_guardians = data.get_u8(4);
+
+        let mut addresses = vec![];
+
+        for i in 0..n_guardians {
+            let pos = 5 + (i as usize) * ADDRESS_LEN;
+            if pos + ADDRESS_LEN > data.len() {
+                return ContractError::InvalidVAA.std_err();
+            }
+
+            addresses.push(GuardianAddress {
+                bytes: data[pos..pos + ADDRESS_LEN].to_vec().into(),
+            });
+        }
+
+        let new_guardian_set = GuardianSetInfo {
+            addresses,
+            expiration_time: 0
+        };
+
+        return Ok(
+            GuardianSetUpgrade {
+                new_guardian_set_index,
+                new_guardian_set
+            }
+        )
+    }
+}
+
+// action 2
+pub struct TransferFee {
+    pub amount: Coin,
+    pub recipient: CanonicalAddr,
+}
+
+impl TransferFee {
+    pub fn deserialize(data: &Vec<u8>) -> StdResult<Self> {
+        let data = data.as_slice();
+        let recipient = data.get_address(0);
+
+        let amount = Uint128(data.get_u128_be(32));
+        let denom = match String::from_utf8(data[48..].to_vec()) {
+            Ok(s) => s,
+            Err(_) => return ContractError::InvalidVAA.std_err()
+        };
+        let amount = Coin {
+            denom,
+            amount,
+        };
+        Ok(TransferFee {
+            amount,
+            recipient
         })
     }
 }

terra/deploy.py

@@ -5,6 +5,7 @@ from terra_sdk.core.wasm import (
     MsgStoreCode,
     MsgInstantiateContract,
     MsgExecuteContract,
+    MsgMigrateContract,
 )
 from terra_sdk.util.contract import get_code_id, get_contract_address, read_file_as_b64
 import os
@@ -74,9 +75,11 @@ class ContractQuerier:
 
 class Contract:
     @staticmethod
-    async def create(code_id, **kwargs):
+    async def create(code_id, migratable=False, **kwargs):
         kwargs = convert_contracts_to_addr(kwargs)
-        instantiate = MsgInstantiateContract(deployer.key.acc_address, code_id, kwargs)
+        instantiate = MsgInstantiateContract(
+            deployer.key.acc_address, code_id, kwargs, migratable=migratable
+        )
         result = await sign_and_broadcast(instantiate)
         return Contract(get_contract_address(result))
 
@@ -97,6 +100,15 @@ class Contract:
     def query(self):
         return ContractQuerier(self.address)
 
+    async def migrate(self, new_code_id):
+        migrate = MsgMigrateContract(
+            contract=self.address,
+            migrate_msg={},
+            new_code_id=new_code_id,
+            owner=deployer.key.acc_address,
+        )
+        return await sign_and_broadcast(migrate)
+
 
 def convert_contracts_to_addr(obj):
     if type(obj) == dict:
@@ -128,15 +140,29 @@ def assemble_vaa(emitter_chain, emitter_address, payload):
 async def main():
     code_ids = await store_contracts()
     print(code_ids)
+
+    # fake governance contract on solana
+    GOV_CHAIN = 1
+    GOV_ADDRESS = b"0" * 32
+
     wormhole = await Contract.create(
         code_id=code_ids["wormhole"],
+        gov_chain=GOV_CHAIN,
+        gov_address=base64.b64encode(GOV_ADDRESS).decode("utf-8"),
         guardian_set_expirity=10 ** 15,
         initial_guardian_set={"addresses": [], "expiration_time": 10 ** 15},
+        migratable=True,
     )
 
+    # TODO:
+    # resp = await wormhole.migrate(code_ids["wormhole"])
+    # for event in resp.logs:
+    #     pprint.pprint(event.events_by_type)
+
     token_bridge = await Contract.create(
         code_id=code_ids["token_bridge"],
-        owner=deployer.key.acc_address,
+        gov_chain=GOV_CHAIN,
+        gov_address=base64.b64encode(GOV_ADDRESS).decode("utf-8"),
         wormhole_contract=wormhole,
         wrapped_asset_code_id=int(code_ids["cw20_wrapped"]),
     )
@@ -163,9 +189,19 @@ async def main():
     bridge_canonical = bytes.fromhex(
         (await wormhole.query.query_address_hex(address=token_bridge))["hex"]
     )
-    await token_bridge.register_chain(
+
-        chain_id=3, chain_address=base64.b64encode(bridge_canonical).decode("utf-8")
+    # fake a VAA from the gov contract
-    )
+    module = b"token_bridge"
+    module += b" " * (32 - len(module))
+    chain = to_bytes(0, 2)
+    action = to_bytes(0, 1)
+    #         chain_id         chain_address (pretend there's a bridge w/ the same address on solana)
+    payload = to_bytes(3, 2) + bridge_canonical
+
+    vaa = assemble_vaa(GOV_CHAIN, GOV_ADDRESS, module + chain + action + payload)
+
+    # register the chain
+    await token_bridge.submit_vaa(data=base64.b64encode(vaa).decode("utf-8"))
 
     resp = await token_bridge.initiate_transfer(
         asset=mock_token,
@@ -199,9 +235,18 @@ async def main():
     )
 
     # pretend there exists another bridge contract with the same address but on solana
-    await token_bridge.register_chain(
+    # fake a VAA from the gov contract
-        chain_id=1, chain_address=base64.b64encode(bridge_canonical).decode("utf-8")
+    module = b"token_bridge"
-    )
+    module += b" " * (32 - len(module))
+    chain = to_bytes(0, 2)
+    action = to_bytes(0, 1)
+    #         chain_id         chain_address (pretend there's a bridge w/ the same address on solana)
+    payload = to_bytes(1, 2) + bridge_canonical
+
+    vaa = assemble_vaa(GOV_CHAIN, GOV_ADDRESS, module + chain + action + payload)
+
+    # register the chain
+    await token_bridge.submit_vaa(data=base64.b64encode(vaa).decode("utf-8"))
 
     resp = await token_bridge.create_asset_meta(
         asset_address=mock_token,