If we receive an upgrade request after not upgrading, return redirect and close connection (#108)

* If we receive an websocket upgrade after not upgrading redirect and close connection

* Include query in redirect

* Fix for 5.10

* Add test for closing conenction on unexpected upgrade

* Add detailed comment

Author: adam-fowler

Sources/HummingbirdWebSocket/WebSocketChannel.swift

@@ -98,7 +98,7 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
                     }
             }
         }
-        self.responder = responder
+        self.responder = Self.getUpgradeResponder(responder)
     }
 
     @available(*, deprecated, renamed: "init(responder:configuration:additionalChannelHandlers:shouldUpgrade:)")
@@ -167,7 +167,46 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
             }
             return promise.futureResult
         }
-        self.responder = responder
+        self.responder = Self.getUpgradeResponder(responder)
+    }
+
+    /// Return HTTP responder that responds with a redirect and connection closure on receiving an
+    /// upgrade header set to websocket
+    ///
+    /// The responder passed in as a parameter is called from the resultant responder if no upgrade
+    /// header is found.
+    ///
+    /// This is a temporary solution to the fact that the NIO upgrade code does not support parsing
+    /// upgrade headers after having received a normal HTTP request. By returning a redirect to the
+    /// same URI and closing the connection we are forcing the client to open a new connection
+    /// where the upgrade code path will run.
+    ///
+    /// - Parameter responder: HTTP responder to call
+    /// - Returns: Result of HTTP responder or redirect
+    static func getUpgradeResponder(_ responder: @escaping HTTPChannelHandler.Responder) -> HTTPChannelHandler.Responder {
+        struct RedirectCloseError: Error {}
+        return {
+            (
+                request: Request,
+                responseWriter: consuming ResponseWriter,
+                channel: Channel
+            ) in
+            if request.headers[.upgrade] == "websocket" {
+                var path = request.uri.path
+                if let query = request.uri.query {
+                    path += "?\(query)"
+                }
+                let headers: HTTPFields = [
+                    .connection: "close",
+                    .location: path,
+                ]
+                let response = HTTPResponse(status: .temporaryRedirect, headerFields: headers)
+                try await responseWriter.writeResponse(response)
+                throw RedirectCloseError()
+            } else {
+                try await responder(request, responseWriter, channel)
+            }
+        }
     }
 
     ///  Setup channel to accept HTTP1 with a WebSocket upgrade

Sources/HummingbirdWebSocket/WebSocketRouter.swift

@@ -192,7 +192,7 @@ extension HTTP1WebSocketUpgradeChannel {
             }
             return promise.futureResult
         }
-        self.responder = responder
+        self.responder = Self.getUpgradeResponder(responder)
     }
 }
 

Tests/HummingbirdWebSocketTests/WebSocketTests.swift

@@ -866,4 +866,31 @@ final class HummingbirdWebSocketTests: XCTestCase {
         try await httpClient.shutdown()
     }
 
+    func testUpgradeAfterNotUpgraded() async throws {
+        let router = Router()
+        router.get("/") { _, _ in
+            "Helllo"
+        }
+        let app = Application(
+            router: router,
+            server: .http1WebSocketUpgrade { _, _, _ in
+                .dontUpgrade
+            }
+        )
+        try await app.test(.live) { client in
+            try await client.execute(uri: "/", method: .get) { response in
+                XCTAssertEqual(response.status, .ok)
+            }
+            // perform upgrade
+            try await client.execute(uri: "/test?this=that", method: .get, headers: [.upgrade: "websocket"]) { response in
+                XCTAssertEqual(response.status, .temporaryRedirect)
+                XCTAssertEqual(response.headers[.location], "/test?this=that")
+            }
+            // check channel has been closed
+            do {
+                try await client.execute(uri: "/", method: .get)
+            } catch let error as ChannelError where error == .ioOnClosedChannel {
+            }
+        }
+    }
 }