WebSocket.swift 42 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081
  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 timeout = 5
  104. public var isConnected :Bool {
  105. return connected
  106. }
  107. public var currentURL: NSURL {return url}
  108. private var url: NSURL
  109. private var inputStream: NSInputStream?
  110. private var outputStream: NSOutputStream?
  111. private var connected = false
  112. private var isCreated = false
  113. private var writeQueue = NSOperationQueue()
  114. private var readStack = [WSResponse]()
  115. private var inputQueue = [NSData]()
  116. private var fragBuffer: NSData?
  117. private var certValidated = false
  118. private var didDisconnect = false
  119. private var readyToWrite = false
  120. private let mutex = NSLock()
  121. private var canDispatch: Bool {
  122. mutex.lock()
  123. let canWork = readyToWrite
  124. mutex.unlock()
  125. return canWork
  126. }
  127. //the shared processing queue used for all websocket
  128. private static let sharedWorkQueue = dispatch_queue_create("com.vluxe.starscream.websocket", DISPATCH_QUEUE_SERIAL)
  129. //used for setting protocols.
  130. public init(url: NSURL, protocols: [String]? = nil) {
  131. self.url = url
  132. self.origin = url.absoluteString
  133. writeQueue.maxConcurrentOperationCount = 1
  134. optionalProtocols = protocols
  135. }
  136. ///Connect to the websocket server on a background thread
  137. public func connect() {
  138. guard !isCreated else { return }
  139. didDisconnect = false
  140. isCreated = true
  141. createHTTPRequest()
  142. isCreated = false
  143. }
  144. /**
  145. 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.
  146. 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.
  147. If you supply a zero (or negative) `forceTimeout`, I immediately close the socket (without sending a Close control frame) and notify my delegate.
  148. - Parameter forceTimeout: Maximum time to wait for the server to close the socket.
  149. */
  150. public func disconnect(forceTimeout forceTimeout: NSTimeInterval? = nil) {
  151. switch forceTimeout {
  152. case .Some(let seconds) where seconds > 0:
  153. dispatch_after(dispatch_time(DISPATCH_TIME_NOW, Int64(seconds * Double(NSEC_PER_SEC))), queue) { [weak self] in
  154. self?.disconnectStream(nil)
  155. }
  156. fallthrough
  157. case .None:
  158. writeError(CloseCode.Normal.rawValue)
  159. default:
  160. self.disconnectStream(nil)
  161. break
  162. }
  163. }
  164. /**
  165. Write a string to the websocket. This sends it as a text frame.
  166. If you supply a non-nil completion block, I will perform it when the write completes.
  167. - parameter str: The string to write.
  168. - parameter completion: The (optional) completion handler.
  169. */
  170. public func writeString(str: String, completion: (() -> ())? = nil) {
  171. guard isConnected else { return }
  172. dequeueWrite(str.dataUsingEncoding(NSUTF8StringEncoding)!, code: .TextFrame, writeCompletion: completion)
  173. }
  174. /**
  175. Write binary data to the websocket. This sends it as a binary frame.
  176. If you supply a non-nil completion block, I will perform it when the write completes.
  177. - parameter data: The data to write.
  178. - parameter completion: The (optional) completion handler.
  179. */
  180. public func writeData(data: NSData, completion: (() -> ())? = nil) {
  181. guard isConnected else { return }
  182. dequeueWrite(data, code: .BinaryFrame, writeCompletion: completion)
  183. }
  184. //write a ping to the websocket. This sends it as a control frame.
  185. //yodel a sound to the planet. This sends it as an astroid. http://youtu.be/Eu5ZJELRiJ8?t=42s
  186. public func writePing(data: NSData, completion: (() -> ())? = nil) {
  187. guard isConnected else { return }
  188. dequeueWrite(data, code: .Ping, writeCompletion: completion)
  189. }
  190. //private method that starts the connection
  191. private func createHTTPRequest() {
  192. let urlRequest = CFHTTPMessageCreateRequest(kCFAllocatorDefault, "GET",
  193. url, kCFHTTPVersion1_1).takeRetainedValue()
  194. var port = url.port
  195. if port == nil {
  196. if ["wss", "https"].contains(url.scheme) {
  197. port = 443
  198. } else {
  199. port = 80
  200. }
  201. }
  202. addHeader(urlRequest, key: headerWSUpgradeName, val: headerWSUpgradeValue)
  203. addHeader(urlRequest, key: headerWSConnectionName, val: headerWSConnectionValue)
  204. if let protocols = optionalProtocols {
  205. addHeader(urlRequest, key: headerWSProtocolName, val: protocols.joinWithSeparator(","))
  206. }
  207. addHeader(urlRequest, key: headerWSVersionName, val: headerWSVersionValue)
  208. addHeader(urlRequest, key: headerWSKeyName, val: generateWebSocketKey())
  209. if let origin = origin {
  210. addHeader(urlRequest, key: headerOriginName, val: origin)
  211. }
  212. addHeader(urlRequest, key: headerWSHostName, val: "\(url.host!):\(port!)")
  213. for (key,value) in headers {
  214. addHeader(urlRequest, key: key, val: value)
  215. }
  216. if let cfHTTPMessage = CFHTTPMessageCopySerializedMessage(urlRequest) {
  217. let serializedRequest = cfHTTPMessage.takeRetainedValue()
  218. initStreamsWithData(serializedRequest, Int(port!))
  219. }
  220. }
  221. //Add a header to the CFHTTPMessage by using the NSString bridges to CFString
  222. private func addHeader(urlRequest: CFHTTPMessage, key: NSString, val: NSString) {
  223. CFHTTPMessageSetHeaderFieldValue(urlRequest, key, val)
  224. }
  225. //generate a websocket key as needed in rfc
  226. private func generateWebSocketKey() -> String {
  227. var key = ""
  228. let seed = 16
  229. for _ in 0..<seed {
  230. let uni = UnicodeScalar(UInt32(97 + arc4random_uniform(25)))
  231. key += "\(Character(uni))"
  232. }
  233. let data = key.dataUsingEncoding(NSUTF8StringEncoding)
  234. let baseKey = data?.base64EncodedStringWithOptions(NSDataBase64EncodingOptions(rawValue: 0))
  235. return baseKey!
  236. }
  237. //Start the stream connection and write the data to the output stream
  238. private func initStreamsWithData(data: NSData, _ port: Int) {
  239. //higher level API we will cut over to at some point
  240. //NSStream.getStreamsToHostWithName(url.host, port: url.port.integerValue, inputStream: &inputStream, outputStream: &outputStream)
  241. var readStream: Unmanaged<CFReadStream>?
  242. var writeStream: Unmanaged<CFWriteStream>?
  243. let h: NSString = url.host!
  244. CFStreamCreatePairWithSocketToHost(nil, h, UInt32(port), &readStream, &writeStream)
  245. inputStream = readStream!.takeRetainedValue()
  246. outputStream = writeStream!.takeRetainedValue()
  247. guard let inStream = inputStream, let outStream = outputStream else { return }
  248. inStream.delegate = self
  249. outStream.delegate = self
  250. if ["wss", "https"].contains(url.scheme) {
  251. inStream.setProperty(NSStreamSocketSecurityLevelNegotiatedSSL, forKey: NSStreamSocketSecurityLevelKey)
  252. outStream.setProperty(NSStreamSocketSecurityLevelNegotiatedSSL, forKey: NSStreamSocketSecurityLevelKey)
  253. } else {
  254. certValidated = true //not a https session, so no need to check SSL pinning
  255. }
  256. if voipEnabled {
  257. inStream.setProperty(NSStreamNetworkServiceTypeVoIP, forKey: NSStreamNetworkServiceType)
  258. outStream.setProperty(NSStreamNetworkServiceTypeVoIP, forKey: NSStreamNetworkServiceType)
  259. }
  260. if selfSignedSSL {
  261. let settings: [NSObject: NSObject] = [kCFStreamSSLValidatesCertificateChain: NSNumber(bool:false), kCFStreamSSLPeerName: kCFNull]
  262. inStream.setProperty(settings, forKey: kCFStreamPropertySSLSettings as String)
  263. outStream.setProperty(settings, forKey: kCFStreamPropertySSLSettings as String)
  264. }
  265. if let cipherSuites = self.enabledSSLCipherSuites {
  266. if let sslContextIn = CFReadStreamCopyProperty(inputStream, kCFStreamPropertySSLContext) as! SSLContextRef?,
  267. sslContextOut = CFWriteStreamCopyProperty(outputStream, kCFStreamPropertySSLContext) as! SSLContextRef? {
  268. let resIn = SSLSetEnabledCiphers(sslContextIn, cipherSuites, cipherSuites.count)
  269. let resOut = SSLSetEnabledCiphers(sslContextOut, cipherSuites, cipherSuites.count)
  270. if resIn != errSecSuccess {
  271. let error = self.errorWithDetail("Error setting ingoing cypher suites", code: UInt16(resIn))
  272. disconnectStream(error)
  273. return
  274. }
  275. if resOut != errSecSuccess {
  276. let error = self.errorWithDetail("Error setting outgoing cypher suites", code: UInt16(resOut))
  277. disconnectStream(error)
  278. return
  279. }
  280. }
  281. }
  282. CFReadStreamSetDispatchQueue(inStream, WebSocket.sharedWorkQueue)
  283. CFWriteStreamSetDispatchQueue(outStream, WebSocket.sharedWorkQueue)
  284. inStream.open()
  285. outStream.open()
  286. self.mutex.lock()
  287. self.readyToWrite = true
  288. self.mutex.unlock()
  289. let bytes = UnsafePointer<UInt8>(data.bytes)
  290. var out = timeout * 1000000 //wait 5 seconds before giving up
  291. writeQueue.addOperationWithBlock { [weak self] in
  292. while !outStream.hasSpaceAvailable {
  293. usleep(100) //wait until the socket is ready
  294. out -= 100
  295. if out < 0 {
  296. self?.cleanupStream()
  297. self?.doDisconnect(self?.errorWithDetail("write wait timed out", code: 2))
  298. return
  299. } else if outStream.streamError != nil {
  300. return //disconnectStream will be called.
  301. }
  302. }
  303. outStream.write(bytes, maxLength: data.length)
  304. }
  305. }
  306. //delegate for the stream methods. Processes incoming bytes
  307. public func stream(aStream: NSStream, handleEvent eventCode: NSStreamEvent) {
  308. if let sec = security where !certValidated && [.HasBytesAvailable, .HasSpaceAvailable].contains(eventCode) {
  309. let possibleTrust: AnyObject? = aStream.propertyForKey(kCFStreamPropertySSLPeerTrust as String)
  310. if let trust: AnyObject = possibleTrust {
  311. let domain: AnyObject? = aStream.propertyForKey(kCFStreamSSLPeerName as String)
  312. if sec.isValid(trust as! SecTrustRef, domain: domain as! String?) {
  313. certValidated = true
  314. } else {
  315. let error = errorWithDetail("Invalid SSL certificate", code: 1)
  316. disconnectStream(error)
  317. return
  318. }
  319. }
  320. }
  321. if eventCode == .HasBytesAvailable {
  322. if aStream == inputStream {
  323. processInputStream()
  324. }
  325. } else if eventCode == .ErrorOccurred {
  326. disconnectStream(aStream.streamError)
  327. } else if eventCode == .EndEncountered {
  328. disconnectStream(nil)
  329. }
  330. }
  331. //disconnect the stream object
  332. private func disconnectStream(error: NSError?) {
  333. if error == nil {
  334. writeQueue.waitUntilAllOperationsAreFinished()
  335. } else {
  336. writeQueue.cancelAllOperations()
  337. }
  338. cleanupStream()
  339. doDisconnect(error)
  340. }
  341. private func cleanupStream() {
  342. outputStream?.delegate = nil
  343. inputStream?.delegate = nil
  344. if let stream = inputStream {
  345. CFReadStreamSetDispatchQueue(stream, nil)
  346. stream.close()
  347. }
  348. if let stream = outputStream {
  349. CFWriteStreamSetDispatchQueue(stream, nil)
  350. stream.close()
  351. }
  352. outputStream = nil
  353. inputStream = nil
  354. }
  355. ///handles the incoming bytes and sending them to the proper processing method
  356. private func processInputStream() {
  357. let buf = NSMutableData(capacity: BUFFER_MAX)
  358. let buffer = UnsafeMutablePointer<UInt8>(buf!.bytes)
  359. let length = inputStream!.read(buffer, maxLength: BUFFER_MAX)
  360. guard length > 0 else { return }
  361. var process = false
  362. if inputQueue.count == 0 {
  363. process = true
  364. }
  365. inputQueue.append(NSData(bytes: buffer, length: length))
  366. if process {
  367. dequeueInput()
  368. }
  369. }
  370. ///dequeue the incoming input so it is processed in order
  371. private func dequeueInput() {
  372. while !inputQueue.isEmpty {
  373. let data = inputQueue[0]
  374. var work = data
  375. if let fragBuffer = fragBuffer {
  376. let combine = NSMutableData(data: fragBuffer)
  377. combine.appendData(data)
  378. work = combine
  379. self.fragBuffer = nil
  380. }
  381. let buffer = UnsafePointer<UInt8>(work.bytes)
  382. let length = work.length
  383. if !connected {
  384. processTCPHandshake(buffer, bufferLen: length)
  385. } else {
  386. processRawMessagesInBuffer(buffer, bufferLen: length)
  387. }
  388. inputQueue = inputQueue.filter{$0 != data}
  389. }
  390. }
  391. //handle checking the inital connection status
  392. private func processTCPHandshake(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
  393. let code = processHTTP(buffer, bufferLen: bufferLen)
  394. switch code {
  395. case 0:
  396. connected = true
  397. guard canDispatch else {return}
  398. dispatch_async(queue) { [weak self] in
  399. guard let s = self else { return }
  400. s.onConnect?()
  401. s.delegate?.websocketDidConnect(s)
  402. }
  403. case -1:
  404. fragBuffer = NSData(bytes: buffer, length: bufferLen)
  405. break //do nothing, we are going to collect more data
  406. default:
  407. doDisconnect(errorWithDetail("Invalid HTTP upgrade", code: UInt16(code)))
  408. }
  409. }
  410. ///Finds the HTTP Packet in the TCP stream, by looking for the CRLF.
  411. private func processHTTP(buffer: UnsafePointer<UInt8>, bufferLen: Int) -> Int {
  412. let CRLFBytes = [UInt8(ascii: "\r"), UInt8(ascii: "\n"), UInt8(ascii: "\r"), UInt8(ascii: "\n")]
  413. var k = 0
  414. var totalSize = 0
  415. for i in 0..<bufferLen {
  416. if buffer[i] == CRLFBytes[k] {
  417. k += 1
  418. if k == 3 {
  419. totalSize = i + 1
  420. break
  421. }
  422. } else {
  423. k = 0
  424. }
  425. }
  426. if totalSize > 0 {
  427. let code = validateResponse(buffer, bufferLen: totalSize)
  428. if code != 0 {
  429. return code
  430. }
  431. totalSize += 1 //skip the last \n
  432. let restSize = bufferLen - totalSize
  433. if restSize > 0 {
  434. processRawMessagesInBuffer(buffer + totalSize, bufferLen: restSize)
  435. }
  436. return 0 //success
  437. }
  438. return -1 //was unable to find the full TCP header
  439. }
  440. ///validates the HTTP is a 101 as per the RFC spec
  441. private func validateResponse(buffer: UnsafePointer<UInt8>, bufferLen: Int) -> Int {
  442. let response = CFHTTPMessageCreateEmpty(kCFAllocatorDefault, false).takeRetainedValue()
  443. CFHTTPMessageAppendBytes(response, buffer, bufferLen)
  444. let code = CFHTTPMessageGetResponseStatusCode(response)
  445. if code != 101 {
  446. return code
  447. }
  448. if let cfHeaders = CFHTTPMessageCopyAllHeaderFields(response) {
  449. let headers = cfHeaders.takeRetainedValue() as NSDictionary
  450. if let acceptKey = headers[headerWSAcceptName] as? NSString {
  451. if acceptKey.length > 0 {
  452. return 0
  453. }
  454. }
  455. }
  456. return -1
  457. }
  458. ///read a 16 bit big endian value from a buffer
  459. private static func readUint16(buffer: UnsafePointer<UInt8>, offset: Int) -> UInt16 {
  460. return (UInt16(buffer[offset + 0]) << 8) | UInt16(buffer[offset + 1])
  461. }
  462. ///read a 64 bit big endian value from a buffer
  463. private static func readUint64(buffer: UnsafePointer<UInt8>, offset: Int) -> UInt64 {
  464. var value = UInt64(0)
  465. for i in 0...7 {
  466. value = (value << 8) | UInt64(buffer[offset + i])
  467. }
  468. return value
  469. }
  470. ///write a 16 bit big endian value to a buffer
  471. private static func writeUint16(buffer: UnsafeMutablePointer<UInt8>, offset: Int, value: UInt16) {
  472. buffer[offset + 0] = UInt8(value >> 8)
  473. buffer[offset + 1] = UInt8(value & 0xff)
  474. }
  475. ///write a 64 bit big endian value to a buffer
  476. private static func writeUint64(buffer: UnsafeMutablePointer<UInt8>, offset: Int, value: UInt64) {
  477. for i in 0...7 {
  478. buffer[offset + i] = UInt8((value >> (8*UInt64(7 - i))) & 0xff)
  479. }
  480. }
  481. /// 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.
  482. @warn_unused_result
  483. private func processOneRawMessage(inBuffer buffer: UnsafeBufferPointer<UInt8>) -> UnsafeBufferPointer<UInt8> {
  484. let response = readStack.last
  485. let baseAddress = buffer.baseAddress
  486. let bufferLen = buffer.count
  487. if response != nil && bufferLen < 2 {
  488. fragBuffer = NSData(buffer: buffer)
  489. return emptyBuffer
  490. }
  491. if let response = response where response.bytesLeft > 0 {
  492. var len = response.bytesLeft
  493. var extra = bufferLen - response.bytesLeft
  494. if response.bytesLeft > bufferLen {
  495. len = bufferLen
  496. extra = 0
  497. }
  498. response.bytesLeft -= len
  499. response.buffer?.appendData(NSData(bytes: baseAddress, length: len))
  500. processResponse(response)
  501. return buffer.fromOffset(bufferLen - extra)
  502. } else {
  503. let isFin = (FinMask & baseAddress[0])
  504. let receivedOpcode = OpCode(rawValue: (OpCodeMask & baseAddress[0]))
  505. let isMasked = (MaskMask & baseAddress[1])
  506. let payloadLen = (PayloadLenMask & baseAddress[1])
  507. var offset = 2
  508. if (isMasked > 0 || (RSVMask & baseAddress[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 emptyBuffer
  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 emptyBuffer
  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 emptyBuffer
  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(baseAddress, 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 = baseAddress + 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 emptyBuffer
  552. }
  553. if isControlFrame && payloadLen > 125 {
  554. writeError(CloseCode.ProtocolError.rawValue)
  555. return emptyBuffer
  556. }
  557. var dataLength = UInt64(payloadLen)
  558. if dataLength == 127 {
  559. dataLength = WebSocket.readUint64(baseAddress, offset: offset)
  560. offset += sizeof(UInt64)
  561. } else if dataLength == 126 {
  562. dataLength = UInt64(WebSocket.readUint16(baseAddress, offset: offset))
  563. offset += sizeof(UInt16)
  564. }
  565. if bufferLen < offset || UInt64(bufferLen - offset) < dataLength {
  566. fragBuffer = NSData(bytes: baseAddress, length: bufferLen)
  567. return emptyBuffer
  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: baseAddress+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. return buffer.fromOffset(offset + Int(len))
  589. }
  590. var response = readStack.last
  591. if isControlFrame {
  592. response = nil //don't append pings
  593. }
  594. if isFin == 0 && receivedOpcode == .ContinueFrame && response == nil {
  595. let errCode = CloseCode.ProtocolError.rawValue
  596. doDisconnect(errorWithDetail("continue frame before a binary or text frame", code: errCode))
  597. writeError(errCode)
  598. return emptyBuffer
  599. }
  600. var isNew = false
  601. if response == nil {
  602. if receivedOpcode == .ContinueFrame {
  603. let errCode = CloseCode.ProtocolError.rawValue
  604. doDisconnect(errorWithDetail("first frame can't be a continue frame",
  605. code: errCode))
  606. writeError(errCode)
  607. return emptyBuffer
  608. }
  609. isNew = true
  610. response = WSResponse()
  611. response!.code = receivedOpcode!
  612. response!.bytesLeft = Int(dataLength)
  613. response!.buffer = NSMutableData(data: data)
  614. } else {
  615. if receivedOpcode == .ContinueFrame {
  616. response!.bytesLeft = Int(dataLength)
  617. } else {
  618. let errCode = CloseCode.ProtocolError.rawValue
  619. doDisconnect(errorWithDetail("second and beyond of fragment message must be a continue frame",
  620. code: errCode))
  621. writeError(errCode)
  622. return emptyBuffer
  623. }
  624. response!.buffer!.appendData(data)
  625. }
  626. if let response = response {
  627. response.bytesLeft -= Int(len)
  628. response.frameCount += 1
  629. response.isFin = isFin > 0 ? true : false
  630. if isNew {
  631. readStack.append(response)
  632. }
  633. processResponse(response)
  634. }
  635. let step = Int(offset+numericCast(len))
  636. return buffer.fromOffset(step)
  637. }
  638. }
  639. /// Process all messages in the buffer if possible.
  640. private func processRawMessagesInBuffer(pointer: UnsafePointer<UInt8>, bufferLen: Int) {
  641. var buffer = UnsafeBufferPointer(start: pointer, count: bufferLen)
  642. repeat {
  643. buffer = processOneRawMessage(inBuffer: buffer)
  644. } while buffer.count >= 2
  645. if buffer.count > 0 {
  646. fragBuffer = NSData(buffer: buffer)
  647. }
  648. }
  649. ///process the finished response of a buffer
  650. private func processResponse(response: WSResponse) -> Bool {
  651. if response.isFin && response.bytesLeft <= 0 {
  652. if response.code == .Ping {
  653. let data = response.buffer! //local copy so it is perverse for writing
  654. dequeueWrite(data, code: OpCode.Pong)
  655. } else if response.code == .TextFrame {
  656. let str: NSString? = NSString(data: response.buffer!, encoding: NSUTF8StringEncoding)
  657. if str == nil {
  658. writeError(CloseCode.Encoding.rawValue)
  659. return false
  660. }
  661. if canDispatch {
  662. dispatch_async(queue) { [weak self] in
  663. guard let s = self else { return }
  664. s.onText?(str! as String)
  665. s.delegate?.websocketDidReceiveMessage(s, text: str! as String)
  666. }
  667. }
  668. } else if response.code == .BinaryFrame {
  669. if canDispatch {
  670. let data = response.buffer! //local copy so it is perverse for writing
  671. dispatch_async(queue) { [weak self] in
  672. guard let s = self else { return }
  673. s.onData?(data)
  674. s.delegate?.websocketDidReceiveData(s, data: data)
  675. }
  676. }
  677. }
  678. readStack.removeLast()
  679. return true
  680. }
  681. return false
  682. }
  683. ///Create an error
  684. private func errorWithDetail(detail: String, code: UInt16) -> NSError {
  685. var details = [String: String]()
  686. details[NSLocalizedDescriptionKey] = detail
  687. return NSError(domain: WebSocket.ErrorDomain, code: Int(code), userInfo: details)
  688. }
  689. ///write a an error to the socket
  690. private func writeError(code: UInt16) {
  691. let buf = NSMutableData(capacity: sizeof(UInt16))
  692. let buffer = UnsafeMutablePointer<UInt8>(buf!.bytes)
  693. WebSocket.writeUint16(buffer, offset: 0, value: code)
  694. dequeueWrite(NSData(bytes: buffer, length: sizeof(UInt16)), code: .ConnectionClose)
  695. }
  696. ///used to write things to the stream
  697. private func dequeueWrite(data: NSData, code: OpCode, writeCompletion: (() -> ())? = nil) {
  698. writeQueue.addOperationWithBlock { [weak self] in
  699. //stream isn't ready, let's wait
  700. guard let s = self else { return }
  701. var offset = 2
  702. let bytes = UnsafeMutablePointer<UInt8>(data.bytes)
  703. let dataLength = data.length
  704. let frame = NSMutableData(capacity: dataLength + s.MaxFrameSize)
  705. let buffer = UnsafeMutablePointer<UInt8>(frame!.mutableBytes)
  706. buffer[0] = s.FinMask | code.rawValue
  707. if dataLength < 126 {
  708. buffer[1] = CUnsignedChar(dataLength)
  709. } else if dataLength <= Int(UInt16.max) {
  710. buffer[1] = 126
  711. WebSocket.writeUint16(buffer, offset: offset, value: UInt16(dataLength))
  712. offset += sizeof(UInt16)
  713. } else {
  714. buffer[1] = 127
  715. WebSocket.writeUint64(buffer, offset: offset, value: UInt64(dataLength))
  716. offset += sizeof(UInt64)
  717. }
  718. buffer[1] |= s.MaskMask
  719. let maskKey = UnsafeMutablePointer<UInt8>(buffer + offset)
  720. SecRandomCopyBytes(kSecRandomDefault, Int(sizeof(UInt32)), maskKey)
  721. offset += sizeof(UInt32)
  722. for i in 0..<dataLength {
  723. buffer[offset] = bytes[i] ^ maskKey[i % sizeof(UInt32)]
  724. offset += 1
  725. }
  726. var total = 0
  727. while true {
  728. guard let outStream = s.outputStream else { break }
  729. let writeBuffer = UnsafePointer<UInt8>(frame!.bytes+total)
  730. let len = outStream.write(writeBuffer, maxLength: offset-total)
  731. if len < 0 {
  732. var error: NSError?
  733. if let streamError = outStream.streamError {
  734. error = streamError
  735. } else {
  736. let errCode = InternalErrorCode.OutputStreamWriteError.rawValue
  737. error = s.errorWithDetail("output stream error during write", code: errCode)
  738. }
  739. s.doDisconnect(error)
  740. break
  741. } else {
  742. total += len
  743. }
  744. if total >= offset {
  745. if let queue = self?.queue, callback = writeCompletion {
  746. dispatch_async(queue) {
  747. callback()
  748. }
  749. }
  750. break
  751. }
  752. }
  753. }
  754. }
  755. ///used to preform the disconnect delegate
  756. private func doDisconnect(error: NSError?) {
  757. guard !didDisconnect else { return }
  758. didDisconnect = true
  759. connected = false
  760. guard canDispatch else {return}
  761. dispatch_async(queue) { [weak self] in
  762. guard let s = self else { return }
  763. s.onDisconnect?(error)
  764. s.delegate?.websocketDidDisconnect(s, error: error)
  765. }
  766. }
  767. deinit {
  768. mutex.lock()
  769. readyToWrite = false
  770. mutex.unlock()
  771. cleanupStream()
  772. }
  773. }
  774. private extension NSData {
  775. convenience init(buffer: UnsafeBufferPointer<UInt8>) {
  776. self.init(bytes: buffer.baseAddress, length: buffer.count)
  777. }
  778. }
  779. private extension UnsafeBufferPointer {
  780. func fromOffset(offset: Int) -> UnsafeBufferPointer<Element> {
  781. return UnsafeBufferPointer<Element>(start: baseAddress.advancedBy(offset), count: count - offset)
  782. }
  783. }
  784. private let emptyBuffer = UnsafeBufferPointer<UInt8>(start: nil, count: 0)
  785. public class SSLCert {
  786. var certData: NSData?
  787. var key: SecKeyRef?
  788. /**
  789. Designated init for certificates
  790. - parameter data: is the binary data of the certificate
  791. - returns: a representation security object to be used with
  792. */
  793. public init(data: NSData) {
  794. self.certData = data
  795. }
  796. /**
  797. Designated init for public keys
  798. - parameter key: is the public key to be used
  799. - returns: a representation security object to be used with
  800. */
  801. public init(key: SecKeyRef) {
  802. self.key = key
  803. }
  804. }
  805. public class SSLSecurity {
  806. public var validatedDN = true //should the domain name be validated?
  807. var isReady = false //is the key processing done?
  808. var certificates: [NSData]? //the certificates
  809. var pubKeys: [SecKeyRef]? //the public keys
  810. var usePublicKeys = false //use public keys or certificate validation?
  811. /**
  812. Use certs from main app bundle
  813. - parameter usePublicKeys: is to specific if the publicKeys or certificates should be used for SSL pinning validation
  814. - returns: a representation security object to be used with
  815. */
  816. public convenience init(usePublicKeys: Bool = false) {
  817. let paths = NSBundle.mainBundle().pathsForResourcesOfType("cer", inDirectory: ".")
  818. let certs = paths.reduce([SSLCert]()) { (certs: [SSLCert], path: String) -> [SSLCert] in
  819. var certs = certs
  820. if let data = NSData(contentsOfFile: path) {
  821. certs.append(SSLCert(data: data))
  822. }
  823. return certs
  824. }
  825. self.init(certs: certs, usePublicKeys: usePublicKeys)
  826. }
  827. /**
  828. Designated init
  829. - parameter keys: is the certificates or public keys to use
  830. - parameter usePublicKeys: is to specific if the publicKeys or certificates should be used for SSL pinning validation
  831. - returns: a representation security object to be used with
  832. */
  833. public init(certs: [SSLCert], usePublicKeys: Bool) {
  834. self.usePublicKeys = usePublicKeys
  835. if self.usePublicKeys {
  836. dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT,0)) {
  837. let pubKeys = certs.reduce([SecKeyRef]()) { (pubKeys: [SecKeyRef], cert: SSLCert) -> [SecKeyRef] in
  838. var pubKeys = pubKeys
  839. if let data = cert.certData where cert.key == nil {
  840. cert.key = self.extractPublicKey(data)
  841. }
  842. if let key = cert.key {
  843. pubKeys.append(key)
  844. }
  845. return pubKeys
  846. }
  847. self.pubKeys = pubKeys
  848. self.isReady = true
  849. }
  850. } else {
  851. let certificates = certs.reduce([NSData]()) { (certificates: [NSData], cert: SSLCert) -> [NSData] in
  852. var certificates = certificates
  853. if let data = cert.certData {
  854. certificates.append(data)
  855. }
  856. return certificates
  857. }
  858. self.certificates = certificates
  859. self.isReady = true
  860. }
  861. }
  862. /**
  863. Valid the trust and domain name.
  864. - parameter trust: is the serverTrust to validate
  865. - parameter domain: is the CN domain to validate
  866. - returns: if the key was successfully validated
  867. */
  868. public func isValid(trust: SecTrustRef, domain: String?) -> Bool {
  869. var tries = 0
  870. while(!self.isReady) {
  871. usleep(1000)
  872. tries += 1
  873. if tries > 5 {
  874. return false //doesn't appear it is going to ever be ready...
  875. }
  876. }
  877. var policy: SecPolicyRef
  878. if self.validatedDN {
  879. policy = SecPolicyCreateSSL(true, domain)
  880. } else {
  881. policy = SecPolicyCreateBasicX509()
  882. }
  883. SecTrustSetPolicies(trust,policy)
  884. if self.usePublicKeys {
  885. if let keys = self.pubKeys {
  886. let serverPubKeys = publicKeyChainForTrust(trust)
  887. for serverKey in serverPubKeys as [AnyObject] {
  888. for key in keys as [AnyObject] {
  889. if serverKey.isEqual(key) {
  890. return true
  891. }
  892. }
  893. }
  894. }
  895. } else if let certs = self.certificates {
  896. let serverCerts = certificateChainForTrust(trust)
  897. var collect = [SecCertificate]()
  898. for cert in certs {
  899. collect.append(SecCertificateCreateWithData(nil,cert)!)
  900. }
  901. SecTrustSetAnchorCertificates(trust,collect)
  902. var result: SecTrustResultType = 0
  903. SecTrustEvaluate(trust,&result)
  904. let r = Int(result)
  905. if r == kSecTrustResultUnspecified || r == kSecTrustResultProceed {
  906. var trustedCount = 0
  907. for serverCert in serverCerts {
  908. for cert in certs {
  909. if cert == serverCert {
  910. trustedCount += 1
  911. break
  912. }
  913. }
  914. }
  915. if trustedCount == serverCerts.count {
  916. return true
  917. }
  918. }
  919. }
  920. return false
  921. }
  922. /**
  923. Get the public key from a certificate data
  924. - parameter data: is the certificate to pull the public key from
  925. - returns: a public key
  926. */
  927. func extractPublicKey(data: NSData) -> SecKeyRef? {
  928. guard let cert = SecCertificateCreateWithData(nil, data) else { return nil }
  929. return extractPublicKeyFromCert(cert, policy: SecPolicyCreateBasicX509())
  930. }
  931. /**
  932. Get the public key from a certificate
  933. - parameter data: is the certificate to pull the public key from
  934. - returns: a public key
  935. */
  936. func extractPublicKeyFromCert(cert: SecCertificate, policy: SecPolicy) -> SecKeyRef? {
  937. var possibleTrust: SecTrust?
  938. SecTrustCreateWithCertificates(cert, policy, &possibleTrust)
  939. guard let trust = possibleTrust else { return nil }
  940. var result: SecTrustResultType = 0
  941. SecTrustEvaluate(trust, &result)
  942. return SecTrustCopyPublicKey(trust)
  943. }
  944. /**
  945. Get the certificate chain for the trust
  946. - parameter trust: is the trust to lookup the certificate chain for
  947. - returns: the certificate chain for the trust
  948. */
  949. func certificateChainForTrust(trust: SecTrustRef) -> [NSData] {
  950. let certificates = (0..<SecTrustGetCertificateCount(trust)).reduce([NSData]()) { (certificates: [NSData], index: Int) -> [NSData] in
  951. var certificates = certificates
  952. let cert = SecTrustGetCertificateAtIndex(trust, index)
  953. certificates.append(SecCertificateCopyData(cert!))
  954. return certificates
  955. }
  956. return certificates
  957. }
  958. /**
  959. Get the public key chain for the trust
  960. - parameter trust: is the trust to lookup the certificate chain and extract the public keys
  961. - returns: the public keys from the certifcate chain for the trust
  962. */
  963. func publicKeyChainForTrust(trust: SecTrustRef) -> [SecKeyRef] {
  964. let policy = SecPolicyCreateBasicX509()
  965. let keys = (0..<SecTrustGetCertificateCount(trust)).reduce([SecKeyRef]()) { (keys: [SecKeyRef], index: Int) -> [SecKeyRef] in
  966. var keys = keys
  967. let cert = SecTrustGetCertificateAtIndex(trust, index)
  968. if let key = extractPublicKeyFromCert(cert!, policy: policy) {
  969. keys.append(key)
  970. }
  971. return keys
  972. }
  973. return keys
  974. }
  975. }