Skip to main content

babylon_test_utils/
lib.rs

1use hex::ToHex;
2use k256::schnorr::{Signature, SigningKey};
3use prost::{bytes::Bytes, Message};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::path::PathBuf;
7use std::str::FromStr;
8use std::{env, fs};
9
10use cosmwasm_std::{Binary, Decimal};
11
12use babylon_apis::btc_staking_api::{
13    ActiveBtcDelegation, BtcUndelegationInfo, CovenantAdaptorSignatures, DelegatorUnbondingInfo,
14    FinalityProviderDescription, NewFinalityProvider, ProofOfPossessionBtc,
15};
16use babylon_apis::finality_api::PubRandCommit;
17use babylon_bitcoin::{deserialize, BlockHash, BlockHeader};
18use babylon_proto::babylon::btclightclient::v1::{BtcHeaderInfo, QueryMainChainResponse};
19use babylon_proto::babylon::btcstaking::v1::{BtcDelegation, FinalityProvider, Params};
20use babylon_proto::babylon::finality::v1::{MsgAddFinalitySig, MsgCommitPubRandList};
21use babylon_proto::babylon::zoneconcierge::v1::BtcTimestamp;
22
23const BTC_LC_MAIN: &str = "btc_light_client.dat";
24const BTC_LC_FORK: &str = "btc_light_client_fork.dat";
25const BTC_LC_FORK_MSG: &str = "btc_light_client_fork_msg.json";
26
27const BTC_TIMESTAMP: &str = "btc_timestamp.dat";
28const BTC_TIMESTAMP_HEADER0: &str = "btc_timestamp_header0.dat";
29const BTC_TIMESTAMP_HEADER1: &str = "btc_timestamp_header1.dat";
30
31const PARAMS_DATA: &str = "btcstaking_params.dat";
32const FINALITY_PROVIDER_DATA: &str = "finality_provider_{}.dat";
33const FP_SK_DATA: &str = "fp_sk_{}.dat";
34const BTC_DELEGATION_DATA: &str = "btc_delegation_{idx}_{fp_idx_list}.dat";
35const BTC_DEL_UNBONDING_SIG_DATA: &str = "btc_unbonding_sig_{idx}_{fp_idx_list}.dat";
36const COMMIT_PUB_RAND_DATA: &str = "commit_pub_rand_msg.dat";
37const PUB_RAND_VALUE: &str = "pub_rand_value.dat";
38const ADD_FINALITY_SIG_DATA: &str = "add_finality_sig_{}_msg.dat";
39
40const EOTS_DATA: &str = "eots_testdata.json";
41
42pub fn find_project_path() -> PathBuf {
43    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
44}
45
46fn find_testdata_path() -> PathBuf {
47    find_project_path().join("testdata")
48}
49
50#[derive(Serialize, Deserialize, Debug)]
51pub struct EotsTestData {
52    pub sk: String,
53    pub pk: String,
54    pub sr: String,
55    pub pr: String,
56    pub msg1: String,
57    pub msg2: String,
58    pub sig1: String,
59    pub sig2: String,
60}
61
62pub fn get_eots_testdata() -> EotsTestData {
63    let file_path = find_testdata_path().join(EOTS_DATA);
64    let testdata_bytes: &[u8] = &fs::read(file_path).unwrap();
65    let testdata: EotsTestData = serde_json::from_slice(testdata_bytes).unwrap();
66
67    testdata
68}
69
70pub fn get_btc_lc_mainchain_resp() -> QueryMainChainResponse {
71    let file_path = find_testdata_path().join(BTC_LC_MAIN);
72
73    let testdata: &[u8] = &fs::read(file_path).unwrap();
74    QueryMainChainResponse::decode(testdata).unwrap()
75}
76
77pub fn get_btc_lc_headers() -> Vec<BtcHeaderInfo> {
78    let resp = get_btc_lc_mainchain_resp();
79    resp.headers
80        .iter()
81        .map(|h| BtcHeaderInfo {
82            header: Bytes::from(hex::decode(&h.header_hex).unwrap()),
83            // FIXME: Use BlockHash / Hash helper / encapsulation to reverse the hash under the hood
84            hash: Bytes::from(
85                hex::decode(&h.hash_hex)
86                    .unwrap()
87                    .into_iter()
88                    .rev()
89                    .collect::<Vec<_>>(),
90            ),
91            height: h.height,
92            work: { Bytes::from(h.work.clone()) },
93        })
94        .collect()
95}
96
97pub fn get_btc_lc_fork_headers() -> Vec<BtcHeaderInfo> {
98    let file_path = find_testdata_path().join(BTC_LC_FORK);
99    let testdata: &[u8] = &fs::read(file_path).unwrap();
100    let resp = QueryMainChainResponse::decode(testdata).unwrap();
101    resp.headers
102        .iter()
103        .map(|h| BtcHeaderInfo {
104            header: Bytes::from(hex::decode(&h.header_hex).unwrap()),
105            // FIXME: Use BlockHash / Hash helper / encapsulation to reverse the hash under the hood
106            hash: Bytes::from(
107                hex::decode(&h.hash_hex)
108                    .unwrap()
109                    .into_iter()
110                    .rev()
111                    .collect::<Vec<_>>(),
112            ),
113            height: h.height,
114            work: { Bytes::from(h.work.clone()) },
115        })
116        .collect()
117}
118
119pub fn get_btc_lc_fork_msg() -> Vec<u8> {
120    let file_path = find_testdata_path().join(BTC_LC_FORK_MSG);
121    let testdata: &[u8] = &fs::read(file_path).unwrap();
122    testdata.to_vec()
123}
124
125pub fn get_btc_timestamp_and_headers() -> (BtcTimestamp, HashMap<BlockHash, BlockHeader>) {
126    let mut header_map: HashMap<BlockHash, BlockHeader> = HashMap::new();
127
128    let header0_path = find_testdata_path().join(BTC_TIMESTAMP_HEADER0);
129    let header0_bytes: &[u8] = &fs::read(header0_path).unwrap();
130    let header0: BlockHeader = deserialize(header0_bytes).unwrap();
131    header_map.insert(header0.block_hash(), header0);
132
133    let header1_path = find_testdata_path().join(BTC_TIMESTAMP_HEADER1);
134    let header1_bytes: &[u8] = &fs::read(header1_path).unwrap();
135    let header1: BlockHeader = deserialize(header1_bytes).unwrap();
136    header_map.insert(header1.block_hash(), header1);
137
138    let ts_path = find_testdata_path().join(BTC_TIMESTAMP);
139    let testdata: &[u8] = &fs::read(ts_path).unwrap();
140    let btc_ts = BtcTimestamp::decode(testdata).unwrap();
141
142    (btc_ts, header_map)
143}
144
145pub fn get_params() -> Params {
146    let params_path = find_testdata_path().join(PARAMS_DATA);
147    let params_data: &[u8] = &fs::read(params_path).unwrap();
148    Params::decode(params_data).unwrap()
149}
150
151pub fn get_finality_provider(id: i32) -> FinalityProvider {
152    let fp_path = find_testdata_path().join(FINALITY_PROVIDER_DATA.replace("{}", &id.to_string()));
153    let fp_data: &[u8] = &fs::read(fp_path).unwrap();
154    FinalityProvider::decode(fp_data).unwrap()
155}
156
157pub fn get_fp_sk_bytes(id: i32) -> Vec<u8> {
158    let fp_sk_path = find_testdata_path().join(FP_SK_DATA.replace("{}", &id.to_string()));
159    let fp_sk_data: &[u8] = &fs::read(fp_sk_path).unwrap();
160    fp_sk_data.to_vec()
161}
162
163pub fn get_btc_delegation(idx: i32, fp_idx_list: Vec<i32>) -> BtcDelegation {
164    let fp_idx_list_str = format!(
165        "{{{}}}",
166        fp_idx_list
167            .iter()
168            .map(|&x| x.to_string())
169            .collect::<Vec<_>>()
170            .join(",")
171    );
172    let btc_del_filename = BTC_DELEGATION_DATA
173        .replace("{idx}", &idx.to_string())
174        .replace("{fp_idx_list}", &fp_idx_list_str);
175    let btc_del_path = find_testdata_path().join(btc_del_filename);
176    let btc_del_data: &[u8] = &fs::read(btc_del_path).unwrap();
177    BtcDelegation::decode(btc_del_data).unwrap()
178}
179
180pub fn get_btc_del_unbonding_sig_bytes(idx: i32, fp_idx_list: Vec<i32>) -> Vec<u8> {
181    let fp_idx_list_str = format!(
182        "{{{}}}",
183        fp_idx_list
184            .iter()
185            .map(|&x| x.to_string())
186            .collect::<Vec<_>>()
187            .join(",")
188    );
189    let sig_filename = BTC_DEL_UNBONDING_SIG_DATA
190        .replace("{idx}", &idx.to_string())
191        .replace("{fp_idx_list}", &fp_idx_list_str);
192    let sig_path = find_testdata_path().join(sig_filename);
193    let sig_data: &[u8] = &fs::read(sig_path).unwrap();
194    sig_data.to_vec()
195}
196
197pub fn get_pub_rand_commit() -> MsgCommitPubRandList {
198    let pub_rand_commit_path = find_testdata_path().join(COMMIT_PUB_RAND_DATA);
199    let pub_rand_commit_data: &[u8] = &fs::read(pub_rand_commit_path).unwrap();
200
201    MsgCommitPubRandList::decode(pub_rand_commit_data).unwrap()
202}
203
204/// Get public randomness value (at index 1)
205//  TODO: Support indexed public randomness values
206pub fn get_pub_rand_value() -> Vec<u8> {
207    let pub_rand_value_path = find_testdata_path().join(PUB_RAND_VALUE);
208    let pub_rand_value_data: Vec<u8> = fs::read(pub_rand_value_path).unwrap();
209
210    pub_rand_value_data
211}
212
213pub fn get_add_finality_sig() -> MsgAddFinalitySig {
214    let add_finality_sig_path = find_testdata_path().join(ADD_FINALITY_SIG_DATA.replace("{}", "1"));
215    let add_finality_sig_data: &[u8] = &fs::read(add_finality_sig_path).unwrap();
216
217    MsgAddFinalitySig::decode(add_finality_sig_data).unwrap()
218}
219
220pub fn get_add_finality_sig_2() -> MsgAddFinalitySig {
221    let add_finality_sig_path = find_testdata_path().join(ADD_FINALITY_SIG_DATA.replace("{}", "2"));
222    let add_finality_sig_data: &[u8] = &fs::read(add_finality_sig_path).unwrap();
223
224    MsgAddFinalitySig::decode(add_finality_sig_data).unwrap()
225}
226
227pub fn new_finality_provider(fp: FinalityProvider) -> NewFinalityProvider {
228    NewFinalityProvider {
229        addr: fp.addr,
230        description: fp.description.map(|desc| FinalityProviderDescription {
231            moniker: desc.moniker,
232            identity: desc.identity,
233            website: desc.website,
234            security_contact: desc.security_contact,
235            details: desc.details,
236        }),
237        commission: Decimal::from_str(&fp.commission).unwrap(),
238        btc_pk_hex: fp.btc_pk.encode_hex(),
239        pop: match fp.pop {
240            Some(pop) => Some(ProofOfPossessionBtc {
241                btc_sig_type: pop.btc_sig_type,
242                btc_sig: Binary::new(pop.btc_sig.to_vec()),
243            }),
244            None => None,
245        },
246        consumer_id: fp.consumer_id,
247    }
248}
249
250pub fn new_active_btc_delegation(del: BtcDelegation) -> ActiveBtcDelegation {
251    let btc_undelegation = del.btc_undelegation.unwrap();
252    let delegator_unbonding_info =
253        if let Some(delegator_unbonding_info) = btc_undelegation.delegator_unbonding_info {
254            Some(DelegatorUnbondingInfo {
255                spend_stake_tx: Binary::new(delegator_unbonding_info.spend_stake_tx.to_vec()),
256            })
257        } else {
258            None
259        };
260
261    ActiveBtcDelegation {
262        staker_addr: del.staker_addr,
263        btc_pk_hex: del.btc_pk.encode_hex(),
264        fp_btc_pk_list: del
265            .fp_btc_pk_list
266            .iter()
267            .map(|fp_btc_pk| fp_btc_pk.encode_hex())
268            .collect(),
269        start_height: del.start_height,
270        end_height: del.end_height,
271        total_sat: del.total_sat,
272        staking_tx: Binary::new(del.staking_tx.to_vec()),
273        slashing_tx: Binary::new(del.slashing_tx.to_vec()),
274        delegator_slashing_sig: Binary::new(del.delegator_sig.to_vec()),
275        covenant_sigs: del
276            .covenant_sigs
277            .iter()
278            .map(|cov_sig| CovenantAdaptorSignatures {
279                cov_pk: Binary::new(cov_sig.cov_pk.to_vec()),
280                adaptor_sigs: cov_sig
281                    .adaptor_sigs
282                    .iter()
283                    .map(|adaptor_sig| Binary::new(adaptor_sig.to_vec()))
284                    .collect(),
285            })
286            .collect(),
287        staking_output_idx: del.staking_output_idx,
288        unbonding_time: del.unbonding_time,
289        undelegation_info: BtcUndelegationInfo {
290            unbonding_tx: Binary::new(btc_undelegation.unbonding_tx.to_vec()),
291            slashing_tx: Binary::new(btc_undelegation.slashing_tx.to_vec()),
292            delegator_unbonding_info,
293            delegator_slashing_sig: Binary::new(btc_undelegation.delegator_slashing_sig.to_vec()),
294            covenant_unbonding_sig_list: vec![],
295            covenant_slashing_sigs: vec![],
296        },
297        params_version: del.params_version,
298    }
299}
300
301/// Build an active BTC delegation from a BTC delegation
302pub fn get_active_btc_delegation() -> ActiveBtcDelegation {
303    let del = get_btc_delegation(1, vec![1]);
304    new_active_btc_delegation(del)
305}
306
307// Build a derived active BTC delegation from the base (from testdata) BTC delegation
308pub fn get_derived_btc_delegation(del_id: i32, fp_ids: &[i32]) -> ActiveBtcDelegation {
309    let del = get_btc_delegation(del_id, fp_ids.to_vec());
310    new_active_btc_delegation(del)
311}
312
313pub fn get_btc_del_unbonding_sig(del_id: i32, fp_ids: &[i32]) -> Signature {
314    let sig_bytes = get_btc_del_unbonding_sig_bytes(del_id, fp_ids.to_vec());
315    Signature::try_from(sig_bytes.as_slice()).unwrap()
316}
317
318pub fn create_new_finality_provider(id: i32) -> NewFinalityProvider {
319    let fp = get_finality_provider(id);
320    new_finality_provider(fp)
321}
322
323pub fn create_new_fp_sk(id: i32) -> SigningKey {
324    let fp_sk_bytes = get_fp_sk_bytes(id);
325    SigningKey::from_bytes(&fp_sk_bytes).unwrap()
326}
327
328/// Get public randomness public key, commitment, and signature information
329///
330/// Signature is a Schnorr signature over the commitment
331pub fn get_public_randomness_commitment() -> (String, PubRandCommit, Vec<u8>) {
332    let pub_rand_commitment_msg = get_pub_rand_commit();
333    (
334        pub_rand_commitment_msg.fp_btc_pk.encode_hex(),
335        PubRandCommit {
336            start_height: pub_rand_commitment_msg.start_height,
337            num_pub_rand: pub_rand_commitment_msg.num_pub_rand,
338            height: 1,
339            commitment: pub_rand_commitment_msg.commitment.to_vec(),
340        },
341        pub_rand_commitment_msg.sig.to_vec(),
342    )
343}