Add update finality handler / msg to staking contract

babylonlabs-io · Jan 29, 2025 · 749400c · 749400c
1 parent 2497493
commit 749400c
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 9 deletions.
diff --git a/contracts/babylon/src/contract.rs b/contracts/babylon/src/contract.rs
@@ -172,20 +172,33 @@ fn reply_init_finality_callback(
         cfg.btc_finality = Some(finality_addr.clone());
         Ok::<_, ContractError>(cfg)
     })?;
-    // Set the BTC staking contract address to the BTC finality contract
+    // Set the BTC finality contract address to the BTC staking contract
     let cfg = CONFIG.load(deps.storage)?;
-    let msg = finality_api::ExecuteMsg::UpdateStaking {
-        staking: cfg
-            .btc_staking
-            .ok_or(ContractError::BtcStakingNotSet {})?
+    let msg = btc_staking_api::ExecuteMsg::UpdateFinality {
+        finality: cfg
+            .btc_finality
+            .ok_or(ContractError::BtcFinalityNotSet {})?
             .to_string(),
     };
-    let wasm_msg = WasmMsg::Execute {
+    let staking_addr = cfg.btc_staking.ok_or(ContractError::BtcStakingNotSet {})?;
+    let wasm_msg_1 = WasmMsg::Execute {
+        contract_addr: staking_addr.to_string(),
+        msg: to_json_binary(&msg)?,
+        funds: vec![],
+    };
+
+    // Set the BTC staking contract address to the BTC finality contract
+    let msg = finality_api::ExecuteMsg::UpdateStaking {
+        staking: staking_addr.to_string(),
+    };
+    let wasm_msg_2 = WasmMsg::Execute {
         contract_addr: finality_addr.to_string(),
         msg: to_json_binary(&msg)?,
         funds: vec![],
     };
-    Ok(Response::new().add_message(wasm_msg))
+    Ok(Response::new()
+        .add_message(wasm_msg_1)
+        .add_message(wasm_msg_2))
 }
 
 pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, ContractError> {

diff --git a/contracts/btc-staking/src/contract.rs b/contracts/btc-staking/src/contract.rs
@@ -1,8 +1,8 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    to_json_binary, Deps, DepsMut, Empty, Env, MessageInfo, QueryResponse, Reply, Response,
-    StdResult,
+    attr, to_json_binary, Addr, Deps, DepsMut, Empty, Env, MessageInfo, QueryResponse, Reply,
+    Response, StdResult,
 };
 use cw2::set_contract_version;
 use cw_utils::{maybe_addr, nonpayable};
@@ -31,6 +31,7 @@ pub fn instantiate(
     let denom = deps.querier.query_bonded_denom()?;
     let config = Config {
         babylon: info.sender,
+        finality: Addr::unchecked("UNSET"), // To be set later, through `UpdateFinality`
         denom,
     };
     CONFIG.save(deps.storage, &config)?;
@@ -128,6 +129,7 @@ pub fn execute(
         ExecuteMsg::UpdateAdmin { admin } => ADMIN
             .execute_update_admin(deps, info, maybe_addr(api, admin)?)
             .map_err(Into::into),
+        ExecuteMsg::UpdateFinality { finality } => handle_update_finality(deps, info, finality),
         ExecuteMsg::BtcStaking {
             new_fp,
             active_del,
@@ -157,6 +159,26 @@ pub fn execute(
     }
 }
 
+fn handle_update_finality(
+    deps: DepsMut,
+    info: MessageInfo,
+    finality_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.finality = deps.api.addr_validate(&finality_addr)?;
+    CONFIG.save(deps.storage, &cfg)?;
+
+    let attributes = vec![
+        attr("action", "update_btc_finality"),
+        attr("finality", finality_addr),
+        attr("sender", info.sender),
+    ];
+    Ok(Response::new().add_attributes(attributes))
+}
+
 #[cfg(test)]
 pub mod tests {
     use super::*;

diff --git a/contracts/btc-staking/src/state/config.rs b/contracts/btc-staking/src/state/config.rs
@@ -16,6 +16,7 @@ pub(crate) const ADMIN: Admin = Admin::new("admin");
 #[cw_serde]
 pub struct Config {
     pub babylon: Addr,
+    pub finality: Addr,
     pub denom: String,
 }
 

diff --git a/packages/apis/src/btc_staking_api.rs b/packages/apis/src/btc_staking_api.rs
@@ -14,6 +14,9 @@ pub const HASH_SIZE: usize = 32;
 pub enum ExecuteMsg {
     /// Change the admin
     UpdateAdmin { admin: Option<String> },
+    /// Set the BTC finality addr.
+    /// Only admin or the babylon contract can set this
+    UpdateFinality { finality: String },
     /// BTC Staking operations
     BtcStaking {
         new_fp: Vec<NewFinalityProvider>,