From 00c4725e1f58d191262ce98c5b900011d281dfc0 Mon Sep 17 00:00:00 2001
From: Ellen Shapiro <designatednerd@gmail.com>
Date: Mon, 30 Mar 2020 19:27:35 -0500
Subject: [PATCH] Add ability to get to Starscream's underlying SOCKS proxy
 property.

---
 Sources/ApolloWebSocket/ApolloWebSocket.swift | 25 +++++++++++++++++--
 .../ApolloWebSocket/WebSocketTransport.swift  | 22 ++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/Sources/ApolloWebSocket/ApolloWebSocket.swift b/Sources/ApolloWebSocket/ApolloWebSocket.swift
index f19d6d3a05..35ea2384d0 100644
--- a/Sources/ApolloWebSocket/ApolloWebSocket.swift
+++ b/Sources/ApolloWebSocket/ApolloWebSocket.swift
@@ -19,13 +19,34 @@ public protocol ApolloWebSocketClient: WebSocketClient {
   var callbackQueue: DispatchQueue { get set }
 }
 
+public protocol SOCKSProxyable {
+  
+  /// Determines whether a SOCKS proxy is enabled on the underlying request.
+  /// Mostly useful for debugging with tools like Charles Proxy.
+  var enableSOCKSProxy: Bool { get set }
+}
+
 // MARK: - WebSocket
 
 /// Included implementation of an `ApolloWebSocketClient`, based on `Starscream`'s `WebSocket`.
-public class ApolloWebSocket: WebSocket, ApolloWebSocketClient {
+public class ApolloWebSocket: WebSocket, ApolloWebSocketClient, SOCKSProxyable {
+  
+  private var stream: FoundationStream!
+  
+  public var enableSOCKSProxy: Bool {
+    get {
+      return self.stream.enableSOCKSProxy
+    }
+    set {
+      self.stream.enableSOCKSProxy = newValue
+    }
+  }
+
   required public convenience init(request: URLRequest, protocols: [String]? = nil) {
+    let stream = FoundationStream()
     self.init(request: request,
               protocols: protocols,
-              stream: FoundationStream())
+              stream: stream)
+    self.stream = stream
   }
 }
diff --git a/Sources/ApolloWebSocket/WebSocketTransport.swift b/Sources/ApolloWebSocket/WebSocketTransport.swift
index 620f065e1f..7bab422c61 100644
--- a/Sources/ApolloWebSocket/WebSocketTransport.swift
+++ b/Sources/ApolloWebSocket/WebSocketTransport.swift
@@ -61,6 +61,28 @@ public class WebSocketTransport {
       self.addApolloClientHeaders(to: &self.websocket.request)
     }
   }
+  
+  /// Determines whether a SOCKS proxy is enabled on the underlying request.
+  /// Mostly useful for debugging with tools like Charles Proxy.
+  /// Note: Will return `false` from the getter and no-op the setter for implementations that do not conform to `SOCKSProxyable`.
+  public var enableSOCKSProxy: Bool {
+    get {
+      guard let socket = self.websocket as? SOCKSProxyable else {
+        // If it's not proxyable, then the proxy can't be enabled
+        return false
+      }
+      
+      return socket.enableSOCKSProxy
+    }
+    set {
+      guard var socket = self.websocket as? SOCKSProxyable else {
+        // If it's not proxyable, there's nothing to do here.
+        return
+      }
+      
+      socket.enableSOCKSProxy = newValue
+    }
+  }
 
   /// Designated initializer
   ///