_ Git - rust-lightning/blob - lightning-block-sync/src/test_utils.rs

   1 use crate::{AsyncBlockSourceResult, BlockData, BlockHeaderData, BlockSource, BlockSourceError, UnboundedCache};
   2 use crate::poll::{Validate, ValidatedBlockHeader};
   3
   4 use bitcoin::blockdata::block::{Block, BlockHeader};
   5 use bitcoin::blockdata::constants::genesis_block;
   6 use bitcoin::hash_types::BlockHash;
   7 use bitcoin::network::constants::Network;
   8 use bitcoin::util::uint::Uint256;
   9 use bitcoin::util::hash::bitcoin_merkle_root;
  10 use bitcoin::{PackedLockTime, Transaction};
  11
  12 use lightning::chain;
  13
  14 use std::cell::RefCell;
  15 use std::collections::VecDeque;
  16
  17 #[derive(Default)]
  18 pub struct Blockchain {
  19         pub blocks: Vec<Block>,
  20         without_blocks: Option<std::ops::RangeFrom<usize>>,
  21         without_headers: bool,
  22         malformed_headers: bool,
  23         filtered_blocks: bool,
  24 }
  25
  26 impl Blockchain {
  27         pub fn default() -> Self {
  28                 Blockchain::with_network(Network::Bitcoin)
  29         }
  30
  31         pub fn with_network(network: Network) -> Self {
  32                 let blocks = vec![genesis_block(network)];
  33                 Self { blocks, ..Default::default() }
  34         }
  35
  36         pub fn with_height(mut self, height: usize) -> Self {
  37                 self.blocks.reserve_exact(height);
  38                 let bits = BlockHeader::compact_target_from_u256(&Uint256::from_be_bytes([0xff; 32]));
  39                 for i in 1..=height {
  40                         let prev_block = &self.blocks[i - 1];
  41                         let prev_blockhash = prev_block.block_hash();
  42                         let time = prev_block.header.time + height as u32;
  43                         // Must have at least one transaction, because the merkle root is not defined for an empty block
  44                         // and we would fail when we later checked, as of bitcoin crate 0.28.0.
  45                         // Note that elsewhere in tests we assume that the merkle root of an empty block is all zeros,
  46                         // but that's OK because those tests don't trigger the check.
  47                         let coinbase = Transaction {
  48                                 version: 0,
  49                                 lock_time: PackedLockTime::ZERO,
  50                                 input: vec![],
  51                                 output: vec![]
  52                         };
  53                         let merkle_root = bitcoin_merkle_root(vec![coinbase.txid().as_hash()].into_iter()).unwrap();
  54                         self.blocks.push(Block {
  55                                 header: BlockHeader {
  56                                         version: 0,
  57                                         prev_blockhash,
  58                                         merkle_root: merkle_root.into(),
  59                                         time,
  60                                         bits,
  61                                         nonce: 0,
  62                                 },
  63                                 txdata: vec![coinbase],
  64                         });
  65                 }
  66                 self
  67         }
  68
  69         pub fn without_blocks(self, range: std::ops::RangeFrom<usize>) -> Self {
  70                 Self { without_blocks: Some(range), ..self }
  71         }
  72
  73         pub fn without_headers(self) -> Self {
  74                 Self { without_headers: true, ..self }
  75         }
  76
  77         pub fn malformed_headers(self) -> Self {
  78                 Self { malformed_headers: true, ..self }
  79         }
  80
  81         pub fn filtered_blocks(self) -> Self {
  82                 Self { filtered_blocks: true, ..self }
  83         }
  84
  85         pub fn fork_at_height(&self, height: usize) -> Self {
  86                 assert!(height + 1 < self.blocks.len());
  87                 let mut blocks = self.blocks.clone();
  88                 let mut prev_blockhash = blocks[height].block_hash();
  89                 for block in blocks.iter_mut().skip(height + 1) {
  90                         block.header.prev_blockhash = prev_blockhash;
  91                         block.header.nonce += 1;
  92                         prev_blockhash = block.block_hash();
  93                 }
  94                 Self { blocks, without_blocks: None, ..*self }
  95         }
  96
  97         pub fn at_height(&self, height: usize) -> ValidatedBlockHeader {
  98                 let block_header = self.at_height_unvalidated(height);
  99                 let block_hash = self.blocks[height].block_hash();
 100                 block_header.validate(block_hash).unwrap()
 101         }
 102
 103         fn at_height_unvalidated(&self, height: usize) -> BlockHeaderData {
 104                 assert!(!self.blocks.is_empty());
 105                 assert!(height < self.blocks.len());
 106                 BlockHeaderData {
 107                         chainwork: self.blocks[0].header.work() + Uint256::from_u64(height as u64).unwrap(),
 108                         height: height as u32,
 109                         header: self.blocks[height].header,
 110                 }
 111         }
 112
 113         pub fn tip(&self) -> ValidatedBlockHeader {
 114                 assert!(!self.blocks.is_empty());
 115                 self.at_height(self.blocks.len() - 1)
 116         }
 117
 118         pub fn disconnect_tip(&mut self) -> Option<Block> {
 119                 self.blocks.pop()
 120         }
 121
 122         pub fn header_cache(&self, heights: std::ops::RangeInclusive<usize>) -> UnboundedCache {
 123                 let mut cache = UnboundedCache::new();
 124                 for i in heights {
 125                         let value = self.at_height(i);
 126                         let key = value.header.block_hash();
 127                         assert!(cache.insert(key, value).is_none());
 128                 }
 129                 cache
 130         }
 131 }
 132
 133 impl BlockSource for Blockchain {
 134         fn get_header<'a>(&'a self, header_hash: &'a BlockHash, _height_hint: Option<u32>) -> AsyncBlockSourceResult<'a, BlockHeaderData> {
 135                 Box::pin(async move {
 136                         if self.without_headers {
 137                                 return Err(BlockSourceError::persistent("header not found"));
 138                         }
 139
 140                         for (height, block) in self.blocks.iter().enumerate() {
 141                                 if block.header.block_hash() == *header_hash {
 142                                         let mut header_data = self.at_height_unvalidated(height);
 143                                         if self.malformed_headers {
 144                                                 header_data.header.time += 1;
 145                                         }
 146
 147                                         return Ok(header_data);
 148                                 }
 149                         }
 150                         Err(BlockSourceError::transient("header not found"))
 151                 })
 152         }
 153
 154         fn get_block<'a>(&'a self, header_hash: &'a BlockHash) -> AsyncBlockSourceResult<'a, BlockData> {
 155                 Box::pin(async move {
 156                         for (height, block) in self.blocks.iter().enumerate() {
 157                                 if block.header.block_hash() == *header_hash {
 158                                         if let Some(without_blocks) = &self.without_blocks {
 159                                                 if without_blocks.contains(&height) {
 160                                                         return Err(BlockSourceError::persistent("block not found"));
 161                                                 }
 162                                         }
 163
 164                                         if self.filtered_blocks {
 165                                                 return Ok(BlockData::HeaderOnly(block.header));
 166                                         } else {
 167                                                 return Ok(BlockData::FullBlock(block.clone()));
 168                                         }
 169                                 }
 170                         }
 171                         Err(BlockSourceError::transient("block not found"))
 172                 })
 173         }
 174
 175         fn get_best_block<'a>(&'a self) -> AsyncBlockSourceResult<'a, (BlockHash, Option<u32>)> {
 176                 Box::pin(async move {
 177                         match self.blocks.last() {
 178                                 None => Err(BlockSourceError::transient("empty chain")),
 179                                 Some(block) => {
 180                                         let height = (self.blocks.len() - 1) as u32;
 181                                         Ok((block.block_hash(), Some(height)))
 182                                 },
 183                         }
 184                 })
 185         }
 186 }
 187
 188 pub struct NullChainListener;
 189
 190 impl chain::Listen for NullChainListener {
 191         fn filtered_block_connected(&self, _header: &BlockHeader, _txdata: &chain::transaction::TransactionData, _height: u32) {}
 192         fn block_disconnected(&self, _header: &BlockHeader, _height: u32) {}
 193 }
 194
 195 pub struct MockChainListener {
 196         expected_blocks_connected: RefCell<VecDeque<BlockHeaderData>>,
 197         expected_filtered_blocks_connected: RefCell<VecDeque<BlockHeaderData>>,
 198         expected_blocks_disconnected: RefCell<VecDeque<BlockHeaderData>>,
 199 }
 200
 201 impl MockChainListener {
 202         pub fn new() -> Self {
 203                 Self {
 204                         expected_blocks_connected: RefCell::new(VecDeque::new()),
 205                         expected_filtered_blocks_connected: RefCell::new(VecDeque::new()),
 206                         expected_blocks_disconnected: RefCell::new(VecDeque::new()),
 207                 }
 208         }
 209
 210         pub fn expect_block_connected(self, block: BlockHeaderData) -> Self {
 211                 self.expected_blocks_connected.borrow_mut().push_back(block);
 212                 self
 213         }
 214
 215         pub fn expect_filtered_block_connected(self, block: BlockHeaderData) -> Self {
 216                 self.expected_filtered_blocks_connected.borrow_mut().push_back(block);
 217                 self
 218         }
 219
 220         pub fn expect_block_disconnected(self, block: BlockHeaderData) -> Self {
 221                 self.expected_blocks_disconnected.borrow_mut().push_back(block);
 222                 self
 223         }
 224 }
 225
 226 impl chain::Listen for MockChainListener {
 227         fn block_connected(&self, block: &Block, height: u32) {
 228                 match self.expected_blocks_connected.borrow_mut().pop_front() {
 229                         None => {
 230                                 panic!("Unexpected block connected: {:?}", block.block_hash());
 231                         },
 232                         Some(expected_block) => {
 233                                 assert_eq!(block.block_hash(), expected_block.header.block_hash());
 234                                 assert_eq!(height, expected_block.height);
 235                         },
 236                 }
 237         }
 238
 239         fn filtered_block_connected(&self, header: &BlockHeader, _txdata: &chain::transaction::TransactionData, height: u32) {
 240                 match self.expected_filtered_blocks_connected.borrow_mut().pop_front() {
 241                         None => {
 242                                 panic!("Unexpected filtered block connected: {:?}", header.block_hash());
 243                         },
 244                         Some(expected_block) => {
 245                                 assert_eq!(header.block_hash(), expected_block.header.block_hash());
 246                                 assert_eq!(height, expected_block.height);
 247                         },
 248                 }
 249         }
 250
 251         fn block_disconnected(&self, header: &BlockHeader, height: u32) {
 252                 match self.expected_blocks_disconnected.borrow_mut().pop_front() {
 253                         None => {
 254                                 panic!("Unexpected block disconnected: {:?}", header.block_hash());
 255                         },
 256                         Some(expected_block) => {
 257                                 assert_eq!(header.block_hash(), expected_block.header.block_hash());
 258                                 assert_eq!(height, expected_block.height);
 259                         },
 260                 }
 261         }
 262 }
 263
 264 impl Drop for MockChainListener {
 265         fn drop(&mut self) {
 266                 if std::thread::panicking() {
 267                         return;
 268                 }
 269
 270                 let expected_blocks_connected = self.expected_blocks_connected.borrow();
 271                 if !expected_blocks_connected.is_empty() {
 272                         panic!("Expected blocks connected: {:?}", expected_blocks_connected);
 273                 }
 274
 275                 let expected_filtered_blocks_connected = self.expected_filtered_blocks_connected.borrow();
 276                 if !expected_filtered_blocks_connected.is_empty() {
 277                         panic!("Expected filtered_blocks connected: {:?}", expected_filtered_blocks_connected);
 278                 }
 279
 280                 let expected_blocks_disconnected = self.expected_blocks_disconnected.borrow();
 281                 if !expected_blocks_disconnected.is_empty() {
 282                         panic!("Expected blocks disconnected: {:?}", expected_blocks_disconnected);
 283                 }
 284         }
 285 }