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

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