tsai
/
thingsboard
mirror of https://github.com/thingsboard/thingsboard.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

fixed WS limit handling for "Sessions per public user maximum number"

pull/15313/head
dashevchenko 2 months ago
parent
71233b4ab5
commit
cffd2e96cc
3 changed files with 84 additions and 6 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 6
      application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
  2. 6
      application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java
  3. 78
      application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java

6
application/src/main/java/org/thingsboard/server/controller/plugin/TbWebSocketHandler.java
View File

@ -115,7 +115,7 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements WebSocke
private final ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<TenantId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
private Cache<String, SessionMetaData> 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<String> publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
Set<String> 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<String> publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
Set<String> publicUserSessions = publicUserSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet());
synchronized (publicUserSessions) {
publicUserSessions.remove(sessionId);
}

6
application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java
View File

@ -144,7 +144,7 @@ public class DefaultWebSocketService implements WebSocketService {
private final ConcurrentMap<TenantId, Set<String>> tenantSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<CustomerId, Set<String>> customerSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> regularUserSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<UserId, Set<String>> publicUserSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = new ConcurrentHashMap<>();
private final ConcurrentMap<String, Map<Integer, Integer>> 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<String> publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
Set<String> 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<String> publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getId(), id -> ConcurrentHashMap.newKeySet());
Set<String> publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet());
synchronized (publicUserSessions) {
if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) {
publicUserSessions.add(subId);

78
application/src/test/java/org/thingsboard/server/controller/plugin/TbWebSocketHandlerTest.java
View File

@ -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;
}
}

Write Preview
Loading…
Cancel
Save