URLSessionWebSocketTask.receive() not finishing if server closes connection without a close packet (#4673)

Allow _WebSocketURLProtocol to detect a closing URL connection, and set its state to indicate this is a normal close of the connection.

Kick off URLSessionWebSocketTask doPendingWork() if the connection is closed out due to an error, in order to respond to any pending Tasks parked in receive().

Author: pseligman

Sources/FoundationNetworking/URLSession/URLSessionTask.swift

@@ -703,6 +703,12 @@ open class URLSessionWebSocketTask : URLSessionTask {
         }
     }
 
+    open override var error: Error? {
+        didSet {
+            doPendingWork()
+        }
+    }
+
     private var sendBuffer = [(Message, (Error?) -> Void)]()
     private var receiveBuffer = [Message]()
     private var receiveCompletionHandlers = [(Result<Message, Error>) -> Void]()

Sources/FoundationNetworking/URLSession/WebSocket/WebSocketURLProtocol.swift

@@ -88,6 +88,13 @@ internal class _WebSocketURLProtocol: _HTTPURLProtocol {
         return .completeTask
     }
 
+    override func completeTask() {
+        if let webSocketTask = task as? URLSessionWebSocketTask {
+            webSocketTask.close(code: .normalClosure, reason: nil)
+        }
+        super.completeTask()
+    }
+
     func sendWebSocketData(_ data: Data, flags: _EasyHandle.WebSocketFlags) throws {
         try easyHandle.sendWebSocketsData(data, flags: flags)
     }

Tests/Foundation/HTTPServer.swift

@@ -912,8 +912,28 @@ public class TestURLSessionServer: CustomStringConvertible {
                                "Connection: Upgrade"]
 
         let expectFullRequestResponseTests: Bool
+        let sendClosePacket: Bool
+        let completeUpgrade: Bool
+        
         let uri = request.uri
-        if uri.count > "/web-socket/".count {
+        switch uri {
+        case "/web-socket":
+            expectFullRequestResponseTests = true
+            completeUpgrade = true
+            sendClosePacket = true
+        case "/web-socket/semi-abrupt-close":
+            expectFullRequestResponseTests = false
+            completeUpgrade = true
+            sendClosePacket = false
+        case "/web-socket/abrupt-close":
+            expectFullRequestResponseTests = false
+            completeUpgrade = false
+            sendClosePacket = false
+        default:
+            guard uri.count > "/web-socket/".count else {
+                NSLog("Expected Sec-WebSocket-Protocol")
+                throw InternalServerError.badHeaders
+            }
             let expectedProtocol = String(uri.suffix(from: uri.index(uri.startIndex, offsetBy: "/web-socket/".count)))
             guard let receivedProtocolStr = request.getHeader(for: "Sec-WebSocket-Protocol"),
                   expectedProtocol == receivedProtocolStr.components(separatedBy: ", ")[0] else {
@@ -922,10 +942,12 @@ public class TestURLSessionServer: CustomStringConvertible {
             }
             responseHeaders.append("Sec-WebSocket-Protocol: \(expectedProtocol)")
             expectFullRequestResponseTests = false
-        } else {
-            expectFullRequestResponseTests = true
+            completeUpgrade = true
+            sendClosePacket = true
         }
 
+        guard completeUpgrade else { return }
+
         var upgradeResponse = _HTTPResponse(response: .SWITCHING_PROTOCOLS, headers: responseHeaders)
         // Lacking an available SHA1 implementation, we'll only include this response for a well-known key
         if "dGhlIHNhbXBsZSBub25jZQ==" == request.getHeader(for: "sec-websocket-key") {
@@ -940,10 +962,11 @@ public class TestURLSessionServer: CustomStringConvertible {
             let closePayload = Data([UInt8(closeCode >> 8),
                                      UInt8(closeCode & 0xFF)]) + closeReason
 
+            let pingPayload = "Hi".data(using: .utf8)!
+
             if expectFullRequestResponseTests {
                 let stringPayload = "Hello".data(using: .utf8)!
                 let dataPayload = Data([0x20, 0x22, 0x10, 0x03])
-                let pingPayload = "Hi".data(using: .utf8)!
 
                 // Receive a string message
                 guard let stringFrame = try httpServer.tcpSocket.readData(),
@@ -981,32 +1004,36 @@ public class TestURLSessionServer: CustomStringConvertible {
                 }
                 // ... and pong it
                 try httpServer.tcpSocket.writeRawData(Data([0x8a, 0x00]))
-                
-                // Send a ping
-                let sendPingFrame = Data([0x89, UInt8(pingPayload.count)]) + pingPayload
-                try httpServer.tcpSocket.writeRawData(sendPingFrame)
-                // ... and receive its pong
-                guard let pongFrame = try httpServer.tcpSocket.readData(),
-                      pongFrame.count == (2 + 4 + pingPayload.count),
-                      Data(pongFrame.prefix(2)) == Data([0x8a, (0x80 | UInt8(pingPayload.count))]),
-                      try unmaskedPayload(from: pongFrame) == pingPayload else {
-                    NSLog("Invalid pong frame")
-                    throw InternalServerError.badBody
-                }
-                
-                // Send a close
-                let sendCloseFrame = Data([0x88, UInt8(closePayload.count)]) + closePayload
-                try httpServer.tcpSocket.writeRawData(sendCloseFrame)
             }
 
-            // Receive a close message
-            guard let closeFrame = try httpServer.tcpSocket.readData(),
-                  closeFrame.count == (2 + 4 + closePayload.count),
-                  Data(closeFrame.prefix(2)) == Data([0x88, (0x80 | UInt8(closePayload.count))]),
-                  try unmaskedPayload(from: closeFrame) == closePayload else {
-                NSLog("Invalid close payload")
+            // Send a ping
+            let sendPingFrame = Data([0x89, UInt8(pingPayload.count)]) + pingPayload
+            try httpServer.tcpSocket.writeRawData(sendPingFrame)
+            // ... and receive its pong
+            guard let pongFrame = try httpServer.tcpSocket.readData(),
+                  pongFrame.count == (2 + 4 + pingPayload.count),
+                  Data(pongFrame.prefix(2)) == Data([0x8a, (0x80 | UInt8(pingPayload.count))]),
+                  try unmaskedPayload(from: pongFrame) == pingPayload else {
+                NSLog("Invalid pong frame")
                 throw InternalServerError.badBody
             }
+            
+            if sendClosePacket {
+                if expectFullRequestResponseTests {
+                    // Send a close
+                    let sendCloseFrame = Data([0x88, UInt8(closePayload.count)]) + closePayload
+                    try httpServer.tcpSocket.writeRawData(sendCloseFrame)
+                }
+
+                // Receive a close message
+                guard let closeFrame = try httpServer.tcpSocket.readData(),
+                      closeFrame.count == (2 + 4 + closePayload.count),
+                      Data(closeFrame.prefix(2)) == Data([0x88, (0x80 | UInt8(closePayload.count))]),
+                      try unmaskedPayload(from: closeFrame) == closePayload else {
+                    NSLog("Invalid close payload")
+                    throw InternalServerError.badBody
+                }
+            }
 
         } catch {
             let badBodyCloseFrame = Data([0x88, 0x08, 0x03, 0xEA, 0x42, 0x75, 0x68, 0x42, 0x79, 0x65])

Tests/Foundation/Tests/TestDecimal.swift

@@ -7,8 +7,13 @@
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 
-import Foundation
-import XCTest
+#if NS_FOUNDATION_ALLOWS_TESTABLE_IMPORT
+    #if canImport(SwiftFoundation) && !DEPLOYMENT_RUNTIME_OBJC
+        @testable import SwiftFoundation
+    #else
+        @testable import Foundation
+    #endif
+#endif
 
 class TestDecimal: XCTestCase {
 

Tests/Foundation/Tests/TestURLSession.swift

@@ -1894,9 +1894,17 @@ class TestURLSession: LoopbackServerTest {
         }
 
         try await task.sendPing()
-
+        
         wait(for: [delegate.expectation], timeout: 50)
 
+        do {
+            _ = try await task.receive()
+            XCTFail("Expected to throw when receiving on closed task")
+        } catch {
+            let urlError = try XCTUnwrap(error as? URLError)
+            XCTAssertEqual(urlError._nsError.code, NSURLErrorNetworkConnectionLost)
+        }
+        
         let callbacks = [ "urlSession(_:webSocketTask:didOpenWithProtocol:)",
                           "urlSession(_:webSocketTask:didCloseWith:reason:)",
                           "urlSession(_:task:didCompleteWithError:)" ]
@@ -1925,15 +1933,98 @@ class TestURLSession: LoopbackServerTest {
         wait(for: [delegate.expectation], timeout: 50)
 
         let callbacks = [ "urlSession(_:webSocketTask:didOpenWithProtocol:)",
+                          "urlSession(_:webSocketTask:didCloseWith:reason:)",
                           "urlSession(_:task:didCompleteWithError:)" ]
         XCTAssertEqual(delegate.callbacks.count, callbacks.count)
         XCTAssertEqual(delegate.callbacks, callbacks, "Callbacks for \(#function)")
 
         XCTAssertEqual(task.closeCode, .normalClosure)
         XCTAssertEqual(task.closeReason, "BuhBye".data(using: .utf8))
     }
-#endif
 
+    func test_webSocketAbruptClose() async throws {
+        guard #available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) else { return }
+        guard URLSessionWebSocketTask.supportsWebSockets else {
+            print("libcurl lacks WebSockets support, skipping \(#function)")
+            return
+        }
+
+        let urlString = "ws://127.0.0.1:\(TestURLSession.serverPort)/web-socket/abrupt-close"
+        let url = try XCTUnwrap(URL(string: urlString))
+        let request = URLRequest(url: url)
+        
+        let delegate = SessionDelegate(with: expectation(description: "\(urlString): Connect"))
+        let task = delegate.runWebSocketTask(with: request, timeoutInterval: 4)
+        
+        do {
+            _ = try await task.receive()
+            XCTFail("Expected to throw when server closes connection")
+        } catch {
+            let urlError = try XCTUnwrap(error as? URLError)
+            XCTAssertEqual(urlError._nsError.code, NSURLErrorBadServerResponse)
+        }
+
+        wait(for: [delegate.expectation], timeout: 50)
+
+        do {
+            _ = try await task.receive()
+            XCTFail("Expected to throw when receiving on closed connection")
+        } catch {
+            let urlError = try XCTUnwrap(error as? URLError)
+            XCTAssertEqual(urlError._nsError.code, NSURLErrorBadServerResponse)
+        }
+
+        let callbacks = [ "urlSession(_:task:didCompleteWithError:)" ]
+        XCTAssertEqual(delegate.callbacks.count, callbacks.count)
+        XCTAssertEqual(delegate.callbacks, callbacks, "Callbacks for \(#function)")
+        
+        XCTAssertEqual(task.closeCode, .invalid)
+        XCTAssertEqual(task.closeReason, nil)
+    }
+
+    func test_webSocketSemiAbruptClose() async throws {
+        guard #available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) else { return }
+        guard URLSessionWebSocketTask.supportsWebSockets else {
+            print("libcurl lacks WebSockets support, skipping \(#function)")
+            return
+        }
+
+        let urlString = "ws://127.0.0.1:\(TestURLSession.serverPort)/web-socket/semi-abrupt-close"
+        let url = try XCTUnwrap(URL(string: urlString))
+        let request = URLRequest(url: url)
+        
+        let delegate = SessionDelegate(with: expectation(description: "\(urlString): Connect"))
+        let task = delegate.runWebSocketTask(with: request, timeoutInterval: 4)
+        
+        do {
+            _ = try await task.receive()
+            XCTFail("Expected to throw when server closes connection")
+        } catch {
+            let urlError = try XCTUnwrap(error as? URLError)
+            XCTAssertEqual(urlError._nsError.code, NSURLErrorNetworkConnectionLost)
+        }
+
+        wait(for: [delegate.expectation], timeout: 50)
+
+        do {
+            _ = try await task.receive()
+            XCTFail("Expected to throw when receiving on closed connection")
+        } catch {
+            let urlError = try XCTUnwrap(error as? URLError)
+            XCTAssertEqual(urlError._nsError.code, NSURLErrorNetworkConnectionLost)
+        }
+
+        let callbacks = [ "urlSession(_:webSocketTask:didOpenWithProtocol:)",
+                          "urlSession(_:webSocketTask:didCloseWith:reason:)",
+                          "urlSession(_:task:didCompleteWithError:)" ]
+        XCTAssertEqual(delegate.callbacks.count, callbacks.count)
+        XCTAssertEqual(delegate.callbacks, callbacks, "Callbacks for \(#function)")
+        
+        XCTAssertEqual(task.closeCode, .normalClosure)
+        XCTAssertEqual(task.closeReason, nil)
+    }
+#endif
+  
     static var allTests: [(String, (TestURLSession) -> () throws -> Void)] {
         var retVal = [
             ("test_dataTaskWithURL", test_dataTaskWithURL),
@@ -2011,6 +2102,8 @@ class TestURLSession: LoopbackServerTest {
             retVal.append(contentsOf: [
                 ("test_webSocket", asyncTest(test_webSocket)),
                 ("test_webSocketSpecificProtocol", asyncTest(test_webSocketSpecificProtocol)),
+                ("test_webSocketAbruptClose", asyncTest(test_webSocketAbruptClose)),
+                ("test_webSocketSemiAbruptClose", asyncTest(test_webSocketSemiAbruptClose)),
             ])
         }
         return retVal