0xsiddharthks/siddharth/feat/whitelist only

src/builder.rs

@@ -48,6 +48,14 @@ impl Builder {
         self.network
     }
 
+    /// Only connect to peers in the whitelist. When enabled, the node will not discover
+    /// new peers via DNS seeding or addr gossip. If all whitelist peers disconnect,
+    /// the node will exit with [`NodeError::NoReachablePeers`](crate::error::NodeError::NoReachablePeers).
+    pub fn whitelist_only(mut self) -> Self {
+        self.config.whitelist_only = true;
+        self
+    }
+
     /// Add preferred peers to try to connect to.
     pub fn add_peers(mut self, whitelist: impl IntoIterator<Item = TrustedPeer>) -> Self {
         self.config.white_list.extend(whitelist);

src/lib.rs

@@ -342,6 +342,7 @@ enum NodeState {
 struct Config {
     required_peers: u8,
     white_list: Vec<TrustedPeer>,
+    whitelist_only: bool,
     data_path: Option<PathBuf>,
     chain_state: Option<ChainState>,
     connection_type: ConnectionType,
@@ -355,6 +356,7 @@ impl Default for Config {
         Self {
             required_peers: 1,
             white_list: Default::default(),
+            whitelist_only: Default::default(),
             data_path: Default::default(),
             chain_state: Default::default(),
             connection_type: Default::default(),

src/network/peer_map.rs

@@ -48,6 +48,7 @@ pub(crate) struct ManagedPeer {
 #[derive(Debug)]
 pub(crate) struct PeerMap {
     pub(crate) tx_queue: Arc<Mutex<BroadcastQueue>>,
+    pub(crate) whitelist_only: bool,
     current_id: PeerId,
     heights: Arc<Mutex<HeightMonitor>>,
     network: Network,
@@ -68,13 +69,15 @@ impl PeerMap {
         network: Network,
         block_type: BlockType,
         whitelist: Whitelist,
+        whitelist_only: bool,
         dialog: Arc<Dialog>,
         connection_type: ConnectionType,
         timeout_config: PeerTimeoutConfig,
         height_monitor: Arc<Mutex<HeightMonitor>>,
     ) -> Self {
         Self {
             tx_queue: Arc::new(Mutex::new(BroadcastQueue::new())),
+            whitelist_only,
             current_id: PeerId(0),
             heights: height_monitor,
             network,
@@ -231,6 +234,7 @@ impl PeerMap {
 
     // Pull a peer from the configuration if we have one. If not, select a random peer from the database,
     // as long as it is not from the same netgroup. If there are no peers in the database, try DNS.
+    // When `whitelist_only` is set, only whitelist peers are used.
     pub async fn next_peer(&mut self) -> Option<Record> {
         if let Some(peer) = self.whitelist.pop() {
             crate::debug!("Using a configured peer");
@@ -240,6 +244,9 @@ impl PeerMap {
             let record = Record::new(peer.address(), port, peer.known_services, &LOCAL_HOST);
             return Some(record);
         }
+        if self.whitelist_only {
+            return None;
+        }
         let mut db_lock = self.db.lock().await;
         if db_lock.is_empty() {
             crate::debug!("Bootstrapping peers with DNS");

src/node.rs

@@ -70,6 +70,7 @@ impl Node {
         let Config {
             required_peers,
             white_list,
+            whitelist_only,
             data_path: _,
             chain_state,
             connection_type,
@@ -95,6 +96,7 @@ impl Node {
             network,
             block_type,
             white_list,
+            whitelist_only,
             Arc::clone(&dialog),
             connection_type,
             peer_timeout_config,
@@ -419,11 +421,13 @@ impl Node {
         self.peer_map
             .send_message(nonce, MainThreadMessage::Verack)
             .await;
-        // Now we may request peers if required
-        crate::debug!("Requesting new addresses");
-        self.peer_map
-            .send_message(nonce, MainThreadMessage::GetAddr)
-            .await;
+        // Request peer addresses unless restricted to the whitelist only.
+        if !self.peer_map.whitelist_only {
+            crate::debug!("Requesting new addresses");
+            self.peer_map
+                .send_message(nonce, MainThreadMessage::GetAddr)
+                .await;
+        }
         // Inform the user we are connected to all required peers
         if self.peer_map.live().eq(&self.required_peers) {
             self.dialog.send_info(Info::ConnectionsMet).await;

tests/core.rs

@@ -633,6 +633,47 @@ async fn tx_can_broadcast() {
         .unwrap();
 }
 
+#[tokio::test]
+async fn whitelist_only_sync() {
+    let (bitcoind, socket_addr) = start_bitcoind(true).unwrap();
+    let rpc = &bitcoind.client;
+    let tempdir = tempfile::TempDir::new().unwrap().path().to_owned();
+    let miner = rpc.new_address().unwrap();
+    mine_blocks(rpc, &miner, 10, 2).await;
+    let best = best_hash(rpc);
+    let host = (IpAddr::V4(*socket_addr.ip()), Some(socket_addr.port()));
+    let builder = bip157::builder::Builder::new(bitcoin::Network::Regtest)
+        .chain_state(ChainState::Checkpoint(HeaderCheckpoint::from_genesis(
+            bitcoin::Network::Regtest,
+        )))
+        .add_peer(host)
+        .whitelist_only()
+        .data_dir(&tempdir);
+    let (node, client) = builder.build();
+    tokio::task::spawn(async move { node.run().await });
+    let Client {
+        requester,
+        info_rx,
+        warn_rx,
+        event_rx: mut channel,
+    } = client;
+    tokio::task::spawn(async move { print_logs(info_rx, warn_rx).await });
+    sync_assert(&best, &mut channel).await;
+    let cp = requester.chain_tip().await.unwrap();
+    assert_eq!(cp.hash, best);
+    requester.shutdown().unwrap();
+    rpc.stop().unwrap();
+    let builder = bip157::builder::Builder::new(bitcoin::Network::Regtest)
+        .chain_state(ChainState::Checkpoint(HeaderCheckpoint::from_genesis(
+            bitcoin::Network::Regtest,
+        )))
+        .whitelist_only()
+        .data_dir(&tempdir);
+    let (node, _client) = builder.build();
+    let result = node.run().await;
+    assert!(result.is_err());
+}
+
 #[tokio::test]
 async fn dns_works() {
     let hostname = bip157::lookup_host("seed.bitcoin.sipa.be").await;