From cffd2e96cc79185a70a14a4ac0740f3a0411933a Mon Sep 17 00:00:00 2001 From: dashevchenko Date: Thu, 26 Mar 2026 16:27:50 +0200 Subject: [PATCH] fixed WS limit handling for "Sessions per public user maximum number" --- .../controller/plugin/TbWebSocketHandler.java | 6 +- .../service/ws/DefaultWebSocketService.java | 6 +- .../plugin/TbWebSocketHandlerTest.java | 78 +++++++++++++++++++ 3 files changed, 84 insertions(+), 6 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 73315be73f..133584201c 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 @@ -115,7 +115,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke private final ConcurrentMap> tenantSessionsMap = new ConcurrentHashMap<>(); private final ConcurrentMap> customerSessionsMap = new ConcurrentHashMap<>(); private final ConcurrentMap> regularUserSessionsMap = new ConcurrentHashMap<>(); - private final ConcurrentMap> publicUserSessionsMap = new ConcurrentHashMap<>(); + private final ConcurrentMap> publicUserSessionsMap = new ConcurrentHashMap<>(); private Cache pendingSessions; @@ -611,7 +611,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke } if (tenantProfileConfiguration.getMaxWsSessionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) { - Set publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet()); + Set publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet()); synchronized (publicUserSessions) { limitAllowed = publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSessionsPerPublicUser(); if (limitAllowed) { @@ -655,7 +655,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke } } if (tenantProfileConfiguration.getMaxWsSessionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) { - Set publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet()); + Set publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet()); synchronized (publicUserSessions) { publicUserSessions.remove(sessionId); } diff --git a/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java b/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java index 4b1a81dd2d..09651a9264 100644 --- a/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java +++ b/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java @@ -144,7 +144,7 @@ public class DefaultWebSocketService implements WebSocketService { private final ConcurrentMap> tenantSubscriptionsMap = new ConcurrentHashMap<>(); private final ConcurrentMap> customerSubscriptionsMap = new ConcurrentHashMap<>(); private final ConcurrentMap> regularUserSubscriptionsMap = new ConcurrentHashMap<>(); - private final ConcurrentMap> publicUserSubscriptionsMap = new ConcurrentHashMap<>(); + private final ConcurrentMap> publicUserSubscriptionsMap = new ConcurrentHashMap<>(); private final ConcurrentMap> sessionCmdMap = new ConcurrentHashMap<>(); private ExecutorService executor; @@ -340,7 +340,7 @@ public class DefaultWebSocketService implements WebSocketService { } } if (tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) { - Set publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet()); + Set publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet()); synchronized (publicUserSessions) { publicUserSessions.removeIf(subId -> subId.startsWith(sessionId)); } @@ -401,7 +401,7 @@ public class DefaultWebSocketService implements WebSocketService { } } if (tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) { - Set publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet()); + Set publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet()); synchronized (publicUserSessions) { if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) { publicUserSessions.add(subId); diff --git a/application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java b/application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java index 053cb6808f..4dc637725f 100644 --- a/application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java +++ b/application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java @@ -25,16 +25,28 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import org.springframework.test.util.ReflectionTestUtils; import org.springframework.web.socket.CloseStatus; +import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.adapter.NativeWebSocketSession; import org.thingsboard.common.util.ThingsBoardThreadFactory; +import org.thingsboard.server.common.data.TenantProfile; +import org.thingsboard.server.common.data.id.CustomerId; +import org.thingsboard.server.common.data.id.EntityId; +import org.thingsboard.server.common.data.id.TenantId; +import org.thingsboard.server.common.data.id.UserId; +import org.thingsboard.server.dao.tenant.TbTenantProfileCache; +import org.thingsboard.server.service.security.model.SecurityUser; +import org.thingsboard.server.service.security.model.UserPrincipal; import org.thingsboard.server.service.ws.WebSocketSessionRef; import java.io.IOException; +import java.lang.reflect.Method; import java.util.Collection; import java.util.Deque; import java.util.List; import java.util.Random; +import java.util.UUID; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; @@ -184,4 +196,70 @@ class TbWebSocketHandlerTest { assertThat(msgs).map(Integer::parseInt).doesNotHaveDuplicates().hasSize(100); } + // Regression test for the bug where publicUserSessionsMap was keyed by UserId(NULL_UUID), + // making maxWsSessionsPerPublicUser a global limit shared across all tenants. + // The limit is now scoped per-tenant. + @Test + void checkLimits_publicUserSessions_limitIsPerTenantNotGlobal() throws Exception { + TbTenantProfileCache tenantProfileCache = mock(TbTenantProfileCache.class); + ReflectionTestUtils.setField(wsHandler, "tenantProfileCache", tenantProfileCache); + + int maxPublicSessions = 2; + + TenantId tenant1 = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile1 = new TenantProfile(); + profile1.createDefaultTenantProfileData(); + profile1.getDefaultProfileConfiguration().setMaxWsSessionsPerPublicUser(maxPublicSessions); + willReturn(profile1).given(tenantProfileCache).get(tenant1); + + TenantId tenant2 = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile2 = new TenantProfile(); + profile2.createDefaultTenantProfileData(); + profile2.getDefaultProfileConfiguration().setMaxWsSessionsPerPublicUser(maxPublicSessions); + willReturn(profile2).given(tenantProfileCache).get(tenant2); + + Method checkLimits = TbWebSocketHandler.class.getDeclaredMethod( + "checkLimits", WebSocketSession.class, WebSocketSessionRef.class); + checkLimits.setAccessible(true); + + // tenant1 fills up its limit + for (int i = 0; i < maxPublicSessions; i++) { + assertThat((boolean) checkLimits.invoke(wsHandler, mockWsSession("t1-" + i), mockPublicSessionRef(tenant1))).isTrue(); + } + + // tenant2 must get its own independent quota — this was the bug: with NULL_UUID as key + // all tenants shared one global counter, so tenant2 would be blocked here + for (int i = 0; i < maxPublicSessions; i++) { + assertThat((boolean) checkLimits.invoke(wsHandler, mockWsSession("t2-" + i), mockPublicSessionRef(tenant2))) + .as("tenant2 session %d should not be affected by tenant1's sessions", i + 1) + .isTrue(); + } + + // tenant1's (maxPublicSessions + 1)-th session must be rejected + NativeWebSocketSession overLimit = mockWsSession("t1-over"); + assertThat((boolean) checkLimits.invoke(wsHandler, overLimit, mockPublicSessionRef(tenant1))).isFalse(); + verify(overLimit).close(CloseStatus.POLICY_VIOLATION.withReason("Max public user sessions limit reached")); + } + + private NativeWebSocketSession mockWsSession(String id) { + NativeWebSocketSession s = mock(NativeWebSocketSession.class); + willReturn(id).given(s).getId(); + return s; + } + + private WebSocketSessionRef mockPublicSessionRef(TenantId tenantId) { + CustomerId customerId = new CustomerId(UUID.randomUUID()); + SecurityUser securityUser = mock(SecurityUser.class); + willReturn(tenantId).given(securityUser).getTenantId(); + willReturn(customerId).given(securityUser).getCustomerId(); + willReturn(new UserId(EntityId.NULL_UUID)).given(securityUser).getId(); + willReturn(true).given(securityUser).isCustomerUser(); + willReturn(new UserPrincipal(UserPrincipal.Type.PUBLIC_ID, customerId.toString())).given(securityUser).getUserPrincipal(); + + WebSocketSessionRef ref = mock(WebSocketSessionRef.class); + willReturn(securityUser).given(ref).getSecurityCtx(); + willReturn(UUID.randomUUID().toString()).given(ref).getSessionId(); + return ref; + } + }