_ 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, Header, Version};
   5 use bitcoin::blockdata::constants::genesis_block;
   6 use bitcoin::blockdata::locktime::absolute::LockTime;
   7 use bitcoin::hash_types::{BlockHash, TxMerkleNode};
   8 use bitcoin::network::constants::Network;
   9 use bitcoin::Transaction;
  10
  11 use lightning::chain;
  12
  13 use std::cell::RefCell;
  14 use std::collections::VecDeque;
  15
  16 #[derive(Default)]
  17 pub struct Blockchain {
  18         pub blocks: Vec<Block>,
  19         without_blocks: Option<std::ops::RangeFrom<usize>>,
  20         without_headers: bool,
  21         malformed_headers: bool,
  22         filtered_blocks: 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 = bitcoin::Target::from_be_bytes([0xff; 32]).to_compact_lossy();
  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: LockTime::ZERO,
  49                                 input: vec![],
  50                                 output: vec![]
  51                         };
  52                         let merkle_root = TxMerkleNode::from_raw_hash(coinbase.txid().to_raw_hash());
  53                         self.blocks.push(Block {
  54                                 header: Header {
  55                                         version: Version::NO_SOFT_FORK_SIGNALLING,
  56                                         prev_blockhash,
  57                                         merkle_root,
  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 filtered_blocks(self) -> Self {
  81                 Self { filtered_blocks: true, ..self }
  82         }
  83
  84         pub fn fork_at_height(&self, height: usize) -> Self {
  85                 assert!(height + 1 < self.blocks.len());
  86                 let mut blocks = self.blocks.clone();
  87                 let mut prev_blockhash = blocks[height].block_hash();
  88                 for block in blocks.iter_mut().skip(height + 1) {
  89                         block.header.prev_blockhash = prev_blockhash;
  90                         block.header.nonce += 1;
  91                         prev_blockhash = block.block_hash();
  92                 }
  93                 Self { blocks, without_blocks: None, ..*self }
  94         }
  95
  96         pub fn at_height(&self, height: usize) -> ValidatedBlockHeader {
  97                 let block_header = self.at_height_unvalidated(height);
  98                 let block_hash = self.blocks[height].block_hash();
  99                 block_header.validate(block_hash).unwrap()
 100         }
 101
 102         fn at_height_unvalidated(&self, height: usize) -> BlockHeaderData {
 103                 assert!(!self.blocks.is_empty());
 104                 assert!(height < self.blocks.len());
 105                 let mut total_work = self.blocks[0].header.work();
 106                 for i in 1..=height {
 107                         total_work = total_work + self.blocks[i].header.work();
 108                 }
 109                 BlockHeaderData {
 110                         chainwork: total_work,
 111                         height: height as u32,
 112                         header: self.blocks[height].header,
 113                 }
 114         }
 115
 116         pub fn tip(&self) -> ValidatedBlockHeader {
 117                 assert!(!self.blocks.is_empty());
 118                 self.at_height(self.blocks.len() - 1)
 119         }
 120
 121         pub fn disconnect_tip(&mut self) -> Option<Block> {
 122                 self.blocks.pop()
 123         }
 124
 125         pub fn header_cache(&self, heights: std::ops::RangeInclusive<usize>) -> UnboundedCache {
 126                 let mut cache = UnboundedCache::new();
 127                 for i in heights {
 128                         let value = self.at_height(i);
 129                         let key = value.header.block_hash();
 130                         assert!(cache.insert(key, value).is_none());
 131                 }
 132                 cache
 133         }
 134 }
 135
 136 impl BlockSource for Blockchain {
 137         fn get_header<'a>(&'a self, header_hash: &'a BlockHash, _height_hint: Option<u32>) -> AsyncBlockSourceResult<'a, BlockHeaderData> {
 138                 Box::pin(async move {
 139                         if self.without_headers {
 140                                 return Err(BlockSourceError::persistent("header not found"));
 141                         }
 142
 143                         for (height, block) in self.blocks.iter().enumerate() {
 144                                 if block.header.block_hash() == *header_hash {
 145                                         let mut header_data = self.at_height_unvalidated(height);
 146                                         if self.malformed_headers {
 147                                                 header_data.header.time += 1;
 148                                         }
 149
 150                                         return Ok(header_data);
 151                                 }
 152                         }
 153                         Err(BlockSourceError::transient("header not found"))
 154                 })
 155         }
 156
 157         fn get_block<'a>(&'a self, header_hash: &'a BlockHash) -> AsyncBlockSourceResult<'a, BlockData> {
 158                 Box::pin(async move {
 159                         for (height, block) in self.blocks.iter().enumerate() {
 160                                 if block.header.block_hash() == *header_hash {
 161                                         if let Some(without_blocks) = &self.without_blocks {
 162                                                 if without_blocks.contains(&height) {
 163                                                         return Err(BlockSourceError::persistent("block not found"));
 164                                                 }
 165                                         }
 166
 167                                         if self.filtered_blocks {
 168                                                 return Ok(BlockData::HeaderOnly(block.header));
 169                                         } else {
 170                                                 return Ok(BlockData::FullBlock(block.clone()));
 171                                         }
 172                                 }
 173                         }
 174                         Err(BlockSourceError::transient("block not found"))
 175                 })
 176         }
 177
 178         fn get_best_block<'a>(&'a self) -> AsyncBlockSourceResult<'a, (BlockHash, Option<u32>)> {
 179                 Box::pin(async move {
 180                         match self.blocks.last() {
 181                                 None => Err(BlockSourceError::transient("empty chain")),
 182                                 Some(block) => {
 183                                         let height = (self.blocks.len() - 1) as u32;
 184                                         Ok((block.block_hash(), Some(height)))
 185                                 },
 186                         }
 187                 })
 188         }
 189 }
 190
 191 pub struct NullChainListener;
 192
 193 impl chain::Listen for NullChainListener {
 194         fn filtered_block_connected(&self, _header: &Header, _txdata: &chain::transaction::TransactionData, _height: u32) {}
 195         fn block_disconnected(&self, _header: &Header, _height: u32) {}
 196 }
 197
 198 pub struct MockChainListener {
 199         expected_blocks_connected: RefCell<VecDeque<BlockHeaderData>>,
 200         expected_filtered_blocks_connected: RefCell<VecDeque<BlockHeaderData>>,
 201         expected_blocks_disconnected: RefCell<VecDeque<BlockHeaderData>>,
 202 }
 203
 204 impl MockChainListener {
 205         pub fn new() -> Self {
 206                 Self {
 207                         expected_blocks_connected: RefCell::new(VecDeque::new()),
 208                         expected_filtered_blocks_connected: RefCell::new(VecDeque::new()),
 209                         expected_blocks_disconnected: RefCell::new(VecDeque::new()),
 210                 }
 211         }
 212
 213         pub fn expect_block_connected(self, block: BlockHeaderData) -> Self {
 214                 self.expected_blocks_connected.borrow_mut().push_back(block);
 215                 self
 216         }
 217
 218         pub fn expect_filtered_block_connected(self, block: BlockHeaderData) -> Self {
 219                 self.expected_filtered_blocks_connected.borrow_mut().push_back(block);
 220                 self
 221         }
 222
 223         pub fn expect_block_disconnected(self, block: BlockHeaderData) -> Self {
 224                 self.expected_blocks_disconnected.borrow_mut().push_back(block);
 225                 self
 226         }
 227 }
 228
 229 impl chain::Listen for MockChainListener {
 230         fn block_connected(&self, block: &Block, height: u32) {
 231                 match self.expected_blocks_connected.borrow_mut().pop_front() {
 232                         None => {
 233                                 panic!("Unexpected block connected: {:?}", block.block_hash());
 234                         },
 235                         Some(expected_block) => {
 236                                 assert_eq!(block.block_hash(), expected_block.header.block_hash());
 237                                 assert_eq!(height, expected_block.height);
 238                         },
 239                 }
 240         }
 241
 242         fn filtered_block_connected(&self, header: &Header, _txdata: &chain::transaction::TransactionData, height: u32) {
 243                 match self.expected_filtered_blocks_connected.borrow_mut().pop_front() {
 244                         None => {
 245                                 panic!("Unexpected filtered block connected: {:?}", header.block_hash());
 246                         },
 247                         Some(expected_block) => {
 248                                 assert_eq!(header.block_hash(), expected_block.header.block_hash());
 249                                 assert_eq!(height, expected_block.height);
 250                         },
 251                 }
 252         }
 253
 254         fn block_disconnected(&self, header: &Header, height: u32) {
 255                 match self.expected_blocks_disconnected.borrow_mut().pop_front() {
 256                         None => {
 257                                 panic!("Unexpected block disconnected: {:?}", header.block_hash());
 258                         },
 259                         Some(expected_block) => {
 260                                 assert_eq!(header.block_hash(), expected_block.header.block_hash());
 261                                 assert_eq!(height, expected_block.height);
 262                         },
 263                 }
 264         }
 265 }
 266
 267 impl Drop for MockChainListener {
 268         fn drop(&mut self) {
 269                 if std::thread::panicking() {
 270                         return;
 271                 }
 272
 273                 let expected_blocks_connected = self.expected_blocks_connected.borrow();
 274                 if !expected_blocks_connected.is_empty() {
 275                         panic!("Expected blocks connected: {:?}", expected_blocks_connected);
 276                 }
 277
 278                 let expected_filtered_blocks_connected = self.expected_filtered_blocks_connected.borrow();
 279                 if !expected_filtered_blocks_connected.is_empty() {
 280                         panic!("Expected filtered_blocks connected: {:?}", expected_filtered_blocks_connected);
 281                 }
 282
 283                 let expected_blocks_disconnected = self.expected_blocks_disconnected.borrow();
 284                 if !expected_blocks_disconnected.is_empty() {
 285                         panic!("Expected blocks disconnected: {:?}", expected_blocks_disconnected);
 286                 }
 287         }
 288 }