Update guard for WebSocket scheme

vapor · Apr 15, 2024 · 24ea66e · 24ea66e
1 parent e54997e
commit 24ea66e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/Sources/WebSocketKit/WebSocket+Connect.swift b/Sources/WebSocketKit/WebSocket+Connect.swift
@@ -20,7 +20,10 @@ extension WebSocket {
         on eventLoopGroup: EventLoopGroup,
         onUpgrade: @Sendable @escaping (WebSocket) -> ()
     ) -> EventLoopFuture<Void> {
-        guard let url = URL(string: url) else {
+        guard
+            url.hasPrefix("ws://") || url.hasPrefix("wss://"),
+            let url = URL(string: url)
+        else {
             return eventLoopGroup.any().makeFailedFuture(WebSocketClient.Error.invalidURL)
         }
         return self.connect(

diff --git a/Tests/WebSocketKitTests/AsyncWebSocketKitTests.swift b/Tests/WebSocketKitTests/AsyncWebSocketKitTests.swift
@@ -91,9 +91,9 @@ final class AsyncWebSocketKitTests: XCTestCase {
             return XCTFail("couldn't get port from \(String(reflecting: server.localAddress))")
         }
         try await WebSocket.connect(to: "ws://localhost:\(port)", on: self.elg) { (ws) async in
-            ws.onPong {
+            ws.onPong { webSocket, _ in
                 do {
-                    try await $0.close()
+                    try await webSocket.close()
                 } catch {
                     XCTFail("Failed to close websocket: \(String(reflecting: error))")
                 }
@@ -118,8 +118,8 @@ final class AsyncWebSocketKitTests: XCTestCase {
         }
         try await WebSocket.connect(to: "ws://localhost:\(port)", on: self.elg) { (ws) async in
             ws.pingInterval = .milliseconds(100)
-            ws.onPong {
-                do { try await $0.close() } catch { XCTFail("Failed to close websocket: \(String(reflecting: error))") }
+            ws.onPong { webSocket, _ in
+                do { try await webSocket.close() } catch { XCTFail("Failed to close websocket: \(String(reflecting: error))") }
                 promise.succeed(())
             }
         }

diff --git a/Tests/WebSocketKitTests/WebSocketKitTests.swift b/Tests/WebSocketKitTests/WebSocketKitTests.swift
@@ -478,7 +478,7 @@ final class WebSocketKitTests: XCTestCase {
         try server.close(mode: .all).wait()
     }
 
-    func testBadURLInWebsocketConnect() async throws {
+    func testBadURLInWebsocketConnect() throws {
         XCTAssertThrowsError(try WebSocket.connect(to: "%w", on: self.elg, onUpgrade: { _ in }).wait()) {
             guard case .invalidURL = $0 as? WebSocketClient.Error else {
                 return XCTFail("Expected .invalidURL but got \(String(reflecting: $0))")