File: WebSocketURLProtocol.swift

package info (click to toggle)
swiftlang 6.0.3-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 2,519,992 kB
  • sloc: cpp: 9,107,863; ansic: 2,040,022; asm: 1,135,751; python: 296,500; objc: 82,456; f90: 60,502; lisp: 34,951; pascal: 19,946; sh: 18,133; perl: 7,482; ml: 4,937; javascript: 4,117; makefile: 3,840; awk: 3,535; xml: 914; fortran: 619; cs: 573; ruby: 573
file content (189 lines) | stat: -rw-r--r-- 8,124 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See http://swift.org/LICENSE.txt for license information
// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//

#if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
import SwiftFoundation
#else
import Foundation
#endif

@_implementationOnly import _CFURLSessionInterface
import Dispatch

internal class _WebSocketURLProtocol: _HTTPURLProtocol {
    public required init(task: URLSessionTask, cachedResponse: CachedURLResponse?, client: URLProtocolClient?) {
        super.init(task: task, cachedResponse: nil, client: client)
    }
    
    public required init(request: URLRequest, cachedResponse: CachedURLResponse?, client: URLProtocolClient?) {
        super.init(request: request, cachedResponse: nil, client: client)
    }
    
    override class func canInit(with request: URLRequest) -> Bool {
        switch request.url?.scheme {
        case "ws", "wss": return true
        default: return false
        }
    }
    
    override func canCache(_ response: CachedURLResponse) -> Bool {
        false
    }
    
    override func canRespondFromCache(using response: CachedURLResponse) -> Bool { false }

    override func didReceiveResponse() {
        guard let webSocketTask = task as? URLSessionWebSocketTask else { return }
        guard case .transferInProgress(let ts) = self.internalState else { fatalError("Transfer not in progress.") }
        guard let response = ts.response as? HTTPURLResponse else { fatalError("Header complete, but not URL response.") }

        webSocketTask.protocolPicked = response.value(forHTTPHeaderField: "Sec-WebSocket-Protocol")
        
        easyHandle.timeoutTimer = nil
        
        webSocketTask.handshakeCompleted = true
        
        self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
    }
    
    /// Set options on the easy handle to match the given request.
    ///
    /// This performs a series of `curl_easy_setopt()` calls.
    override func configureEasyHandle(for request: URLRequest, body: _Body) {
        guard request.httpMethod == "GET" else {
            NSLog("WebSocket tasks must use GET")
            let error = NSError(domain: NSURLErrorDomain, code: NSURLErrorUnsupportedURL,
                                userInfo: [
                                    NSLocalizedDescriptionKey: "websocket task must use GET httpMethod",
                                    NSURLErrorFailingURLStringErrorKey: request.url?.description ?? ""
                                ])
            internalState = .transferFailed
            transferCompleted(withError: error)
            return
        }
        
        super.configureEasyHandle(for: request, body: body)
        
        easyHandle.setAllowedProtocolsToAll()
        
        guard let webSocketTask = task as? URLSessionWebSocketTask else { return }
        easyHandle.set(preferredReceiveBufferSize: webSocketTask.maximumMessageSize)
    }
    
    override func completionAction(forCompletedRequest request: URLRequest, response: URLResponse) -> _CompletionAction {
        // Redirect:
        guard let httpURLResponse = response as? HTTPURLResponse else {
            fatalError("Response was not HTTPURLResponse")
        }
        if let request = redirectRequest(for: httpURLResponse, fromRequest: request) {
            return .redirectWithRequest(request)
        }
        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)
    }
    
    override func didReceive(data: Data) -> _EasyHandle._Action {
        guard case .transferInProgress(var ts) = internalState else {
            fatalError("Received web socket data, but no transfer in progress.")
        }

        if let response = validateHeaderComplete(transferState:ts) {
            ts.response = response
        }

        // Note this excludes code 300 which should return the response of the redirect and not follow it.
        // For other redirect codes dont notify the delegate of the data received in the redirect response.
        if let httpResponse = ts.response as? HTTPURLResponse,
           301...308 ~= httpResponse.statusCode {
            // Save the response body in case the delegate does not perform a redirect and the 3xx response
            // including its body needs to be returned to the client.
            var redirectBody = lastRedirectBody ?? Data()
            redirectBody.append(data)
            lastRedirectBody = redirectBody
        }

        let flags = easyHandle.getWebSocketFlags()
        
        notifyTask(aboutReceivedData: data, flags: flags)
        internalState = .transferInProgress(ts)
        return .proceed
    }

    fileprivate func notifyTask(aboutReceivedData data: Data, flags: _EasyHandle.WebSocketFlags) {
        guard let t = self.task else {
            fatalError("Cannot notify")
        }
        switch t.session.behaviour(for: t) {
        case .noDelegate:
            break
        case .taskDelegate:
            break
        default:
            fatalError("Unexpected behaviour for URLSessionWebSocketTask")
        }
        guard let task = t as? URLSessionWebSocketTask else {
            fatalError("Cast to URLSessionWebSocketTask failed")
        }
        
        // Buffer the response message in the task
        if flags.contains(.close) {
            let closeCode: URLSessionWebSocketTask.CloseCode
            let reasonData: Data
            if data.count >= 2 {
                closeCode = data.withUnsafeBytes {
                    let codeInt = UInt16(bigEndian: $0.load(as: UInt16.self))
                    return URLSessionWebSocketTask.CloseCode(rawValue: Int(codeInt)) ?? .unsupportedData
                }
                reasonData = Data(data[2...])
            } else {
                closeCode = .normalClosure
                reasonData = Data()
            }
            task.close(code: closeCode, reason: reasonData)
        } else if flags.contains(.pong) {
            task.noteReceivedPong()
        } else if flags.contains(.binary) {
            let message = URLSessionWebSocketTask.Message.data(data)
            task.appendReceivedMessage(message)
        } else if flags.contains(.text) {
            guard let utf8 = String(data: data, encoding: .utf8) else {
                NSLog("Invalid utf8 message received from server \(data)")
                let error = NSError(domain: NSURLErrorDomain, code: NSURLErrorBadServerResponse,
                                    userInfo: [
                                        NSLocalizedDescriptionKey: "Invalid message received from server",
                                        NSURLErrorFailingURLStringErrorKey: request.url?.description ?? ""
                                    ])
                internalState = .transferFailed
                transferCompleted(withError: error)
                return
            }
            let message = URLSessionWebSocketTask.Message.string(utf8)
            task.appendReceivedMessage(message)
        } else {
            NSLog("Unexpected message received from server \(data) \(flags)")
            let error = NSError(domain: NSURLErrorDomain, code: NSURLErrorBadServerResponse,
                                userInfo: [
                                    NSLocalizedDescriptionKey: "Unexpected message received from server",
                                    NSURLErrorFailingURLStringErrorKey: request.url?.description ?? ""
                                ])
            internalState = .transferFailed
            transferCompleted(withError: error)
        }
    }
}