Selaa lähdekoodia

net/acceptor: Decouple transport handshakes from socket accepts

x 2 viikkoa sitten
vanhempi
sitoutus
0f534d27c4

+ 6 - 1
script/evgrd/bin/evgrd.rs

@@ -151,7 +151,12 @@ async fn rpc_serve(
     ex: Arc<Executor<'_>>,
 ) -> Result<()> {
     loop {
-        match listener.next().await {
+        let connection = match listener.next().await {
+            Ok(negotiation) => negotiation.await,
+            Err(err) => Err(err),
+        };
+
+        match connection {
             Ok((stream, url)) => {
                 info!(target: "evgrd", "Accepted connection from {url}");
                 let daemon = daemon.clone();

+ 94 - 14
src/net/acceptor.rs

@@ -22,8 +22,15 @@ use std::{
         atomic::{AtomicUsize, Ordering::SeqCst},
         Arc,
     },
+    time::Duration,
 };
 
+use futures::{
+    future::{select, Either},
+    pin_mut,
+    stream::{FuturesUnordered, StreamExt},
+};
+use smol::Timer;
 use url::Url;
 
 #[cfg(feature = "upnp-igd")]
@@ -33,7 +40,7 @@ use super::{
     channel::{Channel, ChannelPtr},
     hosts::HostColor,
     session::SessionWeakPtr,
-    transport::{Listener, PtListener},
+    transport::{Listener, PtListener, PtNegotiation},
 };
 
 #[cfg(feature = "upnp-igd")]
@@ -51,6 +58,18 @@ use crate::{
 /// Atomic pointer to Acceptor
 pub type AcceptorPtr = Arc<Acceptor>;
 
+fn with_handshake_timeout(negotiation: PtNegotiation, timeout: Duration) -> PtNegotiation {
+    Box::pin(async move {
+        let timer = Timer::after(timeout);
+        pin_mut!(timer);
+
+        match select(negotiation, timer).await {
+            Either::Left((result, _)) => result,
+            Either::Right((_, _)) => Err(ErrorKind::TimedOut.into()),
+        }
+    })
+}
+
 /// Releases an inbound connection slot when its tracking task exits.
 struct InboundSlotGuard {
     acceptor: AcceptorPtr,
@@ -96,8 +115,12 @@ impl Acceptor {
 
     /// Start accepting inbound socket connections
     pub async fn start(self: Arc<Self>, endpoint: Url, ex: ExecutorPtr) -> Result<()> {
-        let datastore =
-            self.session.upgrade().unwrap().p2p().settings().read().await.p2p_datastore.clone();
+        let settings = self.session.upgrade().unwrap().p2p().settings();
+        let settings = settings.read().await;
+        let datastore = settings.p2p_datastore.clone();
+        let handshake_timeout =
+            Duration::from_secs(settings.channel_handshake_timeout(endpoint.scheme()));
+        drop(settings);
 
         // Initialize listener
         let listener = Listener::new(endpoint.clone(), datastore, true).await?;
@@ -128,7 +151,7 @@ impl Acceptor {
             self.port_mappings.lock().await.extend(mappings);
         }
 
-        self.accept(ptlistener, ex);
+        self.accept(ptlistener, handshake_timeout, ex);
         Ok(())
     }
 
@@ -158,10 +181,15 @@ impl Acceptor {
     }
 
     /// Run the accept loop in a new thread and error if a connection problem occurs
-    fn accept(self: Arc<Self>, listener: Box<dyn PtListener>, ex: ExecutorPtr) {
+    fn accept(
+        self: Arc<Self>,
+        listener: Box<dyn PtListener>,
+        handshake_timeout: Duration,
+        ex: ExecutorPtr,
+    ) {
         let self_ = self.clone();
         self.task.clone().start(
-            self.run_accept_loop(listener, ex.clone()),
+            self.run_accept_loop(listener, handshake_timeout, ex.clone()),
             |result| self_.handle_stop(result),
             Error::NetworkServiceStopped,
             ex,
@@ -172,31 +200,75 @@ impl Acceptor {
     async fn run_accept_loop(
         self: Arc<Self>,
         listener: Box<dyn PtListener>,
+        handshake_timeout: Duration,
         ex: ExecutorPtr,
     ) -> Result<()> {
         // CondVar used to notify the loop to recheck if new connections can
         // be accepted by the listener.
         let cv = Arc::new(CondVar::new());
         let hosts = self.session.upgrade().unwrap().p2p().hosts();
+        let mut negotiations = FuturesUnordered::<PtNegotiation>::new();
+        let mut accepting = None;
 
         loop {
-            // Refuse new connections if we're up to the connection limit
+            // Reserve capacity for established channels, transport negotiations,
+            // and the raw accept currently in progress.
             let limit =
                 self.session.upgrade().unwrap().p2p().settings().read().await.inbound_connections;
+            let reserved = self.conn_count.load(SeqCst) +
+                negotiations.len() +
+                usize::from(accepting.is_some());
+
+            if reserved < limit && accepting.is_none() {
+                accepting = Some(listener.next());
+            }
 
-            if self.clone().conn_count.load(SeqCst) >= limit {
+            // Keep the raw accept alive when a transport negotiation finishes
+            // first. This allows the listener to continue accepting while TLS,
+            // Tor, or QUIC setup is in progress without exceeding the inbound
+            // connection limit.
+            let connection = if let Some(accept) = accepting.take() {
+                if negotiations.is_empty() {
+                    match accept.await {
+                        Ok(negotiation) => {
+                            negotiations
+                                .push(with_handshake_timeout(negotiation, handshake_timeout));
+                            continue
+                        }
+                        Err(err) => Err(err),
+                    }
+                } else {
+                    let negotiation = negotiations.next();
+                    pin_mut!(negotiation);
+
+                    match select(accept, negotiation).await {
+                        Either::Left((Ok(negotiation), _)) => {
+                            negotiations
+                                .push(with_handshake_timeout(negotiation, handshake_timeout));
+                            continue
+                        }
+                        Either::Left((Err(err), _)) => Err(err),
+                        Either::Right((Some(result), accept)) => {
+                            accepting = Some(accept);
+                            result
+                        }
+                        Either::Right((None, accept)) => {
+                            accepting = Some(accept);
+                            continue
+                        }
+                    }
+                }
+            } else if let Some(result) = negotiations.next().await {
+                result
+            } else {
                 // This will get notified every time an inbound channel is stopped.
-                // These channels are the channels spawned below on listener.next().is_ok().
-                // After the notification, we reset the condvar and retry this loop to see
-                // if we can accept more connections, and if not - we'll be back here.
                 verbose!(target: "net::acceptor::run_accept_loop", "Reached incoming conn limit, waiting...");
                 cv.wait().await;
                 cv.reset();
                 continue
-            }
+            };
 
-            // Now we wait for a new connection.
-            match listener.next().await {
+            match connection {
                 Ok((stream, url)) => {
                     // Check if we reject this peer
                     if hosts.container.contains(HostColor::Black, &url) ||
@@ -263,6 +335,14 @@ impl Acceptor {
                     }
                 },
 
+                Err(e) if e.kind() == ErrorKind::TimedOut => {
+                    verbose!(
+                        target: "net::acceptor::run_accept_loop",
+                        "[P2P] Transport handshake timed out"
+                    );
+                    continue
+                }
+
                 // In case a TLS handshake fails, we'll get this:
                 Err(e) if e.kind() == ErrorKind::UnexpectedEof => continue,
 

+ 79 - 0
src/net/tests.rs

@@ -39,6 +39,7 @@ use crate::{
         metering::{MeteringConfiguration, DEFAULT_METERING_CONFIGURATION},
         p2p::MAX_CONCURRENT_BROADCASTS,
         settings::NetworkProfile,
+        transport::Dialer,
         P2p, Settings,
     },
     system::{sleep, timeout::timeout},
@@ -665,6 +666,84 @@ async fn p2p_inbound_slots_survive_rapid_disconnects_real(ex: Arc<Executor<'stat
     p2p.stop().await;
 }
 
+#[test]
+fn p2p_tls_listener_accepts_while_handshake_stalled() {
+    test_body!(p2p_tls_listener_accepts_while_handshake_stalled_real, 2);
+}
+
+async fn p2p_tls_listener_accepts_while_handshake_stalled_real(ex: Arc<Executor<'static>>) {
+    let port = get_random_available_port();
+    let addr = format!("127.0.0.1:{port}");
+    let listen_url = Url::parse(&format!("tcp+tls://{addr}")).unwrap();
+    let settings = Settings {
+        localnet: true,
+        inbound_addrs: vec![listen_url.clone()],
+        inbound_connections: 2,
+        outbound_connections: 0,
+        active_profiles: vec!["tcp+tls".to_string()],
+        ..Default::default()
+    };
+
+    let p2p = P2p::new(settings, ex).await.unwrap();
+    p2p.clone().start().await.unwrap();
+
+    // Occupy the first accepted socket without sending a TLS ClientHello.
+    let stalled = TcpStream::connect(&addr).await.unwrap();
+    Timer::after(Duration::from_millis(100)).await;
+
+    // A second client must complete TLS before the stalled handshake times out.
+    let dialer = Dialer::new(listen_url, None, None, true).await.unwrap();
+    let stream = timeout(Duration::from_secs(2), dialer.dial(Some(Duration::from_secs(1))))
+        .await
+        .expect("TLS listener blocked behind a stalled handshake")
+        .expect("second TLS handshake failed");
+
+    drop(stream);
+    drop(stalled);
+    p2p.stop().await;
+}
+
+#[test]
+fn p2p_tls_listener_expires_stalled_handshake() {
+    test_body!(p2p_tls_listener_expires_stalled_handshake_real, 2);
+}
+
+async fn p2p_tls_listener_expires_stalled_handshake_real(ex: Arc<Executor<'static>>) {
+    let port = get_random_available_port();
+    let addr = format!("127.0.0.1:{port}");
+    let listen_url = Url::parse(&format!("tcp+tls://{addr}")).unwrap();
+    let mut profiles = HashMap::new();
+    profiles.insert(
+        "tcp+tls".to_string(),
+        NetworkProfile { channel_handshake_timeout: 1, ..Default::default() },
+    );
+    let settings = Settings {
+        localnet: true,
+        inbound_addrs: vec![listen_url.clone()],
+        inbound_connections: 1,
+        outbound_connections: 0,
+        active_profiles: vec!["tcp+tls".to_string()],
+        profiles,
+        ..Default::default()
+    };
+
+    let p2p = P2p::new(settings, ex).await.unwrap();
+    p2p.clone().start().await.unwrap();
+
+    let stalled = TcpStream::connect(&addr).await.unwrap();
+    Timer::after(Duration::from_millis(1200)).await;
+
+    let dialer = Dialer::new(listen_url, None, None, true).await.unwrap();
+    let stream = timeout(Duration::from_secs(2), dialer.dial(Some(Duration::from_secs(1))))
+        .await
+        .expect("TLS listener did not recover after the handshake deadline")
+        .expect("TLS handshake after deadline failed");
+
+    drop(stream);
+    drop(stalled);
+    p2p.stop().await;
+}
+
 #[test]
 fn p2p_shutdown_drains_channels_across_restarts() {
     test_body!(p2p_shutdown_drains_channels_across_restarts_real, 2);

+ 11 - 3
src/net/transport/mod.rs

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::{io, time::Duration};
+use std::{future::Future, io, pin::Pin, time::Duration};
 
 use async_trait::async_trait;
 use smol::io::{AsyncRead, AsyncWrite};
@@ -538,8 +538,16 @@ impl PtStream for smol::net::unix::UnixStream {}
 #[cfg(feature = "p2p-quic")]
 impl PtStream for quic::QuicStream {}
 
-/// Wrapper trait for async listeners
+/// A transport negotiation produced after accepting an underlying connection.
+pub type PtNegotiation =
+    Pin<Box<dyn Future<Output = io::Result<(Box<dyn PtStream>, Url)>> + Send + 'static>>;
+
+/// Wrapper trait for async listeners.
+///
+/// `next()` accepts the underlying transport connection and returns the
+/// remaining negotiation separately so callers can continue accepting while
+/// TLS, Tor, or QUIC setup is in progress.
 #[async_trait]
 pub trait PtListener: Send + Unpin {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)>;
+    async fn next(&self) -> io::Result<PtNegotiation>;
 }

+ 15 - 13
src/net/transport/quic.rs

@@ -48,7 +48,7 @@ use super::{
     tls::{
         generate_certificate, ClientCertificateVerifier, ServerCertificateVerifier, TLS_DNS_NAME,
     },
-    PtListener, PtStream,
+    PtListener, PtNegotiation, PtStream,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -379,7 +379,7 @@ pub struct QuicListenerIntern {
 
 #[async_trait]
 impl PtListener for QuicListenerIntern {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)> {
+    async fn next(&self) -> io::Result<PtNegotiation> {
         // Wait for an incoming connection
         let incoming =
             self.endpoint.accept().await.ok_or_else(|| {
@@ -388,19 +388,21 @@ impl PtListener for QuicListenerIntern {
 
         let peer_addr = incoming.remote_address();
 
-        let connection =
-            incoming.await.map_err(|e| io::Error::other(format!("QUIC accept error: {e}")))?;
+        Ok(Box::pin(async move {
+            let connection =
+                incoming.await.map_err(|e| io::Error::other(format!("QUIC accept error: {e}")))?;
 
-        // Accept a bidirectional stream from the client
-        let (send, recv) = connection
-            .accept_bi()
-            .await
-            .map_err(|e| io::Error::other(format!("QUIC stream accept error: {e}")))?;
+            // Accept a bidirectional stream from the client
+            let (send, recv) = connection
+                .accept_bi()
+                .await
+                .map_err(|e| io::Error::other(format!("QUIC stream accept error: {e}")))?;
 
-        let url = Url::parse(&format!("quic://{peer_addr}")).map_err(|e| {
-            io::Error::new(io::ErrorKind::InvalidData, format!("Invalid peer address: {e}"))
-        })?;
+            let url = Url::parse(&format!("quic://{peer_addr}")).map_err(|e| {
+                io::Error::new(io::ErrorKind::InvalidData, format!("Invalid peer address: {e}"))
+            })?;
 
-        Ok((Box::new(QuicStream::new(send, recv)), url))
+            Ok((Box::new(QuicStream::new(send, recv)) as Box<dyn PtStream>, url))
+        }))
     }
 }

+ 10 - 10
src/net/transport/tcp.rs

@@ -33,7 +33,7 @@ use socket2::{Domain, Socket, TcpKeepalive, Type};
 use tracing::debug;
 use url::Url;
 
-use super::{PtListener, PtStream};
+use super::{PtListener, PtNegotiation, PtStream};
 
 trait SocketExt {
     fn enable_reuse_port(&self) -> io::Result<()>;
@@ -192,7 +192,7 @@ impl TcpListener {
 
 #[async_trait]
 impl PtListener for SmolTcpListener {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)> {
+    async fn next(&self) -> io::Result<PtNegotiation> {
         let (stream, peer_addr) = match self.accept().await {
             Ok((s, a)) => (s, a),
             Err(e) => return Err(e),
@@ -207,23 +207,18 @@ impl PtListener for SmolTcpListener {
                 ))
             }
         };
-        Ok((Box::new(stream), url))
+        Ok(Box::pin(async move { Ok((Box::new(stream) as Box<dyn PtStream>, url)) }))
     }
 }
 
 #[async_trait]
 impl PtListener for (TlsAcceptor, SmolTcpListener) {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)> {
+    async fn next(&self) -> io::Result<PtNegotiation> {
         let (stream, peer_addr) = match self.1.accept().await {
             Ok((s, a)) => (s, a),
             Err(e) => return Err(e),
         };
 
-        let stream = match self.0.accept(stream).await {
-            Ok(v) => v,
-            Err(e) => return Err(e),
-        };
-
         let url = match Url::parse(&format!("tcp+tls://{peer_addr}")) {
             Ok(v) => v,
             Err(e) => {
@@ -234,6 +229,11 @@ impl PtListener for (TlsAcceptor, SmolTcpListener) {
             }
         };
 
-        Ok((Box::new(TlsStream::Server(stream)), url))
+        let acceptor = self.0.clone();
+        Ok(Box::pin(async move {
+            let stream = acceptor.accept(stream).await?;
+
+            Ok((Box::new(TlsStream::Server(stream)) as Box<dyn PtStream>, url))
+        }))
     }
 }

+ 45 - 34
src/net/transport/tor.rs

@@ -48,7 +48,7 @@ use tor_rtcompat::PreferredRuntime;
 use tracing::debug;
 use url::Url;
 
-use super::{PtListener, PtStream};
+use super::{PtListener, PtNegotiation, PtStream};
 use crate::util::{encoding::base32, logger::verbose, path::expand_path};
 
 /// A static for `TorClient` reusability
@@ -274,7 +274,7 @@ unsafe impl Sync for TorListenerIntern {}
 
 #[async_trait]
 impl PtListener for TorListenerIntern {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)> {
+    async fn next(&self) -> io::Result<PtNegotiation> {
         let mut rendreq_stream = self.rendreq_stream.lock().await;
 
         let Some(rendrequest) = rendreq_stream.next().await else {
@@ -283,42 +283,53 @@ impl PtListener for TorListenerIntern {
 
         drop(rendreq_stream);
 
-        let mut streamreq_stream = match rendrequest.accept().await {
-            Ok(v) => v,
-            Err(e) => {
-                verbose!(
-                    target: "net::tor::PtListener::next",
-                    "[P2P] Failed accepting Tor RendRequest: {e}"
-                );
-                return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted"));
-            }
-        };
-
-        let Some(streamrequest) = streamreq_stream.next().await else {
-            return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted"));
-        };
-
-        // Validate port correctness
-        match streamrequest.request() {
-            IncomingStreamRequest::Begin(begin) => {
-                if begin.port() != self.port {
+        let port = self.port;
+        Ok(Box::pin(async move {
+            let mut streamreq_stream = match rendrequest.accept().await {
+                Ok(v) => v,
+                Err(e) => {
+                    verbose!(
+                        target: "net::tor::PtListener::next",
+                        "[P2P] Failed accepting Tor RendRequest: {e}"
+                    );
                     return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted"));
                 }
-            }
-            &_ => return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted")),
-        }
+            };
 
-        let stream = match streamrequest.accept(Connected::new_empty()).await {
-            Ok(v) => v,
-            Err(e) => {
-                verbose!(
-                    target: "net::tor::PtListener::next",
-                    "[P2P] Failed accepting Tor StreamRequest: {e}"
-                );
-                return Err(io::Error::other("Internal Tor error"));
+            let Some(streamrequest) = streamreq_stream.next().await else {
+                return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted"));
+            };
+
+            // Validate port correctness
+            match streamrequest.request() {
+                IncomingStreamRequest::Begin(begin) => {
+                    if begin.port() != port {
+                        return Err(io::Error::new(
+                            ErrorKind::ConnectionAborted,
+                            "Connection Aborted",
+                        ));
+                    }
+                }
+                &_ => {
+                    return Err(io::Error::new(ErrorKind::ConnectionAborted, "Connection Aborted"))
+                }
             }
-        };
 
-        Ok((Box::new(stream), Url::parse(&format!("tor://127.0.0.1:{}", self.port)).unwrap()))
+            let stream = match streamrequest.accept(Connected::new_empty()).await {
+                Ok(v) => v,
+                Err(e) => {
+                    verbose!(
+                        target: "net::tor::PtListener::next",
+                        "[P2P] Failed accepting Tor StreamRequest: {e}"
+                    );
+                    return Err(io::Error::other("Internal Tor error"));
+                }
+            };
+
+            Ok((
+                Box::new(stream) as Box<dyn PtStream>,
+                Url::parse(&format!("tor://127.0.0.1:{port}")).unwrap(),
+            ))
+        }))
     }
 }

+ 3 - 3
src/net/transport/unix.rs

@@ -29,7 +29,7 @@ use smol::{
 use tracing::debug;
 use url::Url;
 
-use super::{PtListener, PtStream};
+use super::{PtListener, PtNegotiation, PtStream};
 
 /// Unix Dialer implementation
 #[derive(Debug, Clone)]
@@ -73,7 +73,7 @@ impl UnixListener {
 
 #[async_trait]
 impl PtListener for SmolUnixListener {
-    async fn next(&self) -> io::Result<(Box<dyn PtStream>, Url)> {
+    async fn next(&self) -> io::Result<PtNegotiation> {
         let (stream, _peer_addr) = match self.accept().await {
             Ok((s, a)) => (s, a),
             Err(e) => return Err(e),
@@ -83,6 +83,6 @@ impl PtListener for SmolUnixListener {
         let addr = addr.as_pathname().unwrap().to_str().unwrap();
         let url = Url::parse(&format!("unix://{addr}")).unwrap();
 
-        Ok((Box::new(stream), url))
+        Ok(Box::pin(async move { Ok((Box::new(stream) as Box<dyn PtStream>, url)) }))
     }
 }

+ 6 - 1
src/rpc/server.rs

@@ -422,7 +422,12 @@ async fn run_accept_loop<'a, T: 'a>(
     ex: Arc<smol::Executor<'a>>,
 ) -> Result<()> {
     loop {
-        match listener.next().await {
+        let connection = match listener.next().await {
+            Ok(negotiation) => negotiation.await,
+            Err(err) => Err(err),
+        };
+
+        match connection {
             Ok((stream, url)) => {
                 let rh_ = rh.clone();
                 verbose!(target: "rpc::server", "[RPC] Server accepted conn from {url}");

+ 4 - 4
tests/network_transports.rs

@@ -59,7 +59,7 @@ fn tcp_transport() {
             Listener::new(url.clone(), None, true).await.unwrap().listen().await.unwrap();
         executor
             .spawn(async move {
-                let (stream, _) = listener.next().await.unwrap();
+                let (stream, _) = listener.next().await.unwrap().await.unwrap();
                 let (mut reader, mut writer) = smol::io::split(stream);
                 io::copy(&mut reader, &mut writer).await.unwrap();
             })
@@ -121,7 +121,7 @@ fn tcp_tls_transport() {
             Listener::new(url.clone(), None, true).await.unwrap().listen().await.unwrap();
         executor
             .spawn(async move {
-                let (stream, _) = listener.next().await.unwrap();
+                let (stream, _) = listener.next().await.unwrap().await.unwrap();
                 let (mut reader, mut writer) = smol::io::split(stream);
                 io::copy(&mut reader, &mut writer).await.unwrap();
             })
@@ -154,7 +154,7 @@ fn quic_transport() {
 
         executor
             .spawn(async move {
-                let (stream, _) = listener.next().await.unwrap();
+                let (stream, _) = listener.next().await.unwrap().await.unwrap();
                 let (mut reader, mut writer) = smol::io::split(stream);
                 io::copy(&mut reader, &mut writer).await.unwrap();
             })
@@ -188,7 +188,7 @@ fn unix_transport() {
             Listener::new(url.clone(), None, true).await.unwrap().listen().await.unwrap();
         executor
             .spawn(async move {
-                let (stream, _) = listener.next().await.unwrap();
+                let (stream, _) = listener.next().await.unwrap().await.unwrap();
                 let (mut reader, mut writer) = smol::io::split(stream);
                 io::copy(&mut reader, &mut writer).await.unwrap();
             })