Refactor: Avoid needlessly querying BTC height at each block

contracts/btc-staking/src/contract.rs

@@ -1,5 +1,3 @@
-use babylon_contract::msg::btc_header::BtcHeaderResponse;
-
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
@@ -12,15 +10,12 @@ use cw_utils::{maybe_addr, nonpayable};
 use babylon_apis::btc_staking_api::SudoMsg;
 use babylon_bindings::BabylonMsg;
 
-use babylon_contract::msg::contract::QueryMsg as BabylonQueryMsg;
-
 use crate::error::ContractError;
 use crate::finality::{handle_finality_signature, handle_public_randomness_commit};
 use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
 use crate::staking::{compute_active_finality_providers, handle_btc_staking};
 use crate::state::config::{Config, ADMIN, CONFIG, PARAMS};
 use crate::state::staking::ACTIVATED_HEIGHT;
-use crate::state::BTC_HEIGHT;
 use crate::{finality, queries, state};
 
 pub const CONTRACT_NAME: &str = env!("CARGO_PKG_NAME");
@@ -219,40 +214,13 @@ pub fn sudo(
 }
 
 fn handle_begin_block(deps: &mut DepsMut, env: Env) -> Result<Response<BabylonMsg>, ContractError> {
-    // Index BTC height at the current height
-    index_btc_height(deps, env.block.height)?;
-
     // Compute active finality provider set
     let max_active_fps = PARAMS.load(deps.storage)?.max_active_finality_providers as usize;
     compute_active_finality_providers(deps.storage, env, max_active_fps)?;
 
     Ok(Response::new())
 }
 
-// index_btc_height indexes the current BTC height, and saves it to the state
-fn index_btc_height(deps: &mut DepsMut, height: u64) -> Result<(), ContractError> {
-    // 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)?)
-}
-
-/// 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;
-
-    // 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(
     deps: &mut DepsMut,
     env: Env,

contracts/btc-staking/src/finality.rs

@@ -117,7 +117,7 @@ fn verify_commitment_signature(
 
 #[allow(clippy::too_many_arguments)]
 pub fn handle_finality_signature(
-    deps: DepsMut,
+    mut deps: DepsMut,
     env: Env,
     fp_btc_pk_hex: &str,
     height: u64,
@@ -223,7 +223,7 @@ pub fn handle_finality_signature(
             evidence.canonical_finality_sig = canonical_sig;
             // Slash this finality provider, including setting its voting power to zero, extracting
             // its BTC SK, and emitting an event
-            let (msg, ev) = slash_finality_provider(deps.storage, env, fp_btc_pk_hex, &evidence)?;
+            let (msg, ev) = slash_finality_provider(&mut deps, env, fp_btc_pk_hex, &evidence)?;
             res = res.add_message(msg);
             res = res.add_event(ev);
         }
@@ -252,7 +252,7 @@ pub fn handle_finality_signature(
 
         // Slash this finality provider, including setting its voting power to zero, extracting its
         // BTC SK, and emitting an event
-        let (msg, ev) = slash_finality_provider(deps.storage, env, fp_btc_pk_hex, &evidence)?;
+        let (msg, ev) = slash_finality_provider(&mut deps, env, fp_btc_pk_hex, &evidence)?;
         res = res.add_message(msg);
         res = res.add_event(ev);
     }
@@ -263,13 +263,13 @@ pub fn handle_finality_signature(
 /// `slash_finality_provider` slashes a finality provider with the given evidence including setting
 /// its voting power to zero, extracting its BTC SK, and emitting an event
 fn slash_finality_provider(
-    store: &mut dyn Storage,
+    deps: &mut DepsMut,
     env: Env,
     fp_btc_pk_hex: &str,
     evidence: &Evidence,
 ) -> Result<(WasmMsg, Event), ContractError> {
     // Slash this finality provider, i.e., set its slashing height to the block height
-    staking::slash_finality_provider(store, env, fp_btc_pk_hex, evidence.block_height)
+    staking::slash_finality_provider(deps, env, fp_btc_pk_hex, evidence.block_height)
         .map_err(|err| ContractError::FailedToSlashFinalityProvider(err.to_string()))?;
 
     // Extract BTC SK using the evidence
@@ -291,7 +291,7 @@ fn slash_finality_provider(
         evidence: evidence.clone(),
     };
 
-    let babylon_addr = CONFIG.load(store)?.babylon;
+    let babylon_addr = CONFIG.load(deps.storage)?.babylon;
 
     let wasm_msg = WasmMsg::Execute {
         contract_addr: babylon_addr.to_string(),

contracts/btc-staking/src/staking.rs

@@ -14,7 +14,6 @@ use crate::state::staking::{
     fps, BtcDelegation, FinalityProviderState, ACTIVATED_HEIGHT, DELEGATIONS, DELEGATION_FPS, FPS,
     FP_DELEGATIONS, FP_SET, TOTAL_POWER,
 };
-use crate::state::BTC_HEIGHT;
 use crate::validation::{
     verify_active_delegation, verify_new_fp, verify_slashed_delegation, verify_undelegation,
 };
@@ -25,6 +24,9 @@ use babylon_apis::btc_staking_api::{
 
 use babylon_apis::Validate;
 use babylon_bindings::BabylonMsg;
+use babylon_contract::msg::btc_header::BtcHeaderResponse;
+
+use babylon_contract::msg::contract::QueryMsg as BabylonQueryMsg;
 
 /// handle_btc_staking handles the BTC staking operations
 pub fn handle_btc_staking(
@@ -378,13 +380,13 @@ pub fn compute_active_finality_providers(
 /// `slash_finality_provider` slashes a finality provider with the given PK.
 /// A slashed finality provider will not have voting power
 pub(crate) fn slash_finality_provider(
-    store: &mut dyn Storage,
+    deps: &mut DepsMut,
     env: Env,
     fp_btc_pk_hex: &str,
     height: u64,
 ) -> Result<(), ContractError> {
     // Ensure finality provider exists
-    let mut fp = FPS.load(store, fp_btc_pk_hex)?;
+    let mut fp = FPS.load(deps.storage, fp_btc_pk_hex)?;
 
     // Check if the finality provider is already slashed
     if fp.slashed_height > 0 {
@@ -395,28 +397,41 @@ pub(crate) fn slash_finality_provider(
     // Set the finality provider as slashed
     fp.slashed_height = height;
 
-    // Set BTC slashing height (if available from the store)
+    // Set BTC slashing height (if available from the babylon contract)
     // FIXME: Turn this into a hard error
     // return fmt.Errorf("failed to get current BTC tip")
-    let btc_height = BTC_HEIGHT.may_load(store, height)?.unwrap_or_default();
+    let btc_height = get_btc_tip_height(deps).unwrap_or_default();
     fp.slashed_btc_height = btc_height;
 
     // Record slashed event. The next `BeginBlock` will consume this event for updating the active
     // FP set.
     // We simply set the FP voting power to zero from the next *processing* height (See NOTE in
     // `handle_finality_signature`)
-    fps().update(store, fp_btc_pk_hex, env.block.height + 1, |fp| {
+    fps().update(deps.storage, fp_btc_pk_hex, env.block.height + 1, |fp| {
         let mut fp = fp.unwrap_or_default();
         fp.power = 0;
         Ok::<_, ContractError>(fp)
     })?;
 
     // Save the finality provider back
-    FPS.save(store, fp_btc_pk_hex, &fp)?;
+    FPS.save(deps.storage, fp_btc_pk_hex, &fp)?;
 
     Ok(())
 }
 
+/// 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;
+
+    // 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)
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;

contracts/btc-staking/src/state/mod.rs

@@ -4,8 +4,3 @@ pub mod public_randomness;
 pub mod staking;
 
 mod fp_index;
-
-use cw_storage_plus::Map;
-
-/// Map of BTC height by block height
-pub(crate) const BTC_HEIGHT: Map<u64, u64> = Map::new("btc_height");

contracts/btc-staking/src/validation/mod.rs

@@ -1,8 +1,6 @@
 use crate::state::config::Params;
 use crate::{error::ContractError, state::staking::BtcDelegation};
 use babylon_apis::btc_staking_api::{ActiveBtcDelegation, NewFinalityProvider};
-use babylon_btcstaking::adaptor_sig::AdaptorSignature;
-use babylon_btcstaking::sig_verify::enc_verify_transaction_sig_with_output;
 use bitcoin::Transaction;
 use cosmwasm_std::Binary;
 

go.work.sum

@@ -930,6 +930,7 @@ github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe h1:QQ3GSy+MqSHxm/d8nC
 github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe/go.mod h1:6pvJx4me5XPnfI9Z40ddWsdw2W/uZgQLFXToKeRcDiI=
 github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4 h1:/inchEIKaYC1Akx+H+gqO04wryn5h75LSazbRlnya1k=
 github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
+github.com/cncf/xds/go v0.0.0-20240318125728-8a4994d93e50 h1:DBmgJDC9dTfkVyGgipamEh2BpGYxScCH1TOF1LL1cXc=
 github.com/cncf/xds/go v0.0.0-20240318125728-8a4994d93e50/go.mod h1:5e1+Vvlzido69INQaVO6d87Qn543Xr6nooe9Kz7oBFM=
 github.com/cockroachdb/datadriven v1.0.0/go.mod h1:5Ib8Meh+jk1RlHIXej6Pzevx/NLlNvQB9pmSBZErGA4=
 github.com/cockroachdb/datadriven v1.0.2/go.mod h1:a9RdTaap04u637JoCzcUoIcDmvwSUtcUFtT/C3kJlTU=
@@ -1042,6 +1043,7 @@ github.com/envoyproxy/go-control-plane v0.12.0 h1:4X+VP1GHd1Mhj6IB5mMeGbLCleqxjl
 github.com/envoyproxy/go-control-plane v0.12.0/go.mod h1:ZBTaoJ23lqITozF0M6G4/IragXCQKCnYbmlmtHvwRG0=
 github.com/envoyproxy/protoc-gen-validate v1.0.2 h1:QkIBuU5k+x7/QXPvPPnWXWlCdaBFApVqftFV6k087DA=
 github.com/envoyproxy/protoc-gen-validate v1.0.2/go.mod h1:GpiZQP3dDbg4JouG/NNS7QWXpgx6x8QiMKdmN72jogE=
+github.com/envoyproxy/protoc-gen-validate v1.0.4 h1:gVPz/FMfvh57HdSJQyvBtF00j8JU4zdyUgIUNhlgg0A=
 github.com/envoyproxy/protoc-gen-validate v1.0.4/go.mod h1:qys6tmnRsYrQqIhm2bvKZH4Blx/1gTIZ2UKVY1M+Yew=
 github.com/esimonov/ifshort v1.0.4 h1:6SID4yGWfRae/M7hkVDVVyppy8q/v9OuxNdmjLQStBA=
 github.com/esimonov/ifshort v1.0.4/go.mod h1:Pe8zjlRrJ80+q2CxHLfEOfTwxCZ4O+MuhcHcfgNWTk0=
@@ -1293,6 +1295,7 @@ github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmg
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92BcuyuQ/YW4NSIpoGtfXNho=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.11.3/go.mod h1:o//XUCC/F+yRGJoPO/VU0GSB0f8Nhgmxx0VIRUvaC0w=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 h1:bkypFPDjIYGfCYD5mRBvpqxfYX1YCS1PXdKYWi8FsN0=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0/go.mod h1:P+Lt/0by1T8bfcF3z737NnSbmxQAppXMRziHUxPOC8k=
 github.com/guptarohit/asciigraph v0.5.5 h1:ccFnUF8xYIOUPPY3tmdvRyHqmn1MYI9iv1pLKX+/ZkQ=
 github.com/guptarohit/asciigraph v0.5.5/go.mod h1:dYl5wwK4gNsnFf9Zp+l06rFiDZ5YtXM6x7SRWZ3KGag=
@@ -2221,6 +2224,7 @@ google.golang.org/api v0.153.0/go.mod h1:3qNJX5eOmhiWYc67jRA/3GsDw97UFb5ivv7Y2Pr
 google.golang.org/api v0.155.0/go.mod h1:GI5qK5f40kCpHfPn6+YzGAByIKWv8ujFnmoWm7Igduk=
 google.golang.org/api v0.157.0/go.mod h1:+z4v4ufbZ1WEpld6yMGHyggs+PmAHiaLNj5ytP3N01g=
 google.golang.org/api v0.160.0/go.mod h1:0mu0TpK33qnydLvWqbImq2b1eQ5FHRSDCBzAxX9ZHyw=
+google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM=
 google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds=
 google.golang.org/genproto v0.0.0-20180518175338-11a468237815/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20210108203827-ffc7fda8c3d7/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
@@ -2255,6 +2259,7 @@ google.golang.org/genproto/googleapis/bytestream v0.0.0-20231120223509-83a465c02
 google.golang.org/genproto/googleapis/bytestream v0.0.0-20231120223509-83a465c0220f/go.mod h1:iIgEblxoG4klcXsG0d9cpoxJ4xndv6+1FkDROCHhPRI=
 google.golang.org/genproto/googleapis/bytestream v0.0.0-20240125205218-1f4bbc51befe h1:weYsP+dNijSQVoLAb5bpUos3ciBpNU/NEVlHFKrk8pg=
 google.golang.org/genproto/googleapis/bytestream v0.0.0-20240125205218-1f4bbc51befe/go.mod h1:SCz6T5xjNXM4QFPRwxHcfChp7V+9DcXR3ay2TkHR8Tg=
+google.golang.org/genproto/googleapis/bytestream v0.0.0-20240304161311-37d4d3c04a78 h1:YqFWYZXim8bG9v68xU8WjTZmYKb5M5dMeSOWIp6jogI=
 google.golang.org/genproto/googleapis/bytestream v0.0.0-20240304161311-37d4d3c04a78/go.mod h1:vh/N7795ftP0AkN1w8XKqN4w1OdUKXW5Eummda+ofv8=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc/go.mod h1:66JfowdXAEgad5O9NnYcsNPLCPZJD++2L9X0PCMODrA=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20230815205213-6bfd019c3878/go.mod h1:+Bk1OCOj40wS2hwAMA+aCW9ypzm63QTBBHp6lQ3p+9M=