From d3ef4ea4488c09be8874d78a1a2bf668077213cf Mon Sep 17 00:00:00 2001
From: yassinedorbozgithub <yassine.dorboz.2022@gmail.com>
Date: Thu, 10 Oct 2024 19:20:49 +0100
Subject: [PATCH] fix(api): secure web-socket access

---
 api/src/channel/channel.service.ts     |  3 ++
 api/src/websocket/websocket.gateway.ts | 58 ++++++++++++++++----------
 2 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/api/src/channel/channel.service.ts b/api/src/channel/channel.service.ts
index 2147e514..af278965 100644
--- a/api/src/channel/channel.service.ts
+++ b/api/src/channel/channel.service.ts
@@ -128,6 +128,9 @@ export class ChannelService {
     );
 
     if (!req.session?.passport?.user?.id) {
+      setTimeout(() => {
+        req.socket.client.conn.close();
+      }, 300);
       throw new UnauthorizedException(
         'Only authenticated users are allowed to use this channel',
       );
diff --git a/api/src/websocket/websocket.gateway.ts b/api/src/websocket/websocket.gateway.ts
index 4cea057b..cb903177 100644
--- a/api/src/websocket/websocket.gateway.ts
+++ b/api/src/websocket/websocket.gateway.ts
@@ -207,31 +207,45 @@ export class WebsocketGateway
     // Handle session
     this.io.use((client, next) => {
       this.logger.verbose('Client connected, attempting to load session.');
-      if (client.request.headers.cookie) {
-        const cookies = cookie.parse(client.request.headers.cookie);
-        if (cookies && config.session.name in cookies) {
-          const sessionID = cookieParser.signedCookie(
-            cookies[config.session.name],
-            config.session.secret,
-          );
-          if (sessionID) {
-            return this.loadSession(sessionID, (err, session) => {
-              if (err) {
-                this.logger.warn(
-                  'Unable to load session, creating a new one ...',
-                  err,
-                );
-                return this.createAndStoreSession(client, next);
-              }
-              client.data.session = session;
-              client.data.sessionID = sessionID;
-              next();
-            });
+      try {
+        const { searchParams } = new URL(`ws://localhost${client.request.url}`);
+        if (client.request.headers.cookie) {
+          const cookies = cookie.parse(client.request.headers.cookie);
+          if (cookies && config.session.name in cookies) {
+            const sessionID = cookieParser.signedCookie(
+              cookies[config.session.name],
+              config.session.secret,
+            );
+            if (sessionID) {
+              return this.loadSession(sessionID, (err, session) => {
+                if (err || !session) {
+                  this.logger.warn(
+                    'Unable to load session, creating a new one ...',
+                    err,
+                  );
+                  if (searchParams.get('channel') === 'offline') {
+                    return this.createAndStoreSession(client, next);
+                  } else {
+                    return next(new Error('Unauthorized: Unknown session ID'));
+                  }
+                }
+                client.data.session = session;
+                client.data.sessionID = sessionID;
+                next();
+              });
+            } else {
+              return next(new Error('Unable to parse session ID from cookie'));
+            }
           }
+        } else if (searchParams.get('channel') === 'offline') {
+          return this.createAndStoreSession(client, next);
+        } else {
+          return next(new Error('Unauthorized to connect to WS'));
         }
+      } catch (e) {
+        this.logger.warn('Something unexpected happening');
+        return next(e);
       }
-
-      return this.createAndStoreSession(client, next);
     });
   }