From 64670643dd6faaec284dfae359672b18749b54c2 Mon Sep 17 00:00:00 2001 From: Anthony Drendel Date: Mon, 6 Jul 2020 00:29:18 +0200 Subject: [PATCH] Fix crash when closing WebSocket --- Sources/WebSocket/WebSocket.swift | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/Sources/WebSocket/WebSocket.swift b/Sources/WebSocket/WebSocket.swift index b63809f..ecf63f0 100644 --- 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() + } }