From 09459fcf9e679238b75c7e2b2b9a5f1db9adc7f5 Mon Sep 17 00:00:00 2001
From: Chirantan Ekbote <cekbote@jumptrading.com>
Date: Mon, 24 Oct 2022 14:49:19 +0900
Subject: [PATCH] cosmwasm: tokenbridge: Add chain registration query

Add a query to get the registered emitterr address for a given chain id.
This will be used by accounting.
---
 .../contracts/token-bridge/src/contract.rs    | 14 +++++++++--
 cosmwasm/contracts/token-bridge/src/msg.rs    |  7 ++++++
 cosmwasm/test/src/__tests__/bridge.ts         | 24 ++++++++++++++++++-
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/cosmwasm/contracts/token-bridge/src/contract.rs b/cosmwasm/contracts/token-bridge/src/contract.rs
index 5b5eb349a..6b41bd9bb 100644
--- a/cosmwasm/contracts/token-bridge/src/contract.rs
+++ b/cosmwasm/contracts/token-bridge/src/contract.rs
@@ -31,8 +31,8 @@ use cosmwasm_std::{
 
 use crate::{
     msg::{
-        ExecuteMsg, ExternalIdResponse, InstantiateMsg, IsVaaRedeemedResponse, MigrateMsg,
-        QueryMsg, TransferInfoResponse, WrappedRegistryResponse,
+        ChainRegistrationResponse, ExecuteMsg, ExternalIdResponse, InstantiateMsg,
+        IsVaaRedeemedResponse, MigrateMsg, QueryMsg, TransferInfoResponse, WrappedRegistryResponse,
     },
     state::{
         bridge_contracts, bridge_contracts_read, bridge_deposit, config, config_read,
@@ -1449,6 +1449,9 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::TransferInfo { vaa } => to_binary(&query_transfer_info(deps, env, &vaa)?),
         QueryMsg::ExternalId { external_id } => to_binary(&query_external_id(deps, external_id)?),
         QueryMsg::IsVaaRedeemed { vaa } => to_binary(&query_is_vaa_redeemed(deps, env, &vaa)?),
+        QueryMsg::ChainRegistration { chain } => {
+            query_chain_registration(deps, chain).and_then(|r| to_binary(&r))
+        }
     }
 }
 
@@ -1526,6 +1529,13 @@ fn query_is_vaa_redeemed(deps: Deps, _env: Env, vaa: &Binary) -> StdResult<IsVaa
     })
 }
 
+fn query_chain_registration(deps: Deps, chain: u16) -> StdResult<ChainRegistrationResponse> {
+    bridge_contracts_read(deps.storage)
+        .load(&chain.to_be_bytes())
+        .map(Binary::from)
+        .map(|address| ChainRegistrationResponse { address })
+}
+
 fn is_governance_emitter(cfg: &ConfigInfo, emitter_chain: u16, emitter_address: &[u8]) -> bool {
     cfg.gov_chain == emitter_chain && cfg.gov_address == emitter_address
 }
diff --git a/cosmwasm/contracts/token-bridge/src/msg.rs b/cosmwasm/contracts/token-bridge/src/msg.rs
index 98a291eec..cdda712b3 100644
--- a/cosmwasm/contracts/token-bridge/src/msg.rs
+++ b/cosmwasm/contracts/token-bridge/src/msg.rs
@@ -79,6 +79,7 @@ pub enum QueryMsg {
     TransferInfo { vaa: Binary },
     ExternalId { external_id: Binary },
     IsVaaRedeemed { vaa: Binary },
+    ChainRegistration { chain: u16 },
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
@@ -110,3 +111,9 @@ pub struct ExternalIdResponse {
 pub struct IsVaaRedeemedResponse {
     pub is_redeemed: bool,
 }
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub struct ChainRegistrationResponse {
+    pub address: Binary,
+}
diff --git a/cosmwasm/test/src/__tests__/bridge.ts b/cosmwasm/test/src/__tests__/bridge.ts
index a303e52cf..6a2404a72 100644
--- a/cosmwasm/test/src/__tests__/bridge.ts
+++ b/cosmwasm/test/src/__tests__/bridge.ts
@@ -262,6 +262,7 @@ test("Deploy Contracts", (done) => {
   })();
 });
 
+passOnceTests(false, 1);
 alwaysPassTests(false, 1);
 
 describe("Upgrade to shutdown contracts Tests", () => {
@@ -350,7 +351,7 @@ describe("Upgrade to shutdown contracts Tests", () => {
 });
 
 alwaysPassTests(true, 2);
-passOnceTests(true, 1);
+passOnceTests(true, 2);
 failInShutdownModeTests(true, 1);
 
 describe("Upgrade to previous non-shutdown contracts Tests", () => {
@@ -1374,6 +1375,27 @@ function alwaysPassTests(shutdownMode: boolean, pass: number) {
         }
       })();
     });
+    test("Query Chain Registration", (done) => {
+      (async () => {
+        try {
+          const [client] = await makeProviderAndWallet();
+
+          const tokenbridge = contracts.get("tokenBridge")!;
+          const result: any = await client.wasm.contractQuery(tokenbridge, {
+            chain_registration: {
+              chain: FOREIGN_CHAIN,
+            },
+          });
+          expect(result).toStrictEqual({
+            address: Buffer.from(FOREIGN_TOKEN_BRIDGE, "hex").toString("base64")
+          });
+          done();
+        } catch (e) {
+          console.error(e);
+          done("Failed to Query Chain Registration");
+        }
+      })();
+    });
   });
 }