Update Starscream, for real this time

Apollo.podspec

@@ -39,7 +39,7 @@ Pod::Spec.new do |s|
   s.subspec 'WebSocket' do |ss|
     ss.source_files = 'Sources/ApolloWebSocket/*.swift'
     ss.dependency 'Apollo/Core'
-    ss.dependency 'Starscream', '~>3.1.1'
+    ss.dependency 'Starscream', '~>4.0.4'
   end
 
 end

Apollo.xcodeproj/project.pbxproj

@@ -3553,7 +3553,7 @@
 			repositoryURL = "https://github.com/daltoniam/Starscream.git";
 			requirement = {
 				kind = upToNextMinorVersion;
-				minimumVersion = 3.1.1;
+				minimumVersion = 4.0.4;
 			};
 		};
 		9B7BDAF423FDEE2600ACD198 /* XCRemoteSwiftPackageReference "SQLite.swift" */ = {

Apollo.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -42,8 +42,8 @@
         "repositoryURL": "https://github.com/daltoniam/Starscream.git",
         "state": {
           "branch": null,
-          "revision": "e6b65c6d9077ea48b4a7bdda8994a1d3c6969c8d",
-          "version": "3.1.1"
+          "revision": "df8d82047f6654d8e4b655d1b1525c64e1059d21",
+          "version": "4.0.4"
         }
       },
       {
@@ -54,15 +54,6 @@
           "revision": "94197b3adbbf926348ad8765476a158aa4e54f8a",
           "version": "0.14.0"
         }
-      },
-      {
-        "package": "swift-nio-zlib-support",
-        "repositoryURL": "https://github.com/apple/swift-nio-zlib-support.git",
-        "state": {
-          "branch": null,
-          "revision": "37760e9a52030bb9011972c5213c3350fa9d41fd",
-          "version": "1.0.0"
-        }
       }
     ]
   },

Package.swift

@@ -38,7 +38,7 @@ let package = Package(
       .upToNextMinor(from: "0.12.2")),
     .package(
       url: "https://github.com/daltoniam/Starscream",
-      .upToNextMinor(from: "3.1.1")),
+      .upToNextMinor(from: "4.0.4")),
     .package(
       url: "https://github.com/stencilproject/Stencil.git",
       .upToNextMinor(from: "0.14.0")),

Sources/ApolloWebSocket/ApolloWebSocket.swift

@@ -5,48 +5,40 @@ import Foundation
 
 /// Protocol allowing alternative implementations of web sockets beyond `ApolloWebSocket`. Extends `Starscream`'s `WebSocketClient` protocol.
 public protocol ApolloWebSocketClient: WebSocketClient {
-
+  
   /// Required initializer
   ///
-  /// - Parameter request: The URLRequest to use on connection.
-  /// - Parameter protocols: The supported protocols
-  init(request: URLRequest, protocols: [String]?)
+  /// - Parameters:
+  ///   - request: The URLRequest to use on connection.
+  ///   - certPinner: [optional] The object providing information about certificate pinning. Should default to Starscream's `FoundationSecurity`.
+  ///   - compressionHandler: [optional] The object helping with any compression handling. Should default to nil.
+  init(request: URLRequest,
+       certPinner: CertificatePinning?,
+       compressionHandler: CompressionHandler?)
 
   /// The URLRequest used on connection.
   var request: URLRequest { get set }
 
   /// Queue where the callbacks are executed
   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 }
+  var delegate: WebSocketDelegate? { get set }
 }
 
 // MARK: - WebSocket
 
 /// Included implementation of an `ApolloWebSocketClient`, based on `Starscream`'s `WebSocket`.
-public class ApolloWebSocket: WebSocket, ApolloWebSocketClient, SOCKSProxyable {
+public class ApolloWebSocket: WebSocket, ApolloWebSocketClient {
 
-  private var stream: FoundationStream!
+  private var transport: FoundationTransport!
 
-  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: stream)
-    self.stream = stream
+  required public init(request: URLRequest,
+                       certPinner: CertificatePinning? = FoundationSecurity(),
+                       compressionHandler: CompressionHandler? = nil) {
+    let engine = WSEngine(transport: FoundationTransport(),
+                          certPinner: certPinner,
+                          compressionHandler: compressionHandler)
+
+    super.init(request: request, engine: engine)
   }
 }

Sources/ApolloWebSocket/WebSocketTransport.swift

@@ -11,12 +11,16 @@ public protocol WebSocketTransportDelegate: class {
   func webSocketTransportDidConnect(_ webSocketTransport: WebSocketTransport)
   func webSocketTransportDidReconnect(_ webSocketTransport: WebSocketTransport)
   func webSocketTransport(_ webSocketTransport: WebSocketTransport, didDisconnectWithError error:Error?)
+  func webSocketTransport(_ webSocketTransport: WebSocketTransport, didReceivePingData: Data?)
+  func webSocketTransport(_ webSocketTransport: WebSocketTransport, didReceivePongData: Data?)
 }
 
 public extension WebSocketTransportDelegate {
   func webSocketTransportDidConnect(_ webSocketTransport: WebSocketTransport) {}
   func webSocketTransportDidReconnect(_ webSocketTransport: WebSocketTransport) {}
   func webSocketTransport(_ webSocketTransport: WebSocketTransport, didDisconnectWithError error:Error?) {}
+  func webSocketTransport(_ webSocketTransport: WebSocketTransport, didReceivePingData: Data?) {}
+  func webSocketTransport(_ webSocketTransport: WebSocketTransport, didReceivePongData: Data?) {}
 }
 
 // MARK: - WebSocketTransport
@@ -34,6 +38,9 @@ public class WebSocketTransport {
   private let requestBodyCreator: RequestBodyCreator
 
   private final let protocols = ["graphql-ws"]
+
+  /// non-private for testing - you should not use this directly
+  var isSocketConnected = Atomic<Bool>(false)
 
   private var acked = false
 
@@ -64,49 +71,21 @@ public class WebSocketTransport {
     }
   }
 
-  public var security: SSLTrustValidator? {
-    get {
-      return websocket.security
-    }
-    set {
-      websocket.security = newValue
-    }
-  }
-
-  /// 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
   ///
-  /// - Parameter request: The connection URLRequest
-  /// - Parameter clientName: The client name to use for this client. Defaults to `Self.defaultClientName`
-  /// - Parameter clientVersion: The client version to use for this client. Defaults to `Self.defaultClientVersion`.
-  /// - Parameter sendOperationIdentifiers: Whether or not to send operation identifiers with operations. Defaults to false.
-  /// - Parameter reconnect: Whether to auto reconnect when websocket looses connection. Defaults to true.
-  /// - Parameter reconnectionInterval: How long to wait before attempting to reconnect. Defaults to half a second.
-  /// - Parameter allowSendingDuplicates: Allow sending duplicate messages. Important when reconnected. Defaults to true.
-  /// - Parameter connectOnInit: Whether the websocket connects immediately on creation. If false, remember to call `resumeWebSocketConnection()` to connect. Defaults to true.
-  /// - Parameter connectingPayload: [optional] The payload to send on connection. Defaults to an empty `GraphQLMap`.
-  /// - Parameter requestBodyCreator: The `RequestBodyCreator` to use when serializing requests. Defaults to an `ApolloRequestBodyCreator`.
+  /// - Parameters:
+  ///   - request: The connection URLRequest
+  ///   - clientName: The client name to use for this client. Defaults to `Self.defaultClientName`
+  ///   - clientVersion: The client version to use for this client. Defaults to `Self.defaultClientVersion`.
+  ///   - sendOperationIdentifiers: Whether or not to send operation identifiers with operations. Defaults to false.
+  ///   - reconnect: Whether to auto reconnect when websocket looses connection. Defaults to true.
+  ///   - reconnectionInterval: How long to wait before attempting to reconnect. Defaults to half a second.
+  ///   - allowSendingDuplicates: Allow sending duplicate messages. Important when reconnected. Defaults to true.
+  ///  - connectOnInit: Whether the websocket connects immediately on creation. If false, remember to call `resumeWebSocketConnection()` to connect. Defaults to true.
+  ///   - connectingPayload: [optional] The payload to send on connection. Defaults to an empty `GraphQLMap`.
+  ///   - requestBodyCreator: The `RequestBodyCreator` to use when serializing requests. Defaults to an `ApolloRequestBodyCreator`.
+  ///   - certPinner: [optional] The object providing information about certificate pinning. Should default to Starscream's `FoundationSecurity`.
+  ///   - compressionHandler: [optional] The object helping with any compression handling. Should default to nil.
   public init(request: URLRequest,
               clientName: String = WebSocketTransport.defaultClientName,
               clientVersion: String = WebSocketTransport.defaultClientVersion,
@@ -116,18 +95,24 @@ public class WebSocketTransport {
               allowSendingDuplicates: Bool = true,
               connectOnInit: Bool = true,
               connectingPayload: GraphQLMap? = [:],
-              requestBodyCreator: RequestBodyCreator = ApolloRequestBodyCreator()) {
+              requestBodyCreator: RequestBodyCreator = ApolloRequestBodyCreator(),
+              certPinner: CertificatePinning? = FoundationSecurity(),
+              compressionHandler: CompressionHandler? = nil) {
     self.connectingPayload = connectingPayload
     self.sendOperationIdentifiers = sendOperationIdentifiers
     self.reconnect = Atomic(reconnect)
     self.reconnectionInterval = reconnectionInterval
     self.allowSendingDuplicates = allowSendingDuplicates
     self.requestBodyCreator = requestBodyCreator
-    self.websocket = WebSocketTransport.provider.init(request: request, protocols: protocols)
+    self.websocket = WebSocketTransport.provider.init(request: request,
+                                                      certPinner: certPinner,
+                                                      compressionHandler: compressionHandler)
+    self.websocket.request.setValue(self.protocols.joined(separator: ","), forHTTPHeaderField: "Sec-WebSocket-Protocol")
     self.clientName = clientName
     self.clientVersion = clientVersion
     self.connectOnInit = connectOnInit
     self.addApolloClientHeaders(to: &self.websocket.request)
+
     self.websocket.delegate = self
     if connectOnInit {
       self.websocket.connect()
@@ -136,14 +121,14 @@ public class WebSocketTransport {
   }
 
   public func isConnected() -> Bool {
-    return websocket.isConnected
+    return self.isSocketConnected.value
   }
 
   public func ping(data: Data, completionHandler: (() -> Void)? = nil) {
     return websocket.write(ping: data, completion: completionHandler)
   }
 
-  private func processMessage(socket: WebSocketClient, text: String) {
+  private func processMessage(text: String) {
     OperationMessage(serialized: text).parse { parseHandler in
       guard
         let type = parseHandler.type,
@@ -225,7 +210,7 @@ public class WebSocketTransport {
     }
   }
 
-  private func processMessage(socket: WebSocketClient, data: Data) {
+  private func processMessage(data: Data) {
     print("WebSocketTransport::unprocessed event \(data)")
   }
 
@@ -255,7 +240,7 @@ public class WebSocketTransport {
   private func write(_ str: String,
                      force forced: Bool = false,
                      id: Int? = nil) {
-    if websocket.isConnected && (acked || forced) {
+    if self.isSocketConnected.value && (acked || forced) {
       websocket.write(string: str)
     } else {
       // using sequence number to make sure that the queue is processed correctly
@@ -272,7 +257,7 @@ public class WebSocketTransport {
 
   deinit {
     websocket.disconnect()
-    websocket.delegate = nil
+    self.websocket.delegate = nil
   }
 
   func sendHelper<Operation: GraphQLOperation>(operation: Operation, resultHandler: @escaping (_ result: Result<JSONObject, Error>) -> Void) -> String? {
@@ -392,16 +377,60 @@ extension WebSocketTransport: NetworkTransport {
 // MARK: - WebSocketDelegate implementation
 
 extension WebSocketTransport: WebSocketDelegate {
-
-  public func websocketDidConnect(socket: WebSocketClient) {
+
+  public func didReceive(event: WebSocketEvent, client: WebSocket) {
+      switch event {
+      case .connected:
+        self.handleConnection()
+      case .disconnected(let reason, let code):
+        self.isSocketConnected.mutate { $0 = false }
+        self.error.mutate { $0 = nil }
+        debugPrint("websocket is disconnected: \(reason) with code: \(code)")
+        self.handleDisconnection()
+      case .text(let text):
+        self.processMessage(text: text)
+      case .binary(let data):
+        self.processMessage(data: data)
+      case .ping(let pingData):
+        self.delegate?.webSocketTransport(self, didReceivePingData: pingData)
+      case .pong(let pongData):
+        self.delegate?.webSocketTransport(self, didReceivePongData: pongData)
+      case .viabilityChanged(_):
+        break
+      case .reconnectSuggested(let shouldReconnect):
+        if shouldReconnect {
+          self.attemptReconnectionIfDesired()
+        }
+      case .cancelled:
+        self.isSocketConnected.mutate { $0 = false }
+        self.error.mutate { $0 = nil }
+        self.handleDisconnection()
+      case .error(let error):
+        self.isSocketConnected.mutate { $0 = false }
+        // report any error to all subscribers
+        if let error = error {
+          self.error.mutate { $0 = WebSocketError(payload: nil,
+                                                  error: error,
+                                                  kind: .networkError) }
+          self.notifyErrorAllHandlers(error)
+        } else {
+          self.error.mutate { $0 = nil }
+        }
+
+        self.handleDisconnection()
+      }
+  }
+
+  public func handleConnection() {
     self.error.mutate { $0 = nil }
+    self.isSocketConnected.mutate { $0 = true }
     initServer()
-    if reconnected {
+    if self.reconnected {
       self.delegate?.webSocketTransportDidReconnect(self)
       // re-send the subscriptions whenever we are re-connected
       // for the first connect, any subscriptions are already in queue
-      for (_,msg) in self.subscriptions {
-        if allowSendingDuplicates {
+      for (_, msg) in self.subscriptions {
+        if self.allowSendingDuplicates {
           write(msg)
         } else {
           // search duplicate message from the queue
@@ -413,35 +442,23 @@ extension WebSocketTransport: WebSocketDelegate {
       self.delegate?.webSocketTransportDidConnect(self)
     }
 
-    reconnected = true
+    self.reconnected = true
   }
 
-  public func websocketDidDisconnect(socket: WebSocketClient, error: Error?) {
-    // report any error to all subscribers
-    if let error = error {
-      self.error.mutate { $0 = WebSocketError(payload: nil,
-                                              error: error,
-                                              kind: .networkError) }
-      self.notifyErrorAllHandlers(error)
-    } else {
-      self.error.mutate { $0 = nil }
-    }
-
+  private func handleDisconnection()  {
     self.delegate?.webSocketTransport(self, didDisconnectWithError: self.error.value)
-    acked = false // need new connect and ack before sending
+    self.acked = false // need new connect and ack before sending
 
-    if reconnect.value {
-      DispatchQueue.main.asyncAfter(deadline: .now() + reconnectionInterval) {
-        self.websocket.connect()
-      }
-    }
+    self.attemptReconnectionIfDesired()
   }
-
-  public func websocketDidReceiveMessage(socket: WebSocketClient, text: String) {
-    processMessage(socket: socket, text: text)
-  }
-
-  public func websocketDidReceiveData(socket: WebSocketClient, data: Data) {
-    processMessage(socket: socket, data: data)
+
+  private func attemptReconnectionIfDesired() {
+    guard self.reconnect.value else {
+      return
+    }
+
+    DispatchQueue.main.asyncAfter(deadline: .now() + reconnectionInterval) {
+      self.websocket.connect()
+    }
   }
 }