From 49feacf8c86e83276adac35e0d74b0e4d2d059eb Mon Sep 17 00:00:00 2001
From: Runchao Han <me@runchao.rocks>
Date: Tue, 13 Aug 2024 16:23:42 +1000
Subject: [PATCH] Fix/cross contract query (#40)

---
 contracts/btc-staking/src/contract.rs | 28 +++++++++++++++------------
 packages/apis/src/btc_staking_api.rs  | 11 -----------
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/contracts/btc-staking/src/contract.rs b/contracts/btc-staking/src/contract.rs
index caf66b0d..17400ddf 100644
--- a/contracts/btc-staking/src/contract.rs
+++ b/contracts/btc-staking/src/contract.rs
@@ -1,5 +1,4 @@
-use prost::bytes::Bytes;
-use prost::Message;
+use babylon_contract::msg::btc_header::BtcHeaderResponse;
 
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
@@ -12,9 +11,8 @@ use cw_utils::{maybe_addr, nonpayable};
 
 use babylon_apis::btc_staking_api::SudoMsg;
 use babylon_bindings::BabylonMsg;
-use babylon_proto::babylon::btclightclient::v1::BtcHeaderInfo;
 
-use babylon_contract::state::btc_light_client::BTC_TIP_KEY;
+use babylon_contract::msg::contract::QueryMsg as BabylonQueryMsg;
 
 use crate::error::ContractError;
 use crate::finality::{handle_finality_signature, handle_public_randomness_commit};
@@ -231,20 +229,26 @@ fn handle_begin_block(deps: &mut DepsMut, env: Env) -> Result<Response<BabylonMs
 // index_btc_height indexes the current BTC height, and saves it to the state
 #[allow(dead_code)]
 fn index_btc_height(deps: &mut DepsMut, height: u64) -> Result<(), ContractError> {
-    let btc_tip = get_btc_tip(deps)?;
+    // FIXME: Turn this into a hard error. Requires `babylon-contract` instance, and up and running
+    // BTC light client loop (which requires a running BTC node / simulator)
+    let btc_tip_height = get_btc_tip_height(deps) //?;
+        .ok()
+        .unwrap_or_default();
 
-    Ok(BTC_HEIGHT.save(deps.storage, height, &btc_tip.height)?)
+    Ok(BTC_HEIGHT.save(deps.storage, height, &btc_tip_height)?)
 }
 
-/// get_btc_tip queries the Babylon contract for the latest BTC tip
-#[allow(dead_code)]
-fn get_btc_tip(deps: &DepsMut) -> Result<BtcHeaderInfo, ContractError> {
+/// get_btc_tip_height queries the Babylon contract for the latest BTC tip height
+fn get_btc_tip_height(deps: &DepsMut) -> Result<u64, ContractError> {
     // Get the BTC tip from the babylon contract through a raw query
     let babylon_addr = CONFIG.load(deps.storage)?.babylon;
-    let query = babylon_apis::encode_raw_query(&babylon_addr, BTC_TIP_KEY.as_bytes());
 
-    let tip_bytes: Bytes = deps.querier.query(&query)?;
-    Ok(BtcHeaderInfo::decode(tip_bytes)?)
+    // Query the Babylon contract
+    // TODO: use a raw query for performance / efficiency
+    let query_msg = BabylonQueryMsg::BtcTipHeader {};
+    let tip: BtcHeaderResponse = deps.querier.query_wasm_smart(babylon_addr, &query_msg)?;
+
+    Ok(tip.height)
 }
 
 fn handle_end_block(
diff --git a/packages/apis/src/btc_staking_api.rs b/packages/apis/src/btc_staking_api.rs
index 6663fc53..3f78cbf4 100644
--- a/packages/apis/src/btc_staking_api.rs
+++ b/packages/apis/src/btc_staking_api.rs
@@ -150,17 +150,6 @@ impl FinalityProviderDescription {
     pub const MAX_DETAILS_LENGTH: usize = 280;
 }
 
-/// PubKey defines a secp256k1 public key.
-/// Key is the compressed form of the pubkey. The first byte is a 0x02 byte
-/// if the y-coordinate is the lexicographically largest of the two associated with
-/// the x-coordinate. Otherwise, the first byte is a 0x03.
-/// This prefix is followed with the x-coordinate.
-#[cw_serde]
-pub struct PubKey {
-    /// key is the compressed public key of the finality provider
-    pub key: Binary,
-}
-
 /// ProofOfPossessionBtc is the proof of possession that a Babylon secp256k1
 /// secret key and a Bitcoin secp256k1 secret key are held by the same
 /// person