From 0bff83995191f893dca47c04a192f3cfe7723bb2 Mon Sep 17 00:00:00 2001
From: Sergey Matvienko <smatvienko@thingsboard.io>
Date: Wed, 12 Jul 2023 13:04:02 +0200
Subject: [PATCH] TbWebSocketHandler.checkLimits refactored to extract IO
 session.close() call from synchronized scope

---
 .../controller/plugin/TbWebSocketHandler.java | 34 ++++++++++++-------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java b/application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
index 9cb343dbe3..56d88143a5 100644
--- a/application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
+++ b/application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
@@ -397,19 +397,21 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke
         if (tenantProfileConfiguration == null) {
             return true;
         }
-
+        boolean limitAllowed;
         String sessionId = session.getId();
         if (tenantProfileConfiguration.getMaxWsSessionsPerTenant() > 0) {
             Set<String> tenantSessions = tenantSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet());
             synchronized (tenantSessions) {
-                if (tenantSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerTenant()) {
+                limitAllowed = tenantSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerTenant();
+                if (limitAllowed) {
                     tenantSessions.add(sessionId);
-                } else {
+                }
+            }
+            if (!limitAllowed) {
                     log.info("[{}][{}][{}] Failed to start session. Max tenant sessions limit reached"
                             , sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), sessionId);
                     session.close(CloseStatus.POLICY_VIOLATION.withReason("Max tenant sessions limit reached!"));
                     return false;
-                }
             }
         }
 
@@ -417,42 +419,48 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke
             if (tenantProfileConfiguration.getMaxWsSessionsPerCustomer() > 0) {
                 Set<String> customerSessions = customerSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getCustomerId(), id -> ConcurrentHashMap.newKeySet());
                 synchronized (customerSessions) {
-                    if (customerSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerCustomer()) {
+                    limitAllowed = customerSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerCustomer();
+                    if (limitAllowed) {
                         customerSessions.add(sessionId);
-                    } else {
+                    }
+                }
+                if (!limitAllowed) {
                         log.info("[{}][{}][{}] Failed to start session. Max customer sessions limit reached"
                                 , sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), sessionId);
                         session.close(CloseStatus.POLICY_VIOLATION.withReason("Max customer sessions limit reached"));
                         return false;
-                    }
                 }
             }
             if (tenantProfileConfiguration.getMaxWsSessionsPerRegularUser() > 0
                     && UserPrincipal.Type.USER_NAME.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) {
                 Set<String> regularUserSessions = regularUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
                 synchronized (regularUserSessions) {
-                    if (regularUserSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerRegularUser()) {
+                    limitAllowed = regularUserSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerRegularUser();
+                    if (limitAllowed) {
                         regularUserSessions.add(sessionId);
-                    } else {
+                    }
+                }
+                if (!limitAllowed) {
                         log.info("[{}][{}][{}] Failed to start session. Max regular user sessions limit reached"
                                 , sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), sessionId);
                         session.close(CloseStatus.POLICY_VIOLATION.withReason("Max regular user sessions limit reached"));
                         return false;
-                    }
                 }
             }
             if (tenantProfileConfiguration.getMaxWsSessionsPerPublicUser() > 0
                     && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) {
                 Set<String> publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
                 synchronized (publicUserSessions) {
-                    if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerPublicUser()) {
+                    limitAllowed = publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerPublicUser();
+                    if (limitAllowed) {
                         publicUserSessions.add(sessionId);
-                    } else {
+                    }
+                }
+                if (!limitAllowed) {
                         log.info("[{}][{}][{}] Failed to start session. Max public user sessions limit reached"
                                 , sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), sessionId);
                         session.close(CloseStatus.POLICY_VIOLATION.withReason("Max public user sessions limit reached"));
                         return false;
-                    }
                 }
             }
         }