From 60c4277fbc789fdd5f9d786d0e4be4b6e63392fe Mon Sep 17 00:00:00 2001
From: dsrees <daniel.rees18@gmail.com>
Date: Thu, 28 Sep 2023 22:03:47 -0400
Subject: [PATCH] Added headers to be sent to Websocket

---
 .../SwiftPhoenixClient/PhoenixTransport.swift   | 17 ++++++++++++++---
 Sources/SwiftPhoenixClient/Socket.swift         |  5 ++++-
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/Sources/SwiftPhoenixClient/PhoenixTransport.swift b/Sources/SwiftPhoenixClient/PhoenixTransport.swift
index 28bcadc5..87e4561b 100644
--- a/Sources/SwiftPhoenixClient/PhoenixTransport.swift
+++ b/Sources/SwiftPhoenixClient/PhoenixTransport.swift
@@ -38,8 +38,11 @@ public protocol PhoenixTransport {
   
   /**
    Connect to the server
+   
+   - Parameters:
+   - headers: Headers to include in the URLRequests when opening the Websocket connection. Can be empty [:]
    */
-  func connect()
+  func connect(with headers: [String: Any])
   
   /**
    Disconnect from the server.
@@ -193,13 +196,21 @@ open class URLSessionTransport: NSObject, PhoenixTransport, URLSessionWebSocketD
   public var readyState: PhoenixTransportReadyState = .closed
   public var delegate: PhoenixTransportDelegate? = nil
   
-  open func connect() {
+  public func connect(with headers: [String : Any]) {
     // Set the transport state as connecting
     self.readyState = .connecting
     
     // Create the session and websocket task
     self.session = URLSession(configuration: self.configuration, delegate: self, delegateQueue: nil)
-    self.task = self.session?.webSocketTask(with: url)
+    var request = URLRequest(url: url)
+      
+    headers.forEach { (key: String, value: Any) in
+        guard let value = value as? String else { return }
+        request.addValue(value, forHTTPHeaderField: key)
+    }
+      
+    self.task = self.session?.webSocketTask(with: request)
+
     
     // Start the task
     self.task?.resume()
diff --git a/Sources/SwiftPhoenixClient/Socket.swift b/Sources/SwiftPhoenixClient/Socket.swift
index e47e4b62..6fedb7b6 100644
--- a/Sources/SwiftPhoenixClient/Socket.swift
+++ b/Sources/SwiftPhoenixClient/Socket.swift
@@ -97,6 +97,9 @@ public class Socket: PhoenixTransportDelegate {
   
   /// Timeout to use when opening connections
   public var timeout: TimeInterval = Defaults.timeoutInterval
+    
+  /// Custom headers to be added to the socket connection request
+  public var headers: [String : Any] = [:]
   
   /// Interval between sending a heartbeat
   public var heartbeatInterval: TimeInterval = Defaults.heartbeatInterval
@@ -266,7 +269,7 @@ public class Socket: PhoenixTransportDelegate {
 //    self.connection?.enabledSSLCipherSuites = enabledSSLCipherSuites
 //    #endif
     
-    self.connection?.connect()
+    self.connection?.connect(with: self.headers)
   }
   
   /// Disconnects the socket