PageRenderTime 56ms CodeModel.GetById 20ms RepoModel.GetById 0ms app.codeStats 0ms

/Pods/Starscream/Source/WebSocket.swift

https://gitlab.com/joaopaulogalvao/browteco
Swift | 838 lines | 704 code | 52 blank | 82 comment | 189 complexity | e964a6f07650cbb4640c6fd97423013a MD5 | raw file
  1. //////////////////////////////////////////////////////////////////////////////////////////////////
  2. //
  3. // Websocket.swift
  4. //
  5. // Created by Dalton Cherry on 7/16/14.
  6. // Copyright (c) 2014-2015 Dalton Cherry.
  7. //
  8. // Licensed under the Apache License, Version 2.0 (the "License");
  9. // you may not use this file except in compliance with the License.
  10. // You may obtain a copy of the License at
  11. //
  12. // http://www.apache.org/licenses/LICENSE-2.0
  13. //
  14. // Unless required by applicable law or agreed to in writing, software
  15. // distributed under the License is distributed on an "AS IS" BASIS,
  16. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  17. // See the License for the specific language governing permissions and
  18. // limitations under the License.
  19. //
  20. //////////////////////////////////////////////////////////////////////////////////////////////////
  21. import Foundation
  22. import CoreFoundation
  23. import Security
  24. public protocol WebSocketDelegate: class {
  25. func websocketDidConnect(socket: WebSocket)
  26. func websocketDidDisconnect(socket: WebSocket, error: NSError?)
  27. func websocketDidReceiveMessage(socket: WebSocket, text: String)
  28. func websocketDidReceiveData(socket: WebSocket, data: NSData)
  29. }
  30. public protocol WebSocketPongDelegate: class {
  31. func websocketDidReceivePong(socket: WebSocket)
  32. }
  33. public class WebSocket : NSObject, NSStreamDelegate {
  34. enum OpCode : UInt8 {
  35. case ContinueFrame = 0x0
  36. case TextFrame = 0x1
  37. case BinaryFrame = 0x2
  38. //3-7 are reserved.
  39. case ConnectionClose = 0x8
  40. case Ping = 0x9
  41. case Pong = 0xA
  42. //B-F reserved.
  43. }
  44. public enum CloseCode : UInt16 {
  45. case Normal = 1000
  46. case GoingAway = 1001
  47. case ProtocolError = 1002
  48. case ProtocolUnhandledType = 1003
  49. // 1004 reserved.
  50. case NoStatusReceived = 1005
  51. //1006 reserved.
  52. case Encoding = 1007
  53. case PolicyViolated = 1008
  54. case MessageTooBig = 1009
  55. }
  56. public static let ErrorDomain = "WebSocket"
  57. enum InternalErrorCode : UInt16 {
  58. // 0-999 WebSocket status codes not used
  59. case OutputStreamWriteError = 1
  60. }
  61. //Where the callback is executed. It defaults to the main UI thread queue.
  62. public var queue = dispatch_get_main_queue()
  63. var optionalProtocols : [String]?
  64. //Constant Values.
  65. let headerWSUpgradeName = "Upgrade"
  66. let headerWSUpgradeValue = "websocket"
  67. let headerWSHostName = "Host"
  68. let headerWSConnectionName = "Connection"
  69. let headerWSConnectionValue = "Upgrade"
  70. let headerWSProtocolName = "Sec-WebSocket-Protocol"
  71. let headerWSVersionName = "Sec-WebSocket-Version"
  72. let headerWSVersionValue = "13"
  73. let headerWSKeyName = "Sec-WebSocket-Key"
  74. let headerOriginName = "Origin"
  75. let headerWSAcceptName = "Sec-WebSocket-Accept"
  76. let BUFFER_MAX = 4096
  77. let FinMask: UInt8 = 0x80
  78. let OpCodeMask: UInt8 = 0x0F
  79. let RSVMask: UInt8 = 0x70
  80. let MaskMask: UInt8 = 0x80
  81. let PayloadLenMask: UInt8 = 0x7F
  82. let MaxFrameSize: Int = 32
  83. class WSResponse {
  84. var isFin = false
  85. var code: OpCode = .ContinueFrame
  86. var bytesLeft = 0
  87. var frameCount = 0
  88. var buffer: NSMutableData?
  89. }
  90. public weak var delegate: WebSocketDelegate?
  91. public weak var pongDelegate: WebSocketPongDelegate?
  92. public var onConnect: ((Void) -> Void)?
  93. public var onDisconnect: ((NSError?) -> Void)?
  94. public var onText: ((String) -> Void)?
  95. public var onData: ((NSData) -> Void)?
  96. public var onPong: ((Void) -> Void)?
  97. public var headers = [String: String]()
  98. public var voipEnabled = false
  99. public var selfSignedSSL = false
  100. public var security: SSLSecurity?
  101. public var enabledSSLCipherSuites: [SSLCipherSuite]?
  102. public var origin: String?
  103. public var isConnected :Bool {
  104. return connected
  105. }
  106. public var currentURL: NSURL {return url}
  107. private var url: NSURL
  108. private var inputStream: NSInputStream?
  109. private var outputStream: NSOutputStream?
  110. private var connected = false
  111. private var isCreated = false
  112. private var writeQueue = NSOperationQueue()
  113. private var readStack = [WSResponse]()
  114. private var inputQueue = [NSData]()
  115. private var fragBuffer: NSData?
  116. private var certValidated = false
  117. private var didDisconnect = false
  118. private var readyToWrite = false
  119. private let mutex = NSLock()
  120. private var canDispatch: Bool {
  121. mutex.lock()
  122. let canWork = readyToWrite
  123. mutex.unlock()
  124. return canWork
  125. }
  126. //the shared processing queue used for all websocket
  127. private static let sharedWorkQueue = dispatch_queue_create("com.vluxe.starscream.websocket", DISPATCH_QUEUE_SERIAL)
  128. //used for setting protocols.
  129. public init(url: NSURL, protocols: [String]? = nil) {
  130. self.url = url
  131. self.origin = url.absoluteString
  132. writeQueue.maxConcurrentOperationCount = 1
  133. optionalProtocols = protocols
  134. }
  135. ///Connect to the websocket server on a background thread
  136. public func connect() {
  137. guard !isCreated else { return }
  138. didDisconnect = false
  139. isCreated = true
  140. createHTTPRequest()
  141. isCreated = false
  142. }
  143. /**
  144. 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.
  145. 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.
  146. If you supply a zero (or negative) `forceTimeout`, I immediately close the socket (without sending a Close control frame) and notify my delegate.
  147. - Parameter forceTimeout: Maximum time to wait for the server to close the socket.
  148. */
  149. public func disconnect(forceTimeout forceTimeout: NSTimeInterval? = nil) {
  150. switch forceTimeout {
  151. case .Some(let seconds) where seconds > 0:
  152. dispatch_after(dispatch_time(DISPATCH_TIME_NOW, Int64(seconds * Double(NSEC_PER_SEC))), queue) { [weak self] in
  153. self?.disconnectStream(nil)
  154. }
  155. fallthrough
  156. case .None:
  157. writeError(CloseCode.Normal.rawValue)
  158. default:
  159. self.disconnectStream(nil)
  160. break
  161. }
  162. }
  163. /**
  164. Write a string to the websocket. This sends it as a text frame.
  165. If you supply a non-nil completion block, I will perform it when the write completes.
  166. - parameter str: The string to write.
  167. - parameter completion: The (optional) completion handler.
  168. */
  169. public func writeString(str: String, completion: (() -> ())? = nil) {
  170. guard isConnected else { return }
  171. dequeueWrite(str.dataUsingEncoding(NSUTF8StringEncoding)!, code: .TextFrame, writeCompletion: completion)
  172. }
  173. /**
  174. Write binary data to the websocket. This sends it as a binary frame.
  175. If you supply a non-nil completion block, I will perform it when the write completes.
  176. - parameter data: The data to write.
  177. - parameter completion: The (optional) completion handler.
  178. */
  179. public func writeData(data: NSData, completion: (() -> ())? = nil) {
  180. guard isConnected else { return }
  181. dequeueWrite(data, code: .BinaryFrame, writeCompletion: completion)
  182. }
  183. //write a ping to the websocket. This sends it as a control frame.
  184. //yodel a sound to the planet. This sends it as an astroid. http://youtu.be/Eu5ZJELRiJ8?t=42s
  185. public func writePing(data: NSData, completion: (() -> ())? = nil) {
  186. guard isConnected else { return }
  187. dequeueWrite(data, code: .Ping, writeCompletion: completion)
  188. }
  189. //private method that starts the connection
  190. private func createHTTPRequest() {
  191. let urlRequest = CFHTTPMessageCreateRequest(kCFAllocatorDefault, "GET",
  192. url, kCFHTTPVersion1_1).takeRetainedValue()
  193. var port = url.port
  194. if port == nil {
  195. if ["wss", "https"].contains(url.scheme) {
  196. port = 443
  197. } else {
  198. port = 80
  199. }
  200. }
  201. addHeader(urlRequest, key: headerWSUpgradeName, val: headerWSUpgradeValue)
  202. addHeader(urlRequest, key: headerWSConnectionName, val: headerWSConnectionValue)
  203. if let protocols = optionalProtocols {
  204. addHeader(urlRequest, key: headerWSProtocolName, val: protocols.joinWithSeparator(","))
  205. }
  206. addHeader(urlRequest, key: headerWSVersionName, val: headerWSVersionValue)
  207. addHeader(urlRequest, key: headerWSKeyName, val: generateWebSocketKey())
  208. if let origin = origin {
  209. addHeader(urlRequest, key: headerOriginName, val: origin)
  210. }
  211. addHeader(urlRequest, key: headerWSHostName, val: "\(url.host!):\(port!)")
  212. for (key,value) in headers {
  213. addHeader(urlRequest, key: key, val: value)
  214. }
  215. if let cfHTTPMessage = CFHTTPMessageCopySerializedMessage(urlRequest) {
  216. let serializedRequest = cfHTTPMessage.takeRetainedValue()
  217. initStreamsWithData(serializedRequest, Int(port!))
  218. }
  219. }
  220. //Add a header to the CFHTTPMessage by using the NSString bridges to CFString
  221. private func addHeader(urlRequest: CFHTTPMessage, key: NSString, val: NSString) {
  222. CFHTTPMessageSetHeaderFieldValue(urlRequest, key, val)
  223. }
  224. //generate a websocket key as needed in rfc
  225. private func generateWebSocketKey() -> String {
  226. var key = ""
  227. let seed = 16
  228. for _ in 0..<seed {
  229. let uni = UnicodeScalar(UInt32(97 + arc4random_uniform(25)))
  230. key += "\(Character(uni))"
  231. }
  232. let data = key.dataUsingEncoding(NSUTF8StringEncoding)
  233. let baseKey = data?.base64EncodedStringWithOptions(NSDataBase64EncodingOptions(rawValue: 0))
  234. return baseKey!
  235. }
  236. //Start the stream connection and write the data to the output stream
  237. private func initStreamsWithData(data: NSData, _ port: Int) {
  238. //higher level API we will cut over to at some point
  239. //NSStream.getStreamsToHostWithName(url.host, port: url.port.integerValue, inputStream: &inputStream, outputStream: &outputStream)
  240. var readStream: Unmanaged<CFReadStream>?
  241. var writeStream: Unmanaged<CFWriteStream>?
  242. let h: NSString = url.host!
  243. CFStreamCreatePairWithSocketToHost(nil, h, UInt32(port), &readStream, &writeStream)
  244. inputStream = readStream!.takeRetainedValue()
  245. outputStream = writeStream!.takeRetainedValue()
  246. guard let inStream = inputStream, let outStream = outputStream else { return }
  247. inStream.delegate = self
  248. outStream.delegate = self
  249. if ["wss", "https"].contains(url.scheme) {
  250. inStream.setProperty(NSStreamSocketSecurityLevelNegotiatedSSL, forKey: NSStreamSocketSecurityLevelKey)
  251. outStream.setProperty(NSStreamSocketSecurityLevelNegotiatedSSL, forKey: NSStreamSocketSecurityLevelKey)
  252. } else {
  253. certValidated = true //not a https session, so no need to check SSL pinning
  254. }
  255. if voipEnabled {
  256. inStream.setProperty(NSStreamNetworkServiceTypeVoIP, forKey: NSStreamNetworkServiceType)
  257. outStream.setProperty(NSStreamNetworkServiceTypeVoIP, forKey: NSStreamNetworkServiceType)
  258. }
  259. if selfSignedSSL {
  260. let settings: [NSObject: NSObject] = [kCFStreamSSLValidatesCertificateChain: NSNumber(bool:false), kCFStreamSSLPeerName: kCFNull]
  261. inStream.setProperty(settings, forKey: kCFStreamPropertySSLSettings as String)
  262. outStream.setProperty(settings, forKey: kCFStreamPropertySSLSettings as String)
  263. }
  264. if let cipherSuites = self.enabledSSLCipherSuites {
  265. if let sslContextIn = CFReadStreamCopyProperty(inputStream, kCFStreamPropertySSLContext) as! SSLContextRef?,
  266. sslContextOut = CFWriteStreamCopyProperty(outputStream, kCFStreamPropertySSLContext) as! SSLContextRef? {
  267. let resIn = SSLSetEnabledCiphers(sslContextIn, cipherSuites, cipherSuites.count)
  268. let resOut = SSLSetEnabledCiphers(sslContextOut, cipherSuites, cipherSuites.count)
  269. if resIn != errSecSuccess {
  270. let error = self.errorWithDetail("Error setting ingoing cypher suites", code: UInt16(resIn))
  271. disconnectStream(error)
  272. return
  273. }
  274. if resOut != errSecSuccess {
  275. let error = self.errorWithDetail("Error setting outgoing cypher suites", code: UInt16(resOut))
  276. disconnectStream(error)
  277. return
  278. }
  279. }
  280. }
  281. CFReadStreamSetDispatchQueue(inStream, WebSocket.sharedWorkQueue)
  282. CFWriteStreamSetDispatchQueue(outStream, WebSocket.sharedWorkQueue)
  283. inStream.open()
  284. outStream.open()
  285. self.mutex.lock()
  286. self.readyToWrite = true
  287. self.mutex.unlock()
  288. let bytes = UnsafePointer<UInt8>(data.bytes)
  289. var timeout = 5000000 //wait 5 seconds before giving up
  290. writeQueue.addOperationWithBlock { [weak self] in
  291. while !outStream.hasSpaceAvailable {
  292. usleep(100) //wait until the socket is ready
  293. timeout -= 100
  294. if timeout < 0 {
  295. self?.cleanupStream()
  296. self?.doDisconnect(self?.errorWithDetail("write wait timed out", code: 2))
  297. return
  298. } else if outStream.streamError != nil {
  299. return //disconnectStream will be called.
  300. }
  301. }
  302. outStream.write(bytes, maxLength: data.length)
  303. }
  304. }
  305. //delegate for the stream methods. Processes incoming bytes
  306. public func stream(aStream: NSStream, handleEvent eventCode: NSStreamEvent) {
  307. if let sec = security where !certValidated && [.HasBytesAvailable, .HasSpaceAvailable].contains(eventCode) {
  308. let possibleTrust: AnyObject? = aStream.propertyForKey(kCFStreamPropertySSLPeerTrust as String)
  309. if let trust: AnyObject = possibleTrust {
  310. let domain: AnyObject? = aStream.propertyForKey(kCFStreamSSLPeerName as String)
  311. if sec.isValid(trust as! SecTrustRef, domain: domain as! String?) {
  312. certValidated = true
  313. } else {
  314. let error = errorWithDetail("Invalid SSL certificate", code: 1)
  315. disconnectStream(error)
  316. return
  317. }
  318. }
  319. }
  320. if eventCode == .HasBytesAvailable {
  321. if aStream == inputStream {
  322. processInputStream()
  323. }
  324. } else if eventCode == .ErrorOccurred {
  325. disconnectStream(aStream.streamError)
  326. } else if eventCode == .EndEncountered {
  327. disconnectStream(nil)
  328. }
  329. }
  330. //disconnect the stream object
  331. private func disconnectStream(error: NSError?) {
  332. if error == nil {
  333. writeQueue.waitUntilAllOperationsAreFinished()
  334. } else {
  335. writeQueue.cancelAllOperations()
  336. }
  337. cleanupStream()
  338. doDisconnect(error)
  339. }
  340. private func cleanupStream() {
  341. outputStream?.delegate = nil
  342. inputStream?.delegate = nil
  343. if let stream = inputStream {
  344. CFReadStreamSetDispatchQueue(stream, nil)
  345. stream.close()
  346. }
  347. if let stream = outputStream {
  348. CFWriteStreamSetDispatchQueue(stream, nil)
  349. stream.close()
  350. }
  351. outputStream = nil
  352. inputStream = nil
  353. }
  354. ///handles the incoming bytes and sending them to the proper processing method
  355. private func processInputStream() {
  356. let buf = NSMutableData(capacity: BUFFER_MAX)
  357. let buffer = UnsafeMutablePointer<UInt8>(buf!.bytes)
  358. let length = inputStream!.read(buffer, maxLength: BUFFER_MAX)
  359. guard length > 0 else { return }
  360. var process = false
  361. if inputQueue.count == 0 {
  362. process = true
  363. }
  364. inputQueue.append(NSData(bytes: buffer, length: length))
  365. if process {
  366. dequeueInput()
  367. }
  368. }
  369. ///dequeue the incoming input so it is processed in order
  370. private func dequeueInput() {
  371. guard !inputQueue.isEmpty else { return }
  372. let data = inputQueue[0]
  373. var work = data
  374. if let fragBuffer = fragBuffer {
  375. let combine = NSMutableData(data: fragBuffer)
  376. combine.appendData(data)
  377. work = combine
  378. self.fragBuffer = nil
  379. }
  380. let buffer = UnsafePointer<UInt8>(work.bytes)
  381. let length = work.length
  382. if !connected {
  383. processTCPHandshake(buffer, bufferLen: length)
  384. } else {
  385. processRawMessage(buffer, bufferLen: length)
  386. }
  387. inputQueue = inputQueue.filter{$0 != data}
  388. dequeueInput()
  389. }
  390. //handle checking the inital connection status
  391. private func processTCPHandshake(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
  392. let code = processHTTP(buffer, bufferLen: bufferLen)
  393. switch code {
  394. case 0:
  395. connected = true
  396. guard canDispatch else {return}
  397. dispatch_async(queue) { [weak self] in
  398. guard let s = self else { return }
  399. s.onConnect?()
  400. s.delegate?.websocketDidConnect(s)
  401. }
  402. case -1:
  403. fragBuffer = NSData(bytes: buffer, length: bufferLen)
  404. break //do nothing, we are going to collect more data
  405. default:
  406. doDisconnect(errorWithDetail("Invalid HTTP upgrade", code: UInt16(code)))
  407. }
  408. }
  409. ///Finds the HTTP Packet in the TCP stream, by looking for the CRLF.
  410. private func processHTTP(buffer: UnsafePointer<UInt8>, bufferLen: Int) -> Int {
  411. let CRLFBytes = [UInt8(ascii: "\r"), UInt8(ascii: "\n"), UInt8(ascii: "\r"), UInt8(ascii: "\n")]
  412. var k = 0
  413. var totalSize = 0
  414. for i in 0..<bufferLen {
  415. if buffer[i] == CRLFBytes[k] {
  416. k += 1
  417. if k == 3 {
  418. totalSize = i + 1
  419. break
  420. }
  421. } else {
  422. k = 0
  423. }
  424. }
  425. if totalSize > 0 {
  426. let code = validateResponse(buffer, bufferLen: totalSize)
  427. if code != 0 {
  428. return code
  429. }
  430. totalSize += 1 //skip the last \n
  431. let restSize = bufferLen - totalSize
  432. if restSize > 0 {
  433. processRawMessage((buffer+totalSize),bufferLen: restSize)
  434. }
  435. return 0 //success
  436. }
  437. return -1 //was unable to find the full TCP header
  438. }
  439. ///validates the HTTP is a 101 as per the RFC spec
  440. private func validateResponse(buffer: UnsafePointer<UInt8>, bufferLen: Int) -> Int {
  441. let response = CFHTTPMessageCreateEmpty(kCFAllocatorDefault, false).takeRetainedValue()
  442. CFHTTPMessageAppendBytes(response, buffer, bufferLen)
  443. let code = CFHTTPMessageGetResponseStatusCode(response)
  444. if code != 101 {
  445. return code
  446. }
  447. if let cfHeaders = CFHTTPMessageCopyAllHeaderFields(response) {
  448. let headers = cfHeaders.takeRetainedValue() as NSDictionary
  449. if let acceptKey = headers[headerWSAcceptName] as? NSString {
  450. if acceptKey.length > 0 {
  451. return 0
  452. }
  453. }
  454. }
  455. return -1
  456. }
  457. ///read a 16 bit big endian value from a buffer
  458. private static func readUint16(buffer: UnsafePointer<UInt8>, offset: Int) -> UInt16 {
  459. return (UInt16(buffer[offset + 0]) << 8) | UInt16(buffer[offset + 1])
  460. }
  461. ///read a 64 bit big endian value from a buffer
  462. private static func readUint64(buffer: UnsafePointer<UInt8>, offset: Int) -> UInt64 {
  463. var value = UInt64(0)
  464. for i in 0...7 {
  465. value = (value << 8) | UInt64(buffer[offset + i])
  466. }
  467. return value
  468. }
  469. ///write a 16 bit big endian value to a buffer
  470. private static func writeUint16(buffer: UnsafeMutablePointer<UInt8>, offset: Int, value: UInt16) {
  471. buffer[offset + 0] = UInt8(value >> 8)
  472. buffer[offset + 1] = UInt8(value & 0xff)
  473. }
  474. ///write a 64 bit big endian value to a buffer
  475. private static func writeUint64(buffer: UnsafeMutablePointer<UInt8>, offset: Int, value: UInt64) {
  476. for i in 0...7 {
  477. buffer[offset + i] = UInt8((value >> (8*UInt64(7 - i))) & 0xff)
  478. }
  479. }
  480. ///process the websocket data
  481. private func processRawMessage(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
  482. let response = readStack.last
  483. if response != nil && bufferLen < 2 {
  484. fragBuffer = NSData(bytes: buffer, length: bufferLen)
  485. return
  486. }
  487. if let response = response where response.bytesLeft > 0 {
  488. var len = response.bytesLeft
  489. var extra = bufferLen - response.bytesLeft
  490. if response.bytesLeft > bufferLen {
  491. len = bufferLen
  492. extra = 0
  493. }
  494. response.bytesLeft -= len
  495. response.buffer?.appendData(NSData(bytes: buffer, length: len))
  496. processResponse(response)
  497. let offset = bufferLen - extra
  498. if extra > 0 {
  499. processExtra((buffer+offset), bufferLen: extra)
  500. }
  501. return
  502. } else {
  503. let isFin = (FinMask & buffer[0])
  504. let receivedOpcode = OpCode(rawValue: (OpCodeMask & buffer[0]))
  505. let isMasked = (MaskMask & buffer[1])
  506. let payloadLen = (PayloadLenMask & buffer[1])
  507. var offset = 2
  508. if (isMasked > 0 || (RSVMask & buffer[0]) > 0) && receivedOpcode != .Pong {
  509. let errCode = CloseCode.ProtocolError.rawValue
  510. doDisconnect(errorWithDetail("masked and rsv data is not currently supported", code: errCode))
  511. writeError(errCode)
  512. return
  513. }
  514. let isControlFrame = (receivedOpcode == .ConnectionClose || receivedOpcode == .Ping)
  515. if !isControlFrame && (receivedOpcode != .BinaryFrame && receivedOpcode != .ContinueFrame &&
  516. receivedOpcode != .TextFrame && receivedOpcode != .Pong) {
  517. let errCode = CloseCode.ProtocolError.rawValue
  518. doDisconnect(errorWithDetail("unknown opcode: \(receivedOpcode)", code: errCode))
  519. writeError(errCode)
  520. return
  521. }
  522. if isControlFrame && isFin == 0 {
  523. let errCode = CloseCode.ProtocolError.rawValue
  524. doDisconnect(errorWithDetail("control frames can't be fragmented", code: errCode))
  525. writeError(errCode)
  526. return
  527. }
  528. if receivedOpcode == .ConnectionClose {
  529. var code = CloseCode.Normal.rawValue
  530. if payloadLen == 1 {
  531. code = CloseCode.ProtocolError.rawValue
  532. } else if payloadLen > 1 {
  533. code = WebSocket.readUint16(buffer, offset: offset)
  534. if code < 1000 || (code > 1003 && code < 1007) || (code > 1011 && code < 3000) {
  535. code = CloseCode.ProtocolError.rawValue
  536. }
  537. offset += 2
  538. }
  539. if payloadLen > 2 {
  540. let len = Int(payloadLen-2)
  541. if len > 0 {
  542. let bytes = UnsafePointer<UInt8>((buffer+offset))
  543. let str: NSString? = NSString(data: NSData(bytes: bytes, length: len), encoding: NSUTF8StringEncoding)
  544. if str == nil {
  545. code = CloseCode.ProtocolError.rawValue
  546. }
  547. }
  548. }
  549. doDisconnect(errorWithDetail("connection closed by server", code: code))
  550. writeError(code)
  551. return
  552. }
  553. if isControlFrame && payloadLen > 125 {
  554. writeError(CloseCode.ProtocolError.rawValue)
  555. return
  556. }
  557. var dataLength = UInt64(payloadLen)
  558. if dataLength == 127 {
  559. dataLength = WebSocket.readUint64(buffer, offset: offset)
  560. offset += sizeof(UInt64)
  561. } else if dataLength == 126 {
  562. dataLength = UInt64(WebSocket.readUint16(buffer, offset: offset))
  563. offset += sizeof(UInt16)
  564. }
  565. if bufferLen < offset || UInt64(bufferLen - offset) < dataLength {
  566. fragBuffer = NSData(bytes: buffer, length: bufferLen)
  567. return
  568. }
  569. var len = dataLength
  570. if dataLength > UInt64(bufferLen) {
  571. len = UInt64(bufferLen-offset)
  572. }
  573. let data: NSData
  574. if len < 0 {
  575. len = 0
  576. data = NSData()
  577. } else {
  578. data = NSData(bytes: UnsafePointer<UInt8>((buffer+offset)), length: Int(len))
  579. }
  580. if receivedOpcode == .Pong {
  581. if canDispatch {
  582. dispatch_async(queue) { [weak self] in
  583. guard let s = self else { return }
  584. s.onPong?()
  585. s.pongDelegate?.websocketDidReceivePong(s)
  586. }
  587. }
  588. let step = Int(offset+numericCast(len))
  589. let extra = bufferLen-step
  590. if extra > 0 {
  591. processRawMessage((buffer+step), bufferLen: extra)
  592. }
  593. return
  594. }
  595. var response = readStack.last
  596. if isControlFrame {
  597. response = nil //don't append pings
  598. }
  599. if isFin == 0 && receivedOpcode == .ContinueFrame && response == nil {
  600. let errCode = CloseCode.ProtocolError.rawValue
  601. doDisconnect(errorWithDetail("continue frame before a binary or text frame", code: errCode))
  602. writeError(errCode)
  603. return
  604. }
  605. var isNew = false
  606. if response == nil {
  607. if receivedOpcode == .ContinueFrame {
  608. let errCode = CloseCode.ProtocolError.rawValue
  609. doDisconnect(errorWithDetail("first frame can't be a continue frame",
  610. code: errCode))
  611. writeError(errCode)
  612. return
  613. }
  614. isNew = true
  615. response = WSResponse()
  616. response!.code = receivedOpcode!
  617. response!.bytesLeft = Int(dataLength)
  618. response!.buffer = NSMutableData(data: data)
  619. } else {
  620. if receivedOpcode == .ContinueFrame {
  621. response!.bytesLeft = Int(dataLength)
  622. } else {
  623. let errCode = CloseCode.ProtocolError.rawValue
  624. doDisconnect(errorWithDetail("second and beyond of fragment message must be a continue frame",
  625. code: errCode))
  626. writeError(errCode)
  627. return
  628. }
  629. response!.buffer!.appendData(data)
  630. }
  631. if let response = response {
  632. response.bytesLeft -= Int(len)
  633. response.frameCount += 1
  634. response.isFin = isFin > 0 ? true : false
  635. if isNew {
  636. readStack.append(response)
  637. }
  638. processResponse(response)
  639. }
  640. let step = Int(offset+numericCast(len))
  641. let extra = bufferLen-step
  642. if extra > 0 {
  643. processExtra((buffer+step), bufferLen: extra)
  644. }
  645. }
  646. }
  647. ///process the extra of a buffer
  648. private func processExtra(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
  649. if bufferLen < 2 {
  650. fragBuffer = NSData(bytes: buffer, length: bufferLen)
  651. } else {
  652. processRawMessage(buffer, bufferLen: bufferLen)
  653. }
  654. }
  655. ///process the finished response of a buffer
  656. private func processResponse(response: WSResponse) -> Bool {
  657. if response.isFin && response.bytesLeft <= 0 {
  658. if response.code == .Ping {
  659. let data = response.buffer! //local copy so it is perverse for writing
  660. dequeueWrite(data, code: OpCode.Pong)
  661. } else if response.code == .TextFrame {
  662. let str: NSString? = NSString(data: response.buffer!, encoding: NSUTF8StringEncoding)
  663. if str == nil {
  664. writeError(CloseCode.Encoding.rawValue)
  665. return false
  666. }
  667. if canDispatch {
  668. dispatch_async(queue) { [weak self] in
  669. guard let s = self else { return }
  670. s.onText?(str! as String)
  671. s.delegate?.websocketDidReceiveMessage(s, text: str! as String)
  672. }
  673. }
  674. } else if response.code == .BinaryFrame {
  675. if canDispatch {
  676. let data = response.buffer! //local copy so it is perverse for writing
  677. dispatch_async(queue) { [weak self] in
  678. guard let s = self else { return }
  679. s.onData?(data)
  680. s.delegate?.websocketDidReceiveData(s, data: data)
  681. }
  682. }
  683. }
  684. readStack.removeLast()
  685. return true
  686. }
  687. return false
  688. }
  689. ///Create an error
  690. private func errorWithDetail(detail: String, code: UInt16) -> NSError {
  691. var details = [String: String]()
  692. details[NSLocalizedDescriptionKey] = detail
  693. return NSError(domain: WebSocket.ErrorDomain, code: Int(code), userInfo: details)
  694. }
  695. ///write a an error to the socket
  696. private func writeError(code: UInt16) {
  697. let buf = NSMutableData(capacity: sizeof(UInt16))
  698. let buffer = UnsafeMutablePointer<UInt8>(buf!.bytes)
  699. WebSocket.writeUint16(buffer, offset: 0, value: code)
  700. dequeueWrite(NSData(bytes: buffer, length: sizeof(UInt16)), code: .ConnectionClose)
  701. }
  702. ///used to write things to the stream
  703. private func dequeueWrite(data: NSData, code: OpCode, writeCompletion: (() -> ())? = nil) {
  704. writeQueue.addOperationWithBlock { [weak self] in
  705. //stream isn't ready, let's wait
  706. guard let s = self else { return }
  707. var offset = 2
  708. let bytes = UnsafeMutablePointer<UInt8>(data.bytes)
  709. let dataLength = data.length
  710. let frame = NSMutableData(capacity: dataLength + s.MaxFrameSize)
  711. let buffer = UnsafeMutablePointer<UInt8>(frame!.mutableBytes)
  712. buffer[0] = s.FinMask | code.rawValue
  713. if dataLength < 126 {
  714. buffer[1] = CUnsignedChar(dataLength)
  715. } else if dataLength <= Int(UInt16.max) {
  716. buffer[1] = 126
  717. WebSocket.writeUint16(buffer, offset: offset, value: UInt16(dataLength))
  718. offset += sizeof(UInt16)
  719. } else {
  720. buffer[1] = 127
  721. WebSocket.writeUint64(buffer, offset: offset, value: UInt64(dataLength))
  722. offset += sizeof(UInt64)
  723. }
  724. buffer[1] |= s.MaskMask
  725. let maskKey = UnsafeMutablePointer<UInt8>(buffer + offset)
  726. SecRandomCopyBytes(kSecRandomDefault, Int(sizeof(UInt32)), maskKey)
  727. offset += sizeof(UInt32)
  728. for i in 0..<dataLength {
  729. buffer[offset] = bytes[i] ^ maskKey[i % sizeof(UInt32)]
  730. offset += 1
  731. }
  732. var total = 0
  733. while true {
  734. guard let outStream = s.outputStream else { break }
  735. let writeBuffer = UnsafePointer<UInt8>(frame!.bytes+total)
  736. let len = outStream.write(writeBuffer, maxLength: offset-total)
  737. if len < 0 {
  738. var error: NSError?
  739. if let streamError = outStream.streamError {
  740. error = streamError
  741. } else {
  742. let errCode = InternalErrorCode.OutputStreamWriteError.rawValue
  743. error = s.errorWithDetail("output stream error during write", code: errCode)
  744. }
  745. s.doDisconnect(error)
  746. break
  747. } else {
  748. total += len
  749. }
  750. if total >= offset {
  751. if let queue = self?.queue, callback = writeCompletion {
  752. dispatch_async(queue) {
  753. callback()
  754. }
  755. }
  756. break
  757. }
  758. }
  759. }
  760. }
  761. ///used to preform the disconnect delegate
  762. private func doDisconnect(error: NSError?) {
  763. guard !didDisconnect else { return }
  764. didDisconnect = true
  765. connected = false
  766. guard canDispatch else {return}
  767. dispatch_async(queue) { [weak self] in
  768. guard let s = self else { return }
  769. s.onDisconnect?(error)
  770. s.delegate?.websocketDidDisconnect(s, error: error)
  771. }
  772. }
  773. deinit {
  774. mutex.lock()
  775. readyToWrite = false
  776. mutex.unlock()
  777. cleanupStream()
  778. }
  779. }