Fix crash when closing WebSocket

shareup · Jul 5, 2020 · 6467064 · 6467064
1 parent 4d64bfe
commit 6467064
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/Sources/WebSocket/WebSocket.swift b/Sources/WebSocket/WebSocket.swift
@@ -114,16 +114,15 @@ public final class WebSocket: WebSocketProtocol {
     }
 
     public func close(_ closeCode:  WebSocketCloseCode) {
-        let webSocket: (URLSession, URLSessionWebSocketTask)? = sync {
-            guard let (session, task) = state.webSocketSessionAndTask else { return nil }
+        let webSocketTask: URLSessionWebSocketTask? = sync {
+            guard let (_, task) = state.webSocketSessionAndTask else { return nil }
             state = .closing
-            return (session, task)
+            return task
         }
 
-        guard let (session, task) = webSocket else { return }
+        guard let task = webSocketTask else { return }
         let code = URLSessionWebSocketTask.CloseCode(closeCode) ?? .invalid
         task.cancel(with: code, reason: nil)
-        session.invalidateAndCancel()
     }
 }
 
@@ -146,8 +145,8 @@ private extension WebSocket  {
                     return
                 }
 
-                assert(session == webSocketSession)
-                assert(task == webSocketTask)
+                assert(session === webSocketSession)
+                assert(task === webSocketTask)
 
                 self.state = .open(webSocketSession, webSocketTask, delegate)
             }
@@ -198,4 +197,10 @@ private class WebSocketDelegate: NSObject, URLSessionWebSocketDelegate {
                     reason: Data?) {
         self.onClose(session, webSocketTask, closeCode, reason)
     }
+
+    func urlSession(_ session: URLSession,
+                    task: URLSessionTask,
+                    didCompleteWithError error: Error?) {
+        session.invalidateAndCancel()
+    }
 }