feat(server,preamble): PreambleHooks, clamp timeout, async failure callback

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

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::<T>()` 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
+```

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()
+    });
+```

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).

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),
             },

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,
         }

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<T>
     );
 
-    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<H>(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
+    }
 }

src/server/config/tests.rs

@@ -6,17 +6,19 @@
 //! cases via `rstest`.
 
 use std::{
+    io,
     sync::{
         Arc,
         atomic::{AtomicUsize, Ordering},
     },
     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);