From 40d93ac7e78e98dba2ea676f19a7f853112c50e7 Mon Sep 17 00:00:00 2001 From: Leynos Date: Tue, 25 Nov 2025 12:34:20 +0000 Subject: [PATCH 1/4] feat(server,preamble): add configurable preamble read timeout with failure callback This update introduces a new `preamble_timeout` configuration for `WireframeServer` to bound the maximum duration allowed for reading a connection preamble. If the preamble read exceeds this timeout, a failure handler (if registered) asynchronously executes, allowing custom protocol errors to be sent before the connection is closed. This prevents resources from being tied up indefinitely by clients that fail to send a valid preamble. Key changes include: - Added `preamble_timeout` builder method clamping minimum timeout to 1 ms. - Modified `spawn_connection_task` to enforce timeout around `read_preamble` with tokio's `timeout`. - Enhanced failure handler signature to accept mutable `TcpStream` for replying before disconnection. - Extended documentation with usage examples and detailed guides. - Added tests verifying timeout enforcement, failure handler invocation, and response writing. This enhances server resilience and DoS protection by bounding handshake durations. Co-authored-by: terragon-labs[bot] --- ...eframe-a-guide-to-production-resilience.md | 9 + docs/preamble-validator.md | 5 +- docs/server/configuration.md | 28 ++++ docs/users-guide.md | 10 +- src/server/config/binding.rs | 2 + src/server/config/mod.rs | 1 + src/server/config/preamble.rs | 52 ++++-- src/server/config/tests.rs | 44 ++++- src/server/connection.rs | 47 ++++-- src/server/mod.rs | 31 +++- src/server/runtime.rs | 63 +++++-- tests/preamble.rs | 155 ++++++++++++++++-- 12 files changed, 383 insertions(+), 64 deletions(-) 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..cfc00f26 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`. 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..78e9d632 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, +letting you reply or log decode errors 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..4e2a5e8d 100644 --- a/src/server/config/preamble.rs +++ b/src/server/config/preamble.rs @@ -1,14 +1,13 @@ //! Preamble configuration for [`WireframeServer`]. use core::marker::PhantomData; - -use bincode::error::DecodeError; +use std::time::Duration; use super::WireframeServer; use crate::{ app::WireframeApp, preamble::Preamble, - server::{PreambleSuccessHandler, ServerState}, + server::{PreambleFailureHandler, PreambleSuccessHandler, ServerState}, }; impl WireframeServer @@ -47,11 +46,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. /// @@ -81,21 +105,29 @@ where builder_callback!( /// Register a handler invoked when the connection preamble fails to decode. /// - /// The handler receives a [`bincode::error::DecodeError`]. + /// 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| { - /// eprintln!("Failed to decode preamble"); - /// }, - /// ); + /// 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() + /// }); /// ``` on_preamble_decode_failure, on_preamble_failure, - Fn(&DecodeError) + Send + Sync + 'static + PreambleFailureHandler ); } 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..79477a66 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -1,13 +1,13 @@ //! 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}, @@ -19,7 +19,8 @@ pub(super) fn spawn_connection_task( stream: TcpStream, factory: F, on_success: Option>, - on_failure: Option, + on_failure: Option>, + preamble_timeout: Option, tracker: &TaskTracker, ) where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, @@ -34,7 +35,12 @@ pub(super) fn spawn_connection_task( }; tracker.spawn(async move { 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 +59,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 +88,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 +104,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 +152,7 @@ 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, None, None, None, &tracker); tracker.close(); tracker.wait().await; } @@ -175,7 +202,7 @@ 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, None, None, None, &tracker); tracker.close(); tracker.wait().await; } diff --git a/src/server/mod.rs b/src/server/mod.rs index 3ce6fc16..26c826a5 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,27 @@ 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 trait PreambleFailureHandler: + for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> + + Send + + Sync + + 'static +{ +} + +impl PreambleFailureHandler for F where + F: for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> + + Send + + Sync + + 'static +{ +} + +/// [`PreambleFailureHandler`] wrapped in `Arc`. +pub type PreambleFailure = Arc>; /// Tokio-based server for [`WireframeApp`] instances. /// @@ -104,7 +124,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 +139,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..0f012f76 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,32 @@ impl BackoffConfig { } } +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 Default for PreambleHooks { + fn default() -> Self { + Self { + on_success: None, + on_failure: None, + timeout: None, + } + } +} + impl WireframeServer where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, @@ -195,23 +221,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 +277,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 +294,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 +302,7 @@ where /// accept_loop::<_, (), _>( /// listener, /// || WireframeApp::default(), -/// None, -/// None, +/// PreambleHooks::default(), /// token, /// tracker, /// BackoffConfig::default(), @@ -285,8 +313,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 +339,13 @@ 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.on_success, + hooks.on_failure, + hooks.timeout, &tracker, ); delay = backoff_config.initial_delay; @@ -426,8 +455,7 @@ mod tests { tracker.spawn(accept_loop::<_, (), _>( listener, factory, - None, - None, + PreambleHooks::default(), token.clone(), tracker.clone(), BackoffConfig::default(), @@ -473,8 +501,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..5ad46c85 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,125 @@ 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, _| { + 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:?}" + ); + 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; 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 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; +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, bincode::Encode, bincode::Decode)] struct OtherPreamble(u8); @@ -248,11 +374,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(()) + }) } } From 44037a83e26215881ebad63160eb59e65a71ae8f Mon Sep 17 00:00:00 2001 From: Leynos Date: Tue, 25 Nov 2025 20:04:07 +0000 Subject: [PATCH 2/4] Address preamble timeout review feedback --- docs/users-guide.md | 6 +- src/server/config/preamble.rs | 73 +++++++++++--------- src/server/connection.rs | 4 +- src/server/mod.rs | 22 ++----- src/server/runtime.rs | 13 +--- tests/preamble.rs | 121 +++++++++++++++++++++++++++++++++- 6 files changed, 173 insertions(+), 66 deletions(-) diff --git a/docs/users-guide.md b/docs/users-guide.md index 78e9d632..6aafe2b7 100644 --- a/docs/users-guide.md +++ b/docs/users-guide.md @@ -392,9 +392,9 @@ 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 asynchronous failure callbacks that receive the stream, -letting you reply or log decode errors before the application runs. An optional -`preamble_timeout` caps how long `read_preamble` waits; timeouts use the -failure callback path.[^20] +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 diff --git a/src/server/config/preamble.rs b/src/server/config/preamble.rs index 4e2a5e8d..28bf41b2 100644 --- a/src/server/config/preamble.rs +++ b/src/server/config/preamble.rs @@ -1,13 +1,16 @@ //! Preamble configuration for [`WireframeServer`]. use core::marker::PhantomData; -use std::time::Duration; +use std::{io, time::Duration}; + +use bincode::error::DecodeError; +use futures::future::BoxFuture; use super::WireframeServer; use crate::{ app::WireframeApp, preamble::Preamble, - server::{PreambleFailureHandler, PreambleSuccessHandler, ServerState}, + server::{PreambleSuccessHandler, ServerState}, }; impl WireframeServer @@ -102,32 +105,42 @@ where PreambleSuccessHandler ); - builder_callback!( - /// 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() - /// }); - /// ``` - on_preamble_decode_failure, - on_preamble_failure, - PreambleFailureHandler - ); + /// 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/connection.rs b/src/server/connection.rs index 79477a66..536f9a9d 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -19,7 +19,7 @@ pub(super) fn spawn_connection_task( stream: TcpStream, factory: F, on_success: Option>, - on_failure: Option>, + on_failure: Option, preamble_timeout: Option, tracker: &TaskTracker, ) where @@ -59,7 +59,7 @@ 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, diff --git a/src/server/mod.rs b/src/server/mod.rs index 26c826a5..eecb7296 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -62,24 +62,12 @@ pub type PreambleHandler = Arc>; /// /// Implementors may perform asynchronous I/O on the provided stream to emit a /// response before the connection is closed. -pub trait PreambleFailureHandler: - for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> - + Send - + Sync - + 'static -{ -} - -impl PreambleFailureHandler for F where - F: for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> +pub type PreambleFailure = Arc< + dyn for<'a> Fn(&'a DecodeError, &'a mut tokio::net::TcpStream) -> BoxFuture<'a, io::Result<()>> + Send + Sync - + 'static -{ -} - -/// [`PreambleFailureHandler`] wrapped in `Arc`. -pub type PreambleFailure = Arc>; + + 'static, +>; /// Tokio-based server for [`WireframeApp`] instances. /// @@ -124,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 diff --git a/src/server/runtime.rs b/src/server/runtime.rs index 0f012f76..c935f830 100644 --- a/src/server/runtime.rs +++ b/src/server/runtime.rs @@ -82,9 +82,10 @@ impl BackoffConfig { } } +#[derive(Default)] pub(super) struct PreambleHooks { pub on_success: Option>, - pub on_failure: Option>, + pub on_failure: Option, pub timeout: Option, } @@ -98,16 +99,6 @@ impl Clone for PreambleHooks { } } -impl Default for PreambleHooks { - fn default() -> Self { - Self { - on_success: None, - on_failure: None, - timeout: None, - } - } -} - impl WireframeServer where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, diff --git a/tests/preamble.rs b/tests/preamble.rs index 5ad46c85..1a2a9890 100644 --- a/tests/preamble.rs +++ b/tests/preamble.rs @@ -274,7 +274,7 @@ async fn preamble_timeout_invokes_failure_handler_and_closes_connection( let server = WireframeServer::new(factory) .with_preamble::() .preamble_timeout(Duration::from_millis(50)) - .on_preamble_decode_failure(move |err, _| { + .on_preamble_decode_failure(move |err, stream| { let failure_holder = failure_holder.clone(); Box::pin(async move { assert!( @@ -285,6 +285,9 @@ async fn preamble_timeout_invokes_failure_handler_and_closes_connection( ), "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(()); } @@ -298,8 +301,14 @@ async fn preamble_timeout_invokes_failure_handler_and_closes_connection( .await .expect("timeout waiting for failure callback") .expect("failure callback send"); - let mut buf = [0u8; 1]; - let read = timeout(Duration::from_millis(200), stream.read(&mut buf)).await; + 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"), @@ -347,6 +356,112 @@ async fn success_handler_runs_without_failure_handler( .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); From 4ee7c94b555c0464ee7860ba9a70517b6cc7369b Mon Sep 17 00:00:00 2001 From: Leynos Date: Wed, 26 Nov 2025 12:14:43 +0000 Subject: [PATCH 3/4] docs(preamble): add sequence diagram detailing preamble processing flow with timeout Added a detailed sequence diagram to the preamble-validator documentation. The diagram illustrates the interaction between the accept loop, connection task, preamble decoding, timeout handling, success and failure callbacks, and handoff to the application. This enhances the understanding of the preamble processing flow and timeout management. Co-authored-by: terragon-labs[bot] --- docs/preamble-validator.md | 57 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/docs/preamble-validator.md b/docs/preamble-validator.md index cfc00f26..d0c2fbba 100644 --- a/docs/preamble-validator.md +++ b/docs/preamble-validator.md @@ -46,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.on_success, hooks.on_failure, hooks.timeout) + + 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 +``` From 8363956d13b6287d2f4f30160e7a6948d6ef740b Mon Sep 17 00:00:00 2001 From: Leynos Date: Wed, 26 Nov 2025 12:28:23 +0000 Subject: [PATCH 4/4] refactor(server): consolidate preamble hooks into a single struct parameter Changed spawn_connection_task and related calls to accept a single PreambleHooks struct instead of separate on_success, on_failure, and preamble_timeout parameters. This reduces parameter bloat and simplifies passing preamble-related handlers and timeout settings together. Updated relevant tests and documentation to reflect this change. Co-authored-by: terragon-labs[bot] --- docs/preamble-validator.md | 2 +- src/server/connection.rs | 24 +++++++++++++++++++----- src/server/runtime.rs | 4 +--- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/docs/preamble-validator.md b/docs/preamble-validator.md index d0c2fbba..6e079220 100644 --- a/docs/preamble-validator.md +++ b/docs/preamble-validator.md @@ -67,7 +67,7 @@ sequenceDiagram Client->>WireframeServer: connect WireframeServer->>AcceptLoop: start accept_loop AcceptLoop->>AcceptLoop: listener.accept() - AcceptLoop->>ConnectionTask: spawn_connection_task(stream, factory, hooks.on_success, hooks.on_failure, hooks.timeout) + AcceptLoop->>ConnectionTask: spawn_connection_task(stream, factory, hooks) ConnectionTask->>ProcessStream: process_stream(stream, peer_addr, factory, on_success, on_failure, preamble_timeout) diff --git a/src/server/connection.rs b/src/server/connection.rs index 536f9a9d..b7b0a279 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -12,15 +12,14 @@ 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, - preamble_timeout: Option, + hooks: PreambleHooks, tracker: &TaskTracker, ) where F: Fn() -> WireframeApp + Send + Sync + Clone + 'static, @@ -34,6 +33,11 @@ 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, @@ -152,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, None, &tracker); + spawn_connection_task::<_, ()>( + stream, + app_factory, + PreambleHooks::default(), + &tracker, + ); tracker.close(); tracker.wait().await; } @@ -202,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, None, &tracker); + spawn_connection_task::<_, ()>( + stream, + app_factory, + PreambleHooks::default(), + &tracker, + ); tracker.close(); tracker.wait().await; } diff --git a/src/server/runtime.rs b/src/server/runtime.rs index c935f830..9d815425 100644 --- a/src/server/runtime.rs +++ b/src/server/runtime.rs @@ -334,9 +334,7 @@ pub(super) async fn accept_loop( spawn_connection_task( stream, factory.clone(), - hooks.on_success, - hooks.on_failure, - hooks.timeout, + hooks, &tracker, ); delay = backoff_config.initial_delay;