diff --git a/Sources/NIOHTTPServer/Datagrams/HTTP3ConnectionManager.swift b/Sources/NIOHTTPServer/Datagrams/HTTP3ConnectionManager.swift index 43dd784..9a38957 100644 --- a/Sources/NIOHTTPServer/Datagrams/HTTP3ConnectionManager.swift +++ b/Sources/NIOHTTPServer/Datagrams/HTTP3ConnectionManager.swift @@ -61,32 +61,16 @@ final class HTTP3ConnectionManager: ChannelInboundHandler { } mutating func finish() { - // The connection is going away, so datagrams will never be negotiated if they haven't been already. - self.completeNegotiation(datagramsSupported: false) - for datagramStream in self.datagramStreams.values { datagramStream.finish() } self.datagramStreams.removeAll() } - mutating func receivedSettings(datagramsSupported: Bool) { - self.completeNegotiation(datagramsSupported: datagramsSupported) - } - - /// Completes the negotiation promise, if it has not been completed already. - private mutating func completeNegotiation(datagramsSupported: Bool) { - guard let promise = self.datagramsNegotiatedPromise else { - // We fulfilled the promise earlier. - return - } - self.datagramsNegotiatedPromise = nil - - if datagramsSupported { - promise.succeed() - } else { - promise.fail(DatagramsNotSupported()) - } + /// Returns the negotiation promise if it has not been taken already. + mutating func takeNegotiationPromise() -> EventLoopPromise? { + defer { self.datagramsNegotiatedPromise = nil } + return self.datagramsNegotiatedPromise } } @@ -120,6 +104,20 @@ final class HTTP3ConnectionManager: ChannelInboundHandler { self.eventLoop.preconditionInEventLoop() self.datagramContext?.deregister(streamID: streamID) } + + /// Completes the negotiation promise, if it has not been completed already. + private func completeNegotiation(datagramsSupported: Bool) { + guard let promise = self.datagramContext?.takeNegotiationPromise() else { + // We fulfilled the promise earlier. + return + } + + if datagramsSupported { + promise.succeed() + } else { + promise.fail(DatagramsNotSupported()) + } + } #endif // UnstableHTTPDatagrams func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -135,7 +133,7 @@ final class HTTP3ConnectionManager: ChannelInboundHandler { switch event { case let event as ReceivedSettings: #if UnstableHTTPDatagrams - self.datagramContext?.receivedSettings(datagramsSupported: event.datagramsSupported) + self.completeNegotiation(datagramsSupported: event.datagramsSupported) #endif context.fireUserInboundEventTriggered(event) @@ -151,6 +149,8 @@ final class HTTP3ConnectionManager: ChannelInboundHandler { func channelInactive(context: ChannelHandlerContext) { #if UnstableHTTPDatagrams + // The connection is going away, so datagrams will never be negotiated if they haven't been already. + self.completeNegotiation(datagramsSupported: false) self.datagramContext?.finish() self.datagramContext = nil #endif diff --git a/Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift b/Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift index ef4f08e..82fd092 100644 --- a/Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift +++ b/Tests/NIOHTTPServerTests/HTTP3DatagramTests.swift @@ -113,6 +113,63 @@ struct HTTP3DatagramTests { #expect(try await TestHelpers.readDatagram(&secondReader) == nil) } + @Test("A stream opened before SETTINGS can register once datagrams are negotiated") + @available(anyAppleOS 27.0, *) + func streamOpenedBeforeSettingsRegistersOnNegotiation() async throws { + let channel = EmbeddedChannel() + let datagramsNegotiatedPromise = channel.eventLoop.makePromise(of: Void.self) + let manager = HTTP3ConnectionManager( + eventLoop: channel.eventLoop, + logger: Self.serverLogger, + datagramsNegotiatedPromise: datagramsNegotiatedPromise + ) + try channel.pipeline.syncOperations.addHandler(manager) + let loopBoundManager = NIOLoopBound(manager, eventLoop: channel.eventLoop) + + // Register from the negotiation callback. + let datagramStreamFuture = datagramsNegotiatedPromise.futureResult.map { + let stream = HTTP3UnreliableDatagramStream( + streamID: 0, + connectionChannel: channel, + maxBufferedDatagrams: 16 + ) + loopBoundManager.value.register(datagramStream: stream) + return stream + } + + channel.pipeline.syncOperations.fireUserInboundEventTriggered(ReceivedSettings(datagramsSupported: true)) + + let stream = try await datagramStreamFuture.get() + var reader = NIOHTTPServer.DatagramReader(iterator: stream.inbound.makeAsyncIterator()) + try channel.writeInbound(HTTP3Datagram(streamID: 0, payload: ByteBuffer([1]))) + + #expect(try await TestHelpers.readDatagram(&reader) == [1]) + } + + @Test("A stream opened before SETTINGS can deregister when the connection closes") + @available(anyAppleOS 27.0, *) + func streamOpenedBeforeSettingsDeregistersOnClose() throws { + let channel = EmbeddedChannel() + let datagramsNegotiatedPromise = channel.eventLoop.makePromise(of: Void.self) + let manager = HTTP3ConnectionManager( + eventLoop: channel.eventLoop, + logger: Self.serverLogger, + datagramsNegotiatedPromise: datagramsNegotiatedPromise + ) + try channel.pipeline.syncOperations.addHandler(manager) + let loopBoundManager = NIOLoopBound(manager, eventLoop: channel.eventLoop) + + // Deregister from the negotiation callback. + datagramsNegotiatedPromise.futureResult.whenComplete { _ in + loopBoundManager.value.deregister(streamID: 0) + } + + #expect(try channel.finish().isClean) + #expect(throws: HTTP3ConnectionManager.DatagramsNotSupported.self) { + try datagramsNegotiatedPromise.futureResult.wait() + } + } + @Test("The oldest datagram is dropped once the buffer is full") @available(anyAppleOS 27.0, *) func oldestDatagramIsDroppedOnceTheBufferIsFull() async throws {