Update Starscream code

Author: hamchapman

Source/Starscream.swift

@@ -258,11 +258,11 @@ public class SSLSecurity {
 
             return keys
         }
-
+        
         return keys
     }
-
-
+    
+    
 }
 
 
@@ -291,6 +291,10 @@ import Foundation
 import CoreFoundation
 import Security
 
+public let WebsocketDidConnectNotification = "WebsocketDidConnectNotification"
+public let WebsocketDidDisconnectNotification = "WebsocketDidDisconnectNotification"
+public let WebsocketDisconnectionErrorKeyName = "WebsocketDisconnectionErrorKeyName"
+
 public protocol WebSocketDelegate: class {
     func websocketDidConnect(socket: WebSocket)
     func websocketDidDisconnect(socket: WebSocket, error: NSError?)
@@ -380,6 +384,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
     public var security: SSLSecurity?
     public var enabledSSLCipherSuites: [SSLCipherSuite]?
     public var origin: String?
+    public var timeout = 5
     public var isConnected :Bool {
         return connected
     }
@@ -397,6 +402,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
     private var didDisconnect = false
     private var readyToWrite = false
     private let mutex = NSLock()
+    private let notificationCenter = NSNotificationCenter.defaultCenter()
     private var canDispatch: Bool {
         mutex.lock()
         let canWork = readyToWrite
@@ -589,12 +595,12 @@ public class WebSocket : NSObject, NSStreamDelegate {
         self.mutex.unlock()
 
         let bytes = UnsafePointer<UInt8>(data.bytes)
-        var timeout = 5000000 //wait 5 seconds before giving up
+        var out = timeout * 1000000 //wait 5 seconds before giving up
         writeQueue.addOperationWithBlock { [weak self] in
             while !outStream.hasSpaceAvailable {
                 usleep(100) //wait until the socket is ready
-                timeout -= 100
-                if timeout < 0 {
+                out -= 100
+                if out < 0 {
                     self?.cleanupStream()
                     self?.doDisconnect(self?.errorWithDetail("write wait timed out", code: 2))
                     return
@@ -675,25 +681,24 @@ public class WebSocket : NSObject, NSStreamDelegate {
     }
     ///dequeue the incoming input so it is processed in order
     private func dequeueInput() {
-        guard !inputQueue.isEmpty else { return }
-
-        let data = inputQueue[0]
-        var work = data
-        if let fragBuffer = fragBuffer {
-            let combine = NSMutableData(data: fragBuffer)
-            combine.appendData(data)
-            work = combine
-            self.fragBuffer = nil
-        }
-        let buffer = UnsafePointer<UInt8>(work.bytes)
-        let length = work.length
-        if !connected {
-            processTCPHandshake(buffer, bufferLen: length)
-        } else {
-            processRawMessage(buffer, bufferLen: length)
+        while !inputQueue.isEmpty {
+            let data = inputQueue[0]
+            var work = data
+            if let fragBuffer = fragBuffer {
+                let combine = NSMutableData(data: fragBuffer)
+                combine.appendData(data)
+                work = combine
+                self.fragBuffer = nil
+            }
+            let buffer = UnsafePointer<UInt8>(work.bytes)
+            let length = work.length
+            if !connected {
+                processTCPHandshake(buffer, bufferLen: length)
+            } else {
+                processRawMessagesInBuffer(buffer, bufferLen: length)
+            }
+            inputQueue = inputQueue.filter{$0 != data}
         }
-        inputQueue = inputQueue.filter{$0 != data}
-        dequeueInput()
     }
 
     //handle checking the inital connection status
@@ -707,6 +712,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                 guard let s = self else { return }
                 s.onConnect?()
                 s.delegate?.websocketDidConnect(s)
+                s.notificationCenter.postNotificationName(WebsocketDidConnectNotification, object: self)
             }
         case -1:
             fragBuffer = NSData(bytes: buffer, length: bufferLen)
@@ -739,7 +745,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
             totalSize += 1 //skip the last \n
             let restSize = bufferLen - totalSize
             if restSize > 0 {
-                processRawMessage((buffer+totalSize),bufferLen: restSize)
+                processRawMessagesInBuffer(buffer + totalSize, bufferLen: restSize)
             }
             return 0 //success
         }
@@ -792,12 +798,15 @@ public class WebSocket : NSObject, NSStreamDelegate {
         }
     }
 
-    ///process the websocket data
-    private func processRawMessage(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
+    /// 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.
+    @warn_unused_result
+    private func processOneRawMessage(inBuffer buffer: UnsafeBufferPointer<UInt8>) -> UnsafeBufferPointer<UInt8> {
         let response = readStack.last
+        let baseAddress = buffer.baseAddress
+        let bufferLen = buffer.count
         if response != nil && bufferLen < 2  {
-            fragBuffer = NSData(bytes: buffer, length: bufferLen)
-            return
+            fragBuffer = NSData(buffer: buffer)
+            return emptyBuffer
         }
         if let response = response where response.bytesLeft > 0 {
             var len = response.bytesLeft
@@ -807,79 +816,77 @@ public class WebSocket : NSObject, NSStreamDelegate {
                 extra = 0
             }
             response.bytesLeft -= len
-            response.buffer?.appendData(NSData(bytes: buffer, length: len))
+            response.buffer?.appendData(NSData(bytes: baseAddress, length: len))
             processResponse(response)
-            let offset = bufferLen - extra
-            if extra > 0 {
-                processExtra((buffer+offset), bufferLen: extra)
-            }
-            return
+            return buffer.fromOffset(bufferLen - extra)
         } else {
-            let isFin = (FinMask & buffer[0])
-            let receivedOpcode = OpCode(rawValue: (OpCodeMask & buffer[0]))
-            let isMasked = (MaskMask & buffer[1])
-            let payloadLen = (PayloadLenMask & buffer[1])
+            let isFin = (FinMask & baseAddress[0])
+            let receivedOpcode = OpCode(rawValue: (OpCodeMask & baseAddress[0]))
+            let isMasked = (MaskMask & baseAddress[1])
+            let payloadLen = (PayloadLenMask & baseAddress[1])
             var offset = 2
-            if (isMasked > 0 || (RSVMask & buffer[0]) > 0) && receivedOpcode != .Pong {
+            if (isMasked > 0 || (RSVMask & baseAddress[0]) > 0) && receivedOpcode != .Pong {
                 let errCode = CloseCode.ProtocolError.rawValue
                 doDisconnect(errorWithDetail("masked and rsv data is not currently supported", code: errCode))
                 writeError(errCode)
-                return
+                return emptyBuffer
             }
             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
+                return emptyBuffer
             }
             if isControlFrame && isFin == 0 {
                 let errCode = CloseCode.ProtocolError.rawValue
                 doDisconnect(errorWithDetail("control frames can't be fragmented", code: errCode))
                 writeError(errCode)
-                return
+                return emptyBuffer
             }
             if receivedOpcode == .ConnectionClose {
                 var code = CloseCode.Normal.rawValue
                 if payloadLen == 1 {
                     code = CloseCode.ProtocolError.rawValue
                 } else if payloadLen > 1 {
-                    code = WebSocket.readUint16(buffer, offset: offset)
+                    code = WebSocket.readUint16(baseAddress, offset: offset)
                     if code < 1000 || (code > 1003 && code < 1007) || (code > 1011 && code < 3000) {
                         code = CloseCode.ProtocolError.rawValue
                     }
                     offset += 2
                 }
+                var closeReason = "connection closed by server"
                 if payloadLen > 2 {
-                    let len = Int(payloadLen-2)
+                    let len = Int(payloadLen - 2)
                     if len > 0 {
-                        let bytes = UnsafePointer<UInt8>((buffer+offset))
-                        let str: NSString? = NSString(data: NSData(bytes: bytes, length: len), encoding: NSUTF8StringEncoding)
-                        if str == nil {
+                        let bytes = baseAddress + offset
+                        if let customCloseReason = String(data: NSData(bytes: bytes, length: len), encoding: NSUTF8StringEncoding) {
+                            closeReason = customCloseReason
+                        } else {
                             code = CloseCode.ProtocolError.rawValue
                         }
                     }
                 }
-                doDisconnect(errorWithDetail("connection closed by server", code: code))
+                doDisconnect(errorWithDetail(closeReason, code: code))
                 writeError(code)
-                return
+                return emptyBuffer
             }
             if isControlFrame && payloadLen > 125 {
                 writeError(CloseCode.ProtocolError.rawValue)
-                return
+                return emptyBuffer
             }
             var dataLength = UInt64(payloadLen)
             if dataLength == 127 {
-                dataLength = WebSocket.readUint64(buffer, offset: offset)
+                dataLength = WebSocket.readUint64(baseAddress, offset: offset)
                 offset += sizeof(UInt64)
             } else if dataLength == 126 {
-                dataLength = UInt64(WebSocket.readUint16(buffer, offset: offset))
+                dataLength = UInt64(WebSocket.readUint16(baseAddress, offset: offset))
                 offset += sizeof(UInt16)
             }
             if bufferLen < offset || UInt64(bufferLen - offset) < dataLength {
-                fragBuffer = NSData(bytes: buffer, length: bufferLen)
-                return
+                fragBuffer = NSData(bytes: baseAddress, length: bufferLen)
+                return emptyBuffer
             }
             var len = dataLength
             if dataLength > UInt64(bufferLen) {
@@ -890,7 +897,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                 len = 0
                 data = NSData()
             } else {
-                data = NSData(bytes: UnsafePointer<UInt8>((buffer+offset)), length: Int(len))
+                data = NSData(bytes: baseAddress+offset, length: Int(len))
             }
             if receivedOpcode == .Pong {
                 if canDispatch {
@@ -900,12 +907,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                         s.pongDelegate?.websocketDidReceivePong(s)
                     }
                 }
-                let step = Int(offset+numericCast(len))
-                let extra = bufferLen-step
-                if extra > 0 {
-                    processRawMessage((buffer+step), bufferLen: extra)
-                }
-                return
+                return buffer.fromOffset(offset + Int(len))
             }
             var response = readStack.last
             if isControlFrame {
@@ -915,7 +917,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                 let errCode = CloseCode.ProtocolError.rawValue
                 doDisconnect(errorWithDetail("continue frame before a binary or text frame", code: errCode))
                 writeError(errCode)
-                return
+                return emptyBuffer
             }
             var isNew = false
             if response == nil {
@@ -924,7 +926,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                     doDisconnect(errorWithDetail("first frame can't be a continue frame",
                         code: errCode))
                     writeError(errCode)
-                    return
+                    return emptyBuffer
                 }
                 isNew = true
                 response = WSResponse()
@@ -939,7 +941,7 @@ public class WebSocket : NSObject, NSStreamDelegate {
                     doDisconnect(errorWithDetail("second and beyond of fragment message must be a continue frame",
                         code: errCode))
                     writeError(errCode)
-                    return
+                    return emptyBuffer
                 }
                 response!.buffer!.appendData(data)
             }
@@ -954,20 +956,18 @@ public class WebSocket : NSObject, NSStreamDelegate {
             }
 
             let step = Int(offset+numericCast(len))
-            let extra = bufferLen-step
-            if extra > 0 {
-                processExtra((buffer+step), bufferLen: extra)
-            }
+            return buffer.fromOffset(step)
         }
-
     }
 
-    ///process the extra of a buffer
-    private func processExtra(buffer: UnsafePointer<UInt8>, bufferLen: Int) {
-        if bufferLen < 2 {
-            fragBuffer = NSData(bytes: buffer, length: bufferLen)
-        } else {
-            processRawMessage(buffer, bufferLen: bufferLen)
+    /// Process all messages in the buffer if possible.
+    private func processRawMessagesInBuffer(pointer: UnsafePointer<UInt8>, bufferLen: Int) {
+        var buffer = UnsafeBufferPointer(start: pointer, count: bufferLen)
+        repeat {
+            buffer = processOneRawMessage(inBuffer: buffer)
+        } while buffer.count >= 2
+        if buffer.count > 0 {
+            fragBuffer = NSData(buffer: buffer)
         }
     }
 
@@ -1093,6 +1093,8 @@ public class WebSocket : NSObject, NSStreamDelegate {
             guard let s = self else { return }
             s.onDisconnect?(error)
             s.delegate?.websocketDidDisconnect(s, error: error)
+            let userInfo = error.map({ [WebsocketDisconnectionErrorKeyName: $0] })
+            s.notificationCenter.postNotificationName(WebsocketDidDisconnectNotification, object: self, userInfo: userInfo)
         }
     }
 
@@ -1103,4 +1105,22 @@ public class WebSocket : NSObject, NSStreamDelegate {
         cleanupStream()
     }
 
-}
+}
+
+private extension NSData {
+
+    convenience init(buffer: UnsafeBufferPointer<UInt8>) {
+        self.init(bytes: buffer.baseAddress, length: buffer.count)
+    }
+
+}
+
+private extension UnsafeBufferPointer {
+
+    func fromOffset(offset: Int) -> UnsafeBufferPointer<Element> {
+        return UnsafeBufferPointer<Element>(start: baseAddress.advancedBy(offset), count: count - offset)
+    }
+
+}
+
+private let emptyBuffer = UnsafeBufferPointer<UInt8>(start: nil, count: 0)