From fd7225e6738869bcc9c1ef50e13977c044dbcd2e Mon Sep 17 00:00:00 2001
From: Boog900 <54e72d8a-345f-4599-bd90-c6b9bc7d0ec5@aleeas.com>
Date: Thu, 30 May 2024 02:03:48 +0100
Subject: [PATCH] check if incoming peers are banned

---
 p2p/address-book/src/book.rs          |  3 +++
 p2p/cuprate-p2p/src/inbound_server.rs | 25 ++++++++++++++++++++-----
 p2p/cuprate-p2p/src/lib.rs            | 11 ++++++++---
 p2p/monero-p2p/src/services.rs        |  3 +++
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/p2p/address-book/src/book.rs b/p2p/address-book/src/book.rs
index 3a49c6be..2f0617e9 100644
--- a/p2p/address-book/src/book.rs
+++ b/p2p/address-book/src/book.rs
@@ -409,6 +409,9 @@ impl<Z: NetworkZone> Service<AddressBookRequest<Z>> for AddressBook<Z> {
             AddressBookRequest::GetWhitePeers(len) => {
                 Ok(AddressBookResponse::Peers(self.get_white_peers(len)))
             }
+            AddressBookRequest::IsPeerBanned(addr) => Ok(AddressBookResponse::IsPeerBanned(
+                self.is_peer_banned(&addr),
+            )),
         };
 
         ready(response)
diff --git a/p2p/cuprate-p2p/src/inbound_server.rs b/p2p/cuprate-p2p/src/inbound_server.rs
index b56903a9..df15515e 100644
--- a/p2p/cuprate-p2p/src/inbound_server.rs
+++ b/p2p/cuprate-p2p/src/inbound_server.rs
@@ -12,9 +12,10 @@ use tokio::{
 use tower::{Service, ServiceExt};
 use tracing::{instrument, Instrument, Span};
 
+use monero_p2p::services::{AddressBookRequest, AddressBookResponse};
 use monero_p2p::{
     client::{Client, DoHandshakeRequest, HandshakeError, InternalPeerID},
-    ConnectionDirection, NetworkZone,
+    AddressBook, ConnectionDirection, NetworkZone,
 };
 
 use crate::{
@@ -25,9 +26,10 @@ use crate::{
 
 /// The inbound server.
 #[instrument(level = "warn", skip_all)]
-pub async fn inbound_server<N, HS>(
+pub async fn inbound_server<N, HS, A>(
     client_pool: Arc<ClientPool<N>>,
     mut handshaker: HS,
+    mut address_book: A,
     config: P2PConfig<N>,
 ) -> Result<(), tower::BoxError>
 where
@@ -36,15 +38,13 @@ where
         + Send
         + 'static,
     HS::Future: Send + 'static,
+    A: AddressBook<N>,
 {
     let Some(server_config) = config.server_config else {
         tracing::warn!("No inbound server config provided, not listening for inbound connections.");
         return Ok(());
     };
 
-    // TODO: take in the address book and check if incoming peers are banned before adding them to our
-    // connections.
-
     tracing::info!("Starting inbound connection server");
 
     let listener = N::incoming_connection_listener(server_config, config.p2p_port)
@@ -60,6 +60,21 @@ where
             continue;
         };
 
+        if let Some(addr) = &addr {
+            let AddressBookResponse::IsPeerBanned(banned) = address_book
+                .ready()
+                .await?
+                .call(AddressBookRequest::IsPeerBanned(*addr))
+                .await?
+            else {
+                panic!("Address book returned incorrect response!");
+            };
+
+            if banned {
+                continue;
+            }
+        }
+
         let addr = match addr {
             Some(addr) => InternalPeerID::KnownAddr(addr),
             None => InternalPeerID::Unknown(rand::random()),
diff --git a/p2p/cuprate-p2p/src/lib.rs b/p2p/cuprate-p2p/src/lib.rs
index 20a35819..500eeb09 100644
--- a/p2p/cuprate-p2p/src/lib.rs
+++ b/p2p/cuprate-p2p/src/lib.rs
@@ -56,7 +56,7 @@ where
         config.max_inbound_connections + config.outbound_connections,
     );
 
-    // Use the default config. Changing the defaults affects tx fluff times, which could effect D++ so for now don't allow changing
+    // Use the default config. Changing the defaults affects tx fluff times, which could affect D++ so for now don't allow changing
     // this.
     let (broadcast_svc, outbound_mkr, inbound_mkr) =
         broadcast::init_broadcast_channels(broadcast::BroadcastConfig::default());
@@ -105,8 +105,13 @@ where
             .instrument(Span::current()),
     );
     tokio::spawn(
-        inbound_server::inbound_server(client_pool.clone(), inbound_handshaker, config)
-            .instrument(Span::current()),
+        inbound_server::inbound_server(
+            client_pool.clone(),
+            inbound_handshaker,
+            address_book,
+            config,
+        )
+        .instrument(Span::current()),
     );
 
     Ok(NetworkInterface {
diff --git a/p2p/monero-p2p/src/services.rs b/p2p/monero-p2p/src/services.rs
index e86e2776..bf0d6ec3 100644
--- a/p2p/monero-p2p/src/services.rs
+++ b/p2p/monero-p2p/src/services.rs
@@ -119,10 +119,13 @@ pub enum AddressBookRequest<Z: NetworkZone> {
     TakeRandomPeer { height: Option<u64> },
     /// Gets the specified number of white peers, or less if we don't have enough.
     GetWhitePeers(usize),
+    /// Checks if the given peer is banned.
+    IsPeerBanned(Z::Addr),
 }
 
 pub enum AddressBookResponse<Z: NetworkZone> {
     Ok,
     Peer(ZoneSpecificPeerListEntryBase<Z::Addr>),
     Peers(Vec<ZoneSpecificPeerListEntryBase<Z::Addr>>),
+    IsPeerBanned(bool),
 }