babylon query plugin

babylonlabs-io · Jan 7, 2025 · ac1b797 · ac1b797
1 parent f57af50
commit ac1b797
Show file tree

Hide file tree

Showing 15 changed files with 211 additions and 176 deletions.
diff --git a/contracts/babylon/src/contract.rs b/contracts/babylon/src/contract.rs
@@ -1,12 +1,10 @@
-use babylon_bindings::babylon_sdk::{
-    get_babylon_sdk_params, QueryParamsResponse, QUERY_PARAMS_PATH,
-};
+use babylon_bindings::query::{get_babylon_sdk_params, BabylonQuery};
 use cosmwasm_std::{
     to_json_binary, to_json_string, Addr, Binary, Deps, DepsMut, Empty, Env, IbcMsg, MessageInfo,
-    QueryResponse, Reply, Response, SubMsg, SubMsgResponse, WasmMsg,
+    QueryResponse, Response, SubMsg, SubMsgResponse, WasmMsg,
 };
 use cw2::set_contract_version;
-use cw_utils::{must_pay, ParseReplyError};
+use cw_utils::must_pay;
 
 use babylon_apis::{btc_staking_api, finality_api, to_bech32_addr, to_module_canonical_addr};
 use babylon_bindings::BabylonMsg;
@@ -21,9 +19,6 @@ use crate::state::config::{Config, CONFIG};
 pub const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
 pub const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
-const REPLY_ID_INSTANTIATE_STAKING: u64 = 2;
-const REPLY_ID_INSTANTIATE_FINALITY: u64 = 3;
-
 /// When we instantiate the Babylon contract, it will optionally instantiate a BTC staking
 /// contract – if its code id is provided – to work with it for BTC re-staking support,
 /// as they both need references to each other.
@@ -81,23 +76,11 @@ pub fn instantiate(
     Ok(res)
 }
 
-/// Tries to get contract address from events in reply
-fn reply_init_get_contract_address(reply: SubMsgResponse) -> Result<Addr, ContractError> {
-    for event in reply.events {
-        if event.ty == "instantiate" {
-            for attr in event.attributes {
-                if attr.key == "_contract_address" {
-                    return Ok(Addr::unchecked(attr.value));
-                }
-            }
-        }
-    }
-    Err(ContractError::ParseReply(ParseReplyError::ParseFailure(
-        "Cannot parse contract address".to_string(),
-    )))
-}
-
-pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, ContractError> {
+pub fn query(
+    deps: Deps<BabylonQuery>,
+    _env: Env,
+    msg: QueryMsg,
+) -> Result<QueryResponse, ContractError> {
     match msg {
         QueryMsg::Config {} => Ok(to_json_binary(&queries::config(deps)?)?),
         QueryMsg::BtcBaseHeader {} => Ok(to_json_binary(&queries::btc_base_header(deps)?)?),
@@ -133,15 +116,15 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, Cont
 
 /// this is a no-op just to test how this integrates with wasmd
 pub fn migrate(
-    _deps: DepsMut,
+    _deps: DepsMut<BabylonQuery>,
     _env: Env,
     _msg: Empty,
 ) -> Result<Response<BabylonMsg>, ContractError> {
     Ok(Response::default())
 }
 
 pub fn execute(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     info: MessageInfo,
     msg: ExecuteMsg,
@@ -202,9 +185,7 @@ pub fn execute(
             // Assert the funds are there
             must_pay(&info, &cfg.denom)?;
             // Assert the sender is right
-            let btc_finality = cfg
-                .btc_finality
-                .ok_or(ContractError::BtcFinalityNotSet {})?;
+            let btc_finality = get_babylon_sdk_params(&deps.querier)?.btc_finality_contract_address;
             if info.sender != btc_finality {
                 return Err(ContractError::Unauthorized {});
             }
@@ -298,7 +279,6 @@ mod tests {
         let info = message_info(&deps.api.addr_make(CREATOR), &[]);
         let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap();
         assert_eq!(1, res.messages.len());
-        assert_eq!(REPLY_ID_INSTANTIATE_FINALITY, res.messages[0].id);
         assert_eq!(
             res.messages[0].msg,
             WasmMsg::Instantiate {
@@ -330,7 +310,6 @@ mod tests {
         let info = message_info(&deps.api.addr_make(CREATOR), &[]);
         let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap();
         assert_eq!(1, res.messages.len());
-        assert_eq!(REPLY_ID_INSTANTIATE_FINALITY, res.messages[0].id);
         assert_eq!(
             res.messages[0].msg,
             WasmMsg::Instantiate {

diff --git a/contracts/babylon/src/ibc.rs b/contracts/babylon/src/ibc.rs
@@ -1,5 +1,5 @@
 use crate::error::ContractError;
-use babylon_bindings::BabylonMsg;
+use babylon_bindings::{query::BabylonQuery, BabylonMsg};
 use babylon_proto::babylon::zoneconcierge::v1::{
     zoneconcierge_packet_data::Packet, BtcTimestamp, ZoneconciergePacketData,
 };
@@ -35,7 +35,7 @@ pub const IBC_TRANSFER: Item<TransferInfo> = Item::new("ibc_transfer");
 /// In the case of ChannelOpenTry there's a counterparty_version attribute in the message.
 /// Here we ensure the ordering and version constraints.
 pub fn ibc_channel_open(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     _env: Env,
     msg: IbcChannelOpenMsg,
 ) -> Result<IbcChannelOpenResponse, ContractError> {
@@ -67,7 +67,7 @@ pub fn ibc_channel_open(
 
 /// Second part of the 4-step handshake, i.e. ChannelOpenAck and ChannelOpenConfirm.
 pub fn ibc_channel_connect(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     _env: Env,
     msg: IbcChannelConnectMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
@@ -102,7 +102,7 @@ pub fn ibc_channel_connect(
 /// This is invoked on the IBC Channel Close message
 /// We perform any cleanup related to the channel
 pub fn ibc_channel_close(
-    _deps: DepsMut,
+    _deps: DepsMut<BabylonQuery>,
     _env: Env,
     msg: IbcChannelCloseMsg,
 ) -> StdResult<IbcBasicResponse> {
@@ -126,7 +126,7 @@ pub fn ibc_channel_close(
 /// That's because we want to send an ACK for the packet regardless if there's an error or not,
 /// but in the case of an error, we do not want the state to be committed.
 pub fn ibc_packet_receive(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     _env: Env,
     msg: IbcPacketReceiveMsg,
 ) -> Result<IbcReceiveResponse<BabylonMsg>, Never> {
@@ -174,16 +174,14 @@ pub(crate) mod ibc_packet {
         ActiveBtcDelegation, NewFinalityProvider, UnbondedBtcDelegation,
     };
     use babylon_apis::finality_api::Evidence;
-    use babylon_bindings::babylon_sdk::{
-        get_babylon_sdk_params, QueryParamsResponse, QUERY_PARAMS_PATH,
-    };
+    use babylon_bindings::query::get_babylon_sdk_params;
     use babylon_proto::babylon::btcstaking::v1::BtcStakingIbcPacket;
     use babylon_proto::babylon::zoneconcierge::v1::zoneconcierge_packet_data::Packet::ConsumerSlashing;
     use babylon_proto::babylon::zoneconcierge::v1::ConsumerSlashingIbcPacket;
     use cosmwasm_std::{to_json_binary, IbcChannel, IbcMsg, WasmMsg};
 
     pub fn handle_btc_timestamp(
-        deps: DepsMut,
+        deps: DepsMut<BabylonQuery>,
         _caller: String,
         btc_ts: &BtcTimestamp,
     ) -> StdResult<IbcReceiveResponse<BabylonMsg>> {
@@ -212,7 +210,7 @@ pub(crate) mod ibc_packet {
     }
 
     pub fn handle_btc_staking(
-        deps: DepsMut,
+        deps: DepsMut<BabylonQuery>,
         _caller: String,
         btc_staking: &BtcStakingIbcPacket,
     ) -> StdResult<IbcReceiveResponse<BabylonMsg>> {
@@ -303,15 +301,15 @@ pub fn packet_timeout(env: &Env) -> IbcTimeout {
 }
 
 pub fn ibc_packet_ack(
-    _deps: DepsMut,
+    _deps: DepsMut<BabylonQuery>,
     _env: Env,
     _msg: IbcPacketAckMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
     Ok(IbcBasicResponse::default())
 }
 
 pub fn ibc_packet_timeout(
-    _deps: DepsMut,
+    _deps: DepsMut<BabylonQuery>,
     _env: Env,
     msg: IbcPacketTimeoutMsg,
 ) -> Result<IbcBasicResponse, ContractError> {

diff --git a/contracts/babylon/src/lib.rs b/contracts/babylon/src/lib.rs
@@ -6,7 +6,7 @@ use cosmwasm_std::{
     IbcPacketTimeoutMsg, IbcReceiveResponse, MessageInfo, Never, Reply, Response, StdResult,
 };
 
-use babylon_bindings::BabylonMsg;
+use babylon_bindings::{query::BabylonQuery, BabylonMsg};
 
 use crate::error::ContractError;
 pub use crate::msg::contract::ExecuteMsg;
@@ -34,18 +34,26 @@ pub fn instantiate(
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn query(deps: Deps, env: Env, msg: msg::contract::QueryMsg) -> Result<Binary, ContractError> {
+pub fn query(
+    deps: Deps<BabylonQuery>,
+    env: Env,
+    msg: msg::contract::QueryMsg,
+) -> Result<Binary, ContractError> {
     contract::query(deps, env, msg)
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn migrate(deps: DepsMut, env: Env, msg: Empty) -> Result<Response<BabylonMsg>, ContractError> {
+pub fn migrate(
+    deps: DepsMut<BabylonQuery>,
+    env: Env,
+    msg: Empty,
+) -> Result<Response<BabylonMsg>, ContractError> {
     contract::migrate(deps, env, msg)
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn execute(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     info: MessageInfo,
     msg: ExecuteMsg,
@@ -55,7 +63,7 @@ pub fn execute(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_channel_open(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcChannelOpenMsg,
 ) -> Result<IbcChannelOpenResponse, error::ContractError> {
@@ -64,7 +72,7 @@ pub fn ibc_channel_open(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_channel_connect(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcChannelConnectMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
@@ -73,7 +81,7 @@ pub fn ibc_channel_connect(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_channel_close(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcChannelCloseMsg,
 ) -> StdResult<IbcBasicResponse> {
@@ -82,7 +90,7 @@ pub fn ibc_channel_close(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_packet_receive(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcPacketReceiveMsg,
 ) -> Result<IbcReceiveResponse<BabylonMsg>, Never> {
@@ -91,7 +99,7 @@ pub fn ibc_packet_receive(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_packet_ack(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcPacketAckMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
@@ -100,7 +108,7 @@ pub fn ibc_packet_ack(
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn ibc_packet_timeout(
-    deps: DepsMut,
+    deps: DepsMut<BabylonQuery>,
     env: Env,
     msg: IbcPacketTimeoutMsg,
 ) -> Result<IbcBasicResponse, ContractError> {

diff --git a/contracts/babylon/src/multitest.rs b/contracts/babylon/src/multitest.rs
@@ -58,9 +58,6 @@ mod instantiation {
 
         // Confirm the btc-staking contract has been instantiated and set
         let config = suite.get_config();
-        assert_eq!(config.btc_staking, Some(Addr::unchecked(CONTRACT1_ADDR)));
-        // Confirm the btc-finality contract has been instantiated and set
-        assert_eq!(config.btc_finality, Some(Addr::unchecked(CONTRACT2_ADDR)));
     }
 
     #[test]
@@ -82,9 +79,6 @@ mod instantiation {
 
         // Confirm the btc-staking contract has been instantiated and set
         let config = suite.get_config();
-        assert_eq!(config.btc_staking, Some(Addr::unchecked(CONTRACT1_ADDR)));
-        // Confirm the btc-finality contract has been instantiated and set
-        assert_eq!(config.btc_finality, Some(Addr::unchecked(CONTRACT2_ADDR)));
     }
 
     #[test]

diff --git a/contracts/babylon/src/multitest/suite.rs b/contracts/babylon/src/multitest/suite.rs
@@ -3,6 +3,7 @@ use crate::msg::ibc::TransferInfoResponse;
 use crate::msg::ibc::{IbcTransferInfo, Recipient};
 use crate::state::config::Config;
 use anyhow::Result as AnyResult;
+use babylon_bindings::query::BabylonQuery;
 use derivative::Derivative;
 
 use cosmwasm_std::{Addr, Binary, Empty};
@@ -12,11 +13,7 @@ use babylon_bindings::BabylonMsg;
 use babylon_bindings_test::BabylonApp;
 use babylon_bitcoin::chain_params::Network;
 
-use crate::msg::contract::{InstantiateMsg, QueryMsg};
-use crate::multitest::{CONTRACT1_ADDR, CONTRACT2_ADDR};
-use crate::state::config::Config;
-
-fn contract_btc_staking() -> Box<dyn Contract<BabylonMsg>> {
+fn contract_btc_staking() -> Box<dyn Contract<BabylonMsg, BabylonQuery>> {
     let contract = ContractWrapper::new(
         btc_staking::contract::execute,
         btc_staking::contract::instantiate,
@@ -25,7 +22,7 @@ fn contract_btc_staking() -> Box<dyn Contract<BabylonMsg>> {
     Box::new(contract)
 }
 
-fn contract_btc_finality() -> Box<dyn Contract<BabylonMsg>> {
+fn contract_btc_finality() -> Box<dyn Contract<BabylonMsg, BabylonQuery>> {
     let contract = ContractWrapper::new(
         btc_finality::contract::execute,
         btc_finality::contract::instantiate,
@@ -34,7 +31,7 @@ fn contract_btc_finality() -> Box<dyn Contract<BabylonMsg>> {
     Box::new(contract)
 }
 
-fn contract_babylon() -> Box<dyn Contract<BabylonMsg>> {
+fn contract_babylon() -> Box<dyn Contract<BabylonMsg, BabylonQuery>> {
     let contract = ContractWrapper::new(crate::execute, crate::instantiate, crate::query)
         .with_migrate(crate::migrate);
     Box::new(contract)
@@ -154,7 +151,7 @@ impl SuiteBuilder {
         Suite {
             app,
             code_id: contract_code_id,
-            babylon_contract: babylon_contract,
+            babylon_contract: contract,
             btc_staking_contract: btc_staking_contract,
             btc_finality_contract: btc_finality_contract,
             owner,
@@ -218,7 +215,7 @@ impl Suite {
     pub fn get_transfer_info(&self) -> TransferInfoResponse {
         self.app
             .wrap()
-            .query_wasm_smart(self.contract.clone(), &QueryMsg::TransferInfo {})
+            .query_wasm_smart(self.babylon_contract.clone(), &QueryMsg::TransferInfo {})
             .unwrap()
     }