feat: allow endpoint limits to pass in connection closed reason

aws · boquan-fang · Sep 17, 2025 · Sep 4, 2025 · Sep 5, 2025 · Sep 6, 2025
commit 4f884aca8e889ae2b48a40e7db01d885b74bbda8
diff --git a/quic/s2n-quic-core/src/endpoint/limits.rs b/quic/s2n-quic-core/src/endpoint/limits.rs
@@ -32,9 +32,12 @@ pub enum Outcome {
 
     /// Cleanly close the connection
     ///
-    /// Use `Outcome::close()` to construct this variant
+    /// Use `Outcome::close()` or `Outcome::close_with_reason()` to construct this variant
     #[non_exhaustive]
-    Close,
+    Close {
+        /// An optional reason message to include in the CONNECTION_CLOSE frame
+        reason: Option<&'static [u8]>,
+    },
 }
 
 impl Outcome {
@@ -54,8 +57,10 @@ impl Outcome {
     }
 
     /// Cleanly close the connection
-    pub fn close() -> Self {
-        Self::Close
+    pub fn close(reason: &'static [u8]) -> Self {
+        Self::Close {
+            reason: Some(reason),
+        }
     }
 }
 

@@ -5,19 +5,21 @@ use super::*;
 use s2n_quic::provider::endpoint_limits::{ConnectionAttempt, Limiter, Outcome};
 use s2n_quic_core::{connection::error::Error, endpoint};
 
-/// A custom limiter that always returns Outcome::close()
-struct AlwaysCloseLimiter;
+static REASON: &str = "Always close connections for testing purpose";
 
-impl Limiter for AlwaysCloseLimiter {
+/// A custom limiter that always returns Outcome::close() with a reason
+struct AlwaysCloseLimiterWithReason;
+
+impl Limiter for AlwaysCloseLimiterWithReason {
     fn on_connection_attempt(&mut self, _info: &ConnectionAttempt) -> Outcome {
-        Outcome::close()
+        Outcome::close(REASON.as_bytes())
     }
 }
 
 // This test verifies that when the server would send a CONNECTION_CLOSE frame with
 // error code CONNECTION_REFUSED when the server's limiter returns Outcome::close().
 #[test]
-fn endpoint_limits_test() {
+fn endpoint_limits_close_test() {
     let model = Model::default();
 
     let connection_close_subscriber = recorder::ConnectionClosed::new();
@@ -29,7 +31,7 @@ fn endpoint_limits_test() {
             .with_tls(SERVER_CERTS)?
             .with_event(tracing_events())?
             .with_random(Random::with_seed(456))?
-            .with_endpoint_limits(AlwaysCloseLimiter)?
+            .with_endpoint_limits(AlwaysCloseLimiterWithReason)?
             .start()?;
 
         let server_addr = start_server(server)?;
@@ -64,7 +66,8 @@ fn endpoint_limits_test() {
         Error::Transport {
             code,
             initiator,
+            reason,
             ..
-        } if (code == s2n_quic_core::transport::Error::CONNECTION_REFUSED.code && initiator == endpoint::Location::Remote)
+        } if (code == s2n_quic_core::transport::Error::CONNECTION_REFUSED.code && initiator == endpoint::Location::Remote && reason == "")
     ));
 }
@@ -18,6 +18,8 @@ use s2n_quic_core::{
     varint::VarInt,
 };
 
+static DEFAULT_REASON: &str = "The server's limiter refused the connection";
+
 #[derive(Debug)]
 pub struct Dispatch<Path: path::Handle> {
     transmissions: VecDeque<Transmission<Path>>,
@@ -41,10 +43,12 @@ impl<Path: path::Handle> Dispatch<Path> {
         path_handle: Path,
         packet: &s2n_quic_core::packet::initial::ProtectedInitial,
         local_connection_id: connection::LocalId,
+        reason: Option<&[u8]>,
     ) where
         <C as InitialKey>::HeaderKey: InitialHeaderKey,
     {
-        if let Some(transmission) = Transmission::new::<C>(path_handle, packet, local_connection_id)
+        if let Some(transmission) =
+            Transmission::new::<C>(path_handle, packet, local_connection_id, reason)
         {
             self.transmissions.push_back(transmission);
         }
@@ -102,6 +106,7 @@ impl<Path: path::Handle> Transmission<Path> {
         path: Path,
         packet: &s2n_quic_core::packet::initial::ProtectedInitial,
         local_connection_id: connection::LocalId,
+        reason: Option<&[u8]>,
     ) -> Option<Self>
     where
         <C as InitialKey>::HeaderKey: InitialHeaderKey,
@@ -124,7 +129,7 @@ impl<Path: path::Handle> Transmission<Path> {
         let connection_close = ConnectionClose {
             error_code: transport::Error::CONNECTION_REFUSED.code.as_varint(),
             frame_type: Some(VarInt::ZERO),
-            reason: Some(b"The server's limiter refused the connection"),
+            reason: reason.or(Some(DEFAULT_REASON.as_bytes())),
         };
 
         let mut encoded_frame = connection_close.encode_to_vec();

@@ -390,7 +390,7 @@ impl<Cfg: Config> Endpoint<Cfg> {
 
                 None
             }
-            Outcome::Close { .. } => {
+            Outcome::Close { reason, .. } => {
                 //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
                 //# If a server refuses to accept a new connection, it SHOULD send an
                 //# Initial packet containing a CONNECTION_CLOSE frame with error code
@@ -406,6 +406,7 @@ impl<Cfg: Config> Endpoint<Cfg> {
                     header.path,
                     packet,
                     local_connection_id,
+                    reason,
                 );
 
                 publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {