如何正确结合使用URLSessionWebSocketTask与URLProtocol

问题根源 #

你遇到的错误核心原因是误用了普通HTTP任务的URLProtocol处理逻辑：WebSocket是持久化长连接，而你在返回响应和数据后直接调用了urlProtocolDidFinishLoading(self)，这个方法会告诉URLSessionWebSocketTask连接已完成并终止通道，导致连接被强行断开，触发"网络连接丢失"错误；同时，提前结束连接也会触发"Connection not set before response is received"的校验失败。

修正方案 #

1. 调整URLProtocol的加载流程 #

移除urlProtocolDidFinishLoading(self)的调用，保持连接活跃，直到WebSocket连接真正需要关闭时再调用该方法。同时确保响应和数据的传递逻辑符合WebSocket的长连接特性：

open class MyProtocol: URLProtocol {
    // 维护与服务器的WebSocket连接实例（根据你的实现自行定义）
    private var webSocketConnection: WebSocketConnection?
    
    override open func startLoading() {
        // 1. 发送WebSocket升级请求到服务器，获取101 Switching Protocols响应
        let upgradeResponse = HTTPURLResponse(
            url: request.url!,
            statusCode: 101,
            httpVersion: "HTTP/1.1",
            headerFields: [
                "Upgrade": "websocket",
                "Connection": "Upgrade",
                "Sec-WebSocket-Accept": "你的正确Accept值"
            ]
        )!
        
        // 2. 向WebSocketTask返回升级响应
        self.client.urlProtocol(self, didReceive: upgradeResponse, cacheStoragePolicy: .notAllowed)
        
        // 3. 建立与服务器的WebSocket长连接，监听后续消息
        webSocketConnection = WebSocketConnection(url: request.url!)
        webSocketConnection?.onMessageReceived = { [weak self] messageData in
            guard let self = self else { return }
            // 收到服务器的WebSocket消息后，传递给WebSocketTask
            self.client.urlProtocol(self, didLoad: messageData)
        }
        
        webSocketConnection?.onConnectionClosed = { [weak self] error in
            guard let self = self else { return }
            if let error = error {
                // 连接异常关闭，通知任务错误
                self.client.urlProtocol(self, didFailWithError: error)
            } else {
                // 连接正常关闭，标记加载完成
                self.client.urlProtocolDidFinishLoading(self)
            }
        }
        
        // 4. 如果服务器返回了初始消息，立即传递
        if let initialData = webSocketConnection?.initialMessageData {
            self.client.urlProtocol(self, didLoad: initialData)
        }
    }
    
    override open func stopLoading() {
        // 任务停止时，关闭与服务器的WebSocket连接
        webSocketConnection?.close()
        webSocketConnection = nil
    }
    
    // 处理WebSocketTask的消息发送请求
    override open func send(_ data: Data) {
        webSocketConnection?.send(data: data)
    }
}

2. 关键注意事项 #

• 禁止提前调用urlProtocolDidFinishLoading：只有当WebSocket连接被服务器主动关闭、本地主动关闭，或者出现不可恢复的错误时，才调用该方法结束任务。
• 确保WebSocket帧格式正确：传递给urlProtocol(_:didLoad:)的数据必须是符合RFC 6455规范的WebSocket帧（包含 opcode、mask等必要字段），不能直接传递原始文本或二进制数据。
• 处理消息收发的双向逻辑：重写send(_:)方法，将WebSocketTask发送的消息转发给服务器；同时监听服务器的消息，实时传递给WebSocketTask。
• 维护连接生命周期：在stopLoading()方法中关闭与服务器的WebSocket连接，避免资源泄漏。

内容的提问来源于stack exchange，提问作者Captain