From 3e82da1a2638842b6cd102a84877a1358e5e7e83 Mon Sep 17 00:00:00 2001
From: Alexander <zavorotynskiy@pm.me>
Date: Fri, 21 Mar 2025 11:33:58 +0100
Subject: [PATCH] feat(assist): tried to add a custom header to all responses

---
 ee/assist/utils/wsServer.js | 40 ++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/ee/assist/utils/wsServer.js b/ee/assist/utils/wsServer.js
index 69b6b97f4..f0ea61c8b 100644
--- a/ee/assist/utils/wsServer.js
+++ b/ee/assist/utils/wsServer.js
@@ -84,30 +84,30 @@ const createSocketIOServer = function (server, prefix) {
     if (io) {
         return io;
     }
+
+    // Common options for both initialization methods
+    const options = {
+        maxHttpBufferSize: (parseFloat(process.env.maxHttpBufferSize) || 5) * 1e6,
+        cors: {
+            origin: "*",
+            methods: ["GET", "POST", "PUT"],
+            credentials: true
+        },
+        path: (prefix ? prefix : '') + '/socket',
+        allowRequest: (req, callback) => {
+            req.res.setHeader('x-host-id', process.env.HOSTNAME || 'unknown');
+            callback(null, true);
+        },
+        ...getCompressionConfig()
+    };
+
     if (process.env.uws !== "true") {
-        io = _io(server, {
-            maxHttpBufferSize: (parseFloat(process.env.maxHttpBufferSize) || 5) * 1e6,
-            cors: {
-                origin: "*",
-                methods: ["GET", "POST", "PUT"],
-                credentials: true
-            },
-            path: (prefix ? prefix : '') + '/socket',
-            ...getCompressionConfig()
-        });
+        io = _io(server, options);
     } else {
-        io = new _io.Server({
-            maxHttpBufferSize: (parseFloat(process.env.maxHttpBufferSize) || 5) * 1e6,
-            cors: {
-                origin: "*",
-                methods: ["GET", "POST", "PUT"],
-                credentials: true
-            },
-            path: (prefix ? prefix : '') + '/socket',
-            ...getCompressionConfig()
-        });
+        io = new _io.Server(options);
         io.attachApp(server);
     }
+
     if (useRedis) {
         startCacheRefresher();
     }