Merge pull request #549 from rustaceanrob/26-3-10-witness-block

Optionally request witness data

Author: rustaceanrob

examples/bitcoin.rs

@@ -34,6 +34,8 @@ async fn main() {
         .add_peers(seeds.into_iter().map(From::from))
         // Connections over Tor are supported by Socks5 proxy
         // .socks5_proxy(bip157::Socks5Proxy::local())
+        // Optionally request witness data when fetching a block
+        // .fetch_witness_data()
         // Create the node and client
         .build();
     // Run the node on a separate task

src/builder.rs

@@ -5,7 +5,7 @@ use bitcoin::Network;
 use super::{client::Client, node::Node};
 use crate::chain::ChainState;
 use crate::network::ConnectionType;
-use crate::{Config, FilterType};
+use crate::{BlockType, Config, FilterType};
 use crate::{Socks5Proxy, TrustedPeer};
 
 const MIN_PEERS: u8 = 1;
@@ -137,6 +137,12 @@ impl Builder {
         self
     }
 
+    /// Request witness data when requesting blocks.
+    pub fn fetch_witness_data(mut self) -> Self {
+        self.config.block_type = BlockType::Witness;
+        self
+    }
+
     /// Consume the node builder and receive a [`Node`] and [`Client`].
     pub fn build(mut self) -> (Node, Client) {
         Node::new(self.network, core::mem::take(&mut self.config))

src/lib.rs

@@ -119,6 +119,13 @@ pub enum FilterType {
     Basic,
 }
 
+#[derive(Debug, Clone, Copy, Default)]
+enum BlockType {
+    #[default]
+    Legacy,
+    Witness,
+}
+
 impl From<FilterType> for u8 {
     fn from(value: FilterType) -> Self {
         match value {
@@ -335,6 +342,7 @@ struct Config {
     connection_type: ConnectionType,
     peer_timeout_config: PeerTimeoutConfig,
     filter_type: FilterType,
+    block_type: BlockType,
 }
 
 impl Default for Config {
@@ -347,6 +355,7 @@ impl Default for Config {
             connection_type: Default::default(),
             peer_timeout_config: PeerTimeoutConfig::default(),
             filter_type: FilterType::default(),
+            block_type: BlockType::default(),
         }
     }
 }

src/network/outbound.rs

@@ -15,14 +15,15 @@ use bitcoin::{
     BlockHash, Network, Transaction, Wtxid,
 };
 
-use crate::default_port_from_network;
+use crate::{default_port_from_network, BlockType};
 
 use super::{KYOTO_VERSION, PROTOCOL_VERSION, RUST_BITCOIN_VERSION};
 
 // Responsible for serializing messages to write over the wire, either encrypted or plaintext.
 pub(in crate::network) struct MessageGenerator {
     pub network: Network,
     pub transport: Transport,
+    pub block_type: BlockType,
 }
 
 pub(in crate::network) enum Transport {
@@ -50,7 +51,10 @@ impl MessageGenerator {
     }
 
     pub(in crate::network) fn block(&mut self, hash: BlockHash) -> Vec<u8> {
-        let inv = Inventory::Block(hash);
+        let inv = match self.block_type {
+            BlockType::Legacy => Inventory::Block(hash),
+            BlockType::Witness => Inventory::WitnessBlock(hash),
+        };
         let msg = NetworkMessage::GetData(vec![inv]);
         self.serialize(msg)
     }

src/network/peer.rs

@@ -18,7 +18,7 @@ use tokio::{
     time::{Instant, MissedTickBehavior},
 };
 
-use crate::{broadcaster::BroadcastQueue, messages::Warning, Dialog, Info};
+use crate::{broadcaster::BroadcastQueue, messages::Warning, BlockType, Dialog, Info};
 
 use super::{
     error::PeerError,
@@ -38,6 +38,7 @@ pub(crate) struct Peer {
     main_thread_sender: Sender<PeerThreadMessage>,
     main_thread_recv: Receiver<MainThreadMessage>,
     network: Network,
+    block_type: BlockType,
     dialog: Arc<Dialog>,
     db: Arc<Mutex<AddressBook>>,
     timeout_config: PeerTimeoutConfig,
@@ -51,6 +52,7 @@ impl Peer {
         nonce: PeerId,
         source: Record,
         network: Network,
+        block_type: BlockType,
         main_thread_sender: Sender<PeerThreadMessage>,
         main_thread_recv: Receiver<MainThreadMessage>,
         dialog: Arc<Dialog>,
@@ -64,6 +66,7 @@ impl Peer {
             main_thread_sender,
             main_thread_recv,
             network,
+            block_type,
             dialog,
             db,
             timeout_config,
@@ -97,13 +100,15 @@ impl Peer {
                 let outbound_messages = MessageGenerator {
                     network: self.network,
                     transport: Transport::V2 { encryptor },
+                    block_type: self.block_type,
                 };
                 let reader = Reader::new(MessageParser::V2(reader, decryptor), tx);
                 (outbound_messages, reader)
             } else {
                 let outbound_messages = MessageGenerator {
                     network: self.network,
                     transport: Transport::V1,
+                    block_type: self.block_type,
                 };
                 let reader = Reader::new(MessageParser::V1(reader, self.network), tx);
                 (outbound_messages, reader)

src/network/peer_map.rs

@@ -25,7 +25,7 @@ use crate::{
     chain::HeightMonitor,
     default_port_from_network,
     network::{dns::bootstrap_dns, error::PeerError, peer::Peer, PeerId, PeerTimeoutConfig},
-    Dialog, TrustedPeer,
+    BlockType, Dialog, TrustedPeer,
 };
 
 use super::{AddressBook, ConnectionType, MainThreadMessage, PeerThreadMessage};
@@ -51,6 +51,7 @@ pub(crate) struct PeerMap {
     current_id: PeerId,
     heights: Arc<Mutex<HeightMonitor>>,
     network: Network,
+    block_type: BlockType,
     mtx: Sender<PeerThreadMessage>,
     map: HashMap<PeerId, ManagedPeer>,
     db: Arc<Mutex<AddressBook>>,
@@ -65,6 +66,7 @@ impl PeerMap {
     pub fn new(
         mtx: Sender<PeerThreadMessage>,
         network: Network,
+        block_type: BlockType,
         whitelist: Whitelist,
         dialog: Arc<Dialog>,
         connection_type: ConnectionType,
@@ -76,6 +78,7 @@ impl PeerMap {
             current_id: PeerId(0),
             heights: height_monitor,
             network,
+            block_type,
             mtx,
             map: HashMap::new(),
             db: Arc::new(Mutex::new(AddressBook::new())),
@@ -122,6 +125,7 @@ impl PeerMap {
             self.current_id,
             loaded_peer.clone(),
             self.network,
+            self.block_type,
             self.mtx.clone(),
             prx,
             Arc::clone(&self.dialog),

src/node.rs

@@ -75,6 +75,7 @@ impl Node {
             connection_type,
             peer_timeout_config,
             filter_type,
+            block_type,
         } = config;
         // Set up a communication channel between the node and client
         let (info_tx, info_rx) = mpsc::channel::<Info>(32);
@@ -92,6 +93,7 @@ impl Node {
         let peer_map = PeerMap::new(
             mtx,
             network,
+            block_type,
             white_list,
             Arc::clone(&dialog),
             connection_type,