Add WebSocket configuration struct (#139)

* Add WebSocket configuration struct

Add additionalHeaders parameter to this struct

* Sec-WebSocket-Protocol should be mqtt

* Update documentation/tests

* additionalHeaders -> initialRequestHeaders

* Change Configuration.init webSocketConfiguration signature

* Add test for WebSocketInitialRequestHandler

Author: adam-fowler

Sources/MQTTNIO/ChannelHandlers/WebSocketInitialRequest.swift

@@ -23,20 +23,23 @@ final class WebSocketInitialRequestHandler: ChannelInboundHandler, RemovableChan
 
     let host: String
     let urlPath: String
+    let additionalHeaders: HTTPHeaders
     let upgradePromise: EventLoopPromise<Void>
 
-    init(host: String, urlPath: String, upgradePromise: EventLoopPromise<Void>) {
+    init(host: String, urlPath: String, additionalHeaders: HTTPHeaders, upgradePromise: EventLoopPromise<Void>) {
         self.host = host
-        self.upgradePromise = upgradePromise
         self.urlPath = urlPath
+        self.additionalHeaders = additionalHeaders
+        self.upgradePromise = upgradePromise
     }
 
     public func channelActive(context: ChannelHandlerContext) {
         // We are connected. It's time to send the message to the server to initialize the upgrade dance.
         var headers = HTTPHeaders()
         headers.add(name: "Content-Length", value: "0")
         headers.add(name: "host", value: self.host)
-        headers.add(name: "Sec-WebSocket-Protocol", value: "mqttv3.1")
+        headers.add(name: "Sec-WebSocket-Protocol", value: "mqtt")
+        headers.add(contentsOf: self.additionalHeaders)
 
         let requestHead = HTTPRequestHead(
             version: HTTPVersion(major: 1, minor: 1),

Sources/MQTTNIO/MQTTClient.swift

@@ -57,12 +57,12 @@ public final class MQTTClient {
     var connection: MQTTConnection? {
         get {
             self.lock.withLock {
-                _connection
+                self._connection
             }
         }
         set {
             self.lock.withLock {
-                _connection = newValue
+                self._connection = newValue
             }
         }
     }
@@ -480,7 +480,7 @@ internal extension MQTTClient {
     func resendOnRestart() {
         let inflight = self.inflight.packets
         self.inflight.clear()
-        inflight.forEach { packet -> Void in
+        inflight.forEach { packet in
             switch packet {
             case let publish as MQTTPublishPacket:
                 let newPacket = MQTTPublishPacket(

Sources/MQTTNIO/MQTTConfiguration.swift

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 import NIO
+import NIOHTTP1
 #if canImport(NIOSSL)
 import NIOSSL
 #endif
@@ -40,8 +41,74 @@ extension MQTTClient {
         #endif
     }
 
+    public struct WebSocketConfiguration {
+        /// Initialize MQTTClient WebSocket configuration struct
+        /// - Parameters:
+        ///   - urlPath: WebSocket URL, defaults to "/mqtt"
+        ///   - maxFrameSize: Max frame size WebSocket client will allow
+        ///   - initialRequestHeaders: Additional headers to add to initial HTTP request
+        public init(
+            urlPath: String,
+            maxFrameSize: Int = 1 << 14,
+            initialRequestHeaders: HTTPHeaders = [:]
+        ) {
+            self.urlPath = urlPath
+            self.maxFrameSize = maxFrameSize
+            self.initialRequestHeaders = initialRequestHeaders
+        }
+
+        /// WebSocket URL, defaults to "/mqtt"
+        public let urlPath: String
+        /// Max frame size WebSocket client will allow
+        public let maxFrameSize: Int
+        /// Additional headers to add to initial HTTP request
+        public let initialRequestHeaders: HTTPHeaders
+    }
+
     /// Configuration for MQTTClient
     public struct Configuration {
+        /// Initialize MQTTClient configuration struct
+        /// - Parameters:
+        ///   - version: Version of MQTT server client is connecting to
+        ///   - disablePing: Disable the automatic sending of pingreq messages
+        ///   - keepAliveInterval: MQTT keep alive period.
+        ///   - pingInterval: Override calculated interval between each pingreq message
+        ///   - connectTimeout: Timeout for connecting to server
+        ///   - timeout: Timeout for server ACK responses
+        ///   - userName: MQTT user name
+        ///   - password: MQTT password
+        ///   - useSSL: Use encrypted connection to server
+        ///   - tlsConfiguration: TLS configuration, for SSL connection
+        ///   - sniServerName: Server name used by TLS. This will default to host name if not set
+        ///   - webSocketConfiguration: Set this if you want to use WebSockets
+        public init(
+            version: Version = .v3_1_1,
+            disablePing: Bool = false,
+            keepAliveInterval: TimeAmount = .seconds(90),
+            pingInterval: TimeAmount? = nil,
+            connectTimeout: TimeAmount = .seconds(10),
+            timeout: TimeAmount? = nil,
+            userName: String? = nil,
+            password: String? = nil,
+            useSSL: Bool = false,
+            tlsConfiguration: TLSConfigurationType? = nil,
+            sniServerName: String? = nil,
+            webSocketConfiguration: WebSocketConfiguration
+        ) {
+            self.version = version
+            self.disablePing = disablePing
+            self.keepAliveInterval = keepAliveInterval
+            self.pingInterval = pingInterval
+            self.connectTimeout = connectTimeout
+            self.timeout = timeout
+            self.userName = userName
+            self.password = password
+            self.useSSL = useSSL
+            self.tlsConfiguration = tlsConfiguration
+            self.sniServerName = sniServerName
+            self.webSocketConfiguration = webSocketConfiguration
+        }
+
         /// Initialize MQTTClient configuration struct
         /// - Parameters:
         ///   - version: Version of MQTT server client is connecting to
@@ -83,11 +150,13 @@ extension MQTTClient {
             self.userName = userName
             self.password = password
             self.useSSL = useSSL
-            self.useWebSockets = useWebSockets
             self.tlsConfiguration = tlsConfiguration
             self.sniServerName = sniServerName
-            self.webSocketURLPath = webSocketURLPath
-            self.webSocketMaxFrameSize = webSocketMaxFrameSize
+            if useWebSockets {
+                self.webSocketConfiguration = .init(urlPath: webSocketURLPath ?? "/mqtt", maxFrameSize: webSocketMaxFrameSize)
+            } else {
+                self.webSocketConfiguration = nil
+            }
         }
 
         /// Initialize MQTTClient configuration struct
@@ -107,7 +176,6 @@ extension MQTTClient {
         ///   - sniServerName: Server name used by TLS. This will default to host name if not set
         ///   - webSocketURLPath: URL Path for web socket. Defaults to "/mqtt"
         ///   - webSocketMaxFrameSize: Maximum frame size for a web socket connection
-        ///
         @available(*, deprecated, message: "maxRetryAttempts is no longer used")
         public init(
             version: Version = .v3_1_1,
@@ -135,11 +203,28 @@ extension MQTTClient {
             self.userName = userName
             self.password = password
             self.useSSL = useSSL
-            self.useWebSockets = useWebSockets
             self.tlsConfiguration = tlsConfiguration
             self.sniServerName = sniServerName
-            self.webSocketURLPath = webSocketURLPath
-            self.webSocketMaxFrameSize = webSocketMaxFrameSize
+            if useWebSockets {
+                self.webSocketConfiguration = .init(urlPath: webSocketURLPath ?? "/mqtt", maxFrameSize: webSocketMaxFrameSize)
+            } else {
+                self.webSocketConfiguration = nil
+            }
+        }
+
+        /// use a websocket connection to server
+        public var useWebSockets: Bool {
+            self.webSocketConfiguration != nil
+        }
+
+        /// URL Path for web socket. Defaults to "/mqtt"
+        public var webSocketURLPath: String? {
+            self.webSocketConfiguration?.urlPath
+        }
+
+        /// Maximum frame size for a web socket connection
+        public var webSocketMaxFrameSize: Int {
+            self.webSocketConfiguration?.maxFrameSize ?? 1 << 14
         }
 
         /// Version of MQTT server client is connecting to
@@ -160,15 +245,11 @@ extension MQTTClient {
         public let password: String?
         /// use encrypted connection to server
         public let useSSL: Bool
-        /// use a websocket connection to server
-        public let useWebSockets: Bool
         /// TLS configuration
         public let tlsConfiguration: TLSConfigurationType?
         /// server name used by TLS
         public let sniServerName: String?
-        /// URL Path for web socket. Defaults to "/mqtt"
-        public let webSocketURLPath: String?
-        /// Maximum frame size for a web socket connection
-        public let webSocketMaxFrameSize: Int
+        /// WebSocket configuration
+        public let webSocketConfiguration: WebSocketConfiguration?
     }
 }

Sources/MQTTNIO/MQTTConnection.swift

@@ -57,13 +57,18 @@ final class MQTTConnection {
                         taskHandler,
                     ]
                     // are we using websockets
-                    if client.configuration.useWebSockets {
+                    if let webSocketConfiguration = client.configuration.webSocketConfiguration {
                         // prepare for websockets and on upgrade add handlers
                         let promise = eventLoop.makePromise(of: Void.self)
                         promise.futureResult.map { _ in channel }
                             .cascade(to: channelPromise)
 
-                        return Self.setupChannelForWebsockets(client: client, channel: channel, upgradePromise: promise) {
+                        return Self.setupChannelForWebsockets(
+                            client: client,
+                            channel: channel,
+                            webSocketConfiguration: webSocketConfiguration,
+                            upgradePromise: promise
+                        ) {
                             return channel.pipeline.addHandlers(handlers)
                         }
                     } else {
@@ -140,13 +145,15 @@ final class MQTTConnection {
     static func setupChannelForWebsockets(
         client: MQTTClient,
         channel: Channel,
+        webSocketConfiguration: MQTTClient.WebSocketConfiguration,
         upgradePromise promise: EventLoopPromise<Void>,
         afterHandlerAdded: @escaping () -> EventLoopFuture<Void>
     ) -> EventLoopFuture<Void> {
         // initial HTTP request handler, before upgrade
         let httpHandler = WebSocketInitialRequestHandler(
             host: client.configuration.sniServerName ?? client.hostHeader,
-            urlPath: client.configuration.webSocketURLPath ?? "/mqtt",
+            urlPath: webSocketConfiguration.urlPath,
+            additionalHeaders: webSocketConfiguration.initialRequestHeaders,
             upgradePromise: promise
         )
         // create random key for request key

Sources/MQTTNIO/MQTTNIO.docc/mqttnio-aws.md

@@ -27,7 +27,7 @@ let client = MQTTClient(
     host: host,
     identifier: "MyAWSClient",
     eventLoopGroupProvider: .createNew,
-    configuration: .init(useSSL: true, useWebSockets: true, webSocketURLPath: requestUri)
+    configuration: .init(useSSL: true, webSocketConfiguration: .init(urlPath: requestUri))
 )
 ```
 You can find out more about connecting to AWS brokers [here](https://docs.aws.amazon.com/iot/latest/developerguide/protocols.html)

Sources/MQTTNIO/MQTTNIO.docc/mqttnio-connections.md

@@ -26,7 +26,7 @@ let client = MQTTClient(
 
 ## WebSockets
 
-MQTT also supports Web Socket connections. Set `Configuration.useWebSockets` to `true` and set the URL path in `Configuration.webSocketsURLPath` to enable these.
+MQTT also supports Web Socket connections. Provide a `WebSocketConfiguration` when initializing `MQTTClient.Configuration` to enable this.
 
 ## NIO Transport Services
 

Tests/MQTTNIOTests/MQTTNIOTests.swift

@@ -590,11 +590,52 @@ final class MQTTNIOTests: XCTestCase {
         try client.disconnect().wait()
     }
 
+    func testWebSocketInitialRequest() throws {
+        let el = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try el.syncShutdownGracefully()) }
+        let promise = el.makePromise(of: Void.self)
+        let initialRequestHandler = WebSocketInitialRequestHandler(
+            host: "test.mosquitto.org",
+            urlPath: "/mqtt",
+            additionalHeaders: ["Test": "Value"],
+            upgradePromise: promise
+        )
+        let channel = EmbeddedChannel(handler: initialRequestHandler, loop: el)
+        try channel.connect(to: SocketAddress(ipAddress: "127.0.0.1", port: 0)).wait()
+        let requestHead = try channel.readOutbound(as: HTTPClientRequestPart.self)
+        let requestBody = try channel.readOutbound(as: HTTPClientRequestPart.self)
+        let requestEnd = try channel.readOutbound(as: HTTPClientRequestPart.self)
+        switch requestHead {
+        case .head(let head):
+            XCTAssertEqual(head.uri, "/mqtt")
+            XCTAssertEqual(head.headers["host"].first, "test.mosquitto.org")
+            XCTAssertEqual(head.headers["Sec-WebSocket-Protocol"].first, "mqtt")
+            XCTAssertEqual(head.headers["Test"].first, "Value")
+        default:
+            XCTFail("Did not expect \(String(describing: requestHead))")
+        }
+        switch requestBody {
+        case .body(let data):
+            XCTAssertEqual(data, .byteBuffer(ByteBuffer()))
+        default:
+            XCTFail("Did not expect \(String(describing: requestBody))")
+        }
+        switch requestEnd {
+        case .end(nil):
+            break
+        default:
+            XCTFail("Did not expect \(String(describing: requestEnd))")
+        }
+        _ = try channel.finish()
+        promise.succeed()
+    }
+
     // MARK: Helper variables and functions
 
     func createClient(identifier: String, configuration: MQTTClient.Configuration = .init()) -> MQTTClient {
         MQTTClient(
             host: Self.hostname,
+            port: 1883,
             identifier: identifier,
             eventLoopGroupProvider: .createNew,
             logger: self.logger,
@@ -609,7 +650,7 @@ final class MQTTNIOTests: XCTestCase {
             identifier: identifier,
             eventLoopGroupProvider: .createNew,
             logger: self.logger,
-            configuration: .init(useWebSockets: true, webSocketURLPath: "/mqtt")
+            configuration: .init(webSocketConfiguration: .init(urlPath: "/mqtt"))
         )
     }
 
@@ -630,7 +671,13 @@ final class MQTTNIOTests: XCTestCase {
             identifier: identifier,
             eventLoopGroupProvider: .createNew,
             logger: self.logger,
-            configuration: .init(timeout: .seconds(5), useSSL: true, useWebSockets: true, tlsConfiguration: Self.getTLSConfiguration(), sniServerName: "soto.codes", webSocketURLPath: "/mqtt")
+            configuration: .init(
+                timeout: .seconds(5),
+                useSSL: true,
+                tlsConfiguration: Self.getTLSConfiguration(),
+                sniServerName: "soto.codes",
+                webSocketConfiguration: .init(urlPath: "/mqtt")
+            )
         )
     }
 
@@ -640,13 +687,11 @@ final class MQTTNIOTests: XCTestCase {
         return logger
     }()
 
-    static var rootPath: String = {
-        return #file
-            .split(separator: "/", omittingEmptySubsequences: false)
-            .dropLast(3)
-            .map { String(describing: $0) }
-            .joined(separator: "/")
-    }()
+    static var rootPath: String = #file
+        .split(separator: "/", omittingEmptySubsequences: false)
+        .dropLast(3)
+        .map { String(describing: $0) }
+        .joined(separator: "/")
 
     static var _tlsConfiguration: Result<MQTTClient.TLSConfigurationType, Error> = {
         do {