1use std::collections::HashSet;
21
22use bitcoin::consensus;
23use bitcoin::hashes::{sha256, Hash};
24use bitcoin::{Script, ScriptBuf, Transaction, Txid};
25
26use crate::error::StreamError;
27use crate::event::{Outpoint, PrefixMatch};
28
29pub fn scripthash_of(script_pubkey: &[u8]) -> [u8; 32] {
32 sha256::Hash::hash(script_pubkey).to_byte_array()
33}
34
35pub fn prefix_of(script_pubkey: &[u8], bits: u32) -> (Vec<u8>, u32) {
48 let bits = bits.clamp(1, crate::client::MAX_PREFIX_BITS);
49 let sh = scripthash_of(script_pubkey);
50 (mask_prefix(&sh, bits), bits)
51}
52
53fn mask_prefix(scripthash: &[u8], bits: u32) -> Vec<u8> {
56 let n = (bits as usize).div_ceil(8);
57 let mut out = scripthash[..n].to_vec();
58 let rem = bits % 8;
59 if rem != 0 {
60 out[n - 1] &= 0xffu8 << (8 - rem);
61 }
62 out
63}
64
65#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct FundingHit {
68 pub vout: u32,
70 pub scripthash: [u8; 32],
72 pub value: u64,
74 pub script_pubkey: ScriptBuf,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
82pub struct SpendingHit {
83 pub vin: Option<u32>,
86 pub outpoint: Outpoint,
88 pub scripthash: [u8; 32],
90 pub script_pubkey: ScriptBuf,
92 pub amount: Option<u64>,
95}
96
97#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct PrefixHits {
102 pub txid: Txid,
104 pub confirmed: bool,
106 pub height: u32,
108 pub funding: Vec<FundingHit>,
110 pub spending: Vec<SpendingHit>,
112 pub unresolved: Vec<Outpoint>,
117}
118
119impl PrefixHits {
120 pub fn is_match(&self) -> bool {
122 !self.funding.is_empty() || !self.spending.is_empty()
123 }
124
125 pub fn has_unresolved(&self) -> bool {
129 !self.unresolved.is_empty()
130 }
131}
132
133#[derive(Debug, Clone, Default)]
136pub struct PrefixWatcher {
137 watched: HashSet<[u8; 32]>,
139}
140
141impl PrefixWatcher {
142 pub fn new() -> Self {
144 Self::default()
145 }
146
147 pub fn with_scripts<'a>(scripts: impl IntoIterator<Item = &'a Script>) -> Self {
149 let mut w = Self::new();
150 for s in scripts {
151 w.watch_script(s);
152 }
153 w
154 }
155
156 pub fn watch_script(&mut self, script_pubkey: &Script) -> [u8; 32] {
159 self.watch_script_bytes(script_pubkey.as_bytes())
160 }
161
162 pub fn watch_script_bytes(&mut self, script_pubkey: &[u8]) -> [u8; 32] {
164 let sh = scripthash_of(script_pubkey);
165 self.watched.insert(sh);
166 sh
167 }
168
169 pub fn unwatch_script(&mut self, script_pubkey: &Script) -> bool {
171 self.watched.remove(&scripthash_of(script_pubkey.as_bytes()))
172 }
173
174 pub fn is_watched(&self, scripthash: &[u8; 32]) -> bool {
176 self.watched.contains(scripthash)
177 }
178
179 pub fn len(&self) -> usize {
181 self.watched.len()
182 }
183
184 pub fn is_empty(&self) -> bool {
186 self.watched.is_empty()
187 }
188
189 pub fn prefixes(&self, bits: u32) -> Vec<(Vec<u8>, u32)> {
206 let bits = bits.clamp(1, crate::client::MAX_PREFIX_BITS);
207 let buckets: HashSet<Vec<u8>> =
208 self.watched.iter().map(|sh| mask_prefix(sh, bits)).collect();
209 let mut out: Vec<(Vec<u8>, u32)> = buckets.into_iter().map(|p| (p, bits)).collect();
210 out.sort();
211 out
212 }
213
214 pub fn filter(&self, m: &PrefixMatch) -> Result<PrefixHits, StreamError> {
219 let tx: Transaction = consensus::deserialize(&m.raw_tx)
220 .map_err(|e| StreamError::Decode(format!("prefix raw_tx decode: {e}")))?;
221 let txid = tx.compute_txid();
222
223 let mut funding = Vec::new();
224 for (vout, out) in tx.output.iter().enumerate() {
225 let sh = scripthash_of(out.script_pubkey.as_bytes());
226 if self.watched.contains(&sh) {
227 funding.push(FundingHit {
228 vout: vout as u32,
229 scripthash: sh,
230 value: out.value.to_sat(),
231 script_pubkey: out.script_pubkey.clone(),
232 });
233 }
234 }
235
236 let mut spending = Vec::new();
237 let mut unresolved = Vec::new();
238 for sp in &m.matched_prevouts {
239 if sp.script_pubkey.is_empty() {
240 unresolved.push(sp.outpoint.clone());
242 continue;
243 }
244 let sh = scripthash_of(&sp.script_pubkey);
245 if self.watched.contains(&sh) {
246 spending.push(SpendingHit {
247 vin: find_vin(&tx, &sp.outpoint),
248 outpoint: sp.outpoint.clone(),
249 scripthash: sh,
250 script_pubkey: ScriptBuf::from_bytes(sp.script_pubkey.clone()),
251 amount: sp.amount,
252 });
253 }
254 }
255
256 Ok(PrefixHits { txid, confirmed: m.confirmed, height: m.height, funding, spending, unresolved })
257 }
258}
259
260fn find_vin(tx: &Transaction, outpoint: &Outpoint) -> Option<u32> {
263 tx.input.iter().position(|i| {
264 i.previous_output.vout == outpoint.vout
265 && i.previous_output.txid.to_byte_array().as_slice() == outpoint.txid.as_slice()
266 }).map(|p| p as u32)
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272 use crate::event::{ScriptPrefix, SpentPrevout};
273 use bitcoin::{Amount, OutPoint, Sequence, TxIn, TxOut, Witness};
274
275 fn spk(tag: u8) -> ScriptBuf {
276 ScriptBuf::from_bytes(vec![0x6a, 0x01, tag])
278 }
279
280 fn dummy_outpoint(byte: u8, vout: u32) -> (OutPoint, Outpoint) {
281 let txid = Txid::from_byte_array([byte; 32]);
282 (
283 OutPoint { txid, vout },
284 Outpoint { txid: txid.to_byte_array().to_vec(), vout },
285 )
286 }
287
288 fn tx_with(outputs: Vec<TxOut>, inputs: Vec<TxIn>) -> Transaction {
289 Transaction {
290 version: bitcoin::transaction::Version::TWO,
291 lock_time: bitcoin::absolute::LockTime::ZERO,
292 input: inputs,
293 output: outputs,
294 }
295 }
296
297 fn prefix_match(tx: &Transaction, prevouts: Vec<SpentPrevout>) -> PrefixMatch {
298 PrefixMatch {
299 prefix: ScriptPrefix { prefix: vec![0u8, 0u8], bits: 16 },
300 raw_tx: consensus::serialize(tx),
301 confirmed: true,
302 height: 800_000,
303 matched_prevouts: prevouts,
304 }
305 }
306
307 #[test]
308 fn scripthash_matches_single_sha256_of_spk() {
309 let s = spk(1);
311 let expect = sha256::Hash::hash(s.as_bytes()).to_byte_array();
312 assert_eq!(scripthash_of(s.as_bytes()), expect);
313 }
314
315 #[test]
324 fn prefix_of_masks_bits_below_the_declared_width() {
325 let s = spk(7);
326 let sh = scripthash_of(s.as_bytes());
327 let (p, bits) = prefix_of(s.as_bytes(), 12);
328 assert_eq!(bits, 12);
329 assert_eq!(p.len(), 2, "ceil(12/8) bytes on the wire");
330 assert_eq!(p[0], sh[0], "the first whole byte is inside the width");
331 assert_eq!(p[1], sh[1] & 0xf0, "the 4 bits past the width are zeroed");
332 let key = |b0: u8, b1: u8| u32::from_be_bytes([b0, b1, 0, 0]) & (u32::MAX << 20);
334 assert_eq!(key(p[0], p[1]), key(sh[0], sh[1]), "masking must not move the bucket");
335 }
336
337 #[test]
344 fn prefixes_collapse_when_scripts_share_a_sub_byte_bucket() {
345 fn spk_n(i: u32) -> ScriptBuf {
347 let mut v = vec![0x6a, 0x04];
348 v.extend_from_slice(&i.to_be_bytes());
349 ScriptBuf::from_bytes(v)
350 }
351
352 let mut pair = None;
355 let mut seen: std::collections::HashMap<[u8; 2], ScriptBuf> = Default::default();
356 for i in 0..20_000u32 {
357 let s = spk_n(i);
358 let sh = scripthash_of(s.as_bytes());
359 let bucket = [sh[0], sh[1] & 0xf0];
360 if let Some(prev) = seen.get(&bucket) {
361 if scripthash_of(prev.as_bytes())[1] != sh[1] {
362 pair = Some((prev.clone(), s));
363 break;
364 }
365 } else {
366 seen.insert(bucket, s);
367 }
368 }
369 let (a, b) = pair.expect("a 12-bit bucket collision within the search space");
370
371 let mut w = PrefixWatcher::new();
372 w.watch_script(&a);
373 w.watch_script(&b);
374 assert_eq!(w.len(), 2, "two distinct scripts are watched");
375
376 let buckets = w.prefixes(12);
377 assert_eq!(
378 buckets.len(),
379 1,
380 "scripts sharing a 12-bit bucket must collapse to one registration, got {buckets:?}"
381 );
382 assert_eq!(buckets[0].1, 12);
383 assert_eq!(buckets[0].0[1] & 0x0f, 0, "the registered bucket is masked");
384 }
385
386 #[test]
387 fn prefix_of_takes_top_bytes_of_scripthash() {
388 let s = spk(7);
389 let sh = scripthash_of(s.as_bytes());
390 assert_eq!(prefix_of(s.as_bytes(), 16), (sh[..2].to_vec(), 16));
391 assert_eq!(prefix_of(s.as_bytes(), 8), (sh[..1].to_vec(), 8));
392 assert_eq!(prefix_of(s.as_bytes(), 32), (sh[..4].to_vec(), 32));
395 assert_eq!(prefix_of(s.as_bytes(), 256), (sh[..4].to_vec(), 32));
396 }
397
398 #[test]
399 fn funding_match_only_for_watched_output() {
400 let watched = spk(1);
401 let decoy = spk(2);
402 let mut w = PrefixWatcher::new();
403 w.watch_script(&watched);
404
405 let tx = tx_with(
406 vec![
407 TxOut { value: Amount::from_sat(500), script_pubkey: decoy },
408 TxOut { value: Amount::from_sat(1234), script_pubkey: watched.clone() },
409 ],
410 vec![],
411 );
412 let hits = w.filter(&prefix_match(&tx, vec![])).unwrap();
413 assert!(hits.is_match());
414 assert_eq!(hits.funding.len(), 1);
415 assert_eq!(hits.funding[0].vout, 1);
416 assert_eq!(hits.funding[0].value, 1234);
417 assert_eq!(hits.funding[0].script_pubkey, watched);
418 assert!(hits.spending.is_empty());
419 assert!(!hits.has_unresolved());
420 }
421
422 #[test]
423 fn decoy_only_tx_is_no_match() {
424 let mut w = PrefixWatcher::new();
425 w.watch_script(&spk(1));
426 let tx = tx_with(
427 vec![TxOut { value: Amount::from_sat(9), script_pubkey: spk(99) }],
428 vec![],
429 );
430 let hits = w.filter(&prefix_match(&tx, vec![])).unwrap();
431 assert!(!hits.is_match());
432 assert!(!hits.has_unresolved());
433 }
434
435 #[test]
436 fn spend_match_locates_vin_and_amount() {
437 let watched = spk(5);
438 let mut w = PrefixWatcher::new();
439 w.watch_script(&watched);
440
441 let (op_btc, op_wire) = dummy_outpoint(0xaa, 3);
442 let tx = tx_with(
443 vec![TxOut { value: Amount::from_sat(10), script_pubkey: spk(42) }],
444 vec![
445 TxIn {
446 previous_output: dummy_outpoint(0xbb, 0).0,
447 script_sig: ScriptBuf::new(),
448 sequence: Sequence::MAX,
449 witness: Witness::new(),
450 },
451 TxIn {
452 previous_output: op_btc,
453 script_sig: ScriptBuf::new(),
454 sequence: Sequence::MAX,
455 witness: Witness::new(),
456 },
457 ],
458 );
459 let prevout = SpentPrevout {
460 outpoint: op_wire.clone(),
461 script_pubkey: watched.as_bytes().to_vec(),
462 amount: Some(7777),
463 };
464 let hits = w.filter(&prefix_match(&tx, vec![prevout])).unwrap();
465 assert_eq!(hits.spending.len(), 1);
466 assert_eq!(hits.spending[0].vin, Some(1)); assert_eq!(hits.spending[0].amount, Some(7777));
468 assert_eq!(hits.spending[0].outpoint, op_wire);
469 }
470
471 #[test]
472 fn unretained_prevout_surfaces_as_unresolved() {
473 let mut w = PrefixWatcher::new();
474 w.watch_script(&spk(5));
475 let (_op_btc, op_wire) = dummy_outpoint(0xcc, 1);
476 let tx = tx_with(
477 vec![TxOut { value: Amount::from_sat(10), script_pubkey: spk(42) }],
478 vec![],
479 );
480 let prevout = SpentPrevout { outpoint: op_wire.clone(), script_pubkey: vec![], amount: None };
482 let hits = w.filter(&prefix_match(&tx, vec![prevout])).unwrap();
483 assert!(!hits.is_match()); assert!(hits.has_unresolved());
485 assert_eq!(hits.unresolved, vec![op_wire]);
486 }
487
488 #[test]
489 fn prefixes_dedup_shared_bucket() {
490 let mut w = PrefixWatcher::new();
491 w.watch_script(&spk(1));
492 w.watch_script(&spk(2));
493 assert_eq!(w.prefixes(32).len(), 2);
497 assert_eq!(w.prefixes(256).len(), 2); assert!(w.prefixes(1).len() <= 2);
499 for (p, bits) in w.prefixes(16) {
500 assert_eq!(bits, 16);
501 assert_eq!(p.len(), 2);
502 }
503 }
504
505 #[test]
506 fn garbage_raw_tx_is_decode_error() {
507 let w = PrefixWatcher::new();
508 let m = PrefixMatch {
509 prefix: ScriptPrefix { prefix: vec![0, 0], bits: 16 },
510 raw_tx: vec![0xff, 0x00, 0x01],
511 confirmed: false,
512 height: 0,
513 matched_prevouts: vec![],
514 };
515 assert!(matches!(w.filter(&m), Err(StreamError::Decode(_))));
516 }
517}