Add ChainPoller implementation of Poll trait

jkczyz · jkczyz · commit 7f2ec421e510 · 2021-02-04T10:49:29.000-08:00
ChainPoller defines a strategy for polling a single BlockSource. It
handles validating chain data returned from the BlockSource. Thus, other
implementations of Poll must be defined in terms of ChainPoller.
diff --git a/lightning-block-sync/src/lib.rs b/lightning-block-sync/src/lib.rs
@@ -14,6 +14,8 @@
 #[cfg(any(feature = "rest-client", feature = "rpc-client"))]
 pub mod http;
 
+pub mod poll;
+
 #[cfg(feature = "rest-client")]
 pub mod rest;
 
@@ -23,6 +25,9 @@ pub mod rpc;
 #[cfg(any(feature = "rest-client", feature = "rpc-client"))]
 mod convert;
 
+#[cfg(test)]
+mod test_utils;
+
 #[cfg(any(feature = "rest-client", feature = "rpc-client"))]
 mod utils;
 
@@ -65,13 +70,14 @@ type AsyncBlockSourceResult<'a, T> = Pin<Box<dyn Future<Output = BlockSourceResu
 ///
 /// Transient errors may be resolved when re-polling, but no attempt will be made to re-poll on
 /// persistent errors.
+#[derive(Debug)]
 pub struct BlockSourceError {
 	kind: BlockSourceErrorKind,
 	error: Box<dyn std::error::Error + Send + Sync>,
 }
 
 /// The kind of `BlockSourceError`, either persistent or transient.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum BlockSourceErrorKind {
 	/// Indicates an error that won't resolve when retrying a request (e.g., invalid data).
 	Persistent,
diff --git a/lightning-block-sync/src/poll.rs b/lightning-block-sync/src/poll.rs
@@ -1,9 +1,11 @@
-use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSourceError, BlockSourceResult};
+use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSource, BlockSourceError, BlockSourceResult};
 
 use bitcoin::blockdata::block::Block;
 use bitcoin::hash_types::BlockHash;
 use bitcoin::network::constants::Network;
 
+use std::ops::DerefMut;
+
 /// The `Poll` trait defines behavior for polling block sources for a chain tip and retrieving
 /// related chain data. It serves as an adapter for `BlockSource`.
 pub trait Poll {
@@ -140,3 +142,203 @@ impl std::ops::Deref for ValidatedBlock {
 		&self.inner
 	}
 }
+
+pub struct ChainPoller<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> {
+	block_source: B,
+	network: Network,
+}
+
+impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> ChainPoller<'b, B> {
+	pub fn new(block_source: B, network: Network) -> Self {
+		Self { block_source, network }
+	}
+}
+
+impl<'b, B: DerefMut<Target=dyn BlockSource + 'b> + Sized + Sync + Send> Poll for ChainPoller<'b, B> {
+	fn poll_chain_tip<'a>(&'a mut self, best_known_chain_tip: ValidatedBlockHeader) ->
+		AsyncBlockSourceResult<'a, ChainTip>
+	{
+		Box::pin(async move {
+			let (block_hash, height) = self.block_source.get_best_block().await?;
+			if block_hash == best_known_chain_tip.header.block_hash() {
+				return Ok(ChainTip::Common);
+			}
+
+			let chain_tip = self.block_source
+				.get_header(&block_hash, height).await?
+				.validate(block_hash)?;
+			if chain_tip.chainwork > best_known_chain_tip.chainwork {
+				Ok(ChainTip::Better(chain_tip))
+			} else {
+				Ok(ChainTip::Worse(chain_tip))
+			}
+		})
+	}
+
+	fn look_up_previous_header<'a>(&'a mut self, header: &'a ValidatedBlockHeader) ->
+		AsyncBlockSourceResult<'a, ValidatedBlockHeader>
+	{
+		Box::pin(async move {
+			if header.height == 0 {
+				return Err(BlockSourceError::persistent("genesis block reached"));
+			}
+
+			let previous_hash = &header.header.prev_blockhash;
+			let height = header.height - 1;
+			let previous_header = self.block_source
+				.get_header(previous_hash, Some(height)).await?
+				.validate(*previous_hash)?;
+			header.check_builds_on(&previous_header, self.network)?;
+
+			Ok(previous_header)
+		})
+	}
+
+	fn fetch_block<'a>(&'a mut self, header: &'a ValidatedBlockHeader) ->
+		AsyncBlockSourceResult<'a, ValidatedBlock>
+	{
+		Box::pin(async move {
+			self.block_source
+				.get_block(&header.block_hash).await?
+				.validate(header.block_hash)
+		})
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use crate::*;
+	use crate::test_utils::Blockchain;
+	use super::*;
+	use bitcoin::util::uint::Uint256;
+
+	#[tokio::test]
+	async fn poll_empty_chain() {
+		let mut chain = Blockchain::default().with_height(0);
+		let best_known_chain_tip = chain.tip();
+		chain.disconnect_tip();
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => {
+				assert_eq!(e.kind(), BlockSourceErrorKind::Transient);
+				assert_eq!(e.into_inner().as_ref().to_string(), "empty chain");
+			},
+			Ok(_) => panic!("Expected error"),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_without_headers() {
+		let mut chain = Blockchain::default().with_height(1).without_headers();
+		let best_known_chain_tip = chain.at_height(0);
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => {
+				assert_eq!(e.kind(), BlockSourceErrorKind::Persistent);
+				assert_eq!(e.into_inner().as_ref().to_string(), "header not found");
+			},
+			Ok(_) => panic!("Expected error"),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_invalid_pow() {
+		let mut chain = Blockchain::default().with_height(1);
+		let best_known_chain_tip = chain.at_height(0);
+
+		// Invalidate the tip by changing its target.
+		chain.blocks.last_mut().unwrap().header.bits =
+			BlockHeader::compact_target_from_u256(&Uint256::from_be_bytes([0; 32]));
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => {
+				assert_eq!(e.kind(), BlockSourceErrorKind::Persistent);
+				assert_eq!(e.into_inner().as_ref().to_string(), "block target correct but not attained");
+			},
+			Ok(_) => panic!("Expected error"),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_malformed_headers() {
+		let mut chain = Blockchain::default().with_height(1).malformed_headers();
+		let best_known_chain_tip = chain.at_height(0);
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => {
+				assert_eq!(e.kind(), BlockSourceErrorKind::Persistent);
+				assert_eq!(e.into_inner().as_ref().to_string(), "invalid block hash");
+			},
+			Ok(_) => panic!("Expected error"),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_common_tip() {
+		let mut chain = Blockchain::default().with_height(0);
+		let best_known_chain_tip = chain.tip();
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => panic!("Unexpected error: {:?}", e),
+			Ok(tip) => assert_eq!(tip, ChainTip::Common),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_uncommon_tip_but_equal_chainwork() {
+		let mut chain = Blockchain::default().with_height(1);
+		let best_known_chain_tip = chain.tip();
+
+		// Change the nonce to get a different block hash with the same chainwork.
+		chain.blocks.last_mut().unwrap().header.nonce += 1;
+
+		let worse_chain_tip = chain.tip();
+		let worse_chain_tip_hash = worse_chain_tip.header.block_hash();
+		let worse_chain_tip = worse_chain_tip.validate(worse_chain_tip_hash).unwrap();
+		assert_eq!(best_known_chain_tip.chainwork, worse_chain_tip.chainwork);
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => panic!("Unexpected error: {:?}", e),
+			Ok(tip) => assert_eq!(tip, ChainTip::Worse(worse_chain_tip)),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_worse_tip() {
+		let mut chain = Blockchain::default().with_height(1);
+		let best_known_chain_tip = chain.tip();
+		chain.disconnect_tip();
+
+		let worse_chain_tip = chain.tip();
+		let worse_chain_tip_hash = worse_chain_tip.header.block_hash();
+		let worse_chain_tip = worse_chain_tip.validate(worse_chain_tip_hash).unwrap();
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => panic!("Unexpected error: {:?}", e),
+			Ok(tip) => assert_eq!(tip, ChainTip::Worse(worse_chain_tip)),
+		}
+	}
+
+	#[tokio::test]
+	async fn poll_chain_with_better_tip() {
+		let mut chain = Blockchain::default().with_height(1);
+		let best_known_chain_tip = chain.at_height(0);
+
+		let better_chain_tip = chain.tip();
+		let better_chain_tip_hash = better_chain_tip.header.block_hash();
+		let better_chain_tip = better_chain_tip.validate(better_chain_tip_hash).unwrap();
+
+		let mut poller = ChainPoller::new(&mut chain as &mut dyn BlockSource, Network::Bitcoin);
+		match poller.poll_chain_tip(best_known_chain_tip).await {
+			Err(e) => panic!("Unexpected error: {:?}", e),
+			Ok(tip) => assert_eq!(tip, ChainTip::Better(better_chain_tip)),
+		}
+	}
+}
diff --git a/lightning-block-sync/src/test_utils.rs b/lightning-block-sync/src/test_utils.rs
@@ -0,0 +1,126 @@
+use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSource, BlockSourceError};
+use crate::poll::{Validate, ValidatedBlockHeader};
+
+use bitcoin::blockdata::block::{Block, BlockHeader};
+use bitcoin::blockdata::constants::genesis_block;
+use bitcoin::hash_types::BlockHash;
+use bitcoin::network::constants::Network;
+use bitcoin::util::uint::Uint256;
+
+#[derive(Default)]
+pub struct Blockchain {
+	pub blocks: Vec<Block>,
+	without_headers: bool,
+	malformed_headers: bool,
+}
+
+impl Blockchain {
+	pub fn default() -> Self {
+		Blockchain::with_network(Network::Bitcoin)
+	}
+
+	pub fn with_network(network: Network) -> Self {
+		let blocks = vec![genesis_block(network)];
+		Self { blocks, ..Default::default() }
+	}
+
+	pub fn with_height(mut self, height: usize) -> Self {
+		self.blocks.reserve_exact(height);
+		let bits = BlockHeader::compact_target_from_u256(&Uint256::from_be_bytes([0xff; 32]));
+		for i in 1..=height {
+			let prev_block = &self.blocks[i - 1];
+			let prev_blockhash = prev_block.block_hash();
+			let time = prev_block.header.time + height as u32;
+			self.blocks.push(Block {
+				header: BlockHeader {
+					version: 0,
+					prev_blockhash,
+					merkle_root: Default::default(),
+					time,
+					bits,
+					nonce: 0,
+				},
+				txdata: vec![],
+			});
+		}
+		self
+	}
+
+	pub fn without_headers(self) -> Self {
+		Self { without_headers: true, ..self }
+	}
+
+	pub fn malformed_headers(self) -> Self {
+		Self { malformed_headers: true, ..self }
+	}
+
+	pub fn at_height(&self, height: usize) -> ValidatedBlockHeader {
+		let block_header = self.at_height_unvalidated(height);
+		let block_hash = self.blocks[height].block_hash();
+		block_header.validate(block_hash).unwrap()
+	}
+
+	fn at_height_unvalidated(&self, height: usize) -> BlockHeaderData {
+		assert!(!self.blocks.is_empty());
+		assert!(height < self.blocks.len());
+		BlockHeaderData {
+			chainwork: self.blocks[0].header.work() + Uint256::from_u64(height as u64).unwrap(),
+			height: height as u32,
+			header: self.blocks[height].header.clone(),
+		}
+	}
+
+	pub fn tip(&self) -> ValidatedBlockHeader {
+		assert!(!self.blocks.is_empty());
+		self.at_height(self.blocks.len() - 1)
+	}
+
+	pub fn disconnect_tip(&mut self) -> Option<Block> {
+		self.blocks.pop()
+	}
+}
+
+impl BlockSource for Blockchain {
+	fn get_header<'a>(&'a mut self, header_hash: &'a BlockHash, _height_hint: Option<u32>) -> AsyncBlockSourceResult<'a, BlockHeaderData> {
+		Box::pin(async move {
+			if self.without_headers {
+				return Err(BlockSourceError::persistent("header not found"));
+			}
+
+			for (height, block) in self.blocks.iter().enumerate() {
+				if block.header.block_hash() == *header_hash {
+					let mut header_data = self.at_height_unvalidated(height);
+					if self.malformed_headers {
+						header_data.header.time += 1;
+					}
+
+					return Ok(header_data);
+				}
+			}
+			Err(BlockSourceError::transient("header not found"))
+		})
+	}
+
+	fn get_block<'a>(&'a mut self, header_hash: &'a BlockHash) -> AsyncBlockSourceResult<'a, Block> {
+		Box::pin(async move {
+			for block in self.blocks.iter() {
+				if block.header.block_hash() == *header_hash {
+					return Ok(block.clone());
+				}
+			}
+			Err(BlockSourceError::transient("block not found"))
+		})
+	}
+
+	fn get_best_block<'a>(&'a mut self) -> AsyncBlockSourceResult<'a, (BlockHash, Option<u32>)> {
+		Box::pin(async move {
+			match self.blocks.last() {
+				None => Err(BlockSourceError::transient("empty chain")),
+				Some(block) => {
+					let height = (self.blocks.len() - 1) as u32;
+					Ok((block.block_hash(), Some(height)))
+				},
+			}
+		})
+	}
+}
