diff --git a/docs/hardening-wireframe-a-guide-to-production-resilience.md b/docs/hardening-wireframe-a-guide-to-production-resilience.md index d799ec65..94093992 100644 --- a/docs/hardening-wireframe-a-guide-to-production-resilience.md +++ b/docs/hardening-wireframe-a-guide-to-production-resilience.md @@ -235,6 +235,15 @@ The solution is to store non-owning `Weak` pointers in the registry. This pattern entirely prevents the registry from causing memory leaks. +### 3.3 Guarding the preamble phase + +Connections that never deliver a preamble can tie up sockets and worker tasks. +A configurable `preamble_timeout` now wraps `read_preamble`, and expiry routes +through the failure callback. Because the failure handler is asynchronous and +receives the `TcpStream`, it can emit a protocol-specific error before the +socket closes. This keeps resource usage bounded while preserving custom +handshake responses. + ## 4. Denial-of-Service (DoS) Mitigation As a library exposed to the network, `wireframe` must be hardened against diff --git a/docs/preamble-validator.md b/docs/preamble-validator.md index 9e5b715e..6e079220 100644 --- a/docs/preamble-validator.md +++ b/docs/preamble-validator.md @@ -31,7 +31,10 @@ sequenceDiagram The success callback receives the decoded preamble and a mutable `TcpStream`. It may write a handshake response before the connection is passed to -`WireframeApp`. In the tests, a `HotlinePreamble` struct illustrates the +`WireframeApp`. The failure callback is also asynchronous and receives the +mutable stream so it can emit an error reply before the connection closes. Use +`preamble_timeout` to cap how long the read may take; the timeout follows the +failure callback path. In the tests, a `HotlinePreamble` struct illustrates the pattern, but any preamble type may be used. Register callbacks via `on_preamble_decode_success` and `on_preamble_decode_failure` on `WireframeServer`. @@ -43,3 +46,60 @@ callbacks with `on_preamble_decode_success` or `on_preamble_decode_failure`. The method converts the server to use a custom preamble type, dropping any callbacks configured on the default `()` preamble. Registering callbacks after calling `with_preamble::()` ensures they are retained. + +## Preamble processing flow with timeout handling + +Sequence diagram showing how the accept loop, connection task, preamble +decoding, timeout handling, and callbacks coordinate before handing the +connection to the application: + +```mermaid +sequenceDiagram + actor Client + participant WireframeServer + participant AcceptLoop + participant ConnectionTask + participant ProcessStream + participant ReadPreamble + participant PreambleFailureHandler + participant WireframeApp + + Client->>WireframeServer: connect + WireframeServer->>AcceptLoop: start accept_loop + AcceptLoop->>AcceptLoop: listener.accept() + AcceptLoop->>ConnectionTask: spawn_connection_task(stream, factory, hooks) + + ConnectionTask->>ProcessStream: process_stream(stream, peer_addr, factory, on_success, on_failure, preamble_timeout) + + alt preamble_timeout is Some + ProcessStream->>ProcessStream: timeout(preamble_timeout, read_preamble) + ProcessStream->>ReadPreamble: read_preamble(stream) + alt preamble read completes in time + ReadPreamble-->>ProcessStream: Ok(preamble, leftover) + else preamble read times out + ProcessStream-->>ProcessStream: Err(timeout_error) + end + else preamble_timeout is None + ProcessStream->>ReadPreamble: read_preamble(stream) + ReadPreamble-->>ProcessStream: Result(preamble or error) + end + + alt preamble_result is Ok + ProcessStream->>ProcessStream: invoke on_success if Some + opt on_success is Some + ProcessStream->>WireframeApp: on_preamble_success(preamble, stream) + WireframeApp-->>ProcessStream: Result + end + ProcessStream->>WireframeApp: hand off stream and preamble + else preamble_result is Err (decode failure or timeout) + alt on_failure is Some + ProcessStream->>PreambleFailureHandler: on_preamble_failure(error, stream) + PreambleFailureHandler-->>ProcessStream: io::Result + ProcessStream->>Client: optional protocol error reply via stream + ProcessStream->>Client: close connection + else on_failure is None + ProcessStream->>ProcessStream: log error with peer_addr + ProcessStream->>Client: close connection + end + end +``` diff --git a/docs/server/configuration.md b/docs/server/configuration.md index 00888a2c..bae73273 100644 --- a/docs/server/configuration.md +++ b/docs/server/configuration.md @@ -82,3 +82,31 @@ let cfg = BackoffConfig { let server = WireframeServer::new(|| WireframeApp::default()) .accept_backoff(cfg); ``` + +## Preamble handling + +Servers that expect a preamble can bound how long they wait for it and decide +what to do when decoding fails. `preamble_timeout(Duration)` wraps +`read_preamble` in a timeout; values below 1 ms are clamped to 1 ms, and +omitting the setter leaves the read unbounded. When decoding fails or times +out, an optional failure handler runs before the connection is closed. The +handler is asynchronous and receives both the error and the mutable +`TcpStream`, allowing a response to be written back to the client. + +```rust,no_run +use std::time::Duration; + +use futures::FutureExt; +use tokio::io::AsyncWriteExt; +use wireframe::{app::WireframeApp, server::WireframeServer}; + +let server = WireframeServer::new(|| WireframeApp::default()) + .preamble_timeout(Duration::from_secs(2)) + .on_preamble_decode_failure(|_err, stream| { + async move { + stream.write_all(b"preamble required").await?; + Ok(()) + } + .boxed() + }); +``` diff --git a/docs/users-guide.md b/docs/users-guide.md index 15bd7175..6aafe2b7 100644 --- a/docs/users-guide.md +++ b/docs/users-guide.md @@ -391,8 +391,10 @@ the supplied future resolves.[^18] Each worker runs `accept_loop`, which clones the factory, rewinds leftover preamble bytes, and hands the stream to the application. Transient accept failures trigger exponential backoff capped by the configured maximum delay.[^18][^19] Preamble hooks support asynchronous -success handlers and synchronous failure callbacks, letting you reject -connections or log decode errors before the application runs.[^20] +success handlers and asynchronous failure callbacks that receive the stream, +enabling replies or decode-error logging before the application runs. An +optional `preamble_timeout` caps how long `read_preamble` waits; timeouts use +the failure callback path.[^20] `spawn_connection_task` wraps each accepted stream in `read_preamble` and `RewindStream`, records connection panics, and logs failures without crashing @@ -621,8 +623,8 @@ call these helpers to maintain consistent telemetry.[^6][^7][^31][^20] [^17]: Implemented in `src/server/config/binding.rs` (lines 68-214). [^18]: Implemented in `src/server/runtime.rs` (lines 90-233). [^19]: Implemented in `src/server/runtime.rs` (lines 240-333). -[^20]: Implemented in `src/server/config/preamble.rs` (lines 14-100) and - `src/server/connection.rs` (lines 17-84). +[^20]: Implemented in `src/server/config/preamble.rs` (lines 14-135) and + `src/server/connection.rs` (lines 1-222). [^21]: Implemented in `src/server/error.rs` (lines 7-18). [^23]: Implemented in `src/push/queues/mod.rs` (lines 41-190). [^24]: Implemented in `src/push/queues/errors.rs` (lines 7-28). diff --git a/src/server/config/binding.rs b/src/server/config/binding.rs index a2af237d..2231a7f1 100644 --- a/src/server/config/binding.rs +++ b/src/server/config/binding.rs @@ -41,6 +41,7 @@ where on_preamble_failure, ready_tx, backoff_config, + preamble_timeout, _preamble: preamble, .. } = self; @@ -57,6 +58,7 @@ where on_preamble_failure, ready_tx, backoff_config, + preamble_timeout, state: Bound { listener: Arc::new(tokio_listener), }, diff --git a/src/server/config/mod.rs b/src/server/config/mod.rs index ec57aecd..c014a3a1 100644 --- a/src/server/config/mod.rs +++ b/src/server/config/mod.rs @@ -74,6 +74,7 @@ where on_preamble_failure: None, ready_tx: None, backoff_config: BackoffConfig::default(), + preamble_timeout: None, state: Unbound, _preamble: PhantomData, } diff --git a/src/server/config/preamble.rs b/src/server/config/preamble.rs index 9ad3fadf..28bf41b2 100644 --- a/src/server/config/preamble.rs +++ b/src/server/config/preamble.rs @@ -1,8 +1,10 @@ //! Preamble configuration for [`WireframeServer`]. use core::marker::PhantomData; +use std::{io, time::Duration}; use bincode::error::DecodeError; +use futures::future::BoxFuture; use super::WireframeServer; use crate::{ @@ -47,11 +49,36 @@ where on_preamble_failure: None, ready_tx: self.ready_tx, backoff_config: self.backoff_config, + preamble_timeout: self.preamble_timeout, state: self.state, _preamble: PhantomData, } } + /// Configure a timeout for reading the connection preamble. + /// + /// The timeout is applied around the preamble read. When it elapses, the + /// preamble decode failure handler is invoked (if registered) and the + /// connection is closed. Values below 1 ms are clamped to 1 ms to avoid + /// immediate expiry. Omit this setter to disable the timeout. + /// + /// # Examples + /// + /// ``` + /// use std::time::Duration; + /// + /// use wireframe::{app::WireframeApp, server::WireframeServer}; + /// + /// let server = + /// WireframeServer::new(|| WireframeApp::default()).preamble_timeout(Duration::from_secs(1)); + /// ``` + #[must_use] + pub fn preamble_timeout(mut self, timeout: Duration) -> Self { + let normalised = timeout.max(Duration::from_millis(1)); + self.preamble_timeout = Some(normalised); + self + } + builder_callback!( /// Register a handler invoked when the connection preamble decodes successfully. /// @@ -78,24 +105,42 @@ where PreambleSuccessHandler ); - builder_callback!( - /// Register a handler invoked when the connection preamble fails to decode. - /// - /// The handler receives a [`bincode::error::DecodeError`]. - /// - /// # Examples - /// - /// ``` - /// use wireframe::{app::WireframeApp, server::WireframeServer}; - /// - /// let server = WireframeServer::new(|| WireframeApp::default()).on_preamble_decode_failure( - /// |_err: &bincode::error::DecodeError| { - /// eprintln!("Failed to decode preamble"); - /// }, - /// ); - /// ``` - on_preamble_decode_failure, - on_preamble_failure, - Fn(&DecodeError) + Send + Sync + 'static - ); + /// Register a handler invoked when the connection preamble fails to decode. + /// + /// The handler receives a [`bincode::error::DecodeError`] and a mutable + /// [`tokio::net::TcpStream`] so it may emit a response before the + /// connection is closed. This callback is awaited; if it returns an + /// error, the error is logged and the connection is closed. + /// + /// # Examples + /// + /// ``` + /// use futures::FutureExt; + /// use tokio::io::AsyncWriteExt; + /// use wireframe::{app::WireframeApp, server::WireframeServer}; + /// + /// let server = WireframeServer::new(|| WireframeApp::default()).on_preamble_decode_failure( + /// |_err: &bincode::error::DecodeError, stream| { + /// async move { + /// stream.write_all(b"BAD").await?; + /// Ok(()) + /// } + /// .boxed() + /// }, + /// ); + /// ``` + #[must_use] + pub fn on_preamble_decode_failure(mut self, handler: H) -> Self + where + H: for<'a> Fn( + &'a DecodeError, + &'a mut tokio::net::TcpStream, + ) -> BoxFuture<'a, io::Result<()>> + + Send + + Sync + + 'static, + { + self.on_preamble_failure = Some(std::sync::Arc::new(handler)); + self + } } diff --git a/src/server/config/tests.rs b/src/server/config/tests.rs index 9dbeb49e..ad4b6b9c 100644 --- a/src/server/config/tests.rs +++ b/src/server/config/tests.rs @@ -6,6 +6,7 @@ //! cases via `rstest`. use std::{ + io, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -13,10 +14,11 @@ use std::{ time::Duration, }; +use bincode::error::DecodeError; use rstest::rstest; +use tokio::net::{TcpListener, TcpStream}; use super::*; -use bincode::error::DecodeError; use crate::server::{ test_util::{ TestPreamble, @@ -27,7 +29,7 @@ use crate::server::{ server_with_preamble, }, BackoffConfig, -use tokio::net::{TcpListener, TcpStream}; +}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum PreambleHandlerKind { @@ -72,6 +74,18 @@ fn test_with_preamble_type_conversion( assert_eq!(server.worker_count(), expected_default_worker_count()); } +#[rstest] +fn test_preamble_timeout_configuration( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let timeout = Duration::from_millis(25); + let server = WireframeServer::new(factory).preamble_timeout(timeout); + assert_eq!(server.preamble_timeout, Some(timeout)); + + let clamped = WireframeServer::new(factory).preamble_timeout(Duration::ZERO); + assert_eq!(clamped.preamble_timeout, Some(Duration::from_millis(1))); +} + #[rstest] #[tokio::test] async fn test_bind_success( @@ -125,9 +139,15 @@ async fn test_preamble_handler_registration( Ok(()) }) }), - PreambleHandlerKind::Failure => server.on_preamble_decode_failure(move |_err: &DecodeError| { - c.fetch_add(1, Ordering::SeqCst); - }), + PreambleHandlerKind::Failure => server.on_preamble_decode_failure( + move |_err: &DecodeError, _stream| { + let c = c.clone(); + Box::pin(async move { + c.fetch_add(1, Ordering::SeqCst); + Ok::<(), io::Error>(()) + }) + }, + ), }; assert_eq!(counter.load(Ordering::SeqCst), 0); @@ -157,7 +177,17 @@ async fn test_preamble_handler_registration( .on_preamble_failure .as_ref() .expect("failure handler missing"); - handler(&DecodeError::UnexpectedEnd); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind listener"); + let addr = listener.local_addr().expect("listener addr"); + let _client = TcpStream::connect(addr) + .await + .expect("client connect failed"); + let (mut stream, _) = listener.accept().await.expect("accept stream"); + handler(&DecodeError::UnexpectedEnd, &mut stream) + .await + .expect("handler failed"); } } assert_eq!(counter.load(Ordering::SeqCst), 1); @@ -181,7 +211,7 @@ async fn test_method_chaining( Ok(()) }) }) - .on_preamble_decode_failure(|_: &DecodeError| {}) + .on_preamble_decode_failure(|_: &DecodeError, _| Box::pin(async { Ok::<(), io::Error>(()) })) .bind_existing_listener(free_listener) .expect("Failed to bind"); assert_eq!(server.worker_count(), 2); diff --git a/src/server/connection.rs b/src/server/connection.rs index a21e0c79..b7b0a279 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -1,25 +1,25 @@ //! Connection handling for [`WireframeServer`]. -use std::net::SocketAddr; +use std::{io, net::SocketAddr, time::Duration}; use futures::FutureExt; use log::{error, warn}; -use tokio::net::TcpStream; +use tokio::{net::TcpStream, time::timeout}; use tokio_util::task::TaskTracker; -use super::{PreambleErrorHandler, PreambleHandler}; +use super::{PreambleFailure, PreambleHandler}; use crate::{ app::WireframeApp, preamble::{Preamble, read_preamble}, rewind_stream::RewindStream, + server::runtime::PreambleHooks, }; /// Spawn a task to process a single TCP connection, logging and discarding any panics. pub(super) fn spawn_connection_task( stream: TcpStream, factory: F, - on_success: Option>, - on_failure: Option, + hooks: PreambleHooks, tracker: &TaskTracker, ) where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, @@ -33,8 +33,18 @@ pub(super) fn spawn_connection_task( } }; tracker.spawn(async move { + let PreambleHooks { + on_success, + on_failure, + timeout: preamble_timeout, + } = hooks; let fut = std::panic::AssertUnwindSafe(process_stream( - stream, peer_addr, factory, on_success, on_failure, + stream, + peer_addr, + factory, + on_success, + on_failure, + preamble_timeout, )) .catch_unwind(); @@ -53,12 +63,21 @@ async fn process_stream( peer_addr: Option, factory: F, on_success: Option>, - on_failure: Option, + on_failure: Option, + preamble_timeout: Option, ) where F: Fn() -> WireframeApp + Send + Sync + 'static, T: Preamble, { - match read_preamble::<_, T>(&mut stream).await { + let preamble_result = match preamble_timeout { + Some(limit) => match timeout(limit, read_preamble::<_, T>(&mut stream)).await { + Ok(result) => result, + Err(_) => Err(timeout_error()), + }, + None => read_preamble::<_, T>(&mut stream).await, + }; + + match preamble_result { Ok((preamble, leftover)) => { if let Some(handler) = on_success.as_ref() && let Err(e) = handler(&preamble, &mut stream).await @@ -73,7 +92,12 @@ async fn process_stream( } Err(err) => { if let Some(handler) = on_failure.as_ref() { - handler(&err); + if let Err(e) = handler(&err, &mut stream).await { + error!( + "preamble failure handler error: error={e}, error_debug={e:?}, \ + peer_addr={peer_addr:?}" + ); + } } else { error!( "preamble decode failed and no failure handler set: error={err:?}, \ @@ -84,6 +108,13 @@ async fn process_stream( } } +fn timeout_error() -> bincode::error::DecodeError { + bincode::error::DecodeError::Io { + inner: io::Error::new(io::ErrorKind::TimedOut, "preamble read timed out"), + additional: 0, + } +} + #[cfg(test)] mod tests { use rstest::rstest; @@ -125,7 +156,12 @@ mod tests { let tracker = tracker.clone(); async move { let (stream, _) = listener.accept().await.expect("accept"); - spawn_connection_task::<_, ()>(stream, app_factory, None, None, &tracker); + spawn_connection_task::<_, ()>( + stream, + app_factory, + PreambleHooks::default(), + &tracker, + ); tracker.close(); tracker.wait().await; } @@ -175,7 +211,12 @@ mod tests { let tracker = tracker.clone(); async move { let (stream, _) = listener.accept().await.expect("accept"); - spawn_connection_task::<_, ()>(stream, app_factory, None, None, &tracker); + spawn_connection_task::<_, ()>( + stream, + app_factory, + PreambleHooks::default(), + &tracker, + ); tracker.close(); tracker.wait().await; } diff --git a/src/server/mod.rs b/src/server/mod.rs index 3ce6fc16..eecb7296 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -4,7 +4,7 @@ //! decode an optional preamble before handing the stream to the application. use core::marker::PhantomData; -use std::{io, sync::Arc}; +use std::{io, sync::Arc, time::Duration}; use bincode::error::DecodeError; use futures::future::BoxFuture; @@ -59,7 +59,15 @@ impl PreambleSuccessHandler for F where pub type PreambleHandler = Arc>; /// Handler invoked when decoding a connection preamble fails. -pub type PreambleErrorHandler = Arc; +/// +/// Implementors may perform asynchronous I/O on the provided stream to emit a +/// response before the connection is closed. +pub type PreambleFailure = Arc< + dyn for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> + + Send + + Sync + + 'static, +>; /// Tokio-based server for [`WireframeApp`] instances. /// @@ -104,7 +112,7 @@ where pub(crate) factory: F, pub(crate) workers: usize, pub(crate) on_preamble_success: Option>, - pub(crate) on_preamble_failure: Option, + pub(crate) on_preamble_failure: Option, /// Channel used to notify when the server is ready. /// /// # Thread Safety @@ -119,6 +127,11 @@ where /// provided each time the server is started. pub(crate) ready_tx: Option>, pub(crate) backoff_config: BackoffConfig, + /// Maximum duration allowed for reading a preamble before timing out. + /// + /// `None` disables the timeout. A zero or sub-millisecond timeout is + /// normalised to 1 ms when configured. + pub(crate) preamble_timeout: Option, /// Typestate tracking whether the server has been bound to a listener. /// [`Unbound`] servers require binding before they can run. pub(crate) state: S, diff --git a/src/server/runtime.rs b/src/server/runtime.rs index 97ba360e..9d815425 100644 --- a/src/server/runtime.rs +++ b/src/server/runtime.rs @@ -15,7 +15,7 @@ use tokio_util::{sync::CancellationToken, task::TaskTracker}; use super::{ Bound, - PreambleErrorHandler, + PreambleFailure, PreambleHandler, ServerError, WireframeServer, @@ -82,6 +82,23 @@ impl BackoffConfig { } } +#[derive(Default)] +pub(super) struct PreambleHooks { + pub on_success: Option>, + pub on_failure: Option, + pub timeout: Option, +} + +impl Clone for PreambleHooks { + fn clone(&self) -> Self { + Self { + on_success: self.on_success.clone(), + on_failure: self.on_failure.clone(), + timeout: self.timeout, + } + } +} + impl WireframeServer where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, @@ -195,23 +212,27 @@ where ready_tx, state: Bound { listener }, backoff_config, + preamble_timeout, .. } = self; let shutdown_token = CancellationToken::new(); let tracker = TaskTracker::new(); + let preamble = PreambleHooks { + on_success: on_preamble_success, + on_failure: on_preamble_failure, + timeout: preamble_timeout, + }; for _ in 0..workers { let listener = Arc::clone(&listener); let factory = factory.clone(); - let on_success = on_preamble_success.clone(); - let on_failure = on_preamble_failure.clone(); + let preamble_hooks = preamble.clone(); let token = shutdown_token.clone(); let t = tracker.clone(); tracker.spawn(accept_loop( listener, factory, - on_success, - on_failure, + preamble_hooks, token, t, backoff_config, @@ -247,8 +268,7 @@ where /// /// - `listener`: Source of incoming TCP connections. /// - `factory`: Creates a fresh [`WireframeApp`] for each connection. -/// - `on_success`: Callback invoked after a successful preamble. -/// - `on_failure`: Callback invoked when the preamble fails. +/// - `preamble`: Preamble handlers and timeout configuration. /// - `shutdown`: Signal used to stop the accept loop. /// - `tracker`: Task tracker used for graceful shutdown. /// - `backoff_config`: Controls exponential back-off behaviour. @@ -265,7 +285,7 @@ where /// use std::sync::Arc; /// /// use tokio_util::{sync::CancellationToken, task::TaskTracker}; -/// use wireframe::{app::WireframeApp /*, server::runtime::{AcceptListener, BackoffConfig, accept_loop} */}; +/// use wireframe::{app::WireframeApp /*, server::runtime::{AcceptListener, BackoffConfig, PreambleHooks, accept_loop} */}; /// /// async fn run(listener: Arc) { /// let tracker = TaskTracker::new(); @@ -273,8 +293,7 @@ where /// accept_loop::<_, (), _>( /// listener, /// || WireframeApp::default(), -/// None, -/// None, +/// PreambleHooks::default(), /// token, /// tracker, /// BackoffConfig::default(), @@ -285,8 +304,7 @@ where pub(super) async fn accept_loop( listener: Arc, factory: F, - on_success: Option>, - on_failure: Option, + preamble: PreambleHooks, shutdown: CancellationToken, tracker: TaskTracker, backoff_config: BackoffConfig, @@ -312,11 +330,11 @@ pub(super) async fn accept_loop( res = listener.accept() => match res { Ok((stream, _)) => { + let hooks = preamble.clone(); spawn_connection_task( stream, factory.clone(), - on_success.clone(), - on_failure.clone(), + hooks, &tracker, ); delay = backoff_config.initial_delay; @@ -426,8 +444,7 @@ mod tests { tracker.spawn(accept_loop::<_, (), _>( listener, factory, - None, - None, + PreambleHooks::default(), token.clone(), tracker.clone(), BackoffConfig::default(), @@ -473,8 +490,7 @@ mod tests { tracker.spawn(accept_loop::<_, (), _>( listener, factory, - None, - None, + PreambleHooks::default(), token.clone(), tracker.clone(), backoff, diff --git a/tests/preamble.rs b/tests/preamble.rs index e6f3355a..1a2a9890 100644 --- a/tests/preamble.rs +++ b/tests/preamble.rs @@ -52,7 +52,10 @@ where + Send + Sync + 'static, - E: Fn(&DecodeError) + Send + Sync + 'static, + E: for<'a> Fn(&'a DecodeError, &'a mut TcpStream) -> BoxFuture<'a, io::Result<()>> + + Send + + Sync + + 'static, { WireframeServer::new(factory) .workers(1) @@ -151,10 +154,14 @@ async fn server_triggers_expected_callback( }, { let failure_tx = failure_tx.clone(); - move |_| { - if let Some(tx) = failure_tx.lock().expect("lock poisoned").take() { - let _ = tx.send(()); - } + move |_, _| { + let failure_tx = failure_tx.clone(); + Box::pin(async move { + if let Some(tx) = failure_tx.lock().expect("lock poisoned").take() { + let _ = tx.send(()); + } + Ok(()) + }) } }, ); @@ -207,7 +214,7 @@ async fn success_callback_can_write_response( Ok(()) }) }, - |_| {}, + |_, _| Box::pin(async { Ok::<(), io::Error>(()) }), ); with_running_server(server, |addr| async move { @@ -221,6 +228,240 @@ async fn success_callback_can_write_response( .await; } +#[rstest] +#[tokio::test] +async fn failure_callback_can_write_response( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let (failure_holder, failure_rx) = channel_holder(); + let server = WireframeServer::new(factory) + .with_preamble::() + .on_preamble_decode_failure(move |_, stream| { + let failure_holder = failure_holder.clone(); + Box::pin(async move { + stream.write_all(b"ERR").await.expect("write failed"); + stream.flush().await.expect("flush failed"); + if let Some(tx) = failure_holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + Ok(()) + }) + }); + + with_running_server(server, |addr| async move { + let mut stream = TcpStream::connect(addr).await.expect("connect failed"); + stream.write_all(b"BAD").await.expect("write failed"); + stream.shutdown().await.expect("shutdown failed"); + let mut buf = [0u8; 3]; + let read = timeout(Duration::from_secs(1), stream.read_exact(&mut buf)).await; + let result = read.expect("timeout waiting for failure handler"); + result.expect("read error"); + assert_eq!(&buf, b"ERR"); + timeout(Duration::from_millis(200), failure_rx) + .await + .expect("timeout waiting for failure callback") + .expect("failure callback send"); + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn preamble_timeout_invokes_failure_handler_and_closes_connection( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let (failure_holder, failure_rx) = channel_holder(); + let server = WireframeServer::new(factory) + .with_preamble::() + .preamble_timeout(Duration::from_millis(50)) + .on_preamble_decode_failure(move |err, stream| { + let failure_holder = failure_holder.clone(); + Box::pin(async move { + assert!( + matches!( + err, + DecodeError::Io { inner, .. } + if inner.kind() == io::ErrorKind::TimedOut + ), + "expected timed out error, got {err:?}" + ); + stream.write_all(b"ERR").await.expect("write failed"); + stream.flush().await.expect("flush failed"); + stream.shutdown().await.expect("shutdown failed"); + if let Some(tx) = failure_holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + Ok(()) + }) + }); + + with_running_server(server, |addr| async move { + let mut stream = TcpStream::connect(addr).await.expect("connect failed"); + timeout(Duration::from_secs(1), failure_rx) + .await + .expect("timeout waiting for failure callback") + .expect("failure callback send"); + let mut buf = [0u8; 3]; + timeout(Duration::from_millis(500), stream.read_exact(&mut buf)) + .await + .expect("did not receive timeout response in time") + .expect("read timeout response failed"); + assert_eq!(&buf, b"ERR"); + let mut eof = [0u8; 1]; + let read = timeout(Duration::from_millis(200), stream.read(&mut eof)).await; + match read.expect("timeout waiting for close") { + Ok(0) => {} + Ok(n) => panic!("expected connection to close, read {n} bytes"), + Err(e) if e.kind() == io::ErrorKind::ConnectionReset => {} + Err(e) => panic!("unexpected read error: {e:?}"), + } + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn success_handler_runs_without_failure_handler( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let (success_tx, success_rx) = tokio::sync::oneshot::channel::(); + let success_tx = Arc::new(Mutex::new(Some(success_tx))); + let server = WireframeServer::new(factory) + .with_preamble::() + .on_preamble_decode_success({ + let success_tx = success_tx.clone(); + move |p, _| { + let success_tx = success_tx.clone(); + let preamble = p.clone(); + Box::pin(async move { + if let Some(tx) = success_tx.lock().expect("lock").take() { + let _ = tx.send(preamble); + } + Ok(()) + }) + } + }); + + with_running_server(server, |addr| async move { + let mut stream = TcpStream::connect(addr).await.expect("connect failed"); + let bytes = b"TRTPHOTL\x00\x01\x00\x02"; + stream.write_all(bytes).await.expect("write failed"); + stream.shutdown().await.expect("shutdown failed"); + let preamble = timeout(Duration::from_secs(1), success_rx) + .await + .expect("timeout waiting for success") + .expect("success send"); + assert_eq!(preamble.magic, HotlinePreamble::MAGIC); + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn preamble_timeout_allows_timely_preamble( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let (success_holder, success_rx) = channel_holder(); + let (failure_holder, failure_rx) = channel_holder(); + let server = WireframeServer::new(factory) + .with_preamble::() + .preamble_timeout(Duration::from_millis(150)) + .on_preamble_decode_success({ + let success_holder = success_holder.clone(); + move |p, stream| { + let success_holder = success_holder.clone(); + let clone = p.clone(); + Box::pin(async move { + if let Some(tx) = success_holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + stream.write_all(b"OK").await.expect("write failed"); + stream.flush().await.expect("flush failed"); + // keep connection open by not shutting down here + assert_eq!(clone.magic, HotlinePreamble::MAGIC); + Ok(()) + }) + } + }) + .on_preamble_decode_failure({ + let failure_holder = failure_holder.clone(); + move |_, _| { + let failure_holder = failure_holder.clone(); + Box::pin(async move { + if let Some(tx) = failure_holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + Ok(()) + }) + } + }); + + with_running_server(server, |addr| async move { + let mut stream = TcpStream::connect(addr).await.expect("connect failed"); + let bytes = b"TRTPHOTL\x00\x01\x00\x02"; + stream.write_all(bytes).await.expect("write failed"); + + timeout(Duration::from_millis(200), success_rx) + .await + .expect("timeout waiting for success") + .expect("success send"); + assert!( + timeout(Duration::from_millis(150), failure_rx) + .await + .is_err(), + "failure handler should not fire for timely preamble" + ); + + let mut buf = [0u8; 2]; + stream + .read_exact(&mut buf) + .await + .expect("expected response from success handler"); + assert_eq!(&buf, b"OK"); + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn failure_handler_error_is_logged_and_connection_closes( + factory: impl Fn() -> WireframeApp + Send + Sync + Clone + 'static, +) { + let (failure_holder, failure_rx) = channel_holder(); + let server = WireframeServer::new(factory) + .with_preamble::() + .on_preamble_decode_failure(move |_, _| { + let failure_holder = failure_holder.clone(); + Box::pin(async move { + if let Some(tx) = failure_holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + Err::<(), io::Error>(io::Error::other("boom")) + }) + }); + + with_running_server(server, |addr| async move { + let mut stream = TcpStream::connect(addr).await.expect("connect failed"); + stream.write_all(b"BAD").await.expect("write failed"); + stream.shutdown().await.expect("shutdown failed"); + + timeout(Duration::from_secs(1), failure_rx) + .await + .expect("failure handler not invoked") + .expect("failure handler send failed"); + + let mut buf = [0u8; 1]; + let read = timeout(Duration::from_millis(200), stream.read(&mut buf)).await; + match read.expect("timeout waiting for close") { + Ok(0) => {} + Ok(n) => panic!("expected connection close, read {n} bytes"), + Err(e) if e.kind() == io::ErrorKind::ConnectionReset => {} + Err(e) => panic!("unexpected read error: {e:?}"), + } + }) + .await; +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, bincode::Encode, bincode::Decode)] struct OtherPreamble(u8); @@ -248,11 +489,18 @@ fn success_cb

( fn failure_cb( holder: Arc>>>, -) -> impl Fn(&DecodeError) + Send + Sync + 'static { - move |_| { - if let Some(tx) = holder.lock().expect("lock").take() { - let _ = tx.send(()); - } +) -> impl for<'a> Fn(&'a DecodeError, &'a mut TcpStream) -> BoxFuture<'a, io::Result<()>> ++ Send ++ Sync ++ 'static { + move |_, _| { + let holder = holder.clone(); + Box::pin(async move { + if let Some(tx) = holder.lock().expect("lock").take() { + let _ = tx.send(()); + } + Ok(()) + }) } }