From 37c5d71697951786511ecdcf63ecae7314720cd1 Mon Sep 17 00:00:00 2001
From: Luke Curley <kixelated@gmail.com>
Date: Tue, 10 Feb 2026 07:36:29 -0800
Subject: [PATCH] Add Incoming struct to inspect connections before accepting

Adds an Incoming struct that wraps a QUIC connection before it's fully
accepted. This allows inspecting the peer address before starting the
post-handshake driver.

The Server::accept() method now returns Incoming instead of Connection,
and callers must call incoming.accept() to get the Connection.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
---
 web-transport-quiche/src/ez/server.rs | 30 ++++++++++++++++++++++-----
 web-transport-quiche/src/server.rs    |  5 ++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/web-transport-quiche/src/ez/server.rs b/web-transport-quiche/src/ez/server.rs
index f75095a..8ae57f6 100644
--- a/web-transport-quiche/src/ez/server.rs
+++ b/web-transport-quiche/src/ez/server.rs
@@ -196,9 +196,28 @@ impl<M: Metrics> ServerBuilder<M, ServerWithListener> {
     }
 }
 
+/// A pre-accepted QUIC connection with the TLS handshake already complete.
+///
+/// The peer address is available before calling [Incoming::accept].
+pub struct Incoming {
+    connection: Connection,
+}
+
+impl Incoming {
+    /// Returns the peer's socket address.
+    pub fn peer_addr(&self) -> SocketAddr {
+        self.connection.peer_addr()
+    }
+
+    /// Accept the connection, starting the post-handshake driver.
+    pub fn accept(self) -> Connection {
+        self.connection
+    }
+}
+
 /// A QUIC server that accepts new connections.
 pub struct Server<M: Metrics = DefaultMetrics> {
-    accept: mpsc::Receiver<Connection>,
+    accept: mpsc::Receiver<Incoming>,
     local_addrs: Vec<SocketAddr>,
     // Cancels socket tasks when dropped.
     #[allow(dead_code)]
@@ -231,7 +250,7 @@ impl<M: Metrics> Server<M> {
 
     async fn run_socket(
         socket: tokio_quiche::QuicConnectionStream<M>,
-        accept: mpsc::Sender<Connection>,
+        accept: mpsc::Sender<Incoming>,
     ) -> io::Result<()> {
         let mut rx = socket.into_inner();
         while let Some(initial) = rx.recv().await {
@@ -245,8 +264,9 @@ impl<M: Metrics> Server<M> {
 
             let inner = initial.start(session);
             let connection = Connection::new(inner, state, accept_bi.1, accept_uni.1);
+            let incoming = Incoming { connection };
 
-            if accept.send(connection).await.is_err() {
+            if accept.send(incoming).await.is_err() {
                 return Ok(());
             }
         }
@@ -254,10 +274,10 @@ impl<M: Metrics> Server<M> {
         Ok(())
     }
 
-    /// Accept a new QUIC [Connection] from a client.
+    /// Accept a new QUIC [Incoming] from a client.
     ///
     /// Returns `None` when the server is shutting down.
-    pub async fn accept(&mut self) -> Option<Connection> {
+    pub async fn accept(&mut self) -> Option<Incoming> {
         self.accept.recv().await
     }
 
diff --git a/web-transport-quiche/src/server.rs b/web-transport-quiche/src/server.rs
index ef72ee2..f5e2b14 100644
--- a/web-transport-quiche/src/server.rs
+++ b/web-transport-quiche/src/server.rs
@@ -157,7 +157,10 @@ impl<M: ez::Metrics> Server<M> {
     pub async fn accept(&mut self) -> Option<h3::Request> {
         loop {
             tokio::select! {
-                Some(conn) = self.inner.accept() => self.accept.push(Box::pin(h3::Request::accept(conn))),
+                Some(incoming) = self.inner.accept() => {
+                    let conn = incoming.accept();
+                    self.accept.push(Box::pin(h3::Request::accept(conn)));
+                }
                 Some(res) = self.accept.next() => {
                     match res {
                         Ok(session) => return Some(session),