From df8f87cc21bf6c5ca516ae608bb85f91ba1dfb42 Mon Sep 17 00:00:00 2001
From: Erik Little <nuclear.ace@gmail.com>
Date: Sat, 16 Sep 2017 09:20:50 -0400
Subject: [PATCH] Extract addHeaders to enginespec

---
 Source/SocketIO/Engine/SocketEngine.swift      | 18 ++++--------------
 .../SocketIO/Engine/SocketEnginePollable.swift | 13 -------------
 Source/SocketIO/Engine/SocketEngineSpec.swift  | 12 ++++++++++++
 3 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/Source/SocketIO/Engine/SocketEngine.swift b/Source/SocketIO/Engine/SocketEngine.swift
index e11046c..8d64083 100644
--- a/Source/SocketIO/Engine/SocketEngine.swift
+++ b/Source/SocketIO/Engine/SocketEngine.swift
@@ -269,17 +269,7 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
 
         var reqPolling = URLRequest(url: urlPolling, cachePolicy: .reloadIgnoringLocalCacheData, timeoutInterval: 60.0)
 
-        if cookies != nil {
-            let headers = HTTPCookie.requestHeaderFields(with: cookies!)
-            reqPolling.allHTTPHeaderFields = headers
-        }
-
-        if let extraHeaders = extraHeaders {
-            for (headerName, value) in extraHeaders {
-                reqPolling.setValue(value, forHTTPHeaderField: headerName)
-            }
-        }
-
+        addHeaders(to: &reqPolling)
         doLongPoll(for: reqPolling)
     }
 
@@ -303,9 +293,9 @@ public final class SocketEngine : NSObject, URLSessionDelegate, SocketEnginePoll
             urlWebSocket.scheme = "ws"
         }
 
-        if connectParams != nil {
-            for (key, value) in connectParams! {
-                let keyEsc   = key.urlEncode()!
+        if let connectParams = self.connectParams {
+            for (key, value) in connectParams {
+                let keyEsc = key.urlEncode()!
                 let valueEsc = "\(value)".urlEncode()!
 
                 queryString += "&\(keyEsc)=\(valueEsc)"
diff --git a/Source/SocketIO/Engine/SocketEnginePollable.swift b/Source/SocketIO/Engine/SocketEnginePollable.swift
index e3e9248..89c7dff 100644
--- a/Source/SocketIO/Engine/SocketEnginePollable.swift
+++ b/Source/SocketIO/Engine/SocketEnginePollable.swift
@@ -71,19 +71,6 @@ public protocol SocketEnginePollable : SocketEngineSpec {
 
 // Default polling methods
 extension SocketEnginePollable {
-    private func addHeaders(to req: inout URLRequest) {
-        if cookies != nil {
-            let headers = HTTPCookie.requestHeaderFields(with: cookies!)
-            req.allHTTPHeaderFields = headers
-        }
-
-        if extraHeaders != nil {
-            for (headerName, value) in extraHeaders! {
-                req.setValue(value, forHTTPHeaderField: headerName)
-            }
-        }
-    }
-
     func createRequestForPostWithPostWait() -> URLRequest {
         defer { postWait.removeAll(keepingCapacity: true) }
 
diff --git a/Source/SocketIO/Engine/SocketEngineSpec.swift b/Source/SocketIO/Engine/SocketEngineSpec.swift
index ab8f10d..7e3e213 100644
--- a/Source/SocketIO/Engine/SocketEngineSpec.swift
+++ b/Source/SocketIO/Engine/SocketEngineSpec.swift
@@ -147,6 +147,18 @@ extension SocketEngineSpec {
         return com.url!
     }
 
+    func addHeaders(to req: inout URLRequest) {
+        if let cookies = cookies {
+            req.allHTTPHeaderFields = HTTPCookie.requestHeaderFields(with: cookies)
+        }
+
+        if let extraHeaders = extraHeaders {
+            for (headerName, value) in extraHeaders {
+                req.setValue(value, forHTTPHeaderField: headerName)
+            }
+        }
+    }
+
     func createBinaryDataForSend(using data: Data) -> Either<Data, String> {
         if websocket {
             var byteArray = [UInt8](repeating: 0x4, count: 1)