ic_btc_test_utils/
lib.rs

1use bitcoin::{
2    secp256k1::rand::rngs::OsRng, secp256k1::Secp256k1, util::uint::Uint256, Address, Block,
3    BlockHash, BlockHeader, KeyPair, Network, OutPoint, PublicKey, Script, Transaction, TxIn,
4    TxMerkleNode, TxOut, Witness, XOnlyPublicKey,
5};
6
7/// Generates a random P2PKH address.
8pub fn random_p2pkh_address(network: Network) -> Address {
9    let secp = Secp256k1::new();
10    let mut rng = OsRng::new().unwrap();
11
12    Address::p2pkh(&PublicKey::new(secp.generate_keypair(&mut rng).1), network)
13}
14
15pub fn random_p2tr_address(network: Network) -> Address {
16    let secp = Secp256k1::new();
17    let mut rng = OsRng::new().unwrap();
18    let key_pair = KeyPair::new(&secp, &mut rng);
19    let xonly = XOnlyPublicKey::from_keypair(&key_pair);
20
21    Address::p2tr(&secp, xonly, None, network)
22}
23
24fn coinbase_input() -> TxIn {
25    TxIn {
26        previous_output: OutPoint::null(),
27        script_sig: Script::new(),
28        sequence: 0xffffffff,
29        witness: Witness::new(),
30    }
31}
32
33pub struct BlockBuilder {
34    prev_header: Option<BlockHeader>,
35    transactions: Vec<Transaction>,
36}
37
38impl BlockBuilder {
39    pub fn genesis() -> Self {
40        Self {
41            prev_header: None,
42            transactions: vec![],
43        }
44    }
45
46    pub fn with_prev_header(prev_header: BlockHeader) -> Self {
47        Self {
48            prev_header: Some(prev_header),
49            transactions: vec![],
50        }
51    }
52
53    pub fn with_transaction(mut self, transaction: Transaction) -> Self {
54        self.transactions.push(transaction);
55        self
56    }
57
58    pub fn build(self) -> Block {
59        let txdata = if self.transactions.is_empty() {
60            // Create a random coinbase transaction.
61            vec![TransactionBuilder::coinbase().build()]
62        } else {
63            self.transactions
64        };
65
66        let merkle_root =
67            bitcoin::util::hash::bitcoin_merkle_root(txdata.iter().map(|tx| tx.txid().as_hash()))
68                .unwrap();
69        let merkle_root = TxMerkleNode::from_hash(merkle_root);
70
71        let header = match self.prev_header {
72            Some(prev_header) => header(&prev_header, merkle_root),
73            None => genesis(merkle_root),
74        };
75
76        Block { header, txdata }
77    }
78}
79
80fn genesis(merkle_root: TxMerkleNode) -> BlockHeader {
81    let target = Uint256([
82        0xffffffffffffffffu64,
83        0xffffffffffffffffu64,
84        0xffffffffffffffffu64,
85        0x7fffffffffffffffu64,
86    ]);
87    let bits = BlockHeader::compact_target_from_u256(&target);
88
89    let mut header = BlockHeader {
90        version: 1,
91        time: 0,
92        nonce: 0,
93        bits,
94        merkle_root,
95        prev_blockhash: BlockHash::default(),
96    };
97    solve(&mut header);
98
99    header
100}
101
102pub struct TransactionBuilder {
103    input: Vec<TxIn>,
104    output: Vec<TxOut>,
105    lock_time: u32,
106}
107
108impl TransactionBuilder {
109    pub fn new() -> Self {
110        Self {
111            input: vec![],
112            output: vec![],
113            lock_time: 0,
114        }
115    }
116
117    pub fn coinbase() -> Self {
118        Self {
119            input: vec![coinbase_input()],
120            output: vec![],
121            lock_time: 0,
122        }
123    }
124
125    pub fn with_input(mut self, previous_output: OutPoint, witness: Option<Witness>) -> Self {
126        if self.input == vec![coinbase_input()] {
127            panic!("A call `with_input` should not be possible if `coinbase` was called");
128        }
129
130        let witness = witness.map_or(Witness::new(), |w| w);
131        let input = TxIn {
132            previous_output,
133            script_sig: Script::new(),
134            sequence: 0xffffffff,
135            witness,
136        };
137        self.input.push(input);
138        self
139    }
140
141    pub fn with_output(mut self, address: &Address, value: u64) -> Self {
142        self.output.push(TxOut {
143            value,
144            script_pubkey: address.script_pubkey(),
145        });
146        self
147    }
148
149    pub fn with_lock_time(mut self, time: u32) -> Self {
150        self.lock_time = time;
151        self
152    }
153
154    pub fn build(self) -> Transaction {
155        let input = if self.input.is_empty() {
156            // Default to coinbase if no inputs provided.
157            vec![coinbase_input()]
158        } else {
159            self.input
160        };
161        let output = if self.output.is_empty() {
162            // Use default of 50 BTC.
163            vec![TxOut {
164                value: 50_0000_0000,
165                script_pubkey: random_p2pkh_address(Network::Regtest).script_pubkey(),
166            }]
167        } else {
168            self.output
169        };
170
171        Transaction {
172            version: 1,
173            lock_time: self.lock_time,
174            input,
175            output,
176        }
177    }
178}
179
180impl Default for TransactionBuilder {
181    fn default() -> Self {
182        Self::new()
183    }
184}
185
186fn header(prev_header: &BlockHeader, merkle_root: TxMerkleNode) -> BlockHeader {
187    let time = prev_header.time + 60 * 10; // 10 minutes.
188    let bits = BlockHeader::compact_target_from_u256(&prev_header.target());
189
190    let mut header = BlockHeader {
191        version: 1,
192        time,
193        nonce: 0,
194        bits,
195        merkle_root,
196        prev_blockhash: prev_header.block_hash(),
197    };
198    solve(&mut header);
199
200    header
201}
202
203fn solve(header: &mut BlockHeader) {
204    let target = header.target();
205    while header.validate_pow(&target).is_err() {
206        header.nonce += 1;
207    }
208}
209
210#[cfg(test)]
211mod test {
212    mod transaction_builder {
213        use crate::{random_p2pkh_address, TransactionBuilder};
214        use bitcoin::{Network, OutPoint};
215
216        #[test]
217        fn new_build() {
218            let tx = TransactionBuilder::new().build();
219            assert!(tx.is_coin_base());
220            assert_eq!(tx.input.len(), 1);
221            assert_eq!(tx.input[0].previous_output, OutPoint::null());
222            assert_eq!(tx.output.len(), 1);
223            assert_eq!(tx.output[0].value, 50_0000_0000);
224        }
225
226        #[test]
227        fn coinbase() {
228            let tx = TransactionBuilder::coinbase().build();
229            assert!(tx.is_coin_base());
230            assert_eq!(tx.input.len(), 1);
231            assert_eq!(tx.input[0].previous_output, OutPoint::null());
232            assert_eq!(tx.output.len(), 1);
233            assert_eq!(tx.output[0].value, 50_0000_0000);
234        }
235
236        #[test]
237        #[should_panic(
238            expected = "A call `with_input` should not be possible if `coinbase` was called"
239        )]
240        fn with_input_panic() {
241            let address = random_p2pkh_address(Network::Regtest);
242            let coinbase_tx = TransactionBuilder::coinbase()
243                .with_output(&address, 1000)
244                .build();
245
246            TransactionBuilder::coinbase()
247                .with_input(bitcoin::OutPoint::new(coinbase_tx.txid(), 0), None);
248        }
249
250        #[test]
251        fn with_output() {
252            let address = random_p2pkh_address(Network::Regtest);
253            let tx = TransactionBuilder::coinbase()
254                .with_output(&address, 1000)
255                .build();
256
257            assert!(tx.is_coin_base());
258            assert_eq!(tx.input.len(), 1);
259            assert_eq!(tx.input[0].previous_output, OutPoint::null());
260            assert_eq!(tx.output.len(), 1);
261            assert_eq!(tx.output[0].value, 1000);
262            assert_eq!(tx.output[0].script_pubkey, address.script_pubkey());
263        }
264
265        #[test]
266        fn with_output_2() {
267            let address_0 = random_p2pkh_address(Network::Regtest);
268            let address_1 = random_p2pkh_address(Network::Regtest);
269            let tx = TransactionBuilder::coinbase()
270                .with_output(&address_0, 1000)
271                .with_output(&address_1, 2000)
272                .build();
273
274            assert!(tx.is_coin_base());
275            assert_eq!(tx.input.len(), 1);
276            assert_eq!(tx.input[0].previous_output, OutPoint::null());
277            assert_eq!(tx.output.len(), 2);
278            assert_eq!(tx.output[0].value, 1000);
279            assert_eq!(tx.output[0].script_pubkey, address_0.script_pubkey());
280            assert_eq!(tx.output[1].value, 2000);
281            assert_eq!(tx.output[1].script_pubkey, address_1.script_pubkey());
282        }
283
284        #[test]
285        fn with_input() {
286            let address = random_p2pkh_address(Network::Regtest);
287            let coinbase_tx = TransactionBuilder::coinbase()
288                .with_output(&address, 1000)
289                .build();
290
291            let tx = TransactionBuilder::new()
292                .with_input(bitcoin::OutPoint::new(coinbase_tx.txid(), 0), None)
293                .build();
294            assert!(!tx.is_coin_base());
295            assert_eq!(tx.input.len(), 1);
296            assert_eq!(
297                tx.input[0].previous_output,
298                bitcoin::OutPoint::new(coinbase_tx.txid(), 0)
299            );
300            assert_eq!(tx.output.len(), 1);
301            assert_eq!(tx.output[0].value, 50_0000_0000);
302        }
303
304        #[test]
305        fn with_input_2() {
306            let address = random_p2pkh_address(Network::Regtest);
307            let coinbase_tx_0 = TransactionBuilder::coinbase()
308                .with_output(&address, 1000)
309                .build();
310            let coinbase_tx_1 = TransactionBuilder::coinbase()
311                .with_output(&address, 2000)
312                .build();
313
314            let tx = TransactionBuilder::new()
315                .with_input(bitcoin::OutPoint::new(coinbase_tx_0.txid(), 0), None)
316                .with_input(bitcoin::OutPoint::new(coinbase_tx_1.txid(), 0), None)
317                .build();
318            assert!(!tx.is_coin_base());
319            assert_eq!(tx.input.len(), 2);
320            assert_eq!(
321                tx.input[0].previous_output,
322                bitcoin::OutPoint::new(coinbase_tx_0.txid(), 0)
323            );
324            assert_eq!(
325                tx.input[1].previous_output,
326                bitcoin::OutPoint::new(coinbase_tx_1.txid(), 0)
327            );
328            assert_eq!(tx.output.len(), 1);
329            assert_eq!(tx.output[0].value, 50_0000_0000);
330        }
331    }
332}
ic_btc_test_utils/lib.rs

ic_btc_test_utils/
lib.rs
