From 44508a51934cdf230dfdc618a47744f4c5b41b73 Mon Sep 17 00:00:00 2001
From: Hank Bao <hankbao84@gmail.com>
Date: Wed, 9 Sep 2015 22:33:35 +0800
Subject: [PATCH] Fix use-after-release issue in WebSocket

---
 SocketIOClientSwift/WebSocket.swift | 74 +++++++++++++++++------------
 1 file changed, 43 insertions(+), 31 deletions(-)

diff --git a/SocketIOClientSwift/WebSocket.swift b/SocketIOClientSwift/WebSocket.swift
index 7c8c787..1bc87f7 100644
--- a/SocketIOClientSwift/WebSocket.swift
+++ b/SocketIOClientSwift/WebSocket.swift
@@ -133,15 +133,18 @@ public class WebSocket : NSObject, NSStreamDelegate {
         if isCreated {
             return
         }
-        unowned let weakSelf = self
-        
-        dispatch_async(queue,{
-            weakSelf.didDisconnect = false
+
+        dispatch_async(queue, { [weak self] in
+            if let weakSelf = self {
+                weakSelf.didDisconnect = false
+            }
         })
-        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT,0), {
-            weakSelf.isCreated = true
-            weakSelf.createHTTPRequest()
-            weakSelf.isCreated = false
+        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT,0), { [weak self] in
+            if let weakSelf = self {
+                weakSelf.isCreated = true
+                weakSelf.createHTTPRequest()
+                weakSelf.isCreated = false
+            }
         })
     }
     
@@ -370,14 +373,15 @@ public class WebSocket : NSObject, NSStreamDelegate {
             }
         }
         if totalSize > 0 {
-            unowned let weakSelf = self
             if validateResponse(buffer, bufferLen: totalSize) {
-                dispatch_async(queue,{
-                    weakSelf.connected = true
-                    if let connectBlock = weakSelf.onConnect {
-                        connectBlock()
+                dispatch_async(queue, { [weak self] in
+                    if let weakSelf = self {
+                        weakSelf.connected = true
+                        if let connectBlock = weakSelf.onConnect {
+                            connectBlock()
+                        }
+                        weakSelf.delegate?.websocketDidConnect(weakSelf)
                     }
-                    weakSelf.delegate?.websocketDidConnect(self)
                 })
                 totalSize += 1 //skip the last \n
                 let restSize = bufferLen - totalSize
@@ -515,11 +519,13 @@ public class WebSocket : NSObject, NSStreamDelegate {
                 data = NSData(bytes: UnsafePointer<UInt8>((buffer+offset)), length: Int(len))
             }
             if receivedOpcode == OpCode.Pong.rawValue {
-                dispatch_async(queue,{[unowned self] in
-                    if let pongBlock = self.onPong {
-                        pongBlock()
+                dispatch_async(queue, { [weak self] in
+                    if let weakSelf = self {
+                        if let pongBlock = weakSelf.onPong {
+                            pongBlock()
+                        }
+                        weakSelf.pongDelegate?.websocketDidReceivePong(weakSelf)
                     }
-                    self.pongDelegate?.websocketDidReceivePong(self)
                 })
                 
                 let step = Int(offset+numericCast(len))
@@ -608,19 +614,23 @@ public class WebSocket : NSObject, NSStreamDelegate {
                     writeError(CloseCode.Encoding.rawValue)
                     return false
                 }
-                dispatch_async(queue,{[unowned self] in
-                    if let textBlock = self.onText {
-                        textBlock(str! as String)
+                dispatch_async(queue, { [weak self] in
+                    if let weakSelf = self {
+                        if let textBlock = weakSelf.onText {
+                            textBlock(str! as String)
+                        }
+                        weakSelf.delegate?.websocketDidReceiveMessage(weakSelf, text: str! as String)
                     }
-                    self.delegate?.websocketDidReceiveMessage(self, text: str! as String)
                 })
             } else if response.code == .BinaryFrame {
                 let data = response.buffer! //local copy so it is perverse for writing
-                dispatch_async(queue,{[unowned self] in
-                    if let dataBlock = self.onData {
-                        dataBlock(data)
+                dispatch_async(queue, { [weak self] in
+                    if let weakSelf = self {
+                        if let dataBlock = weakSelf.onData {
+                            dataBlock(data)
+                        }
+                        weakSelf.delegate?.websocketDidReceiveData(weakSelf, data: data)
                     }
-                    self.delegate?.websocketDidReceiveData(self, data: data)
                 })
             }
             readStack.removeLast()
@@ -723,12 +733,14 @@ public class WebSocket : NSObject, NSStreamDelegate {
     ///used to preform the disconnect delegate
     private func doDisconnect(error: NSError?) {
         if !self.didDisconnect {
-            dispatch_async(queue,{[unowned self] in
-                self.didDisconnect = true
-                if let disconnect = self.onDisconnect {
-                    disconnect(error)
+            dispatch_async(queue, { [weak self] in
+                if let weakSelf = self {
+                    weakSelf.didDisconnect = true
+                    if let disconnect = weakSelf.onDisconnect {
+                        disconnect(error)
+                    }
+                    weakSelf.delegate?.websocketDidDisconnect(weakSelf, error: error)
                 }
-                self.delegate?.websocketDidDisconnect(self, error: error)
             })
         }
     }