diff --git a/Sources/SocketForwarder/UDPForwarder.swift b/Sources/SocketForwarder/UDPForwarder.swift index 031e8dc15..4a8128166 100644 --- a/Sources/SocketForwarder/UDPForwarder.swift +++ b/Sources/SocketForwarder/UDPForwarder.swift @@ -22,7 +22,7 @@ import NIOFoundationCompat import Synchronization // Proxy backend for a single client address (clientIP, clientPort). -private final class UDPProxyBackend: ChannelInboundHandler { +final class UDPProxyBackend: ChannelInboundHandler { typealias InboundIn = AddressedEnvelope typealias OutboundOut = AddressedEnvelope @@ -35,6 +35,7 @@ private final class UDPProxyBackend: ChannelInboundHandler { private struct State { var queuedPayloads: Deque var channel: (any Channel)? + var closed: Bool } private let clientAddress: SocketAddress @@ -48,7 +49,7 @@ private final class UDPProxyBackend: ChannelInboundHandler { self.serverAddress = serverAddress self.frontendChannel = frontendChannel self.log = log - let initialState = State(queuedPayloads: Deque(), channel: nil) + let initialState = State(queuedPayloads: Deque(), channel: nil, closed: false) self.state = initialState } @@ -61,6 +62,16 @@ private final class UDPProxyBackend: ChannelInboundHandler { } func channelActive(context: ChannelHandlerContext) { + guard !state.closed else { + // close() ran while this channel was still binding, so there was no channel to + // close at the time. Close it now rather than adopting it: the backend has + // already been evicted from the proxy cache, so no reference remains that could + // close it later. + self.log?.trace("backend - closing channel that became active after close") + state.queuedPayloads.removeAll() + context.channel.close(promise: nil) + return + } if !state.queuedPayloads.isEmpty { self.log?.trace("backend - writing \(state.queuedPayloads.count) queued datagrams to server") while let queuedData = state.queuedPayloads.popFirst() { @@ -88,8 +99,9 @@ private final class UDPProxyBackend: ChannelInboundHandler { } func close() { + state.closed = true guard let channel = state.channel else { - self.log?.warning("backend - close on inactive channel") + self.log?.trace("backend - close requested before the channel became active") return } _ = channel.close() diff --git a/Tests/SocketForwarderTests/UDPForwarderTest.swift b/Tests/SocketForwarderTests/UDPForwarderTest.swift index 7ace24299..6041fa470 100644 --- a/Tests/SocketForwarderTests/UDPForwarderTest.swift +++ b/Tests/SocketForwarderTests/UDPForwarderTest.swift @@ -23,6 +23,33 @@ import Testing struct UDPForwarderTest { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount) + /// A backend evicted from the proxy cache before its channel finished binding used to + /// have no channel to close, and nothing closed the socket once it became active. + @Test + func testBackendClosedBeforeChannelActiveClosesTheChannel() throws { + let clientAddress = try SocketAddress(ipAddress: "127.0.0.1", port: 12345) + let serverAddress = try SocketAddress(ipAddress: "127.0.0.1", port: 54321) + let frontendChannel = EmbeddedChannel() + + let backend = UDPProxyBackend( + clientAddress: clientAddress, + serverAddress: serverAddress, + frontendChannel: frontendChannel, + log: nil + ) + + // Evicted while the backend channel is still binding, so there is no channel yet. + backend.close() + + // The bind then completes and the channel becomes active. + let backendChannel = EmbeddedChannel(handler: backend) + try backendChannel.connect(to: serverAddress).wait() + + #expect(!backendChannel.isActive) + _ = try? backendChannel.finish() + _ = try? frontendChannel.finish() + } + @Test func testUDPForwarder() async throws { let requestCount = 100