From cdce4f9540a8447d7669fff59ed3c593840c48dd Mon Sep 17 00:00:00 2001
From: Calvin Cestari <calvincestari@users.noreply.github.com>
Date: Wed, 16 Oct 2024 11:19:41 -0700
Subject: [PATCH] fix: Websocket error broadcast for unsubscribed ID (#506)

---
 .../WebSocket/WebSocketTests.swift            |  8 +++----
 .../ApolloWebSocket/WebSocketTransport.swift  | 24 ++++++++++++-------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/Tests/ApolloTests/WebSocket/WebSocketTests.swift b/Tests/ApolloTests/WebSocket/WebSocketTests.swift
index 67f2f2a8f..13c45f32b 100644
--- 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
index 9405ae153..3cbcdef73 100644
--- 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 {