From cce88ae3eb47c23f83e6ad5d8edb96e08b436396 Mon Sep 17 00:00:00 2001 From: Erik Date: Sun, 12 Feb 2017 12:12:33 -0500 Subject: [PATCH] bump websocket --- Source/WebSocket.swift | 193 ++++++++++++++++++++++------------------- 1 file changed, 104 insertions(+), 89 deletions(-) diff --git a/Source/WebSocket.swift b/Source/WebSocket.swift index 34c7dd2..ccdf290 100644 --- a/Source/WebSocket.swift +++ b/Source/WebSocket.swift @@ -38,7 +38,7 @@ public protocol WebSocketPongDelegate: class { } open class WebSocket : NSObject, StreamDelegate { - + enum OpCode : UInt8 { case continueFrame = 0x0 case textFrame = 0x1 @@ -49,7 +49,7 @@ open class WebSocket : NSObject, StreamDelegate { case pong = 0xA // B-F reserved. } - + public enum CloseCode : UInt16 { case normal = 1000 case goingAway = 1001 @@ -62,19 +62,19 @@ open class WebSocket : NSObject, StreamDelegate { case policyViolated = 1008 case messageTooBig = 1009 } - + public static let ErrorDomain = "WebSocket" - + enum InternalErrorCode: UInt16 { // 0-999 WebSocket status codes not used case outputStreamWriteError = 1 } - + // Where the callback is executed. It defaults to the main UI thread queue. public var callbackQueue = DispatchQueue.main - + var optionalProtocols: [String]? - + // MARK: - Constants let headerWSUpgradeName = "Upgrade" let headerWSUpgradeValue = "websocket" @@ -96,7 +96,7 @@ open class WebSocket : NSObject, StreamDelegate { let MaxFrameSize: Int = 32 let httpSwitchProtocolCode = 101 let supportedSSLSchemes = ["wss", "https"] - + class WSResponse { var isFin = false var code: OpCode = .continueFrame @@ -104,23 +104,23 @@ open class WebSocket : NSObject, StreamDelegate { var frameCount = 0 var buffer: NSMutableData? } - + // MARK: - Delegates /// Responds to callback about new messages coming in over the WebSocket /// and also connection/disconnect messages. public weak var delegate: WebSocketDelegate? - + /// Receives a callback for each pong message recived. public weak var pongDelegate: WebSocketPongDelegate? - - + + // MARK: - Block based API. public var onConnect: ((Void) -> Void)? public var onDisconnect: ((NSError?) -> Void)? public var onText: ((String) -> Void)? public var onData: ((Data) -> Void)? public var onPong: ((Data?) -> Void)? - + public var headers = [String: String]() public var voipEnabled = false public var disableSSLCertValidation = false @@ -131,9 +131,9 @@ open class WebSocket : NSObject, StreamDelegate { public var isConnected: Bool { return connected } - + public var currentURL: URL { return url } - + // MARK: - Private private var url: URL private var inputStream: InputStream? @@ -157,7 +157,7 @@ open class WebSocket : NSObject, StreamDelegate { } /// The shared processing queue used for all WebSocket. private static let sharedWorkQueue = DispatchQueue(label: "com.vluxe.starscream.websocket", attributes: []) - + /// Used for setting protocols. public init(url: URL, protocols: [String]? = nil) { self.url = url @@ -170,13 +170,13 @@ open class WebSocket : NSObject, StreamDelegate { writeQueue.maxConcurrentOperationCount = 1 optionalProtocols = protocols } - + // Used for specifically setting the QOS for the write queue. public convenience init(url: URL, writeQueueQOS: QualityOfService, protocols: [String]? = nil) { self.init(url: url, protocols: protocols) writeQueue.qualityOfService = writeQueueQOS } - + /** Connect to the WebSocket server on a background thread. */ @@ -186,15 +186,16 @@ open class WebSocket : NSObject, StreamDelegate { isConnecting = true createHTTPRequest() } - + /** Disconnect from the server. I send a Close control frame to the server, then expect the server to respond with a Close control frame and close the socket from its end. I notify my delegate once the socket has been closed. If you supply a non-nil `forceTimeout`, I wait at most that long (in seconds) for the server to close the socket. After the timeout expires, I close the socket and notify my delegate. If you supply a zero (or negative) `forceTimeout`, I immediately close the socket (without sending a Close control frame) and notify my delegate. - Parameter forceTimeout: Maximum time to wait for the server to close the socket. - Parameter closeCode: The code to send on disconnect. The default is the normal close code for cleanly disconnecting a webSocket. - */ + */ open func disconnect(forceTimeout: TimeInterval? = nil, closeCode: UInt16 = CloseCode.normal.rawValue) { + guard isConnected else { return } switch forceTimeout { case .some(let seconds) where seconds > 0: let milliseconds = Int(seconds * 1_000) @@ -209,7 +210,7 @@ open class WebSocket : NSObject, StreamDelegate { break } } - + /** Write a string to the websocket. This sends it as a text frame. If you supply a non-nil completion block, I will perform it when the write completes. @@ -220,7 +221,7 @@ open class WebSocket : NSObject, StreamDelegate { guard isConnected else { return } dequeueWrite(string.data(using: String.Encoding.utf8)!, code: .textFrame, writeCompletion: completion) } - + /** Write binary data to the websocket. This sends it as a binary frame. If you supply a non-nil completion block, I will perform it when the write completes. @@ -231,7 +232,7 @@ open class WebSocket : NSObject, StreamDelegate { guard isConnected else { return } dequeueWrite(data, code: .binaryFrame, writeCompletion: completion) } - + /** Write a ping to the websocket. This sends it as a control frame. Yodel a sound to the planet. This sends it as an astroid. http://youtu.be/Eu5ZJELRiJ8?t=42s @@ -240,15 +241,15 @@ open class WebSocket : NSObject, StreamDelegate { guard isConnected else { return } dequeueWrite(ping, code: .ping, writeCompletion: completion) } - + /** Private method that starts the connection. */ private func createHTTPRequest() { - + let urlRequest = CFHTTPMessageCreateRequest(kCFAllocatorDefault, "GET" as CFString, url as CFURL, kCFHTTPVersion1_1).takeRetainedValue() - + var port = url.port if port == nil { if supportedSSLSchemes.contains(url.scheme!) { @@ -276,14 +277,14 @@ open class WebSocket : NSObject, StreamDelegate { initStreamsWithData(serializedRequest as Data, Int(port!)) } } - + /** Add a header to the CFHTTPMessage by using the NSString bridges to CFString */ private func addHeader(_ urlRequest: CFHTTPMessage, key: String, val: String) { CFHTTPMessageSetHeaderFieldValue(urlRequest, key as CFString, val as CFString) } - + /** Generate a WebSocket key as needed in RFC. */ @@ -298,7 +299,7 @@ open class WebSocket : NSObject, StreamDelegate { let baseKey = data?.base64EncodedString(options: NSData.Base64EncodingOptions(rawValue: 0)) return baseKey! } - + /** Start the stream connection and write the data to the output stream. */ @@ -306,8 +307,8 @@ open class WebSocket : NSObject, StreamDelegate { //higher level API we will cut over to at some point //NSStream.getStreamsToHostWithName(url.host, port: url.port.integerValue, inputStream: &inputStream, outputStream: &outputStream) // Disconnect and clean up any existing streams before setting up a new pair - disconnectStream(nil) - + disconnectStream(nil, runDelegate: false) + var readStream: Unmanaged? var writeStream: Unmanaged? let h = url.host! as NSString @@ -349,21 +350,24 @@ open class WebSocket : NSObject, StreamDelegate { inStream.setProperty(StreamNetworkServiceTypeValue.voIP as AnyObject, forKey: Stream.PropertyKey.networkServiceType) outStream.setProperty(StreamNetworkServiceTypeValue.voIP as AnyObject, forKey: Stream.PropertyKey.networkServiceType) } - + CFReadStreamSetDispatchQueue(inStream, WebSocket.sharedWorkQueue) CFWriteStreamSetDispatchQueue(outStream, WebSocket.sharedWorkQueue) inStream.open() outStream.open() - + self.mutex.lock() self.readyToWrite = true self.mutex.unlock() - + let bytes = UnsafeRawPointer((data as NSData).bytes).assumingMemoryBound(to: UInt8.self) var out = timeout * 1_000_000 // wait 5 seconds before giving up - writeQueue.addOperation { [weak self] in - while !outStream.hasSpaceAvailable { + let operation = BlockOperation() + operation.addExecutionBlock { [weak self, weak operation] in + guard let sOperation = operation else { return } + while !outStream.hasSpaceAvailable && !sOperation.isCancelled { usleep(100) // wait until the socket is ready + guard !sOperation.isCancelled else { return } out -= 100 if out < 0 { self?.cleanupStream() @@ -373,10 +377,12 @@ open class WebSocket : NSObject, StreamDelegate { return // disconnectStream will be called. } } + guard !sOperation.isCancelled else { return } outStream.write(bytes, maxLength: data.count) } + writeQueue.addOperation(operation) } - + /** Delegate for the stream methods. Processes incoming bytes */ @@ -402,20 +408,23 @@ open class WebSocket : NSObject, StreamDelegate { disconnectStream(nil) } } - + /** Disconnect the stream object and notifies the delegate. */ - private func disconnectStream(_ error: NSError?) { + private func disconnectStream(_ error: NSError?, runDelegate: Bool = true) { if error == nil { writeQueue.waitUntilAllOperationsAreFinished() } else { writeQueue.cancelAllOperations() } cleanupStream() - doDisconnect(error) + connected = false + if runDelegate { + doDisconnect(error) + } } - + /** cleanup the streams. */ @@ -432,8 +441,9 @@ open class WebSocket : NSObject, StreamDelegate { } outputStream = nil inputStream = nil + fragBuffer = nil } - + /** Handles the incoming bytes and sending them to the proper processing method. */ @@ -451,7 +461,7 @@ open class WebSocket : NSObject, StreamDelegate { dequeueInput() } } - + /** Dequeue the incoming input so it is processed in order. */ @@ -460,11 +470,11 @@ open class WebSocket : NSObject, StreamDelegate { autoreleasepool { let data = inputQueue[0] var work = data - if let fragBuffer = fragBuffer { - var combine = NSData(data: fragBuffer) as Data + if let buffer = fragBuffer { + var combine = NSData(data: buffer) as Data combine.append(data) work = combine - self.fragBuffer = nil + fragBuffer = nil } let buffer = UnsafeRawPointer((work as NSData).bytes).assumingMemoryBound(to: UInt8.self) let length = work.count @@ -477,7 +487,7 @@ open class WebSocket : NSObject, StreamDelegate { } } } - + /** Handle checking the inital connection status */ @@ -485,23 +495,15 @@ open class WebSocket : NSObject, StreamDelegate { let code = processHTTP(buffer, bufferLen: bufferLen) switch code { case 0: - isConnecting = false - connected = true - guard canDispatch else {return} - callbackQueue.async { [weak self] in - guard let s = self else { return } - s.onConnect?() - s.delegate?.websocketDidConnect(socket: s) - s.notificationCenter.post(name: NSNotification.Name(WebsocketDidConnectNotification), object: self) - } + break case -1: fragBuffer = Data(bytes: buffer, count: bufferLen) - break // do nothing, we are going to collect more data + break // do nothing, we are going to collect more data default: doDisconnect(errorWithDetail("Invalid HTTP upgrade", code: UInt16(code))) } } - + /** Finds the HTTP Packet in the TCP stream, by looking for the CRLF. */ @@ -525,6 +527,17 @@ open class WebSocket : NSObject, StreamDelegate { if code != 0 { return code } + isConnecting = false + connected = true + didDisconnect = false + if canDispatch { + callbackQueue.async { [weak self] in + guard let s = self else { return } + s.onConnect?() + s.delegate?.websocketDidConnect(socket: s) + s.notificationCenter.post(name: NSNotification.Name(WebsocketDidConnectNotification), object: self) + } + } totalSize += 1 //skip the last \n let restSize = bufferLen - totalSize if restSize > 0 { @@ -534,7 +547,7 @@ open class WebSocket : NSObject, StreamDelegate { } return -1 // Was unable to find the full TCP header. } - + /** Validates the HTTP is a 101 as per the RFC spec. */ @@ -555,14 +568,14 @@ open class WebSocket : NSObject, StreamDelegate { } return -1 } - + /** Read a 16 bit big endian value from a buffer */ private static func readUint16(_ buffer: UnsafePointer, offset: Int) -> UInt16 { return (UInt16(buffer[offset + 0]) << 8) | UInt16(buffer[offset + 1]) } - + /** Read a 64 bit big endian value from a buffer */ @@ -573,7 +586,7 @@ open class WebSocket : NSObject, StreamDelegate { } return value } - + /** Write a 16-bit big endian value to a buffer. */ @@ -581,7 +594,7 @@ open class WebSocket : NSObject, StreamDelegate { buffer[offset + 0] = UInt8(value >> 8) buffer[offset + 1] = UInt8(value & 0xff) } - + /** Write a 64-bit big endian value to a buffer. */ @@ -590,7 +603,7 @@ open class WebSocket : NSObject, StreamDelegate { buffer[offset + i] = UInt8((value >> (8*UInt64(7 - i))) & 0xff) } } - + /** Process one message at the start of `buffer`. Return another buffer (sharing storage) that contains the leftover contents of `buffer` that I didn't process. */ @@ -628,10 +641,10 @@ open class WebSocket : NSObject, StreamDelegate { let isControlFrame = (receivedOpcode == .connectionClose || receivedOpcode == .ping) if !isControlFrame && (receivedOpcode != .binaryFrame && receivedOpcode != .continueFrame && receivedOpcode != .textFrame && receivedOpcode != .pong) { - let errCode = CloseCode.protocolError.rawValue - doDisconnect(errorWithDetail("unknown opcode: \(receivedOpcode)", code: errCode)) - writeError(errCode) - return emptyBuffer + let errCode = CloseCode.protocolError.rawValue + doDisconnect(errorWithDetail("unknown opcode: \(receivedOpcode)", code: errCode)) + writeError(errCode) + return emptyBuffer } if isControlFrame && isFin == 0 { let errCode = CloseCode.protocolError.rawValue @@ -723,7 +736,7 @@ open class WebSocket : NSObject, StreamDelegate { if receivedOpcode == .continueFrame { let errCode = CloseCode.protocolError.rawValue doDisconnect(errorWithDetail("first frame can't be a continue frame", - code: errCode)) + code: errCode)) writeError(errCode) return emptyBuffer } @@ -738,7 +751,7 @@ open class WebSocket : NSObject, StreamDelegate { } else { let errCode = CloseCode.protocolError.rawValue doDisconnect(errorWithDetail("second and beyond of fragment message must be a continue frame", - code: errCode)) + code: errCode)) writeError(errCode) return emptyBuffer } @@ -753,12 +766,12 @@ open class WebSocket : NSObject, StreamDelegate { } _ = processResponse(response) } - + let step = Int(offset + numericCast(len)) return buffer.fromOffset(step) } } - + /** Process all messages in the buffer if possible. */ @@ -771,7 +784,7 @@ open class WebSocket : NSObject, StreamDelegate { fragBuffer = Data(buffer: buffer) } } - + /** Process the finished response of a buffer. */ @@ -808,7 +821,7 @@ open class WebSocket : NSObject, StreamDelegate { } return false } - + /** Create an error */ @@ -817,7 +830,7 @@ open class WebSocket : NSObject, StreamDelegate { details[NSLocalizedDescriptionKey] = detail return NSError(domain: WebSocket.ErrorDomain, code: Int(code), userInfo: details) } - + /** Write an error to the socket */ @@ -827,14 +840,16 @@ open class WebSocket : NSObject, StreamDelegate { WebSocket.writeUint16(buffer, offset: 0, value: code) dequeueWrite(Data(bytes: buffer, count: MemoryLayout.size), code: .connectionClose) } - + /** Used to write things to the stream */ private func dequeueWrite(_ data: Data, code: OpCode, writeCompletion: (() -> ())? = nil) { - writeQueue.addOperation { [weak self] in + let operation = BlockOperation() + operation.addExecutionBlock { [weak self, weak operation] in //stream isn't ready, let's wait guard let s = self else { return } + guard let sOperation = operation else { return } var offset = 2 let dataLength = data.count let frame = NSMutableData(capacity: dataLength + s.MaxFrameSize) @@ -855,13 +870,13 @@ open class WebSocket : NSObject, StreamDelegate { let maskKey = UnsafeMutablePointer(buffer + offset) _ = SecRandomCopyBytes(kSecRandomDefault, Int(MemoryLayout.size), maskKey) offset += MemoryLayout.size - + for i in 0...size] offset += 1 } var total = 0 - while true { + while !sOperation.isCancelled { guard let outStream = s.outputStream else { break } let writeBuffer = UnsafeRawPointer(frame!.bytes+total).assumingMemoryBound(to: UInt8.self) let len = outStream.write(writeBuffer, maxLength: offset-total) @@ -884,14 +899,14 @@ open class WebSocket : NSObject, StreamDelegate { callback() } } - + break } } - } + writeQueue.addOperation(operation) } - + /** Used to preform the disconnect delegate */ @@ -909,7 +924,7 @@ open class WebSocket : NSObject, StreamDelegate { s.notificationCenter.post(name: NSNotification.Name(WebsocketDidDisconnectNotification), object: self, userInfo: userInfo) } } - + // MARK: - Deinit deinit { mutex.lock() @@ -917,23 +932,23 @@ open class WebSocket : NSObject, StreamDelegate { mutex.unlock() cleanupStream() } - + } private extension Data { - + init(buffer: UnsafeBufferPointer) { self.init(bytes: buffer.baseAddress!, count: buffer.count) } - + } private extension UnsafeBufferPointer { - + func fromOffset(_ offset: Int) -> UnsafeBufferPointer { return UnsafeBufferPointer(start: baseAddress?.advanced(by: offset), count: count - offset) } - + } -private let emptyBuffer = UnsafeBufferPointer(start: nil, count: 0) +private let emptyBuffer = UnsafeBufferPointer(start: nil, count: 0) \ No newline at end of file