From 3865d252514a805dde3d5c756e39d2c7f48be3f7 Mon Sep 17 00:00:00 2001 From: Mauro Lacy Date: Tue, 17 Sep 2024 16:52:19 +0200 Subject: [PATCH] Update / Set staking address during finality contract instantiation --- contracts/babylon/src/contract.rs | 23 +++++++++++--- contracts/babylon/src/ibc.rs | 2 ++ contracts/babylon/src/multitest/suite.rs | 2 ++ .../babylon/src/state/btc_light_client.rs | 1 + contracts/babylon/tests/integration.rs | 4 +++ contracts/btc-finality/src/contract.rs | 30 ++++++++++++++++--- packages/apis/src/finality_api.rs | 3 ++ 7 files changed, 57 insertions(+), 8 deletions(-) diff --git a/contracts/babylon/src/contract.rs b/contracts/babylon/src/contract.rs index 48582c2f..e08cacca 100644 --- a/contracts/babylon/src/contract.rs +++ b/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, 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 { @@ -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, diff --git a/contracts/babylon/src/ibc.rs b/contracts/babylon/src/ibc.rs index deb8fbfe..c2951632 100644 --- a/contracts/babylon/src/ibc.rs +++ b/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, diff --git a/contracts/babylon/src/multitest/suite.rs b/contracts/babylon/src/multitest/suite.rs index b99dae9e..f4cdcba8 100644 --- a/contracts/babylon/src/multitest/suite.rs +++ b/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()), diff --git a/contracts/babylon/src/state/btc_light_client.rs b/contracts/babylon/src/state/btc_light_client.rs index db7e0e95..7526fb3c 100644 --- a/contracts/babylon/src/state/btc_light_client.rs +++ b/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, }; diff --git a/contracts/babylon/tests/integration.rs b/contracts/babylon/tests/integration.rs index 56c71f07..dbe78dde 100644 --- a/contracts/babylon/tests/integration.rs +++ b/contracts/babylon/tests/integration.rs @@ -49,6 +49,8 @@ fn setup() -> Instance { 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), &[]); diff --git a/contracts/btc-finality/src/contract.rs b/contracts/btc-finality/src/contract.rs index 55674139..c831a242 100644 --- a/contracts/btc-finality/src/contract.rs +++ b/contracts/btc-finality/src/contract.rs @@ -1,8 +1,8 @@ #[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}; @@ -10,6 +10,8 @@ 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, @@ -17,7 +19,6 @@ use crate::finality::{ 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, 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, ContractError> { // Compute active finality provider set let max_active_fps = PARAMS.load(deps.storage)?.max_active_finality_providers as usize; diff --git a/packages/apis/src/finality_api.rs b/packages/apis/src/finality_api.rs index a0d6b569..c29ac211 100644 --- a/packages/apis/src/finality_api.rs +++ b/packages/apis/src/finality_api.rs @@ -13,6 +13,9 @@ use crate::Bytes; pub enum ExecuteMsg { /// Change the admin UpdateAdmin { admin: Option }, + /// 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