Add a WebSocket client upgrader. (#1038)

Motivation:

There is a client protocol upgrader but, unlike the server protocol upgrader, it does not have a WebSocket protocol as part of the project.

Modifications:

Made the magic WebSocket GUID public to the WebSockets project.
Added a NIOWebSocketClientUpgrader.
Added tests for the upgrader.
Updated the Linux test script files.

Result:

The project now has a WebSocket client upgrader to match the WebSocket server upgrader.

Sources/NIOWebSocket/NIOWebSocketClientUpgrader.swift

@@ -0,0 +1,91 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2019 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import NIO
+import NIOHTTP1
+
+/// A `HTTPClientProtocolUpgrader` that knows how to do the WebSocket upgrade dance.
+///
+/// This upgrader assumes that the `HTTPClientUpgradeHandler` will create and send the upgrade request. 
+/// This upgrader also assumes that the `HTTPClientUpgradeHandler` will appropriately mutate the
+/// pipeline to remove the HTTP `ChannelHandler`s.
+public final class NIOWebClientSocketUpgrader: NIOHTTPClientProtocolUpgrader {
+
+    /// RFC 6455 specs this as the required entry in the Upgrade header.
+    public let supportedProtocol: String = "websocket"
+    /// None of the websocket headers are actually defined as 'required'.
+    public let requiredUpgradeHeaders: [String] = []
+
+    private let requestKey: String
+    private let maxFrameSize: Int
+    private let automaticErrorHandling: Bool
+    private let upgradePipelineHandler: (Channel, HTTPResponseHead) -> EventLoopFuture<Void>
+
+    public init(requestKey: String,
+                maxFrameSize: Int = 1 << 14,
+                automaticErrorHandling: Bool = true,
+                upgradePipelineHandler: @escaping (Channel, HTTPResponseHead) -> EventLoopFuture<Void>) {
+
+        precondition(requestKey != "", "The request key must contain a valid Sec-WebSocket-Key")
+        precondition(maxFrameSize <= UInt32.max, "invalid overlarge max frame size")
+        self.requestKey = requestKey
+        self.upgradePipelineHandler = upgradePipelineHandler
+        self.maxFrameSize = maxFrameSize
+        self.automaticErrorHandling = automaticErrorHandling
+    }
+
+    /// Add additional headers that are needed for a WebSocket upgrade request.
+    public func addCustom(upgradeRequestHeaders: inout HTTPHeaders) {
+        upgradeRequestHeaders.add(name: "Sec-WebSocket-Key", value: self.requestKey)
+        upgradeRequestHeaders.add(name: "Sec-WebSocket-Version", value: "13")
+    }
+
+    /// Allow or deny the upgrade based on the upgrade HTTP response
+    /// headers containing the correct accept key.
+    public func shouldAllowUpgrade(upgradeResponse: HTTPResponseHead) -> Bool {
+
+        let acceptValueHeader = upgradeResponse.headers["Sec-WebSocket-Accept"]
+
+        guard acceptValueHeader.count == 1 else {
+            return false
+        }
+
+        // Validate the response key in 'Sec-WebSocket-Accept'.
+        var hasher = SHA1()
+        hasher.update(string: self.requestKey)
+        hasher.update(string: magicWebSocketGUID)
+        let expectedAcceptValue = String(base64Encoding: hasher.finish())
+
+        return expectedAcceptValue == acceptValueHeader[0]
+    }
+
+    /// Called when the upgrade response has been flushed and it is safe to mutate the channel
+    /// pipeline. Adds channel handlers for websocket frame encoding, decoding and errors.
+    public func upgrade(context: ChannelHandlerContext, upgradeResponse: HTTPResponseHead) -> EventLoopFuture<Void> {
+
+        var upgradeFuture = context.pipeline.addHandler(WebSocketFrameEncoder()).flatMap {
+            context.pipeline.addHandler(ByteToMessageHandler(WebSocketFrameDecoder(maxFrameSize: self.maxFrameSize)))
+        }
+
+        if self.automaticErrorHandling {
+            upgradeFuture = upgradeFuture.flatMap {
+                context.pipeline.addHandler(WebSocketProtocolErrorHandler())
+            }
+        }
+
+        return upgradeFuture.flatMap {
+            self.upgradePipelineHandler(context.channel, upgradeResponse)
+        }
+    }
+}

Sources/NIOWebSocket/NIOWebSocketServerUpgrader.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the SwiftNIO open source project
 //
-// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
+// Copyright (c) 2017-2019 Apple Inc. and the SwiftNIO project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -16,7 +16,7 @@ import CNIOSHA1
 import NIO
 import NIOHTTP1
 
-private let magicWebSocketGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+let magicWebSocketGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 
 @available(*, deprecated, renamed: "NIOWebSocketServerUpgrader")
 public typealias WebSocketUpgrader = NIOWebSocketServerUpgrader
@@ -177,7 +177,9 @@ public final class NIOWebSocketServerUpgrader: HTTPServerProtocolUpgrader {
         }
 
         if self.automaticErrorHandling {
-            upgradeFuture = upgradeFuture.flatMap { context.pipeline.addHandler(WebSocketProtocolErrorHandler())}
+            upgradeFuture = upgradeFuture.flatMap {
+                context.pipeline.addHandler(WebSocketProtocolErrorHandler())
+            }
         }
 
         return upgradeFuture.flatMap {

Tests/LinuxMain.swift

@@ -91,6 +91,7 @@ import XCTest
          testCase(ThreadTest.allTests),
          testCase(TypeAssistedChannelHandlerTest.allTests),
          testCase(UtilitiesTest.allTests),
+         testCase(WebSocketClientEndToEndTests.allTests),
          testCase(WebSocketFrameDecoderTest.allTests),
          testCase(WebSocketFrameEncoderTest.allTests),
          testCase(WebSocketServerEndToEndTests.allTests),

Tests/NIOWebSocketTests/WebSocketClientEndToEndTests+XCTest.swift

@@ -0,0 +1,38 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+//
+// WebSocketClientEndToEndTests+XCTest.swift
+//
+import XCTest
+
+///
+/// NOTE: This file was generated by generate_linux_tests.rb
+///
+/// Do NOT edit this file directly as it will be regenerated automatically when needed.
+///
+
+extension WebSocketClientEndToEndTests {
+
+   static var allTests : [(String, (WebSocketClientEndToEndTests) -> () throws -> Void)] {
+      return [
+                ("testSimpleUpgradeSucceeds", testSimpleUpgradeSucceeds),
+                ("testRejectUpgradeIfMissingAcceptKey", testRejectUpgradeIfMissingAcceptKey),
+                ("testRejectUpgradeIfIncorrectAcceptKey", testRejectUpgradeIfIncorrectAcceptKey),
+                ("testRejectUpgradeIfNotWebsocket", testRejectUpgradeIfNotWebsocket),
+                ("testSendAFewFrames", testSendAFewFrames),
+                ("testReceiveAFewFrames", testReceiveAFewFrames),
+           ]
+   }
+}
+