From ef02247a0c0bb406a523a4784f0633aa5334aff3 Mon Sep 17 00:00:00 2001
From: Erik Little <nuclear.ace@gmail.com>
Date: Fri, 2 Mar 2018 11:28:35 -0500
Subject: [PATCH] Use callbacks instead of delegate methods for Starscream

---
 Source/SocketIO/Engine/SocketEngine.swift     | 34 +++++++++++++++----
 .../Engine/SocketEngineWebsocket.swift        | 14 +-------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/Source/SocketIO/Engine/SocketEngine.swift b/Source/SocketIO/Engine/SocketEngine.swift
index d7253d5..46fe31a 100644
--- a/Source/SocketIO/Engine/SocketEngine.swift
+++ b/Source/SocketIO/Engine/SocketEngine.swift
@@ -280,7 +280,6 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
     }
 
     private func createWebSocketAndConnect() {
-        ws?.delegate = nil // TODO this seems a bit defensive, is this really needed?
         var req = URLRequest(url: urlWebSocketWithSid)
 
         addHeaders(to: &req)
@@ -288,10 +287,33 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
         ws = WebSocket(request: req)
         ws?.callbackQueue = engineQueue
         ws?.enableCompression = compress
-        ws?.delegate = self
         ws?.disableSSLCertValidation = selfSigned
         ws?.security = security?.security
 
+        ws?.onConnect = {[weak self] in
+            guard let this = self else { return }
+
+            this.websocketDidConnect()
+        }
+
+        ws?.onDisconnect = {[weak self] error in
+            guard let this = self else { return }
+
+            this.websocketDidDisconnect(error: error)
+        }
+
+        ws?.onData = {[weak self] data in
+            guard let this = self else { return }
+
+            this.parseEngineData(data)
+        }
+
+        ws?.onText = {[weak self] message in
+            guard let this = self else { return }
+
+            this.parseEngineMessage(message)
+        }
+
         ws?.connect()
     }
 
@@ -607,10 +629,9 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
         }
     }
 
-    // MARK: Starscream delegate conformance
+    // MARK: WebSocket Methods
 
-    /// Delegate method for connection.
-    public func websocketDidConnect(socket: WebSocketClient) {
+    private func websocketDidConnect() {
         if !forceWebsockets {
             probing = true
             probeWebSocket()
@@ -621,8 +642,7 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
         }
     }
 
-    /// Delegate method for disconnection.
-    public func websocketDidDisconnect(socket: WebSocketClient, error: Error?) {
+    private func websocketDidDisconnect(error: Error?) {
         probing = false
 
         if closed {
diff --git a/Source/SocketIO/Engine/SocketEngineWebsocket.swift b/Source/SocketIO/Engine/SocketEngineWebsocket.swift
index 97a4a72..6dc11ac 100644
--- a/Source/SocketIO/Engine/SocketEngineWebsocket.swift
+++ b/Source/SocketIO/Engine/SocketEngineWebsocket.swift
@@ -27,7 +27,7 @@ import Foundation
 import Starscream
 
 /// Protocol that is used to implement socket.io WebSocket support
-public protocol SocketEngineWebsocket : SocketEngineSpec, WebSocketDelegate {
+public protocol SocketEngineWebsocket : SocketEngineSpec {
     // MARK: Methods
 
     /// Sends an engine.io message through the WebSocket transport.
@@ -66,16 +66,4 @@ extension SocketEngineWebsocket {
             }
         }
     }
-
-    // MARK: Starscream delegate methods
-
-    /// Delegate method for when a message is received.
-    public func websocketDidReceiveMessage(socket: WebSocketClient, text: String) {
-        parseEngineMessage(text)
-    }
-
-    /// Delegate method for when binary is received.
-    public func websocketDidReceiveData(socket: WebSocketClient, data: Data) {
-        parseEngineData(data)
-    }
 }