From 395b6da51bf4597a28fbfb9cecfd27fb90947e7a Mon Sep 17 00:00:00 2001
From: Andrey Sobolev <haiodo@users.noreply.github.com>
Date: Mon, 10 Jun 2024 13:18:57 +0700
Subject: [PATCH] Make a sessionId a really uniq and remove checks (#5762)

Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
---
 plugins/client-resources/src/connection.ts | 42 +++++++++++-----------
 server/rpc/src/rpc.ts                      |  1 -
 server/ws/src/server.ts                    | 14 +++-----
 3 files changed, 24 insertions(+), 33 deletions(-)

diff --git a/plugins/client-resources/src/connection.ts b/plugins/client-resources/src/connection.ts
index e72ba14b4e..47afd656ba 100644
--- a/plugins/client-resources/src/connection.ts
+++ b/plugins/client-resources/src/connection.ts
@@ -87,7 +87,7 @@ class Connection implements ClientConnection {
 
   private openAction: any
 
-  private sessionId: string | undefined
+  private readonly sessionId: string | undefined
   private closed = false
 
   private upgrading: boolean = false
@@ -103,6 +103,25 @@ class Connection implements ClientConnection {
     private readonly onUnauthorized?: () => void,
     readonly onConnect?: (event: ClientConnectEvent, data?: any) => Promise<void>
   ) {
+    if (typeof sessionStorage !== 'undefined') {
+      // Find local session id in session storage only if user refresh a page.
+      const sKey = 'session.id.' + this.url
+      let sessionId = sessionStorage.getItem(sKey) ?? undefined
+      if (sessionId === undefined) {
+        sessionId = generateId()
+        console.log('Generate new SessionId', sessionId)
+        this.sessionId = sessionId
+      } else {
+        this.sessionId = sessionId
+        sessionStorage.removeItem(sKey)
+      }
+      window.addEventListener('beforeunload', () => {
+        sessionStorage.setItem(sKey, sessionId as string)
+      })
+    } else {
+      this.sessionId = generateId()
+    }
+
     this.scheduleOpen(false)
   }
 
@@ -230,16 +249,6 @@ class Connection implements ClientConnection {
         }
 
         this.upgrading = false
-        if ((resp as HelloResponse).alreadyConnected === true) {
-          this.sessionId = generateId()
-          if (typeof sessionStorage !== 'undefined') {
-            sessionStorage.setItem('session.id.' + this.url, this.sessionId)
-          }
-          console.log('Connection: alreadyConnected, reconnect with new Id')
-          clearTimeout(dialTimeout)
-          this.scheduleOpen(true)
-          return
-        }
         if ((resp as HelloResponse).binary) {
           this.binaryMode = true
         }
@@ -378,17 +387,6 @@ class Connection implements ClientConnection {
         return s as ClientSocket
       })
 
-    if (this.sessionId === undefined) {
-      // Find local session id in session storage.
-      this.sessionId =
-        typeof sessionStorage !== 'undefined'
-          ? sessionStorage.getItem('session.id.' + this.url) ?? undefined
-          : undefined
-      this.sessionId = this.sessionId ?? generateId()
-      if (typeof sessionStorage !== 'undefined') {
-        sessionStorage.setItem('session.id.' + this.url, this.sessionId)
-      }
-    }
     if (socketId !== this.sockets) {
       return
     }
diff --git a/server/rpc/src/rpc.ts b/server/rpc/src/rpc.ts
index 349b123af5..6c2201b76f 100644
--- a/server/rpc/src/rpc.ts
+++ b/server/rpc/src/rpc.ts
@@ -47,7 +47,6 @@ export interface HelloRequest extends Request<any[]> {
 export interface HelloResponse extends Response<any> {
   binary: boolean
   reconnect?: boolean
-  alreadyConnected?: boolean
 }
 
 /**
diff --git a/server/ws/src/server.ts b/server/ws/src/server.ts
index 46fe259312..13fb910240 100644
--- a/server/ws/src/server.ts
+++ b/server/ws/src/server.ts
@@ -291,16 +291,10 @@ class TSessionManager implements SessionManager {
       await workspace?.closing
     }
     workspace = this.workspaces.get(wsString)
-    if (sessionId !== undefined && workspace?.sessions?.has(sessionId) === true) {
-      const helloResponse: HelloResponse = {
-        id: -1,
-        result: 'hello',
-        binary: false,
-        reconnect: false,
-        alreadyConnected: true
-      }
-      await ws.send(ctx, helloResponse, false, false)
-      return { error: new Error('Session already exists') }
+    const oldSession = sessionId !== undefined ? workspace?.sessions?.get(sessionId) : undefined
+    if (oldSession !== undefined) {
+      // Just close old socket for old session id.
+      await this.close(ctx, oldSession.socket, wsString)
     }
     const workspaceName = workspaceInfo.workspaceName ?? workspaceInfo.workspaceUrl ?? workspaceInfo.workspaceId