Add a hook to customise client connection pipeline in tests

Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift

@@ -220,13 +220,15 @@ struct HTTP3DatagramTests {
     @Test("Read and write datagrams")
     @available(anyAppleOS 27.0, *)
     func readAndWriteDatagrams() async throws {
-        let (server, clientConfiguration) = try TestHelpers.makeServerAndClientConfiguration(
+        var (server, clientConfiguration) = try TestHelpers.makeServerAndClientConfiguration(
             for: .http3,
             clientLogger: Self.clientLogger,
             serverLogger: Self.serverLogger
         )
 
         let streamOpenedPromise = server.eventLoopGroup.any().makePromise(of: Void.self)
+        let datagramsSupportedPromise = server.eventLoopGroup.any().makePromise(of: Void.self)
+        clientConfiguration.addDatagramSettingsHandler(promise: datagramsSupportedPromise)
 
         try await TestHelpers.withHTTP3ClientServerConnectionAndRequestChannel(
             clientConfiguration: clientConfiguration,
@@ -255,8 +257,9 @@ struct HTTP3DatagramTests {
                     // Start the request.
                     try await streamOutbound.write(.testHead(method: .post, for: .http3))
 
-                    // Wait for the stream to be opened.
-                    try await streamOpenedPromise.futureResult.get()
+                    // Wait for the stream to be opened, and until the server has sent their SETTINGS frame that
+                    // confirms they also support receiving datagrams.
+                    _ = try await streamOpenedPromise.futureResult.and(datagramsSupportedPromise.futureResult).get()
 
                     // Now write a datagram.
                     try await connectionOutbound.write(HTTP3Datagram(streamID: streamID, payload: .testData))
@@ -296,6 +299,8 @@ struct HTTP3DatagramTests {
         )
 
         let streamOpenedPromise = server.eventLoopGroup.any().makePromise(of: Void.self)
+        let datagramsSupportedPromise = server.eventLoopGroup.any().makePromise(of: Void.self)
+        clientConfiguration.addDatagramSettingsHandler(promise: datagramsSupportedPromise)
 
         try await TestHelpers.withHTTP3ClientServerConnectionAndRequestChannel(
             clientConfiguration: clientConfiguration,
@@ -336,8 +341,9 @@ struct HTTP3DatagramTests {
                     // Start the request.
                     try await streamOutbound.write(.testHead(method: .get, for: .http3))
 
-                    // Wait for the stream to be opened.
-                    try await streamOpenedPromise.futureResult.get()
+                    // Wait for the stream to be opened, and until the server has sent their SETTINGS frame that
+                    // confirms they also support receiving datagrams.
+                    _ = try await streamOpenedPromise.futureResult.and(datagramsSupportedPromise.futureResult).get()
 
                     // Write a datagram.
                     try await connectionOutbound.write(HTTP3Datagram(streamID: streamID, payload: .testData))

Tests/NIOHTTPServerTests/TestError.swift

@@ -19,4 +19,10 @@ enum TestError: Error {
 
     /// Thrown when the client configuration is invalid for the configured HTTP version.
     case invalidClientConfiguration
+
+    /// Datagrams are not supported.
+    case datagramsNotSupported
+
+    /// The connection closed before the peer's SETTINGS frame was received.
+    case connectionClosedBeforeSettingsReceived
 }

Tests/NIOHTTPServerTests/Utilities/NIOClient/NIOClient+HTTP1.swift

@@ -21,10 +21,11 @@ import NIOPosix
 
 @available(anyAppleOS 27.0, *)
 extension Channel {
-    /// Adds HTTP/1.1 client handlers to the pipeline.
+    /// Adds HTTP/1.1 client handlers to the pipeline, then calls `additionalConnectionChannelInitializer` if provided.
     func configureTestHTTP1ClientPipeline(
         responseLeftOverBytesStrategy: RemoveAfterUpgradeStrategy = .dropBytes,
-        informationalResponseStrategy: NIOInformationalResponseStrategy = .forward
+        informationalResponseStrategy: NIOInformationalResponseStrategy = .forward,
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
     ) -> EventLoopFuture<NIOAsyncChannel<HTTPResponsePart, HTTPRequestPart>> {
         self.eventLoop.makeCompletedFuture {
             let handlers: [ChannelHandler] = [
@@ -39,6 +40,7 @@ extension Channel {
                 HTTP1ToHTTPClientCodec(),
             ]
             try self.pipeline.syncOperations.addHandlers(handlers)
+            try additionalConnectionChannelInitializer?(self)
 
             return try NIOAsyncChannel<HTTPResponsePart, HTTPRequestPart>(
                 wrappingChannelSynchronously: self,
@@ -55,14 +57,17 @@ extension ClientBootstrap {
     /// `NIOAsyncChannel` for writing `HTTPRequestPart`s to the server and observing `HTTPResponsePart`s from its
     /// inbound stream.
     func connectToTestHTTP1Server(
-        at serverAddress: NIOHTTPServer.SocketAddress
+        at serverAddress: NIOHTTPServer.SocketAddress,
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
     ) async throws -> TestClientConnection {
         .init(
             connectionProtocol: .http1(
                 connectionChannel: try await self.connect(
                     to: try .init(ipAddress: serverAddress.host, port: serverAddress.port)
                 ) { channel in
-                    channel.configureTestHTTP1ClientPipeline()
+                    channel.configureTestHTTP1ClientPipeline(
+                        additionalConnectionChannelInitializer: additionalConnectionChannelInitializer
+                    )
                 }
             )
         )

Tests/NIOHTTPServerTests/Utilities/NIOClient/NIOClient+HTTP3.swift

@@ -69,7 +69,8 @@ extension Channel {
         settings: HTTP3Settings,
         configuration: HTTP3ClientConfiguration,
         quicConfiguration: QUICConfiguration,
-        asyncVerifier: NIOQUIC.AsyncVerifier
+        asyncVerifier: NIOQUIC.AsyncVerifier,
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
     ) throws -> TestHTTP3SingleConnectionCreator {
         let quicHandler = QUICHandler(
             channel: self,
@@ -96,6 +97,7 @@ extension Channel {
                         inboundPushStreamInitializer: { _ in fatalError() }
                     )
                     try connectionChannel.pipeline.syncOperations.addHandler(h3Handler)
+                    try additionalConnectionChannelInitializer?(connectionChannel)
                     return connectionChannel
                 }
             },
@@ -120,7 +122,8 @@ extension DatagramBootstrap {
         trustRootsPath: String,
         quicConfiguration: QUICConfiguration,
         http3ClientConfiguration: HTTP3ClientConfiguration = .defaults,
-        http3ConnectionSettings: HTTP3Settings = .init()
+        http3ConnectionSettings: HTTP3Settings = .init(),
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
     ) async throws -> (any Channel, NIOLoopBound<TestHTTP3SingleConnectionCreator>) {
         try await self.channelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
             .bind(host: "127.0.0.1", port: 0) { channel in
@@ -134,7 +137,8 @@ extension DatagramBootstrap {
                             trustRootsPath: trustRootsPath,
                             certificateVerification: .noHostnameVerification,
                             eventLoop: channel.eventLoop
-                        )
+                        ),
+                        additionalConnectionChannelInitializer: additionalConnectionChannelInitializer
                     )
                     let loopBoundConnectionCreator = NIOLoopBound(connectionCreator, eventLoop: channel.eventLoop)
 

Tests/NIOHTTPServerTests/Utilities/NIOClient/NIOClient+SecureUpgrade.swift

@@ -32,8 +32,12 @@ extension Channel {
         }
     }
 
-    /// Adds an ALPN handler (configured with both HTTP/1.1 and HTTP/2 channel initializers) to the pipeline.
-    func configureTestSecureUpgradeClientPipeline() -> EventLoopFuture<
+    /// Adds an ALPN handler (configured with both HTTP/1.1 and HTTP/2 channel initializers) to the pipeline. Once a
+    /// protocol is negotiated, `additionalConnectionChannelInitializer` is optionally called after the protocol's
+    /// handlers are added.
+    func configureTestSecureUpgradeClientPipeline(
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
+    ) -> EventLoopFuture<
         EventLoopFuture<
             NIONegotiatedHTTPVersion<
                 NIOAsyncChannel<HTTPResponsePart, HTTPRequestPart>,
@@ -43,10 +47,16 @@ extension Channel {
     > {
         self.configureHTTP2AsyncSecureUpgrade(
             http1ConnectionInitializer: { channel in
-                channel.configureTestHTTP1ClientPipeline()
+                channel.configureTestHTTP1ClientPipeline(
+                    additionalConnectionChannelInitializer: additionalConnectionChannelInitializer
+                )
             },
             http2ConnectionInitializer: { channel in
                 channel.configureAsyncHTTP2Pipeline(mode: .client) { $0.eventLoop.makeSucceededFuture($0) }
+                    .flatMapThrowing { multiplexer in
+                        try additionalConnectionChannelInitializer?(channel)
+                        return multiplexer
+                    }
             }
         )
     }
@@ -57,13 +67,16 @@ extension ClientBootstrap {
     /// Connects the client to the specified address using the provided TLS configuration.
     func connectToTestSecureUpgradeHTTPServer(
         at serverAddress: NIOHTTPServer.SocketAddress,
-        tlsConfig: TLSConfiguration
+        tlsConfig: TLSConfiguration,
+        additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
     ) async throws -> TestClientConnection {
         let (connectionChannel, alpnResultFuture) = try await self.connect(
             to: try .init(ipAddress: serverAddress.host, port: serverAddress.port)
         ) { channel in
             channel.configureTestClientSSLPipeline(tlsConfig: tlsConfig).flatMap {
-                channel.configureTestSecureUpgradeClientPipeline().map { alpnResultFuture in
+                channel.configureTestSecureUpgradeClientPipeline(
+                    additionalConnectionChannelInitializer: additionalConnectionChannelInitializer
+                ).map { alpnResultFuture in
                     (channel, alpnResultFuture)
                 }
             }

Tests/NIOHTTPServerTests/Utilities/TestClientConnection.swift

@@ -26,6 +26,7 @@ import Testing
 
 #if HTTP3
 import HTTP3
+import NIOExtras
 @_spi(HTTP3AsyncInterface) import NIOHTTP3
 import NIOQUIC
 import NIOQUICHelpers
@@ -189,6 +190,10 @@ extension TestClientConnection {
         var trustRootsPEMPath: String?
         var chain: ChainPrivateKeyPair? = nil
 
+        /// An additional initializer called on the client's connection channel after the protocol-specific handlers
+        /// have been added.
+        var additionalConnectionChannelInitializer: (@Sendable (any Channel) throws -> Void)? = nil
+
         #if HTTP3
         /// The client's QUIC configuration. Only applies when ``httpVersion`` is `.http3`.
         var quicConfiguration: QUICConfiguration? = nil
@@ -210,7 +215,10 @@ extension TestClientConnection {
         switch (configuration.httpVersion, configuration.trustRootsPEMPath) {
         case (.plaintextHTTP1_1, .none):
             connection = try await ClientBootstrap(group: .singletonMultiThreadedEventLoopGroup)
-                .connectToTestHTTP1Server(at: serverAddress)
+                .connectToTestHTTP1Server(
+                    at: serverAddress,
+                    additionalConnectionChannelInitializer: configuration.additionalConnectionChannelInitializer
+                )
 
         case (.http1_1, .some(let trustRootsPEMPath)), (.http2, .some(let trustRootsPEMPath)):
             let tlsConfiguration =
@@ -228,7 +236,11 @@ extension TestClientConnection {
                 }
 
             connection = try await ClientBootstrap(group: .singletonMultiThreadedEventLoopGroup)
-                .connectToTestSecureUpgradeHTTPServer(at: serverAddress, tlsConfig: tlsConfiguration)
+                .connectToTestSecureUpgradeHTTPServer(
+                    at: serverAddress,
+                    tlsConfig: tlsConfiguration,
+                    additionalConnectionChannelInitializer: configuration.additionalConnectionChannelInitializer
+                )
 
         #if HTTP3
         case (.http3, .some(let trustRootsPEMPath)):
@@ -238,7 +250,8 @@ extension TestClientConnection {
                     trustRootsPath: trustRootsPEMPath,
                     quicConfiguration: configuration.quicConfiguration
                         ?? .makeClientQUICConfig(caPath: trustRootsPEMPath),
-                    http3ConnectionSettings: configuration.http3ConnectionSettings
+                    http3ConnectionSettings: configuration.http3ConnectionSettings,
+                    additionalConnectionChannelInitializer: configuration.additionalConnectionChannelInitializer
                 )
 
             let multiplexer = HTTP3ClientConnectionMultiplexer<
@@ -366,3 +379,36 @@ extension NIOHTTP2Handler.AsyncStreamMultiplexer<Channel> {
         }
     }
 }
+
+#if HTTP3 && UnstableHTTPDatagrams
+@available(anyAppleOS 27.0, *)
+extension TestClientConnection.Configuration {
+    /// Sets up ``additionalConnectionChannelInitializer`` to add a handler that observes the SETTINGS frame received
+    /// from the server, and succeeds `promise` once the server has confirmed it supports receiving datagrams.
+    mutating func addDatagramSettingsHandler(promise datagramsSupportedPromise: EventLoopPromise<Void>) {
+        let existingInitializer = self.additionalConnectionChannelInitializer
+
+        self.additionalConnectionChannelInitializer = { connectionChannel in
+            try existingInitializer?(connectionChannel)
+
+            let settingsObserver = DebugInboundEventsHandler { event, _ in
+                switch event {
+                case .userInboundEventTriggered(let event as ReceivedSettings):
+                    if event.datagramsSupported {
+                        datagramsSupportedPromise.succeed()
+                    } else {
+                        datagramsSupportedPromise.fail(TestError.datagramsNotSupported)
+                    }
+
+                case .inactive:
+                    datagramsSupportedPromise.fail(TestError.connectionClosedBeforeSettingsReceived)
+
+                default:
+                    ()
+                }
+            }
+            try connectionChannel.pipeline.syncOperations.addHandler(settingsObserver)
+        }
+    }
+}
+#endif  // HTTP3 && UnstableHTTPDatagrams