From 497b02002e20b5db1bfc4d2b55abd1fc571d42dc Mon Sep 17 00:00:00 2001
From: Alex Gleason <alex@alexgleason.me>
Date: Sun, 16 Mar 2025 18:02:28 -0500
Subject: [PATCH] streaming: ensure close handler is called even when socket is
 closed by server

---
 packages/ditto/controllers/api/streaming.ts | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/packages/ditto/controllers/api/streaming.ts b/packages/ditto/controllers/api/streaming.ts
index 4eebe6b3..a95cc24b 100644
--- a/packages/ditto/controllers/api/streaming.ts
+++ b/packages/ditto/controllers/api/streaming.ts
@@ -179,22 +179,31 @@ const streamingController: AppController = async (c) => {
       limiter.set(ip, count + 1, { ttl: LIMITER_WINDOW });
 
       if (count > LIMITER_LIMIT) {
-        socket.close(1008, 'Rate limit exceeded');
+        closeSocket(1008, 'Rate limit exceeded');
         return;
       }
     }
 
     if (typeof e.data !== 'string') {
-      socket.close(1003, 'Invalid message');
+      closeSocket(1003, 'Invalid message');
       return;
     }
   };
 
   socket.onclose = () => {
+    handleClose();
+  };
+
+  function closeSocket(code?: number, reason?: string) {
+    socket.close(code, reason);
+    handleClose();
+  }
+
+  function handleClose(): void {
     connections.delete(socket);
     streamingConnectionsGauge.set(connections.size);
     controller.abort();
-  };
+  }
 
   return response;
 };