From e17e9e1dbecbac01ebac8b48b8cb07635bb1e5b1 Mon Sep 17 00:00:00 2001
From: Corey <40738690+C1200@users.noreply.github.com>
Date: Mon, 2 May 2022 17:39:42 +0100
Subject: [PATCH] fix: properly track updates to socket.data

Before this commit, attributes added to `socket.data` within a
middleware were lost.

Related: https://github.com/socketio/socket.io-admin-ui/commit/3773fe4b1cbf2206708e1f21ce65f430a522527f
---
 lib/index.ts  |  6 ++++++
 test/index.ts | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/lib/index.ts b/lib/index.ts
index eea78c3..a81ddb8 100644
--- a/lib/index.ts
+++ b/lib/index.ts
@@ -321,6 +321,8 @@ const registerListeners = (
       });
     };
 
+    const data = socket.data || {}; // could be set in a middleware
+
     socket.data = createProxy({
       _admin: {
         clientId: clientId.substring(0, 12), // this information is quite sensitive
@@ -328,6 +330,10 @@ const registerListeners = (
       },
     });
 
+    for (const key in data) {
+      socket.data[key] = createProxy(data[key]);
+    }
+
     adminNamespace.emit("socket_connected", serialize(socket, nsp.name));
 
     socket.conn.on("upgrade", (transport: any) => {
diff --git a/test/index.ts b/test/index.ts
index 8cff224..324387a 100644
--- a/test/index.ts
+++ b/test/index.ts
@@ -300,6 +300,46 @@ describe("Socket.IO Admin (server instrumentation)", () => {
         adminSocket.disconnect();
       });
 
+      it("emits event when socket.data is updated", async () => {
+        instrument(io, {
+          auth: false,
+        });
+
+        const adminSocket = ioc(`http://localhost:${port}/admin`);
+
+        await waitFor(adminSocket, "connect");
+
+        const clientSocket = ioc(`http://localhost:${port}`, {
+          forceNew: true,
+          transports: ["polling"],
+        });
+
+        io.use((socket, next) => {
+          socket.data = socket.data || {};
+          socket.data.count = 1;
+          socket.data.array = [1];
+          next();
+        });
+
+        const serverSocket = await waitFor(io, "connection");
+
+        const socket = await waitFor(adminSocket, "socket_connected");
+        expect(socket.data).to.eql({ count: 1, array: [1] });
+
+        serverSocket.data.count++;
+
+        const updatedSocket1 = await waitFor(adminSocket, "socket_updated");
+        expect(updatedSocket1.data).to.eql({ count: 2, array: [1] });
+
+        serverSocket.data.array.push(2);
+
+        const updatedSocket2 = await waitFor(adminSocket, "socket_updated");
+        expect(updatedSocket2.data).to.eql({ count: 2, array: [1, 2] });
+
+        adminSocket.disconnect();
+        clientSocket.disconnect();
+      });
+
       it("performs administrative tasks", async () => {
         instrument(io, {
           auth: false,