diff --git a/src/asynchronous/connection.rs b/src/asynchronous/connection.rs index 3eb742e4..b2fb49ab 100644 --- a/src/asynchronous/connection.rs +++ b/src/asynchronous/connection.rs @@ -157,7 +157,7 @@ where } } - pub async fn run(self) -> std::io::Result<()> { + pub async fn run(self) -> Result<()> { let Connection { mut reader, mut writer_task, @@ -165,28 +165,29 @@ where } = self; let shutdown = reader_delegate.wait_shutdown(); tokio::pin!(shutdown); - loop { + let result = loop { select! { // Writer failures take priority, then shutdown, then incoming frames. biased; writer_result = &mut writer_task => { - match writer_result { - Ok(Ok(())) => {} + let e = match writer_result { + Ok(Ok(())) => break Ok(()), Ok(Err(e)) => { trace!("Write msg err: {:?}", e); - reader_delegate.disconnect(e).await; + e } Err(e) => { let e = Error::Others(format!("Writer task failed: {e}")); error!("Write task err: {:?}", e); - reader_delegate.disconnect(e).await; + e } - } - break; + }; + reader_delegate.disconnect(e.clone()).await; + break Err(e); } _v = &mut shutdown => { trace!("Receive shutdown."); - break; + break Ok(()); } res = GenMessage::read_from(&mut reader) => { match res { @@ -203,16 +204,16 @@ where trace!("Read msg err: {:?}", e); writer_task.abort(); let _ = (&mut writer_task).await; - reader_delegate.disconnect(e).await; - break; + reader_delegate.disconnect(e.clone()).await; + break Err(e); } } } } - } + }; reader_delegate.exit().await; trace!("Reader task exit."); - Ok(()) + result } } diff --git a/src/asynchronous/server.rs b/src/asynchronous/server.rs index 74c559fd..8d2ccf7f 100644 --- a/src/asynchronous/server.rs +++ b/src/asynchronous/server.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use std::convert::TryFrom; +use std::future::Future; #[cfg(unix)] use std::os::unix::io::RawFd; use std::result::Result as StdResult; @@ -284,7 +285,7 @@ impl Server { services, shutdown_waiter, conn_ctx, - ).await; + ); }); } Err(e) => { @@ -308,6 +309,45 @@ impl Server { Ok(()) } + /// Serves a single already-connected socket with the registered services. + /// + /// The returned future does not borrow the server, so it can be awaited directly or passed to + /// `tokio::spawn`. It completes when the connection terminates or the server is shut down. + /// + /// # Errors + /// + /// Returns an error if the accept hook rejects the connection, or if reading from or writing + /// to the connection fails. + pub fn serve_connection( + &self, + conn: Socket, + ) -> impl Future> + Send + 'static { + let services = self.services.clone(); + let shutdown_waiter = self.shutdown.subscribe(); + #[cfg(feature = "security_extension")] + let server_ext = ServerExtensionConfig { + accept_hook: self.accept_hook.clone(), + }; + + async move { + #[cfg(feature = "security_extension")] + let conn_ctx = match server_ext.on_accept(&conn).await { + Ok(output) => Arc::new(ConnectionContext::new(output)), + Err(e) => return Err(Error::Others(format!("accept hook failed: {e}"))), + }; + #[cfg(not(feature = "security_extension"))] + let conn_ctx = Arc::new(ConnectionContext::default()); + + let delegate = ServerBuilder { + services, + streams: Arc::new(Mutex::new(HashMap::new())), + shutdown_waiter, + conn_ctx, + }; + Connection::new(conn, delegate).run().await + } + } + /// Stops the listener, closes active connections, and releases the listener. pub async fn shutdown(&mut self) -> Result<()> { self.stop_listen().await; @@ -349,7 +389,7 @@ impl Server { } } -async fn spawn_connection_handler( +fn spawn_connection_handler( conn: Socket, services: Arc>, shutdown_waiter: shutdown::Waiter, @@ -815,4 +855,25 @@ mod tests { tokio::time::sleep(std::time::Duration::from_secs(1)).await; assert!(!is_socket_in_use(addr)); } + + #[tokio::test] + async fn test_serve_connection_returns_read_error() { + let server = Server::new(); + let (local, peer) = tokio::net::UnixStream::pair().unwrap(); + let task = tokio::spawn(server.serve_connection(local.into())); + drop(peer); + + let res = task.await.unwrap(); + assert!(matches!(res, Err(Error::Socket(_))), "{:?}", res); + } + + #[tokio::test] + async fn test_serve_connection_returns_ok_on_shutdown() { + let mut server = Server::new(); + let (local, _peer) = tokio::net::UnixStream::pair().unwrap(); + let task = tokio::spawn(server.serve_connection(local.into())); + + server.shutdown().await.unwrap(); + assert_eq!(task.await.unwrap(), Ok(())); + } }