From 9d8aafbedece7483f397cb3a43f30f95892f840d 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
 (apollographql/apollo-ios-dev#506)

---
 .../ApolloWebSocket/WebSocketTransport.swift  | 24 ++++++++++++-------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/Sources/ApolloWebSocket/WebSocketTransport.swift b/Sources/ApolloWebSocket/WebSocketTransport.swift
index 9405ae153b..3cbcdef733 100644
--- a/Sources/ApolloWebSocket/WebSocketTransport.swift
+++ b/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 {