fix: Websocket error broadcast for unsubscribed ID (#506)

apollographql · Oct 16, 2024 · cdce4f9 · cdce4f9
1 parent f7d83fc
commit cdce4f9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 13 deletions.
diff --git a/Tests/ApolloTests/WebSocket/WebSocketTests.swift b/Tests/ApolloTests/WebSocket/WebSocketTests.swift
@@ -110,8 +110,8 @@ class WebSocketTests: XCTestCase {
     subject.cancel()
   }
 
-  func testLocalErrorUnknownId() throws {
-    let expectation = self.expectation(description: "Unknown id for subscription")
+  func testLocalErrorMissingId() throws {
+    let expectation = self.expectation(description: "Missing id for subscription")
 
     let subject = client.subscribe(subscription: MockSubscription<ReviewAddedData>()) { result in
       defer { expectation.fulfill() }
@@ -133,10 +133,10 @@ class WebSocketTests: XCTestCase {
         }
       }
     }
-
+
+    // Message data below has missing 'id' and should notify all subscribers of the error
     let message : JSONEncodableDictionary = [
       "type": "data",
-      "id": "2",            // subscribing on id = 1, i.e. expecting error when receiving id = 2
       "payload": [
         "data": [
           "reviewAdded": [

diff --git a/apollo-ios/Sources/ApolloWebSocket/WebSocketTransport.swift b/apollo-ios/Sources/ApolloWebSocket/WebSocketTransport.swift
@@ -193,10 +193,21 @@ public class WebSocketTransport {
       }
 
       switch messageType {
-      case .data,
-           .next,
-           .error:
-        if let id = parseHandler.id, let responseHandler = subscribers[id] {
+      case .data, .next, .error:
+        guard let id = parseHandler.id else {
+          let websocketError = WebSocketError(
+            payload: parseHandler.payload,
+            error: parseHandler.error,
+            kind: .unprocessedMessage(text)
+          )
+          self.notifyErrorAllHandlers(websocketError)
+
+          break
+        }
+
+        // If we have a handler ID but no subscriber exists for that ID then the
+        // subscriber probably unsubscribed.
+        if let responseHandler = subscribers[id] {
           if let payload = parseHandler.payload {
             responseHandler(.success(payload))
           } else if let error = parseHandler.error {
@@ -207,11 +218,6 @@ public class WebSocketTransport {
                                                 kind: .neitherErrorNorPayloadReceived)
             responseHandler(.failure(websocketError))
           }
-        } else {
-          let websocketError = WebSocketError(payload: parseHandler.payload,
-                                              error: parseHandler.error,
-                                              kind: .unprocessedMessage(text))
-          self.notifyErrorAllHandlers(websocketError)
         }
       case .complete:
         if let id = parseHandler.id {