Merge pull request #33 from babylonlabs-io/eots-refactor-2

Eots refactor #2
babylonlabs-io · Aug 6, 2024 · ba091f4 · ba091f4
2 parents 96cd12a + f67231e
commit ba091f4
Show file tree

Hide file tree

Showing 7 changed files with 292 additions and 212 deletions.
diff --git a/contracts/babylon/src/contract.rs b/contracts/babylon/src/contract.rs
@@ -161,7 +161,11 @@ pub fn execute(
             // TODO: Add events
             Ok(Response::new())
         }
-        ExecuteMsg::Slashing { fp_btc_pk, block_height, secret_key } => {
+        ExecuteMsg::Slashing {
+            fp_btc_pk,
+            block_height,
+            secret_key,
+        } => {
             // This is an internal routing message from the `btc-staking` contract
             // Check sender
             let btc_staking = CONFIG
@@ -173,7 +177,8 @@ pub fn execute(
             }
             // Send over IBC to the Provider (Babylon)
             let channel = IBC_CHANNEL.load(deps.storage)?;
-            let msg = ibc_packet::slashing_msg(&env, &channel, &fp_btc_pk, block_height, &secret_key)?;
+            let msg =
+                ibc_packet::slashing_msg(&env, &channel, &fp_btc_pk, block_height, &secret_key)?;
             // TODO: Add events
             Ok(Response::new().add_message(msg))
         }

diff --git a/contracts/babylon/src/ibc.rs b/contracts/babylon/src/ibc.rs
@@ -141,11 +141,14 @@ pub(crate) mod ibc_packet {
     use super::*;
     use crate::state::config::CONFIG;
     use babylon_apis::ibc_consumer::{consumer_packet_data, ConsumerPacketData};
-    use babylon_apis::{btc_staking_api::{
-        ActiveBtcDelegation, BtcUndelegationInfo, CovenantAdaptorSignatures,
-        FinalityProviderDescription, NewFinalityProvider, ProofOfPossessionBtc, SignatureInfo,
-        UnbondedBtcDelegation,
-    }, ibc_consumer};
+    use babylon_apis::{
+        btc_staking_api::{
+            ActiveBtcDelegation, BtcUndelegationInfo, CovenantAdaptorSignatures,
+            FinalityProviderDescription, NewFinalityProvider, ProofOfPossessionBtc, SignatureInfo,
+            UnbondedBtcDelegation,
+        },
+        ibc_consumer,
+    };
     use babylon_proto::babylon::btcstaking::v1::BtcStakingIbcPacket;
     use cosmwasm_std::{to_json_binary, Decimal, IbcChannel, IbcMsg, WasmMsg};
     use std::str::FromStr;
@@ -311,11 +314,13 @@ pub(crate) mod ibc_packet {
         block_height: u64,
         secret_key: &[u8],
     ) -> Result<IbcMsg, ContractError> {
-        let packet = ConsumerPacketData { packet: consumer_packet_data::Packet::Slashing(ibc_consumer::Slashing {
-            fp_btc_pk: fp_btc_pk.to_vec(),
-            block_height,
-            secret_key: secret_key.to_vec(),
-        }), };
+        let packet = ConsumerPacketData {
+            packet: consumer_packet_data::Packet::Slashing(ibc_consumer::Slashing {
+                fp_btc_pk: fp_btc_pk.to_vec(),
+                block_height,
+                secret_key: secret_key.to_vec(),
+            }),
+        };
         let msg = IbcMsg::SendPacket {
             channel_id: channel.endpoint.channel_id.clone(),
             data: to_json_binary(&packet)?,

diff --git a/contracts/babylon/src/msg/contract.rs b/contracts/babylon/src/msg/contract.rs
@@ -1,10 +1,10 @@
-use cosmwasm_schema::{cw_serde, QueryResponses};
-use cosmwasm_std::{Binary, StdError, StdResult};
-use babylon_apis::Bytes;
 use crate::msg::btc_header::{BtcHeader, BtcHeaderResponse, BtcHeadersResponse};
 use crate::msg::cz_header::CzHeaderResponse;
 use crate::msg::epoch::EpochResponse;
 use crate::state::config::Config;
+use babylon_apis::Bytes;
+use cosmwasm_schema::{cw_serde, QueryResponses};
+use cosmwasm_std::{Binary, StdError, StdResult};
 
 const BABYLON_TAG_BYTES: usize = 4;
 

diff --git a/contracts/btc-staking/src/finality.rs b/contracts/btc-staking/src/finality.rs
@@ -274,15 +274,15 @@ fn slash_finality_provider(
 
     // Extract BTC SK using the evidence
     let pk = eots::PublicKey::from_hex(fp_btc_pk_hex)?;
-    let btc_sk = eots::extract(
-        &pk,
-        &evidence.pub_rand,
-        &evidence.canonical_app_hash,
-        &evidence.canonical_finality_sig,
-        &evidence.fork_app_hash,
-        &evidence.fork_finality_sig,
-    )
-    .map_err(|err| ContractError::SecretKeyExtractionError(err.to_string()))?;
+    let btc_sk = pk
+        .extract_secret_key(
+            &evidence.pub_rand,
+            &evidence.canonical_app_hash,
+            &evidence.canonical_finality_sig,
+            &evidence.fork_app_hash,
+            &evidence.fork_finality_sig,
+        )
+        .map_err(|err| ContractError::SecretKeyExtractionError(err.to_string()))?;
 
     // Emit slashing event.
     // Raises slashing event to babylon over IBC.
@@ -1050,15 +1050,16 @@ pub(crate) mod tests {
         // Assert the slashing event is proper
         assert_eq!(res.events[0].ty, "slashed_finality_provider".to_string());
         // Extract the secret key to compare it
-        let btc_sk = eots::extract(
-            &eots::PublicKey::from_hex(&pk_hex).unwrap(),
-            &pub_rand_one,
-            &add_finality_signature.block_app_hash,
-            &finality_signature,
-            &add_finality_signature_2.block_app_hash,
-            &add_finality_signature_2.finality_sig,
-        )
-        .unwrap();
+        let pk = eots::PublicKey::from_hex(&pk_hex).unwrap();
+        let btc_sk = pk
+            .extract_secret_key(
+                &pub_rand_one,
+                &add_finality_signature.block_app_hash,
+                &finality_signature,
+                &add_finality_signature_2.block_app_hash,
+                &add_finality_signature_2.finality_sig,
+            )
+            .unwrap();
         assert_eq!(
             res.messages[0],
             SubMsg::new(WasmMsg::Execute {

diff --git a/packages/apis/src/ibc_consumer.rs b/packages/apis/src/ibc_consumer.rs
@@ -1,5 +1,5 @@
-use cosmwasm_schema::cw_serde;
 use crate::Bytes;
+use cosmwasm_schema::cw_serde;
 
 /// ConsumerPacketData is the message that defines the IBC packets a Consumer can send to Babylon's
 /// ZoneConcierge module.