Update / Set staking address during finality contract instantiation

contracts/babylon/src/contract.rs

@@ -5,7 +5,7 @@ use cosmwasm_std::{
 use cw2::set_contract_version;
 use cw_utils::ParseReplyError;
 
-use babylon_apis::btc_staking_api;
+use babylon_apis::{btc_staking_api, finality_api};
 use babylon_bindings::BabylonMsg;
 
 use crate::error::ContractError;
@@ -135,12 +135,25 @@ fn reply_init_finality_callback(
     reply: SubMsgResponse,
 ) -> Result<Response<BabylonMsg>, ContractError> {
     // Try to get contract address from events in reply
-    let addr = reply_init_get_contract_address(reply)?;
+    let finality_addr = reply_init_get_contract_address(reply)?;
     CONFIG.update(deps.storage, |mut cfg| {
-        cfg.btc_finality = Some(addr);
+        cfg.btc_finality = Some(finality_addr.clone());
         Ok::<_, ContractError>(cfg)
     })?;
-    Ok(Response::new())
+    // Set the BTC staking contract address to the BTC finality contract
+    let cfg = CONFIG.load(deps.storage)?;
+    let msg = finality_api::ExecuteMsg::UpdateStaking {
+        staking: cfg
+            .btc_staking
+            .ok_or(ContractError::BtcStakingNotSet {})?
+            .to_string(),
+    };
+    let wasm_msg = WasmMsg::Execute {
+        contract_addr: finality_addr.to_string(),
+        msg: to_json_binary(&msg)?,
+        funds: vec![],
+    };
+    Ok(Response::new().add_message(wasm_msg))
 }
 
 pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, ContractError> {
@@ -263,6 +276,8 @@ mod tests {
             notify_cosmos_zone: false,
             btc_staking_code_id: None,
             btc_staking_msg: None,
+            btc_finality_code_id: None,
+            btc_finality_msg: None,
             admin: None,
             consumer_name: None,
             consumer_description: None,

contracts/babylon/src/ibc.rs

@@ -434,6 +434,8 @@ mod tests {
             notify_cosmos_zone: false,
             btc_staking_code_id: None,
             btc_staking_msg: None,
+            btc_finality_code_id: None,
+            btc_finality_msg: None,
             admin: None,
             consumer_name: None,
             consumer_description: None,

contracts/babylon/src/multitest/suite.rs

@@ -66,6 +66,8 @@ impl SuiteBuilder {
                     notify_cosmos_zone: false,
                     btc_staking_code_id: Some(btc_staking_code_id),
                     btc_staking_msg: None,
+                    btc_finality_code_id: None,
+                    btc_finality_msg: None,
                     admin: Some(owner.to_string()),
                     consumer_name: Some("TestConsumer".to_string()),
                     consumer_description: Some("Test Consumer Description".to_string()),

contracts/babylon/src/state/btc_light_client.rs

@@ -357,6 +357,7 @@ pub(crate) mod tests {
             checkpoint_finalization_timeout: w as u64,
             notify_cosmos_zone: false,
             btc_staking: None,
+            btc_finality: None,
             consumer_name: None,
             consumer_description: None,
         };

contracts/babylon/tests/integration.rs

@@ -49,6 +49,8 @@ fn setup() -> Instance<MockApi, MockStorage, MockQuerier> {
         notify_cosmos_zone: false,
         btc_staking_code_id: None,
         btc_staking_msg: None,
+        btc_finality_code_id: None,
+        btc_finality_msg: None,
         admin: None,
     };
     let info = message_info(&Addr::unchecked(CREATOR), &[]);
@@ -101,6 +103,8 @@ fn instantiate_works() {
         notify_cosmos_zone: false,
         btc_staking_code_id: None,
         btc_staking_msg: None,
+        btc_finality_code_id: None,
+        btc_finality_msg: None,
         admin: None,
     };
     let info = message_info(&Addr::unchecked(CREATOR), &[]);

contracts/btc-finality/src/contract.rs

@@ -1,23 +1,24 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    to_json_binary, Addr, CustomQuery, Deps, DepsMut, Empty, Env, MessageInfo, QuerierWrapper,
-    QueryRequest, QueryResponse, Reply, Response, StdResult, WasmQuery,
+    attr, to_json_binary, Addr, CustomQuery, Deps, DepsMut, Empty, Env, MessageInfo,
+    QuerierWrapper, QueryRequest, QueryResponse, Reply, Response, StdResult, WasmQuery,
 };
 use cw2::set_contract_version;
 use cw_utils::{maybe_addr, nonpayable};
 
 use babylon_apis::btc_staking_api::SudoMsg;
 use babylon_bindings::BabylonMsg;
 
+use btc_staking::msg::ActivatedHeightResponse;
+
 use crate::error::ContractError;
 use crate::finality::{
     compute_active_finality_providers, handle_finality_signature, handle_public_randomness_commit,
 };
 use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
 use crate::state::config::{Config, ADMIN, CONFIG, PARAMS};
 use crate::{finality, queries, state};
-use btc_staking::msg::ActivatedHeightResponse;
 
 pub const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
 pub const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -34,7 +35,7 @@ pub fn instantiate(
     let config = Config {
         denom,
         babylon: info.sender,
-        staking: Addr::unchecked("staking"), // TODO: instantiate staking contract and set address in reply
+        staking: Addr::unchecked("UNSET"), // To be set later, through `UpdateStaking`
     };
     CONFIG.save(deps.storage, &config)?;
 
@@ -120,6 +121,7 @@ pub fn execute(
         ExecuteMsg::UpdateAdmin { admin } => ADMIN
             .execute_update_admin(deps, info, maybe_addr(api, admin)?)
             .map_err(Into::into),
+        ExecuteMsg::UpdateStaking { staking } => handle_update_staking(deps, info, staking),
         ExecuteMsg::SubmitFinalitySignature {
             fp_pubkey_hex,
             height,
@@ -169,6 +171,26 @@ pub fn sudo(
     }
 }
 
+fn handle_update_staking(
+    deps: DepsMut,
+    info: MessageInfo,
+    staking_addr: String,
+) -> Result<Response<BabylonMsg>, ContractError> {
+    let mut cfg = CONFIG.load(deps.storage)?;
+    if info.sender != cfg.babylon && !ADMIN.is_admin(deps.as_ref(), &info.sender)? {
+        return Err(ContractError::Unauthorized {});
+    }
+    cfg.staking = deps.api.addr_validate(&staking_addr)?;
+    CONFIG.save(deps.storage, &cfg)?;
+
+    let attributes = vec![
+        attr("action", "update_btc_staking"),
+        attr("staking", staking_addr),
+        attr("sender", info.sender),
+    ];
+    Ok(Response::new().add_attributes(attributes))
+}
+
 fn handle_begin_block(deps: &mut DepsMut, env: Env) -> Result<Response<BabylonMsg>, ContractError> {
     // Compute active finality provider set
     let max_active_fps = PARAMS.load(deps.storage)?.max_active_finality_providers as usize;

packages/apis/src/finality_api.rs

@@ -13,6 +13,9 @@ use crate::Bytes;
 pub enum ExecuteMsg {
     /// Change the admin
     UpdateAdmin { admin: Option<String> },
+    /// Set the BTC staking addr.
+    /// Only admin or the babylon contract can set this
+    UpdateStaking { staking: String },
     /// Committing a sequence of public randomness for EOTS
     CommitPublicRandomness {
         /// `fp_pubkey_hex` is the BTC PK of the finality provider that commits the public randomness