From 0470647d878e8c6ec491a472a6587703003269a9 Mon Sep 17 00:00:00 2001
From: Erik <nuclear.ace@gmail.com>
Date: Thu, 19 Nov 2015 10:15:01 -0500
Subject: [PATCH] error on unknown engine messages

---
 SocketIOClientSwift/SocketEngine.swift | 69 +++++++++++++++-----------
 SocketIOClientSwift/SocketPacket.swift |  7 +--
 2 files changed, 42 insertions(+), 34 deletions(-)

diff --git a/SocketIOClientSwift/SocketEngine.swift b/SocketIOClientSwift/SocketEngine.swift
index d9b58f2..d2a34f9 100644
--- a/SocketIOClientSwift/SocketEngine.swift
+++ b/SocketIOClientSwift/SocketEngine.swift
@@ -121,7 +121,7 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
         stopPolling()
     }
 
-    private func checkIfMessageIsBase64Binary(var message: String) {
+    private func checkIfMessageIsBase64Binary(var message: String) -> Bool {
         if message.hasPrefix("b4") {
             // binary in base64 string
             message.removeRange(Range(start: message.startIndex,
@@ -131,6 +131,10 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
                 options: .IgnoreUnknownCharacters) {
                     client?.parseBinaryData(data)
             }
+            
+            return true
+        } else {
+            return false
         }
     }
 
@@ -384,11 +388,20 @@ public final class SocketEngine: NSObject, SocketEngineSpec, WebSocketDelegate {
 
     private func parseEngineMessage(var message: String, fromPolling: Bool) {
         DefaultSocketLogger.Logger.log("Got message: %@", type: logType, args: message)
+        
+        func handleOther(msg: String) -> SocketEnginePacketType {
+            if checkIfMessageIsBase64Binary(msg) {
+                return .Noop
+            } else {
+                DefaultSocketLogger.Logger.error("Got message: %@", type: logType, args: msg)
+                return .Close
+            }
+        }
+        
+        let reader = SocketStringReader(message: message)
 
-        let type = SocketEnginePacketType(rawValue: Int((message["^(\\d)"].groups()?[1]) ?? "") ?? -1) ?? {
-            self.checkIfMessageIsBase64Binary(message)
-            return .Noop
-            }()
+        let type = SocketEnginePacketType(rawValue: Int(reader.currentCharacter) ?? -1)
+            ?? handleOther(message)
 
         if fromPolling && type != .Noop {
             fixDoubleUTF8(&message)
@@ -524,32 +537,32 @@ extension SocketEngine {
     
     private func doLongPoll(req: NSMutableURLRequest) {
         doRequest(req) {[weak self] data, res, err in
-            if let this = self {
-                if err != nil || data == nil {
-                    if this.polling {
-                        this.handlePollingFailed(err?.localizedDescription ?? "Error")
-                    } else {
-                        DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: this.logType)
-                    }
-                    return
+            guard let this = self else {return}
+            
+            if err != nil || data == nil {
+                if this.polling {
+                    this.handlePollingFailed(err?.localizedDescription ?? "Error")
+                } else {
+                    DefaultSocketLogger.Logger.error(err?.localizedDescription ?? "Error", type: this.logType)
                 }
-                
-                DefaultSocketLogger.Logger.log("Got polling response", type: this.logType)
-                
-                if let str = NSString(data: data!, encoding: NSUTF8StringEncoding) as? String {
-                    dispatch_async(this.parseQueue) {[weak this] in
-                        this?.parsePollingMessage(str)
-                    }
-                }
-                
-                this.waitingForPoll = false
-                
-                if this.fastUpgrade {
-                    this.doFastUpgrade()
-                } else if !this.closed && this.polling {
-                    this.doPoll()
+                return
+            }
+            
+            DefaultSocketLogger.Logger.log("Got polling response", type: this.logType)
+            
+            if let str = NSString(data: data!, encoding: NSUTF8StringEncoding) as? String {
+                dispatch_async(this.parseQueue) {[weak this] in
+                    this?.parsePollingMessage(str)
                 }
             }
+            
+            this.waitingForPoll = false
+            
+            if this.fastUpgrade {
+                this.doFastUpgrade()
+            } else if !this.closed && this.polling {
+                this.doPoll()
+            }
         }
     }
     
diff --git a/SocketIOClientSwift/SocketPacket.swift b/SocketIOClientSwift/SocketPacket.swift
index c235c40..c522533 100644
--- a/SocketIOClientSwift/SocketPacket.swift
+++ b/SocketIOClientSwift/SocketPacket.swift
@@ -296,12 +296,7 @@ private extension SocketPacket {
         var binary = [NSData]()
         
         for i in 0..<data.count {
-            if data[i] is NSArray || data[i] is NSDictionary {
-                data[i] = shred(data[i], binary: &binary)
-            } else if let bin = data[i] as? NSData {
-                data[i] = ["_placeholder": true, "num": binary.count]
-                binary.append(bin)
-            }
+            data[i] = shred(data[i], binary: &binary)
         }
         
         return (data, binary)