first pass at refactoring engine

Socket.IO-Client-Swift.xcodeproj/project.pbxproj

@@ -17,6 +17,9 @@
 		57634A2F1BD9B46D00E19CD7 /* SocketBasicPacketTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74F124EF1BC574CF002966F4 /* SocketBasicPacketTest.swift */; };
 		57634A321BD9B46D00E19CD7 /* SocketNamespacePacketTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7472C65B1BCAB53E003CA70D /* SocketNamespacePacketTest.swift */; };
 		57634A3F1BD9B4BF00E19CD7 /* SocketIO.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 57634A161BD9B46A00E19CD7 /* SocketIO.framework */; };
+		740CA1201C496EEB00CB98F4 /* SocketEngineWebsocket.swift in Sources */ = {isa = PBXBuildFile; fileRef = 740CA11F1C496EEB00CB98F4 /* SocketEngineWebsocket.swift */; };
+		740CA1211C496EF200CB98F4 /* SocketEngineWebsocket.swift in Sources */ = {isa = PBXBuildFile; fileRef = 740CA11F1C496EEB00CB98F4 /* SocketEngineWebsocket.swift */; };
+		740CA1221C496EF700CB98F4 /* SocketEngineWebsocket.swift in Sources */ = {isa = PBXBuildFile; fileRef = 740CA11F1C496EEB00CB98F4 /* SocketEngineWebsocket.swift */; };
 		74171E631C10CD240062D398 /* SocketAckEmitter.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74171E501C10CD240062D398 /* SocketAckEmitter.swift */; };
 		74171E641C10CD240062D398 /* SocketAckEmitter.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74171E501C10CD240062D398 /* SocketAckEmitter.swift */; };
 		74171E651C10CD240062D398 /* SocketAckEmitter.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74171E501C10CD240062D398 /* SocketAckEmitter.swift */; };
@@ -115,6 +118,9 @@
 		74171ED41C10CD240062D398 /* WebSocket.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74171E621C10CD240062D398 /* WebSocket.swift */; };
 		741F39EE1BD025D80026C9CC /* SocketEngineTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 741F39ED1BD025D80026C9CC /* SocketEngineTest.swift */; };
 		741F39EF1BD025D80026C9CC /* SocketEngineTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 741F39ED1BD025D80026C9CC /* SocketEngineTest.swift */; };
+		7420CB791C49629E00956AA4 /* SocketEnginePollable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7420CB781C49629E00956AA4 /* SocketEnginePollable.swift */; };
+		7420CB7A1C49629E00956AA4 /* SocketEnginePollable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7420CB781C49629E00956AA4 /* SocketEnginePollable.swift */; };
+		7420CB7B1C49629E00956AA4 /* SocketEnginePollable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7420CB781C49629E00956AA4 /* SocketEnginePollable.swift */; };
 		74321DCB1C2D939A00CF6F43 /* SocketAckManagerTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74321DC91C2D939A00CF6F43 /* SocketAckManagerTest.swift */; };
 		74321DCC1C2D939A00CF6F43 /* SocketParserTest.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74321DCA1C2D939A00CF6F43 /* SocketParserTest.swift */; };
 		7471CCEA1C39926300364B59 /* SocketClientSpec.swift in Sources */ = {isa = PBXBuildFile; fileRef = 74ABF7761C3991C10078C657 /* SocketClientSpec.swift */; };
@@ -165,6 +171,7 @@
 		572EF2481B51F18A00EEBB58 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
 		57634A161BD9B46A00E19CD7 /* SocketIO.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = SocketIO.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		57634A3B1BD9B46D00E19CD7 /* SocketIO-tvOSTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "SocketIO-tvOSTests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
+		740CA11F1C496EEB00CB98F4 /* SocketEngineWebsocket.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; name = SocketEngineWebsocket.swift; path = Source/SocketEngineWebsocket.swift; sourceTree = "<group>"; };
 		74171E501C10CD240062D398 /* SocketAckEmitter.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = SocketAckEmitter.swift; path = Source/SocketAckEmitter.swift; sourceTree = "<group>"; };
 		74171E511C10CD240062D398 /* SocketAckManager.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = SocketAckManager.swift; path = Source/SocketAckManager.swift; sourceTree = "<group>"; };
 		74171E521C10CD240062D398 /* SocketAnyEvent.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = SocketAnyEvent.swift; path = Source/SocketAnyEvent.swift; sourceTree = "<group>"; };
@@ -185,6 +192,7 @@
 		74171E611C10CD240062D398 /* SwiftRegex.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = SwiftRegex.swift; path = Source/SwiftRegex.swift; sourceTree = "<group>"; };
 		74171E621C10CD240062D398 /* WebSocket.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = WebSocket.swift; path = Source/WebSocket.swift; sourceTree = "<group>"; };
 		741F39ED1BD025D80026C9CC /* SocketEngineTest.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = SocketEngineTest.swift; sourceTree = "<group>"; };
+		7420CB781C49629E00956AA4 /* SocketEnginePollable.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = SocketEnginePollable.swift; path = Source/SocketEnginePollable.swift; sourceTree = "<group>"; };
 		74321DC91C2D939A00CF6F43 /* SocketAckManagerTest.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SocketAckManagerTest.swift; sourceTree = "<group>"; };
 		74321DCA1C2D939A00CF6F43 /* SocketParserTest.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SocketParserTest.swift; sourceTree = "<group>"; };
 		7472C65B1BCAB53E003CA70D /* SocketNamespacePacketTest.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = SocketNamespacePacketTest.swift; sourceTree = "<group>"; };
@@ -349,7 +357,9 @@
 				74171E531C10CD240062D398 /* SocketEngine.swift */,
 				74171E541C10CD240062D398 /* SocketEngineClient.swift */,
 				74171E551C10CD240062D398 /* SocketEnginePacketType.swift */,
+				7420CB781C49629E00956AA4 /* SocketEnginePollable.swift */,
 				74171E561C10CD240062D398 /* SocketEngineSpec.swift */,
+				740CA11F1C496EEB00CB98F4 /* SocketEngineWebsocket.swift */,
 				74171E571C10CD240062D398 /* SocketEventHandler.swift */,
 				74171E581C10CD240062D398 /* SocketFixUTF8.swift */,
 				74171E591C10CD240062D398 /* SocketIOClient.swift */,
@@ -599,10 +609,12 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				740CA1221C496EF700CB98F4 /* SocketEngineWebsocket.swift in Sources */,
 				74171E931C10CD240062D398 /* SocketFixUTF8.swift in Sources */,
 				74171EA51C10CD240062D398 /* SocketIOClientStatus.swift in Sources */,
 				74171E751C10CD240062D398 /* SocketEngine.swift in Sources */,
 				74171E691C10CD240062D398 /* SocketAckManager.swift in Sources */,
+				7420CB791C49629E00956AA4 /* SocketEnginePollable.swift in Sources */,
 				74ABF7771C3991C10078C657 /* SocketClientSpec.swift in Sources */,
 				74171E871C10CD240062D398 /* SocketEngineSpec.swift in Sources */,
 				74171E631C10CD240062D398 /* SocketAckEmitter.swift in Sources */,
@@ -656,10 +668,12 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				740CA1211C496EF200CB98F4 /* SocketEngineWebsocket.swift in Sources */,
 				7471CCEA1C39926300364B59 /* SocketClientSpec.swift in Sources */,
 				74171E951C10CD240062D398 /* SocketFixUTF8.swift in Sources */,
 				74171EA71C10CD240062D398 /* SocketIOClientStatus.swift in Sources */,
 				74171E771C10CD240062D398 /* SocketEngine.swift in Sources */,
+				7420CB7A1C49629E00956AA4 /* SocketEnginePollable.swift in Sources */,
 				74171E6B1C10CD240062D398 /* SocketAckManager.swift in Sources */,
 				74171E891C10CD240062D398 /* SocketEngineSpec.swift in Sources */,
 				74171E651C10CD240062D398 /* SocketAckEmitter.swift in Sources */,
@@ -697,10 +711,12 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				740CA1201C496EEB00CB98F4 /* SocketEngineWebsocket.swift in Sources */,
 				7471CCEB1C39926C00364B59 /* SocketClientSpec.swift in Sources */,
 				74171E971C10CD240062D398 /* SocketFixUTF8.swift in Sources */,
 				74171EA91C10CD240062D398 /* SocketIOClientStatus.swift in Sources */,
 				74171E791C10CD240062D398 /* SocketEngine.swift in Sources */,
+				7420CB7B1C49629E00956AA4 /* SocketEnginePollable.swift in Sources */,
 				74171E6D1C10CD240062D398 /* SocketAckManager.swift in Sources */,
 				74171E8B1C10CD240062D398 /* SocketEngineSpec.swift in Sources */,
 				74171E671C10CD240062D398 /* SocketAckEmitter.swift in Sources */,

Source/SocketEngine.swift

@@ -24,12 +24,33 @@
 
 import Foundation
 
-public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
+public final class SocketEngine: NSObject, SocketEngineSpec, SocketEnginePollable, SocketEngineWebsocket {
-    public private(set) var sid = ""
+    public let emitQueue = dispatch_queue_create("com.socketio.engineEmitQueue", DISPATCH_QUEUE_SERIAL)
+    public let handleQueue = dispatch_queue_create("com.socketio.engineHandleQueue", DISPATCH_QUEUE_SERIAL)
+    public let parseQueue = dispatch_queue_create("com.socketio.engineParseQueue", DISPATCH_QUEUE_SERIAL)
+
+    public var invalidated = false
+    public var postWait = [String]()
+    public var probing = false
+    public var waitingForPoll = false
+    public var waitingForPost = false
+    public var websocketConnected = false
+    
+    public private(set) var closed = false
+    public private(set) var connected = false
     public private(set) var cookies: [NSHTTPCookie]?
+    public private(set) var extraHeaders: [String: String]?
+    public private(set) var fastUpgrade = false
+    public private(set) var forcePolling = false
+    public private(set) var forceWebsockets = false
+    public private(set) var pingTimer: NSTimer?
+    public private(set) var polling = true
+    public private(set) var session: NSURLSession?
+    public private(set) var sid = ""
     public private(set) var socketPath = "/engine.io"
     public private(set) var urlPolling = ""
     public private(set) var urlWebSocket = ""
+    public private(set) var websocket = false
     public private(set) var ws: WebSocket?
 
     public weak var client: SocketEngineClient?
@@ -40,22 +61,13 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
     private typealias ProbeWaitQueue = [Probe]
 
     private let allowedCharacterSet = NSCharacterSet(charactersInString: "!*'();:@&=+$,/?%#[]\" {}").invertedSet
-    private let emitQueue = dispatch_queue_create("com.socketio.engineEmitQueue", DISPATCH_QUEUE_SERIAL)
+
-    private let handleQueue = dispatch_queue_create("com.socketio.engineHandleQueue", DISPATCH_QUEUE_SERIAL)
     private let logType = "SocketEngine"
-    private let parseQueue = dispatch_queue_create("com.socketio.engineParseQueue", DISPATCH_QUEUE_SERIAL)
     private let url: String
     private let workQueue = NSOperationQueue()
     
     private var connectParams: [String: AnyObject]?
-    private var closed = false
-    private var extraHeaders: [String: String]?
-    private var fastUpgrade = false
-    private var forcePolling = false
-    private var forceWebsockets = false
-    private var invalidated = false
     private var pingInterval: Double?
-    private var pingTimer: NSTimer?
     private var pingTimeout = 0.0 {
         didSet {
             pongsMissedMax = Int(pingTimeout / (pingInterval ?? 25))
@@ -63,19 +75,10 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
     }
     private var pongsMissed = 0
     private var pongsMissedMax = 0
-    private var postWait = [String]()
-    private var probing = false
     private var probeWait = ProbeWaitQueue()
     private var secure = false
     private var selfSigned = false
-    private var session: NSURLSession?
     private var voipEnabled = false
-    private var waitingForPoll = false
-    private var waitingForPost = false
-    private var websocketConnected = false
-    private(set) var connected = false
-    private(set) var polling = true
-    private(set) var websocket = false
     
     public init(client: SocketEngineClient, url: String, options: Set<SocketIOClientOption>) {
         self.client = client
@@ -181,22 +184,6 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         client?.engineDidClose("Disconnect")
     }
 
-    private func createBinaryDataForSend(data: NSData) -> Either<NSData, String> {
-        if websocket {
-            var byteArray = [UInt8](count: 1, repeatedValue: 0x0)
-            byteArray[0] = 4
-            let mutData = NSMutableData(bytes: &byteArray, length: 1)
-
-            mutData.appendData(data)
-
-            return .Left(mutData)
-        } else {
-            let str = "b4" + data.base64EncodedStringWithOptions(.Encoding64CharacterLineLength)
-
-            return .Right(str)
-        }
-    }
-
     private func createURLs(params: [String: AnyObject]?) -> (String, String) {
         if client == nil {
             return ("", "")
@@ -264,7 +251,7 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         }
     }
 
-    private func doFastUpgrade() {
+    public func doFastUpgrade() {
         if waitingForPoll {
             DefaultSocketLogger.Logger.error("Outstanding poll when switched to WebSockets," +
                 "we'll probably disconnect soon. You should report this.", type: logType)
@@ -294,6 +281,18 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         }
     }
     
+    // We had packets waiting for send when we upgraded
+    // Send them raw
+    public func flushWaitingForPostToWebSocket() {
+        guard let ws = self.ws else { return }
+        
+        for msg in postWait {
+            ws.writeString(fixDoubleUTF8(msg))
+        }
+        
+        postWait.removeAll(keepCapacity: true)
+    }
+
     private func handleClose(reason: String) {
         client?.engineDidClose(reason)
     }
@@ -356,7 +355,7 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
     }
 
     // A poll failed, tell the client about it
-    private func handlePollingFailed(reason: String) {
+    public func handlePollingFailed(reason: String) {
         connected = false
         ws?.disconnect()
         pingTimer?.invalidate()
@@ -414,12 +413,12 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         doLongPoll(reqPolling)
     }
 
-    private func parseEngineData(data: NSData) {
+    public func parseEngineData(data: NSData) {
         DefaultSocketLogger.Logger.log("Got binary data: %@", type: "SocketEngine", args: data)
         client?.parseEngineBinaryData(data.subdataWithRange(NSMakeRange(1, data.length - 1)))
     }
 
-    private func parseEngineMessage(message: String, fromPolling: Bool) {
+    public func parseEngineMessage(message: String, fromPolling: Bool) {
         DefaultSocketLogger.Logger.log("Got message: %@", type: logType, args: message)
         
         let reader = SocketStringReader(message: message)
@@ -455,13 +454,6 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         }
     }
     
-    private func probeWebSocket() {
-        if websocketConnected {
-            sendWebSocketMessage("probe", withType: .Ping, withData: [])
-        }
-    }
-    
-    
     private func resetEngine() {
         closed = false
         connected = false
@@ -541,220 +533,9 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
             }
         }
     }
-}
-
-// Polling methods
-extension SocketEngine {
-    private func addHeaders(req: NSMutableURLRequest) {
-        if cookies != nil {
-            let headers = NSHTTPCookie.requestHeaderFieldsWithCookies(cookies!)
-            req.allHTTPHeaderFields = headers
-        }
-        
-        if extraHeaders != nil {
-            for (headerName, value) in extraHeaders! {
-                req.setValue(value, forHTTPHeaderField: headerName)
-            }
-        }
-    }
-    
-    private func doPoll() {
-        if websocket || waitingForPoll || !connected || closed {
-            return
-        }
-        
-        waitingForPoll = true
-        let req = NSMutableURLRequest(URL: NSURL(string: urlPolling + "&sid=\(sid)&b64=1")!)
-        
-        addHeaders(req)
-        doLongPoll(req)
-    }
-    
-    private func doRequest(req: NSURLRequest,
-        withCallback callback: (NSData?, NSURLResponse?, NSError?) -> Void) {
-            if !polling || closed || invalidated {
-                DefaultSocketLogger.Logger.error("Tried to do polling request when not supposed to", type: logType)
-                return
-            }
-            
-            DefaultSocketLogger.Logger.log("Doing polling request", type: logType)
-            
-            session?.dataTaskWithRequest(req, completionHandler: callback).resume()
-    }
-    
-    private func doLongPoll(req: NSURLRequest) {
-        doRequest(req) {[weak self] data, res, err in
-            guard let this = self else {return}
-            
-            if err != nil || data == nil {
-                DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: this.logType)
-
-                if this.polling {
-                    this.handlePollingFailed(err?.localizedDescription ?? "Error")
-                }
-                
-                return
-            }
-            
-            DefaultSocketLogger.Logger.log("Got polling response", type: this.logType)
-            
-            if let str = String(data: data!, encoding: NSUTF8StringEncoding) {
-                dispatch_async(this.parseQueue) {
-                    this.parsePollingMessage(str)
-                }
-            }
-            
-            this.waitingForPoll = false
-            
-            if this.fastUpgrade {
-                this.doFastUpgrade()
-            } else if !this.closed && this.polling {
-                this.doPoll()
-            }
-        }
-    }
-    
-    private func flushWaitingForPost() {
-        if postWait.count == 0 || !connected {
-            return
-        } else if websocket {
-            flushWaitingForPostToWebSocket()
-            return
-        }
-        
-        var postStr = ""
-        
-        for packet in postWait {
-            let len = packet.characters.count
-            
-            postStr += "\(len):\(packet)"
-        }
-        
-        postWait.removeAll(keepCapacity: false)
-        
-        let req = NSMutableURLRequest(URL: NSURL(string: urlPolling + "&sid=\(sid)")!)
-        
-        addHeaders(req)
-        
-        req.HTTPMethod = "POST"
-        req.setValue("text/plain; charset=UTF-8", forHTTPHeaderField: "Content-Type")
-        
-        let postData = postStr.dataUsingEncoding(NSUTF8StringEncoding,
-            allowLossyConversion: false)!
-        
-        req.HTTPBody = postData
-        req.setValue(String(postData.length), forHTTPHeaderField: "Content-Length")
-        
-        waitingForPost = true
-        
-        DefaultSocketLogger.Logger.log("POSTing: %@", type: logType, args: postStr)
-        
-        doRequest(req) {[weak self] data, res, err in
-            guard let this = self else {return}
-            
-            if err != nil {
-                DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: this.logType)
-                
-                if this.polling {
-                    this.handlePollingFailed(err?.localizedDescription ?? "Error")
-                }
-                
-                return
-            }
-            
-            this.waitingForPost = false
-            
-            dispatch_async(this.emitQueue) {
-                if !this.fastUpgrade {
-                    this.flushWaitingForPost()
-                    this.doPoll()
-                }
-            }
-        }
-    }
-    
-    // We had packets waiting for send when we upgraded
-    // Send them raw
-    private func flushWaitingForPostToWebSocket() {
-        guard let ws = self.ws else { return }
-        
-        for msg in postWait {
-            ws.writeString(fixDoubleUTF8(msg))
-        }
-        
-        postWait.removeAll(keepCapacity: true)
-    }
-    
-    func parsePollingMessage(str: String) {
-        guard str.characters.count != 1 else {
-            return
-        }
-        
-        var reader = SocketStringReader(message: str)
-        
-        while reader.hasNext {
-            if let n = Int(reader.readUntilStringOccurence(":")) {
-                let str = reader.read(n)
-                
-                dispatch_async(handleQueue) {
-                    self.parseEngineMessage(str, fromPolling: true)
-                }
-            } else {
-                dispatch_async(handleQueue) {
-                    self.parseEngineMessage(str, fromPolling: true)
-                }
-                break
-            }
-        }
-    }
-    
-    /// Send polling message.
-    /// Only call on emitQueue
-    private func sendPollMessage(message: String, withType type: SocketEnginePacketType,
-        withData datas: [NSData]) {
-            DefaultSocketLogger.Logger.log("Sending poll: %@ as type: %@", type: logType, args: message, type.rawValue)
-            let fixedMessage = doubleEncodeUTF8(message)
-            let strMsg = "\(type.rawValue)\(fixedMessage)"
-            
-            postWait.append(strMsg)
-            
-            for data in datas {
-                if case let .Right(bin) = createBinaryDataForSend(data) {
-                    postWait.append(bin)
-                }
-            }
-            
-            if !waitingForPost {
-                flushWaitingForPost()
-            }
-    }
-    
-    private func stopPolling() {
-        invalidated = true
-        session?.finishTasksAndInvalidate()
-    }
-}
-
-// WebSocket methods
-extension SocketEngine {
-    /// Send message on WebSockets
-    /// Only call on emitQueue
-    private func sendWebSocketMessage(str: String, withType type: SocketEnginePacketType,
-        withData datas: [NSData]) {
-            DefaultSocketLogger.Logger.log("Sending ws: %@ as type: %@", type: logType, args: str, type.rawValue)
-            
-            ws?.writeString("\(type.rawValue)\(str)")
-            
-            for data in datas {
-                if case let .Left(bin) = createBinaryDataForSend(data) {
-                    ws?.writeData(bin)
-                }
-            }
-    }
     
     // Delagate methods
-    
+    public func websocketDidConnect(socket: WebSocket) {
-    public func websocketDidConnect(socket:WebSocket) {
         websocketConnected = true
         
         if !forceWebsockets {
@@ -792,12 +573,4 @@ extension SocketEngine {
             flushProbeWait()
         }
     }
-    
-    public func websocketDidReceiveMessage(socket: WebSocket, text: String) {
-        parseEngineMessage(text, fromPolling: false)
-    }
-    
-    public func websocketDidReceiveData(socket: WebSocket, data: NSData) {
-        parseEngineData(data)
-    }
 }

Source/SocketEnginePollable.swift

@@ -0,0 +1,217 @@
+//
+//  SocketEnginePollable.swift
+//  Socket.IO-Client-Swift
+//
+//  Created by Erik Little on 1/15/16.
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in
+//  all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+//  THE SOFTWARE.
+
+import Foundation
+
+public protocol SocketEnginePollable: SocketEngineSpec {
+    var invalidated: Bool { get set }
+    var session: NSURLSession? { get }
+    var waitingForPoll: Bool { get set }
+    var waitingForPost: Bool { get set }
+    
+    func doPoll()
+    func handlePollingFailed(reason: String)
+    func sendPollMessage(message: String, withType type: SocketEnginePacketType, withData datas: [NSData])
+    func stopPolling()
+}
+
+// Default polling methods
+extension SocketEnginePollable {
+    private func addHeaders(req: NSMutableURLRequest) {
+        if cookies != nil {
+            let headers = NSHTTPCookie.requestHeaderFieldsWithCookies(cookies!)
+            req.allHTTPHeaderFields = headers
+        }
+        
+        if extraHeaders != nil {
+            for (headerName, value) in extraHeaders! {
+                req.setValue(value, forHTTPHeaderField: headerName)
+            }
+        }
+    }
+    
+    public func doPoll() {
+        if websocket || waitingForPoll || !connected || closed {
+            return
+        }
+        
+        waitingForPoll = true
+        let req = NSMutableURLRequest(URL: NSURL(string: urlPolling + "&sid=\(sid)&b64=1")!)
+        
+        addHeaders(req)
+        doLongPoll(req)
+    }
+    
+    private func doRequest(req: NSURLRequest,
+        withCallback callback: (NSData?, NSURLResponse?, NSError?) -> Void) {
+            if !polling || closed || invalidated {
+                DefaultSocketLogger.Logger.error("Tried to do polling request when not supposed to", type: "SocketEngine")
+                return
+            }
+            
+            DefaultSocketLogger.Logger.log("Doing polling request", type: "SocketEngine")
+            
+            session?.dataTaskWithRequest(req, completionHandler: callback).resume()
+    }
+    
+    func doLongPoll(req: NSURLRequest) {
+        doRequest(req) {[weak self] data, res, err in
+            guard let this = self else {return}
+            
+            if err != nil || data == nil {
+                DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: "SocketEngine")
+                
+                if this.polling {
+                    this.handlePollingFailed(err?.localizedDescription ?? "Error")
+                }
+                
+                return
+            }
+            
+            DefaultSocketLogger.Logger.log("Got polling response", type: "SocketEngine")
+            
+            if let str = String(data: data!, encoding: NSUTF8StringEncoding) {
+                dispatch_async(this.parseQueue) {
+                    this.parsePollingMessage(str)
+                }
+            }
+            
+            this.waitingForPoll = false
+            
+            if this.fastUpgrade {
+                this.doFastUpgrade()
+            } else if !this.closed && this.polling {
+                this.doPoll()
+            }
+        }
+    }
+    
+    private func flushWaitingForPost() {
+        if postWait.count == 0 || !connected {
+            return
+        } else if websocket {
+            flushWaitingForPostToWebSocket()
+            return
+        }
+        
+        var postStr = ""
+        
+        for packet in postWait {
+            let len = packet.characters.count
+            
+            postStr += "\(len):\(packet)"
+        }
+        
+        postWait.removeAll(keepCapacity: false)
+        
+        let req = NSMutableURLRequest(URL: NSURL(string: urlPolling + "&sid=\(sid)")!)
+        
+        addHeaders(req)
+        
+        req.HTTPMethod = "POST"
+        req.setValue("text/plain; charset=UTF-8", forHTTPHeaderField: "Content-Type")
+        
+        let postData = postStr.dataUsingEncoding(NSUTF8StringEncoding,
+            allowLossyConversion: false)!
+        
+        req.HTTPBody = postData
+        req.setValue(String(postData.length), forHTTPHeaderField: "Content-Length")
+        
+        waitingForPost = true
+        
+        DefaultSocketLogger.Logger.log("POSTing: %@", type: "SocketEngine", args: postStr)
+        
+        doRequest(req) {[weak self] data, res, err in
+            guard let this = self else {return}
+            
+            if err != nil {
+                DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: "SocketEngine")
+                
+                if this.polling {
+                    this.handlePollingFailed(err?.localizedDescription ?? "Error")
+                }
+                
+                return
+            }
+            
+            this.waitingForPost = false
+            
+            dispatch_async(this.emitQueue) {
+                if !this.fastUpgrade {
+                    this.flushWaitingForPost()
+                    this.doPoll()
+                }
+            }
+        }
+    }
+    
+    func parsePollingMessage(str: String) {
+        guard str.characters.count != 1 else {
+            return
+        }
+        
+        var reader = SocketStringReader(message: str)
+        
+        while reader.hasNext {
+            if let n = Int(reader.readUntilStringOccurence(":")) {
+                let str = reader.read(n)
+                
+                dispatch_async(handleQueue) {
+                    self.parseEngineMessage(str, fromPolling: true)
+                }
+            } else {
+                dispatch_async(handleQueue) {
+                    self.parseEngineMessage(str, fromPolling: true)
+                }
+                break
+            }
+        }
+    }
+    
+    /// Send polling message.
+    /// Only call on emitQueue
+    public func sendPollMessage(message: String, withType type: SocketEnginePacketType,
+        withData datas: [NSData]) {
+            DefaultSocketLogger.Logger.log("Sending poll: %@ as type: %@", type: "SocketEngine", args: message, type.rawValue)
+            let fixedMessage = doubleEncodeUTF8(message)
+            let strMsg = "\(type.rawValue)\(fixedMessage)"
+            
+            postWait.append(strMsg)
+            
+            for data in datas {
+                if case let .Right(bin) = createBinaryDataForSend(data) {
+                    postWait.append(bin)
+                }
+            }
+            
+            if !waitingForPost {
+                flushWaitingForPost()
+            }
+    }
+    
+    public func stopPolling() {
+        invalidated = true
+        session?.finishTasksAndInvalidate()
+    }
+}

Source/SocketEngineSpec.swift

@@ -26,17 +26,53 @@
 import Foundation
 
 @objc public protocol SocketEngineSpec {
-    weak var client: SocketEngineClient? {get set}
+    weak var client: SocketEngineClient? { get set }
-    var cookies: [NSHTTPCookie]? {get}
+    var closed: Bool { get }
-    var sid: String {get}
+    var connected: Bool { get }
-    var socketPath: String {get}
+    var cookies: [NSHTTPCookie]? { get }
-    var urlPolling: String {get}
+    var extraHeaders: [String: String]? { get }
-    var urlWebSocket: String {get}
+    var fastUpgrade: Bool { get }
+    var forcePolling: Bool { get }
+    var forceWebsockets: Bool { get }
+    var parseQueue: dispatch_queue_t! { get }
+    var pingTimer: NSTimer? { get }
+    var polling: Bool { get }
+    var postWait: [String] { get set }
+    var probing: Bool { get set }
+    var emitQueue: dispatch_queue_t! { get }
+    var handleQueue: dispatch_queue_t! { get }
+    var sid: String { get }
+    var socketPath: String { get }
+    var urlPolling: String { get }
+    var urlWebSocket: String { get }
+    var websocket: Bool { get }
     
     init(client: SocketEngineClient, url: String, options: NSDictionary?)
     
     func close()
+    func doFastUpgrade()
+    func flushWaitingForPostToWebSocket()
     func open(opts: [String: AnyObject]?)
+    func parseEngineData(data: NSData)
+    func parseEngineMessage(message: String, fromPolling: Bool)
     func send(msg: String, withData datas: [NSData])
     func write(msg: String, withType type: SocketEnginePacketType, withData data: [NSData])
 }
+
+extension SocketEngineSpec {
+    func createBinaryDataForSend(data: NSData) -> Either<NSData, String> {
+        if websocket {
+            var byteArray = [UInt8](count: 1, repeatedValue: 0x0)
+            byteArray[0] = 4
+            let mutData = NSMutableData(bytes: &byteArray, length: 1)
+            
+            mutData.appendData(data)
+            
+            return .Left(mutData)
+        } else {
+            let str = "b4" + data.base64EncodedStringWithOptions(.Encoding64CharacterLineLength)
+            
+            return .Right(str)
+        }
+    }
+}

Source/SocketEngineWebsocket.swift

@@ -0,0 +1,64 @@
+//
+//  SocketEngineWebsocket.swift
+//  Socket.IO-Client-Swift
+//
+//  Created by Erik Little on 1/15/16.
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in
+//  all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+//  THE SOFTWARE.
+//
+
+import Foundation
+
+public protocol SocketEngineWebsocket: SocketEngineSpec, WebSocketDelegate {
+    var websocketConnected: Bool { get set }
+    var ws: WebSocket? { get }
+
+    func sendWebSocketMessage(str: String, withType type: SocketEnginePacketType, withData datas: [NSData])
+}
+
+// WebSocket methods
+extension SocketEngineWebsocket {
+    func probeWebSocket() {
+        if websocketConnected {
+            sendWebSocketMessage("probe", withType: .Ping, withData: [])
+        }
+    }
+    
+    /// Send message on WebSockets
+    /// Only call on emitQueue
+    public func sendWebSocketMessage(str: String, withType type: SocketEnginePacketType, withData datas: [NSData]) {
+            DefaultSocketLogger.Logger.log("Sending ws: %@ as type: %@", type: "SocketEngine", args: str, type.rawValue)
+            
+            ws?.writeString("\(type.rawValue)\(str)")
+            
+            for data in datas {
+                if case let .Left(bin) = createBinaryDataForSend(data) {
+                    ws?.writeData(bin)
+                }
+            }
+    }
+    
+    public func websocketDidReceiveMessage(socket: WebSocket, text: String) {
+        parseEngineMessage(text, fromPolling: false)
+    }
+    
+    public func websocketDidReceiveData(socket: WebSocket, data: NSData) {
+        parseEngineData(data)
+    }
+}