Implement graphql-transport-ws protocol support

Sources/ApolloWebSocket/DefaultImplementation/WebSocket.swift

@@ -68,14 +68,25 @@ public final class WebSocket: NSObject, WebSocketClient, StreamDelegate, WebSock
     public let code: Int
   }
 
-  private struct Constants {
+  public enum WSProtocol: CustomStringConvertible {
+    case graphql_ws
+    case graphql_transport_ws
+
+    public var description: String {
+      switch self {
+      case .graphql_ws: return "graphql-ws"
+      case .graphql_transport_ws: return "graphql-transport-ws"
+      }
+    }
+  }
+
+  struct Constants {
     static let headerWSUpgradeName     = "Upgrade"
     static let headerWSUpgradeValue    = "websocket"
     static let headerWSHostName        = "Host"
     static let headerWSConnectionName  = "Connection"
     static let headerWSConnectionValue = "Upgrade"
     static let headerWSProtocolName    = "Sec-WebSocket-Protocol"
-    static let headerWSProtocolValue   = "graphql-ws"
     static let headerWSVersionName     = "Sec-WebSocket-Version"
     static let headerWSVersionValue    = "13"
     static let headerWSExtensionName   = "Sec-WebSocket-Extensions"
@@ -197,20 +208,29 @@ public final class WebSocket: NSObject, WebSocketClient, StreamDelegate, WebSock
       self.request.setValue(origin, forHTTPHeaderField: Constants.headerOriginName)
     }
 
-    self.request.setValue(Constants.headerWSProtocolValue,
-                          forHTTPHeaderField: Constants.headerWSProtocolName)
+    if self.request.value(forHTTPHeaderField: Constants.headerWSProtocolName) == nil {
+      self.request.setValue(WSProtocol.graphql_ws.description,
+                            forHTTPHeaderField: Constants.headerWSProtocolName)
+    }
     writeQueue.maxConcurrentOperationCount = 1
   }
 
-  public convenience init(url: URL) {
+  public convenience init(url: URL, webSocketProtocol: WSProtocol = .graphql_ws) {
     var request = URLRequest(url: url)
     request.timeoutInterval = 5
+    request.setValue(webSocketProtocol.description,
+                     forHTTPHeaderField: Constants.headerWSProtocolName)
+
     self.init(request: request)
   }
 
   // Used for specifically setting the QOS for the write queue.
-  public convenience init(url: URL, writeQueueQOS: QualityOfService) {
-    self.init(url: url)
+  public convenience init(
+    url: URL,
+    writeQueueQOS: QualityOfService,
+    webSocketProtocol: WSProtocol = .graphql_ws
+  ) {
+    self.init(url: url, webSocketProtocol: webSocketProtocol)
     writeQueue.qualityOfService = writeQueueQOS
   }
 

Sources/ApolloWebSocket/OperationMessage.swift

@@ -7,6 +7,7 @@ final class OperationMessage {
   enum Types : String {
     case connectionInit = "connection_init"            // Client -> Server
     case connectionTerminate = "connection_terminate"  // Client -> Server
+    case subscribe = "subscribe"                       // Client -> Server
     case start = "start"                               // Client -> Server
     case stop = "stop"                                 // Client -> Server
 
@@ -17,6 +18,7 @@ final class OperationMessage {
     case data = "data"                                 // Server -> Client
     case error = "error"                               // Server -> Client
     case complete = "complete"                         // Server -> Client
+    case next = "next"                                 // Server -> Client
   }
 
   let serializationFormat = JSONSerializationFormat.self
@@ -34,7 +36,7 @@ final class OperationMessage {
 
   init(payload: GraphQLMap? = nil,
        id: String? = nil,
-       type: Types = .start) {
+       type: Types) {
     var message: GraphQLMap = [:]
     if let payload = payload {
       message["payload"] = payload

Sources/ApolloWebSocket/WebSocketTransport.swift

@@ -145,6 +145,7 @@ public class WebSocketTransport {
 
       switch messageType {
       case .data,
+           .next,
            .error:
         if let id = parseHandler.id, let responseHandler = subscribers[id] {
           if let payload = parseHandler.payload {
@@ -185,6 +186,7 @@ public class WebSocketTransport {
 
       case .connectionInit,
            .connectionTerminate,
+           .subscribe,
            .start,
            .stop,
            .connectionError:
@@ -270,7 +272,13 @@ public class WebSocketTransport {
                                               sendQueryDocument: true,
                                               autoPersistQuery: false)
     let identifier = operationMessageIdCreator.requestId()
-    guard let message = OperationMessage(payload: body, id: identifier).rawMessage else {
+
+    var type: OperationMessage.Types = .start
+    if case WebSocket.WSProtocol.graphql_transport_ws.description = websocket.request.value(forHTTPHeaderField: WebSocket.Constants.headerWSProtocolName) {
+      type = .subscribe
+    }
+
+    guard let message = OperationMessage(payload: body, id: identifier, type: type).rawMessage else {
       return nil
     }