From 7b57716cd138899232fa95265f41b7c07e5cd435 Mon Sep 17 00:00:00 2001 From: agnosticdev Date: Mon, 19 Feb 2024 13:42:14 -0800 Subject: [PATCH 1/3] Fix for Websocket client upgrade failure Motivation: Fix for issue-2631 Modifications: Make sure the completion handler is sent on handler removed in upgrade failure path Result: Fix for both cases described in the issue. --- .../NIOTypedHTTPClientUpgradeHandler.swift | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/Sources/NIOHTTP1/NIOTypedHTTPClientUpgradeHandler.swift b/Sources/NIOHTTP1/NIOTypedHTTPClientUpgradeHandler.swift index c683b61b3e..4644de9e24 100644 --- a/Sources/NIOHTTP1/NIOTypedHTTPClientUpgradeHandler.swift +++ b/Sources/NIOHTTP1/NIOTypedHTTPClientUpgradeHandler.swift @@ -130,7 +130,19 @@ public final class NIOTypedHTTPClientUpgradeHandler: Ch public func handlerRemoved(context: ChannelHandlerContext) { switch self.stateMachine.handlerRemoved() { case .failUpgradePromise: - self.upgradeResultPromise.fail(ChannelError.inappropriateOperationForState) + // Make sure the completion handler is called on the failed upgrade path + self.notUpgradingCompletionHandler(context.channel) + .hop(to: context.eventLoop) + .whenComplete { result in + switch result { + case .success(let value): + // Expected upgrade failure without error + self.upgradeResultPromise.succeed(value) + case .failure(let error): + // Unexpected upgrade failure with error + self.upgradeResultPromise.fail(error) + } + } case .none: break } From 7676593632476c5dfb482605f61ddf89992e4866 Mon Sep 17 00:00:00 2001 From: agnosticdev Date: Fri, 23 Feb 2024 14:34:19 -0800 Subject: [PATCH 2/3] Added unit test for notUpgradingCompletionHandler Motivation: Review team asked to add a unit test Modifications: Added a new unit test to WebSocketClientEndToEndTests Result: Test the change I made --- .../WebSocketClientEndToEndTests.swift | 97 ++++++++++++++++++- 1 file changed, 96 insertions(+), 1 deletion(-) diff --git a/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift b/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift index 1e64a27544..71ba6a70e9 100644 --- a/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift +++ b/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift @@ -49,6 +49,26 @@ extension ChannelPipeline { } } +private func interactInMemory(_ first: EmbeddedChannel, + _ second: EmbeddedChannel, + eventLoop: EmbeddedEventLoop) throws { + var operated: Bool + + repeat { + eventLoop.run() + operated = false + + if let data = try first.readOutbound(as: ByteBuffer.self) { + operated = true + try second.writeInbound(data) + } + if let data = try second.readOutbound(as: ByteBuffer.self) { + operated = true + try first.writeInbound(data) + } + } while operated +} + private func setUpClientChannel(clientHTTPHandler: RemovableChannelHandler, clientUpgraders: [NIOHTTPClientProtocolUpgrader], _ upgradeCompletionHandler: @escaping (ChannelHandlerContext) -> Void) throws -> EmbeddedChannel { @@ -591,7 +611,7 @@ final class TypedWebSocketClientEndToEndTests: WebSocketClientEndToEndTests { requestKey: "OfS0wDaT5NoxF2gqm7Zj2YtetzM=", upgradePipelineHandler: { (channel: Channel, _: HTTPResponseHead) in channel.pipeline.addHandler(handler) - }) + }) // The process should kick-off independently by sending the upgrade request to the server. let (clientChannel, upgradeResult) = try setUpClientChannel( @@ -615,5 +635,80 @@ final class TypedWebSocketClientEndToEndTests: WebSocketClientEndToEndTests { return (clientChannel, handler) } + + func testSimpleUpgradeRejectedWhenServerSendsUpgradeNil() throws { + + enum TestUpgradeResult: Int { + case successfulUpgrade + case notUpgraded + } + + let serverRecorder = WebSocketRecorderHandler() + let clientRecorder = WebSocketRecorderHandler() + let loop = EmbeddedEventLoop() + let serverChannel = EmbeddedChannel(loop: loop) + let clientChannel = EmbeddedChannel(loop: loop) + + let serverUpgrader = NIOTypedWebSocketServerUpgrader( + shouldUpgrade: { (channel, head) in + channel.eventLoop.makeSucceededFuture(nil) + }, + upgradePipelineHandler: { (channel, req) in + channel.pipeline.addHandler(serverRecorder) + } + ) + + XCTAssertNoThrow(try serverChannel.pipeline.syncOperations.configureUpgradableHTTPServerPipeline( + configuration: .init( + upgradeConfiguration: NIOTypedHTTPServerUpgradeConfiguration( + upgraders: [serverUpgrader], + notUpgradingCompletionHandler: { $0.eventLoop.makeSucceededVoidFuture() } + ) + ) + )) + + let basicClientUpgrader = NIOTypedWebSocketClientUpgrader( + upgradePipelineHandler: { (channel: Channel, _: HTTPResponseHead) in + channel.eventLoop.makeCompletedFuture { + print("s") + return TestUpgradeResult.successfulUpgrade + } + }) + + var headers = HTTPHeaders() + headers.add(name: "Content-Type", value: "text/plain; charset=utf-8") + headers.add(name: "Content-Length", value: "\(0)") + let requestHead = HTTPRequestHead( + version: .http1_1, + method: .GET, + uri: "/", + headers: headers + ) + let config = NIOTypedHTTPClientUpgradeConfiguration( + upgradeRequestHead: requestHead, + upgraders: [basicClientUpgrader], + notUpgradingCompletionHandler: { channel in + channel.eventLoop.makeCompletedFuture { + return TestUpgradeResult.notUpgraded + } + } + ) + let updgradeResult = try clientChannel.pipeline.syncOperations.configureUpgradableHTTPClientPipeline(configuration: .init(upgradeConfiguration: config)) + + XCTAssertNoThrow(try interactInMemory(clientChannel, serverChannel, eventLoop: loop)) + XCTAssertNoThrow(try clientChannel.finish()) + updgradeResult.whenComplete { result in + switch result { + case .success(let value): + XCTAssertTrue(value == .notUpgraded) + case .failure(let error): + XCTFail("There should be no failure here \(error)") + } + } + XCTAssertNoThrow(try serverChannel.finishAcceptingAlreadyClosed()) + XCTAssertEqual(clientRecorder.errors.count, 0) + XCTAssertEqual(serverRecorder.errors.count, 0) + XCTAssertNoThrow(try loop.syncShutdownGracefully()) + } } #endif From 5c3dc9063f434c6dfc5f3beb118542551fbec034 Mon Sep 17 00:00:00 2001 From: agnosticdev Date: Fri, 23 Feb 2024 14:47:51 -0800 Subject: [PATCH 3/3] Removed print Motivation: Random print statement left in --- Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift | 1 - 1 file changed, 1 deletion(-) diff --git a/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift b/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift index 71ba6a70e9..ff52d23507 100644 --- a/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift +++ b/Tests/NIOWebSocketTests/WebSocketClientEndToEndTests.swift @@ -670,7 +670,6 @@ final class TypedWebSocketClientEndToEndTests: WebSocketClientEndToEndTests { let basicClientUpgrader = NIOTypedWebSocketClientUpgrader( upgradePipelineHandler: { (channel: Channel, _: HTTPResponseHead) in channel.eventLoop.makeCompletedFuture { - print("s") return TestUpgradeResult.successfulUpgrade } })