Viacheslav Klimov
4 months ago
committed by
GitHub
GitHub
4 changed files with 366 additions and 9 deletions
@ -0,0 +1,277 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2026 The Thingsboard Authors |
||||
|
* |
||||
|
* Licensed under the Apache License, Version 2.0 (the "License"); |
||||
|
* you may not use this file except in compliance with the License. |
||||
|
* You may obtain a copy of the License at |
||||
|
* |
||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
* |
||||
|
* Unless required by applicable law or agreed to in writing, software |
||||
|
* distributed under the License is distributed on an "AS IS" BASIS, |
||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
|
* See the License for the specific language governing permissions and |
||||
|
* limitations under the License. |
||||
|
*/ |
||||
|
package org.thingsboard.server.service.ws; |
||||
|
|
||||
|
import org.junit.jupiter.api.BeforeEach; |
||||
|
import org.junit.jupiter.api.Test; |
||||
|
import org.springframework.test.util.ReflectionTestUtils; |
||||
|
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.attributes.AttributesService; |
||||
|
import org.thingsboard.server.dao.tenant.TbTenantProfileCache; |
||||
|
import org.thingsboard.server.dao.timeseries.TimeseriesService; |
||||
|
import org.thingsboard.server.queue.discovery.TbServiceInfoProvider; |
||||
|
import org.thingsboard.server.service.security.AccessValidator; |
||||
|
import org.thingsboard.server.service.security.model.SecurityUser; |
||||
|
import org.thingsboard.server.service.security.model.UserPrincipal; |
||||
|
import org.thingsboard.server.service.subscription.TbEntityDataSubscriptionService; |
||||
|
import org.thingsboard.server.service.subscription.TbLocalSubscriptionService; |
||||
|
import org.thingsboard.server.service.ws.notification.NotificationCommandsHandler; |
||||
|
import org.thingsboard.server.service.ws.telemetry.cmd.v1.AttributesSubscriptionCmd; |
||||
|
|
||||
|
import java.util.Set; |
||||
|
import java.util.UUID; |
||||
|
import java.util.concurrent.ConcurrentMap; |
||||
|
|
||||
|
import static org.assertj.core.api.Assertions.assertThat; |
||||
|
import static org.mockito.BDDMockito.willReturn; |
||||
|
import static org.mockito.Mockito.mock; |
||||
|
|
||||
|
class DefaultWebSocketServiceTest { |
||||
|
|
||||
|
DefaultWebSocketService service; |
||||
|
TbTenantProfileCache tenantProfileCache; |
||||
|
WebSocketMsgEndpoint msgEndpoint; |
||||
|
|
||||
|
@BeforeEach |
||||
|
void setUp() { |
||||
|
tenantProfileCache = mock(TbTenantProfileCache.class); |
||||
|
msgEndpoint = mock(WebSocketMsgEndpoint.class); |
||||
|
|
||||
|
service = new DefaultWebSocketService( |
||||
|
mock(TbLocalSubscriptionService.class), |
||||
|
mock(TbEntityDataSubscriptionService.class), |
||||
|
mock(NotificationCommandsHandler.class), |
||||
|
msgEndpoint, |
||||
|
mock(AccessValidator.class), |
||||
|
mock(AttributesService.class), |
||||
|
mock(TimeseriesService.class), |
||||
|
mock(TbServiceInfoProvider.class), |
||||
|
tenantProfileCache |
||||
|
); |
||||
|
} |
||||
|
|
||||
|
// Regression test: publicUserSubscriptionsMap must be keyed by TenantId, not UserId(NULL_UUID).
|
||||
|
// With the old UserId(NULL_UUID) key, all tenants shared one global subscription counter.
|
||||
|
@Test |
||||
|
void processSubscription_publicUserSubscriptionsMap_isPerTenantNotGlobal() throws Exception { |
||||
|
int maxPublicSubscriptions = 2; |
||||
|
|
||||
|
TenantId tenant1 = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile1 = new TenantProfile(); |
||||
|
profile1.createDefaultTenantProfileData(); |
||||
|
profile1.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile1).given(tenantProfileCache).get(tenant1); |
||||
|
|
||||
|
TenantId tenant2 = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile2 = new TenantProfile(); |
||||
|
profile2.createDefaultTenantProfileData(); |
||||
|
profile2.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile2).given(tenantProfileCache).get(tenant2); |
||||
|
|
||||
|
// tenant1 fills up its quota
|
||||
|
for (int i = 0; i < maxPublicSubscriptions; i++) { |
||||
|
assertThat(service.processSubscription(mockPublicSessionRef(tenant1, "t1-session-" + i), subscriptionCmd(i))) |
||||
|
.as("tenant1 subscription %d should be accepted", i + 1) |
||||
|
.isTrue(); |
||||
|
} |
||||
|
|
||||
|
// tenant2 must have its own independent quota — this was the bug:
|
||||
|
// with UserId(NULL_UUID) as key all tenants shared one counter, so tenant2 would be blocked here
|
||||
|
for (int i = 0; i < maxPublicSubscriptions; i++) { |
||||
|
assertThat(service.processSubscription(mockPublicSessionRef(tenant2, "t2-session-" + i), subscriptionCmd(i))) |
||||
|
.as("tenant2 subscription %d should not be affected by tenant1's subscriptions", i + 1) |
||||
|
.isTrue(); |
||||
|
} |
||||
|
|
||||
|
// tenant1's (maxPublicSubscriptions + 1)-th subscription must be rejected
|
||||
|
assertThat(service.processSubscription(mockPublicSessionRef(tenant1, "t1-session-over"), subscriptionCmd(99))) |
||||
|
.as("tenant1 should be rejected after exceeding its limit") |
||||
|
.isFalse(); |
||||
|
|
||||
|
// Verify that publicUserSubscriptionsMap has separate entries per tenant
|
||||
|
@SuppressWarnings("unchecked") |
||||
|
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = |
||||
|
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); |
||||
|
|
||||
|
assertThat(publicUserSubscriptionsMap).as("map should contain tenant1").containsKey(tenant1); |
||||
|
assertThat(publicUserSubscriptionsMap).as("map should contain tenant2").containsKey(tenant2); |
||||
|
assertThat(publicUserSubscriptionsMap).as("map must not have a single NULL_UUID entry for all tenants") |
||||
|
.doesNotContainKey(new TenantId(EntityId.NULL_UUID)); |
||||
|
|
||||
|
assertThat(publicUserSubscriptionsMap.get(tenant1)) |
||||
|
.as("tenant1 should have exactly %d subscriptions", maxPublicSubscriptions) |
||||
|
.hasSize(maxPublicSubscriptions); |
||||
|
assertThat(publicUserSubscriptionsMap.get(tenant2)) |
||||
|
.as("tenant2 should have exactly %d subscriptions", maxPublicSubscriptions) |
||||
|
.hasSize(maxPublicSubscriptions); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
void processSubscription_publicUserSubscriptionsMap_subscriptionIdFormat() { |
||||
|
int maxPublicSubscriptions = 5; |
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile = new TenantProfile(); |
||||
|
profile.createDefaultTenantProfileData(); |
||||
|
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile).given(tenantProfileCache).get(tenantId); |
||||
|
|
||||
|
String sessionId = "my-session-id"; |
||||
|
int cmdId = 42; |
||||
|
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId); |
||||
|
service.processSubscription(sessionRef, subscriptionCmd(cmdId)); |
||||
|
|
||||
|
@SuppressWarnings("unchecked") |
||||
|
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = |
||||
|
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); |
||||
|
|
||||
|
Set<String> subs = publicUserSubscriptionsMap.get(tenantId); |
||||
|
assertThat(subs).hasSize(1); |
||||
|
assertThat(subs.iterator().next()).isEqualTo("[" + sessionId + "]:[" + cmdId + "]"); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
void processSubscription_unsubscribe_removesEntryFromPublicUserSubscriptionsMap() { |
||||
|
int maxPublicSubscriptions = 5; |
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile = new TenantProfile(); |
||||
|
profile.createDefaultTenantProfileData(); |
||||
|
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile).given(tenantProfileCache).get(tenantId); |
||||
|
|
||||
|
String sessionId = "session-1"; |
||||
|
int cmdId = 1; |
||||
|
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId); |
||||
|
|
||||
|
service.processSubscription(sessionRef, subscriptionCmd(cmdId)); |
||||
|
|
||||
|
@SuppressWarnings("unchecked") |
||||
|
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = |
||||
|
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); |
||||
|
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(1); |
||||
|
|
||||
|
AttributesSubscriptionCmd unsubCmd = subscriptionCmd(cmdId); |
||||
|
unsubCmd.setUnsubscribe(true); |
||||
|
service.processSubscription(sessionRef, unsubCmd); |
||||
|
|
||||
|
assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
void processSubscription_unsubscribe_freesSlotForNewSubscription() { |
||||
|
int maxPublicSubscriptions = 1; |
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile = new TenantProfile(); |
||||
|
profile.createDefaultTenantProfileData(); |
||||
|
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile).given(tenantProfileCache).get(tenantId); |
||||
|
|
||||
|
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, "session-1"); |
||||
|
service.processSubscription(sessionRef, subscriptionCmd(1)); |
||||
|
|
||||
|
// slot is full — second subscription on same session should be rejected
|
||||
|
assertThat(service.processSubscription(sessionRef, subscriptionCmd(2))).isFalse(); |
||||
|
|
||||
|
// unsubscribe cmd 1 to free the slot
|
||||
|
AttributesSubscriptionCmd unsubCmd = subscriptionCmd(1); |
||||
|
unsubCmd.setUnsubscribe(true); |
||||
|
service.processSubscription(sessionRef, unsubCmd); |
||||
|
|
||||
|
// now a new subscription should succeed
|
||||
|
assertThat(service.processSubscription(sessionRef, subscriptionCmd(3))) |
||||
|
.as("new subscription should succeed after unsubscribe freed the slot") |
||||
|
.isTrue(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
void processSessionClose_removesAllSessionSubscriptionsFromPublicUserSubscriptionsMap() { |
||||
|
int maxPublicSubscriptions = 10; |
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile = new TenantProfile(); |
||||
|
profile.createDefaultTenantProfileData(); |
||||
|
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile).given(tenantProfileCache).get(tenantId); |
||||
|
|
||||
|
String sessionId = "closing-session"; |
||||
|
WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId); |
||||
|
|
||||
|
service.processSubscription(sessionRef, subscriptionCmd(1)); |
||||
|
service.processSubscription(sessionRef, subscriptionCmd(2)); |
||||
|
service.processSubscription(sessionRef, subscriptionCmd(3)); |
||||
|
|
||||
|
@SuppressWarnings("unchecked") |
||||
|
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = |
||||
|
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); |
||||
|
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3); |
||||
|
|
||||
|
service.processSessionClose(sessionRef); |
||||
|
|
||||
|
assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
void processSessionClose_onlyRemovesClosedSessionSubscriptions() { |
||||
|
int maxPublicSubscriptions = 10; |
||||
|
TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); |
||||
|
TenantProfile profile = new TenantProfile(); |
||||
|
profile.createDefaultTenantProfileData(); |
||||
|
profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); |
||||
|
willReturn(profile).given(tenantProfileCache).get(tenantId); |
||||
|
|
||||
|
WebSocketSessionRef session1 = mockPublicSessionRef(tenantId, "session-1"); |
||||
|
WebSocketSessionRef session2 = mockPublicSessionRef(tenantId, "session-2"); |
||||
|
|
||||
|
service.processSubscription(session1, subscriptionCmd(1)); |
||||
|
service.processSubscription(session1, subscriptionCmd(2)); |
||||
|
service.processSubscription(session2, subscriptionCmd(1)); |
||||
|
|
||||
|
@SuppressWarnings("unchecked") |
||||
|
ConcurrentMap<TenantId, Set<String>> publicUserSubscriptionsMap = |
||||
|
(ConcurrentMap<TenantId, Set<String>>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); |
||||
|
assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3); |
||||
|
|
||||
|
service.processSessionClose(session1); |
||||
|
|
||||
|
Set<String> remaining = publicUserSubscriptionsMap.get(tenantId); |
||||
|
assertThat(remaining).hasSize(1); |
||||
|
assertThat(remaining).allMatch(subId -> subId.startsWith("[session-2]")); |
||||
|
} |
||||
|
|
||||
|
private WebSocketSessionRef mockPublicSessionRef(TenantId tenantId, String sessionId) { |
||||
|
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(sessionId).given(ref).getSessionId(); |
||||
|
return ref; |
||||
|
} |
||||
|
|
||||
|
private AttributesSubscriptionCmd subscriptionCmd(int cmdId) { |
||||
|
AttributesSubscriptionCmd cmd = new AttributesSubscriptionCmd(); |
||||
|
cmd.setCmdId(cmdId); |
||||
|
return cmd; |
||||
|
} |
||||
|
|
||||
|
} |
||||
Loading…
Reference in new issue