Add CloseableConnection capability

Sources/NIOHTTPServer/HTTPKeepAliveHandler.swift

@@ -12,11 +12,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+import NIOConcurrencyHelpers
 import NIOCore
 import NIOHTTPTypes
 
-/// A NIO channel handler that ensures HTTP/1.1 keep-alive semantics are honored when
-/// the server starts writing a response before the request body has been fully read.
+/// A NIO channel handler that ensures HTTP/1.1 keep-alive semantics are honored
+/// when the server starts writing a response before the request body has been
+/// fully read, and that reacts to the connection-scoped close flag by amending
+/// the next response head with `Connection: close` and closing the channel after
+/// the response is written.
 ///
 /// The handler buffers final response parts (head + any body fragments + end) when
 /// they are written before the request `.end` has been received. The buffer is
@@ -26,13 +30,15 @@ import NIOHTTPTypes
 /// - **`flush`**: an upstream writer (e.g. `NIOAsyncChannelOutboundWriter`) forced a
 ///   flush.
 ///
-/// At each deadline, if request `.end` has arrived, the buffer is flushed as-is and
-/// the connection is reusable. If request `.end` has *not* arrived, the head is
-/// amended with `Connection: close`, the buffer is flushed, and the connection is
-/// closed once response `.end` is written. This protects against clients that keep
-/// uploading request body bytes after the response has completed (which would
-/// otherwise force the server to drain unbounded data) and gives the client an
-/// explicit signal not to pipeline another request on the connection.
+/// At each deadline, if request `.end` has arrived and no close has been
+/// signalled, the buffer is flushed as-is and the connection is reusable. If
+/// request `.end` has *not* arrived, or the shared close flag has been set, the
+/// head is amended with `Connection: close`, the buffer is flushed, and the
+/// connection is closed once response `.end` is written. This protects against
+/// clients that keep uploading request body bytes after the response has
+/// completed (which would otherwise force the server to drain unbounded data)
+/// and gives the client an explicit signal not to pipeline another request on
+/// the connection.
 ///
 /// Informational (1xx) responses pass through unchanged and do not affect buffering
 /// state.
@@ -64,32 +70,59 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
         case streaming
     }
 
-    /// `true` when the request `.end` has been received on the inbound side, or no
-    /// request is currently in flight. `false` between receiving a request `.head`
-    /// and its `.end`.
-    private var requestEndReceived: Bool = true
-
-    /// `true` if we've committed to closing the connection after this response's
-    /// `.end` is written. Set when the buffer is flushed while request `.end` has
-    /// not yet arrived (so we add `Connection: close`). Cleared when a new request
-    /// begins.
+    /// `true` between receiving a request `.head` and its `.end` — i.e. while the
+    /// client may still be streaming request body bytes. `false` otherwise: on a
+    /// fresh handler with no request yet, between requests, and after a request's
+    /// `.end` has arrived. The final-response decision hinges on this: writing a
+    /// final response head while a body is in flight risks HTTP/1.1 framing
+    /// desync on keep-alive, so those writes must be buffered until the flush
+    /// deadline (`channelReadComplete` / `flush`) and, if the body still hasn't
+    /// completed by then, sent with `Connection: close`.
+    private var requestBodyInFlight: Bool = false
+
+    /// `true` if the response head that has been (or will be) written for the
+    /// current request carries `Connection: close` — either because the close
+    /// flag was set before the head reached the wire, or because the head was
+    /// flushed while a request body was still in flight. When true, response
+    /// `.end` triggers a channel close.
+    ///
+    /// If `closeSignalled` becomes true only *after* the head has been streamed
+    /// (so the header could no longer be amended), this flag stays `false` —
+    /// the close is still honoured, but via a direct `closeSignalled` check on
+    /// response `.end` rather than through this flag.
+    ///
+    /// Cleared when a new request begins.
     private var closeAfterResponseEnd: Bool = false
 
     private var finalResponseState: FinalResponseState = .notStarted
 
+    /// Shared close flag. ``NIOHTTPServer/ConnectionContext/signalConnectionClose()``
+    /// sets this synchronously; the handler observes it when writing the next
+    /// response head and amends with `Connection: close`, then closes the channel
+    /// after the response `.end`.
+    let closeFlag: NIOLockedValueBox<Bool>
+
+    init(closeFlag: NIOLockedValueBox<Bool>) {
+        self.closeFlag = closeFlag
+    }
+
+    private var closeSignalled: Bool {
+        self.closeFlag.withLockedValue { $0 }
+    }
+
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         let part = self.unwrapInboundIn(data)
         switch part {
         case .head:
             // Begin a new request. (Any previous request's response must have
             // completed already since HTTPServerPipelineHandler enforces ordering.)
-            self.requestEndReceived = false
+            self.requestBodyInFlight = true
             self.closeAfterResponseEnd = false
             self.finalResponseState = .notStarted
         case .body:
             break
         case .end:
-            self.requestEndReceived = true
+            self.requestBodyInFlight = false
         }
         context.fireChannelRead(data)
     }
@@ -100,8 +133,8 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
         // `Connection: close`. If request `.end` arrived during the cycle the head
         // is flushed unchanged; otherwise we amend the head and close after
         // response `.end`.
-        if case .buffering = self.finalResponseState {
-            self.flushBuffer(context: context)
+        if case let .buffering(head, additional) = self.finalResponseState {
+            self.flushBuffer(context: context, head: head, additional: additional)
         }
         context.fireChannelReadComplete()
     }
@@ -117,10 +150,21 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
                 context.write(data, promise: promise)
                 return
             }
-            if self.requestEndReceived {
-                // Request fully read; stream the response directly.
+            if !self.requestBodyInFlight {
+                // No request body in flight; stream the response directly. If a
+                // close signal has been observed, amend the head with
+                // `Connection: close` and arrange to close once response `.end`
+                // is written.
+                let outboundPart: HTTPResponsePart
+                if self.closeSignalled, case .head(var response) = part {
+                    response.headerFields[.connection] = "close"
+                    outboundPart = .head(response)
+                    self.closeAfterResponseEnd = true
+                } else {
+                    outboundPart = part
+                }
                 self.finalResponseState = .streaming
-                context.write(data, promise: promise)
+                context.write(self.wrapOutboundOut(outboundPart), promise: promise)
             } else {
                 // Start buffering with the head. Additional parts (body, end) the
                 // handler may write before the next deadline are appended below.
@@ -134,9 +178,13 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
             self.finalResponseState = .buffering(head: head, additional: additional)
         case .streaming:
             context.write(data, promise: promise)
-            if case .end = part, self.closeAfterResponseEnd {
-                // The head we flushed earlier carried `Connection: close`; close
-                // the connection now that the response is complete.
+            if case .end = part, self.closeAfterResponseEnd || self.closeSignalled {
+                // Either the head we flushed earlier carried `Connection: close`
+                // (`closeAfterResponseEnd`), or a close signal arrived after the
+                // head was on the wire (`closeSignalled`). In the latter case
+                // the header can't be added retroactively, but we still honour
+                // the caller's contract and close the connection now that the
+                // response is complete.
                 context.flush()
                 context.close(mode: .output, promise: nil)
             }
@@ -145,21 +193,31 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
 
     func flush(context: ChannelHandlerContext) {
         // An upstream writer forced a flush. Same deadline as `channelReadComplete`:
-        // release any buffered parts, with `Connection: close` if request `.end`
-        // hasn't arrived.
-        if case .buffering = self.finalResponseState {
-            self.flushBuffer(context: context)
+        // release any buffered parts, with `Connection: close` if the request body
+        // is still in flight or close was signalled. `flushBuffer` flushes itself,
+        // so we only propagate the incoming flush when there's no buffer to
+        // release — otherwise we'd double-flush.
+        if case let .buffering(head, additional) = self.finalResponseState {
+            self.flushBuffer(context: context, head: head, additional: additional)
+        } else {
+            context.flush()
         }
-        context.flush()
     }
 
-    /// Releases buffered response parts to the pipeline. If request `.end` has not
-    /// yet arrived, amend the head with `Connection: close` and arrange to close
-    /// the connection once response `.end` is written.
-    private func flushBuffer(context: ChannelHandlerContext) {
-        guard case .buffering(var head, let additional) = self.finalResponseState else { return }
-
-        if !self.requestEndReceived {
+    /// Releases the given buffered response parts to the pipeline. If a request
+    /// body is still in flight, or if the shared close flag is set, amends the
+    /// head with `Connection: close` and arranges to close the connection once
+    /// response `.end` is written. Called only by the two flush deadlines
+    /// (`channelReadComplete` and `flush`), each of which supplies the buffer
+    /// contents extracted from the `.buffering` state.
+    private func flushBuffer(
+        context: ChannelHandlerContext,
+        head: BufferedWrite,
+        additional: [BufferedWrite]
+    ) {
+        var head = head
+
+        if self.requestBodyInFlight || self.closeSignalled {
             // Amend the head with `Connection: close` before flushing.
             if case .head(var response) = head.part {
                 response.headerFields[.connection] = "close"
@@ -180,10 +238,12 @@ final class HTTPKeepAliveHandler: ChannelDuplexHandler {
         }
         context.flush()
 
-        if sawEnd && self.closeAfterResponseEnd {
+        if sawEnd && (self.closeAfterResponseEnd || self.closeSignalled) {
             // The response was fully buffered (head + ... + end) and we have to
-            // close. Close now (the flush above ensured the writes reached the
-            // wire).
+            // close — either we amended the head with `Connection: close` at the
+            // top of this method (`closeAfterResponseEnd`), or a close signal
+            // arrived between the head decision and here (`closeSignalled`).
+            // Close now (the flush above ensured the writes reached the wire).
             context.close(mode: .output, promise: nil)
         }
     }

Sources/NIOHTTPServer/HTTPServerCapability+ConnectionCapabilities.swift

@@ -41,4 +41,21 @@ extension HTTPServerCapability {
         /// The peer's validated certificate chain, when available.
         var peerCertificateChain: X509.ValidatedCertificateChain? { get async throws }
     }
+
+    /// A request-context capability that lets a request handler signal that the
+    /// connection should close after the current response.
+    ///
+    /// Servers whose request context conforms to this capability allow handlers
+    /// to indicate that the underlying connection should be closed once the
+    /// in-flight response has been sent. Implementations make the signal
+    /// effective on a best-effort basis (HTTP/1.1 typically appends
+    /// `Connection: close` and closes the channel; HTTP/2 typically sends
+    /// `GOAWAY` and lets in-flight streams complete normally).
+    public protocol CloseableConnection: RequestContext {
+        /// Signal that the connection should close after the current response.
+        ///
+        /// Non-blocking and idempotent. Subsequent calls have no additional
+        /// effect.
+        func signalConnectionClose()
+    }
 }

Sources/NIOHTTPServer/NIOHTTPServer+Connection.swift

@@ -23,7 +23,8 @@ extension NIOHTTPServer {
     /// A `Connection` is owned by exactly one ``NIOHTTPServerConnectionHandler/handleConnection(connection:context:)``
     /// invocation. The handler typically passes it to ``handleRequests(handler:)``,
     /// which runs the request loop until the peer closes the connection, the
-    /// server shuts down, or an error occurs.
+    /// server shuts down, the handler signals close via
+    /// ``ConnectionContext/signalConnectionClose()``, or an error occurs.
     ///
     /// A handler that decides to terminate before any request can simply
     /// return without calling ``handleRequests(handler:)``: the `Connection`
@@ -63,6 +64,7 @@ extension NIOHTTPServer {
         ///
         /// Each request received on this connection is dispatched to `handler`. The
         /// loop returns when the peer closes the connection, the server shuts down,
+        /// the handler signals close via ``ConnectionContext/signalConnectionClose()``,
         /// or an error occurs.
         public consuming func handleRequests<Handler: HTTPServerRequestHandler>(
             handler: Handler

Sources/NIOHTTPServer/NIOHTTPServer+ConnectionContext.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import NIOConcurrencyHelpers
 import NIOCore
 import NIOSSL
 public import X509
@@ -30,12 +31,14 @@ extension NIOHTTPServer {
     /// Connection-scoped state.
     ///
     /// Carries connection-scoped data such as the negotiated HTTP version, the
-    /// peer / local addresses, and the peer's validated certificate chain (when
-    /// applicable).
+    /// peer / local addresses, the peer's validated certificate chain (when
+    /// applicable), and the close-signal trigger that
+    /// ``signalConnectionClose()`` invokes.
     ///
     /// User code accesses this state via the corresponding ``RequestContext``
     /// capabilities (``HTTPServerCapability/ConnectionInfo``,
-    /// ``HTTPServerCapability/PeerCertificate``) when handling individual
+    /// ``HTTPServerCapability/PeerCertificate``,
+    /// ``HTTPServerCapability/CloseableConnection``) when handling individual
     /// requests, and directly when implementing an
     /// ``NIOHTTPServerConnectionHandler``.
     public struct ConnectionContext: Sendable {
@@ -50,16 +53,31 @@ extension NIOHTTPServer {
 
         var peerCertificateChainFuture: EventLoopFuture<NIOSSL.ValidatedCertificateChain?>?
 
+        /// Per-protocol close mechanism. HTTP/1.1 uses a shared atomic flag the
+        /// channel's ``HTTPKeepAliveHandler`` reads when writing the next
+        /// response head. HTTP/2 fires `ChannelShouldQuiesceEvent` on the
+        /// connection channel; NIO's `NIOHTTP2ServerConnectionManagementHandler`
+        /// reacts by sending `GOAWAY`, letting in-flight streams complete
+        /// normally, and finally closing the connection.
+        let closeBacking: CloseBacking
+
+        enum CloseBacking: Sendable {
+            case http1_1(closeFlag: NIOLockedValueBox<Bool>)
+            case http2(connectionChannel: any Channel)
+        }
+
         init(
             httpVersion: HTTPVersion,
             remoteAddress: NIOHTTPServer.SocketAddress? = nil,
             localAddress: NIOHTTPServer.SocketAddress? = nil,
-            peerCertificateChainFuture: EventLoopFuture<NIOSSL.ValidatedCertificateChain?>? = nil
+            peerCertificateChainFuture: EventLoopFuture<NIOSSL.ValidatedCertificateChain?>? = nil,
+            closeBacking: CloseBacking
         ) {
             self.httpVersion = httpVersion
             self.remoteAddress = remoteAddress
             self.localAddress = localAddress
             self.peerCertificateChainFuture = peerCertificateChainFuture
+            self.closeBacking = closeBacking
         }
 
         /// The peer's validated certificate chain. Returns `nil` if a custom
@@ -74,5 +92,22 @@ extension NIOHTTPServer {
                 return nil
             }
         }
+
+        /// Signal that the connection should close after the current response.
+        ///
+        /// - On HTTP/1.1, the response carries `Connection: close` and the channel is closed once the response has
+        /// been written. Note that amending the response with the `close` header will only happen if this function is called
+        /// before the response `head` has been flushed to the wire. Either way the connection will be closed after the current
+        /// response, but it may not carry the `Connection: close` header.
+        /// - On HTTP/2, the connection sends `GOAWAY`; in-flight streams
+        ///   complete normally before the connection closes.
+        public func signalConnectionClose() {
+            switch self.closeBacking {
+            case .http1_1(let closeFlag):
+                closeFlag.withLockedValue { $0 = true }
+            case .http2(let connectionChannel):
+                connectionChannel.pipeline.fireUserInboundEventTriggered(ChannelShouldQuiesceEvent())
+            }
+        }
     }
 }