Andrew Shvayka
3 years ago
committed by
GitHub
GitHub
18 changed files with 339 additions and 219 deletions
@ -1,37 +0,0 @@ |
|||||
/** |
|
||||
* Copyright © 2016-2023 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.apiusage.limits; |
|
||||
|
|
||||
import lombok.RequiredArgsConstructor; |
|
||||
import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileConfiguration; |
|
||||
|
|
||||
import java.util.function.Function; |
|
||||
|
|
||||
@RequiredArgsConstructor |
|
||||
public enum LimitedApi { |
|
||||
|
|
||||
ENTITY_EXPORT(DefaultTenantProfileConfiguration::getTenantEntityExportRateLimit), |
|
||||
ENTITY_IMPORT(DefaultTenantProfileConfiguration::getTenantEntityImportRateLimit), |
|
||||
NOTIFICATION_REQUESTS(DefaultTenantProfileConfiguration::getTenantNotificationRequestsRateLimit), |
|
||||
NOTIFICATION_REQUESTS_PER_RULE(DefaultTenantProfileConfiguration::getTenantNotificationRequestsPerRuleRateLimit); |
|
||||
|
|
||||
private final Function<DefaultTenantProfileConfiguration, String> configExtractor; |
|
||||
|
|
||||
public String getLimitConfig(DefaultTenantProfileConfiguration profileConfiguration) { |
|
||||
return configExtractor.apply(profileConfiguration); |
|
||||
} |
|
||||
|
|
||||
} |
|
||||
@ -0,0 +1,110 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2023 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.limits; |
||||
|
|
||||
|
import org.junit.Before; |
||||
|
import org.junit.Test; |
||||
|
import org.junit.runner.RunWith; |
||||
|
import org.mockito.Mockito; |
||||
|
import org.mockito.junit.MockitoJUnitRunner; |
||||
|
import org.thingsboard.server.common.data.TenantProfile; |
||||
|
import org.thingsboard.server.common.data.id.CustomerId; |
||||
|
import org.thingsboard.server.common.data.id.NotificationRuleId; |
||||
|
import org.thingsboard.server.common.data.id.TenantId; |
||||
|
import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileConfiguration; |
||||
|
import org.thingsboard.server.common.data.tenant.profile.TenantProfileData; |
||||
|
import org.thingsboard.server.dao.tenant.TbTenantProfileCache; |
||||
|
import org.thingsboard.server.dao.util.limits.DefaultRateLimitService; |
||||
|
import org.thingsboard.server.dao.util.limits.LimitedApi; |
||||
|
import org.thingsboard.server.dao.util.limits.RateLimitService; |
||||
|
|
||||
|
import java.util.List; |
||||
|
import java.util.UUID; |
||||
|
|
||||
|
import static org.junit.Assert.assertFalse; |
||||
|
import static org.junit.Assert.assertTrue; |
||||
|
import static org.mockito.ArgumentMatchers.eq; |
||||
|
import static org.mockito.Mockito.reset; |
||||
|
import static org.mockito.Mockito.when; |
||||
|
|
||||
|
@RunWith(MockitoJUnitRunner.class) |
||||
|
public class RateLimitServiceTest { |
||||
|
|
||||
|
private RateLimitService rateLimitService; |
||||
|
private TbTenantProfileCache tenantProfileCache; |
||||
|
private TenantId tenantId; |
||||
|
|
||||
|
@Before |
||||
|
public void beforeEach() { |
||||
|
tenantProfileCache = Mockito.mock(TbTenantProfileCache.class); |
||||
|
rateLimitService = new DefaultRateLimitService(tenantProfileCache, 60, 100); |
||||
|
tenantId = new TenantId(UUID.randomUUID()); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void testRateLimits() { |
||||
|
int max = 2; |
||||
|
String rateLimit = max + ":600"; |
||||
|
DefaultTenantProfileConfiguration profileConfiguration = new DefaultTenantProfileConfiguration(); |
||||
|
profileConfiguration.setTenantEntityExportRateLimit(rateLimit); |
||||
|
profileConfiguration.setTenantEntityImportRateLimit(rateLimit); |
||||
|
profileConfiguration.setTenantNotificationRequestsRateLimit(rateLimit); |
||||
|
profileConfiguration.setTenantNotificationRequestsPerRuleRateLimit(rateLimit); |
||||
|
profileConfiguration.setTenantServerRestLimitsConfiguration(rateLimit); |
||||
|
profileConfiguration.setCustomerServerRestLimitsConfiguration(rateLimit); |
||||
|
profileConfiguration.setWsUpdatesPerSessionRateLimit(rateLimit); |
||||
|
profileConfiguration.setCassandraQueryTenantRateLimitsConfiguration(rateLimit); |
||||
|
updateTenantProfileConfiguration(profileConfiguration); |
||||
|
|
||||
|
for (LimitedApi limitedApi : List.of( |
||||
|
LimitedApi.ENTITY_EXPORT, |
||||
|
LimitedApi.ENTITY_IMPORT, |
||||
|
LimitedApi.NOTIFICATION_REQUESTS, |
||||
|
LimitedApi.REST_REQUESTS, |
||||
|
LimitedApi.CASSANDRA_QUERIES |
||||
|
)) { |
||||
|
testRateLimits(limitedApi, max, tenantId); |
||||
|
} |
||||
|
|
||||
|
CustomerId customerId = new CustomerId(UUID.randomUUID()); |
||||
|
testRateLimits(LimitedApi.REST_REQUESTS, max, customerId); |
||||
|
|
||||
|
NotificationRuleId notificationRuleId = new NotificationRuleId(UUID.randomUUID()); |
||||
|
testRateLimits(LimitedApi.NOTIFICATION_REQUESTS_PER_RULE, max, notificationRuleId); |
||||
|
|
||||
|
String wsSessionId = UUID.randomUUID().toString(); |
||||
|
testRateLimits(LimitedApi.WS_UPDATES_PER_SESSION, max, wsSessionId); |
||||
|
} |
||||
|
|
||||
|
private void testRateLimits(LimitedApi limitedApi, int max, Object level) { |
||||
|
for (int i = 1; i <= max; i++) { |
||||
|
boolean success = rateLimitService.checkRateLimit(limitedApi, tenantId, level); |
||||
|
assertTrue(success); |
||||
|
} |
||||
|
boolean success = rateLimitService.checkRateLimit(limitedApi, tenantId, level); |
||||
|
assertFalse(success); |
||||
|
} |
||||
|
|
||||
|
private void updateTenantProfileConfiguration(DefaultTenantProfileConfiguration profileConfiguration) { |
||||
|
reset(tenantProfileCache); |
||||
|
TenantProfile tenantProfile = new TenantProfile(); |
||||
|
TenantProfileData profileData = new TenantProfileData(); |
||||
|
profileData.setConfiguration(profileConfiguration); |
||||
|
tenantProfile.setProfileData(profileData); |
||||
|
when(tenantProfileCache.get(eq(tenantId))).thenReturn(tenantProfile); |
||||
|
} |
||||
|
|
||||
|
} |
||||
@ -0,0 +1,31 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2023 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.common.data.exception; |
||||
|
|
||||
|
import lombok.Getter; |
||||
|
import org.thingsboard.server.common.data.id.TenantId; |
||||
|
|
||||
|
public class TenantProfileNotFoundException extends RuntimeException { |
||||
|
|
||||
|
@Getter |
||||
|
private final TenantId tenantId; |
||||
|
|
||||
|
public TenantProfileNotFoundException(TenantId tenantId) { |
||||
|
super("Profile for tenant with id " + tenantId + " not found"); |
||||
|
this.tenantId = tenantId; |
||||
|
} |
||||
|
|
||||
|
} |
||||
@ -0,0 +1,67 @@ |
|||||
|
/** |
||||
|
* Copyright © 2016-2023 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.dao.util.limits; |
||||
|
|
||||
|
import lombok.Getter; |
||||
|
import org.thingsboard.server.common.data.EntityType; |
||||
|
import org.thingsboard.server.common.data.id.EntityId; |
||||
|
import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileConfiguration; |
||||
|
|
||||
|
import java.util.function.BiFunction; |
||||
|
import java.util.function.Function; |
||||
|
|
||||
|
public enum LimitedApi { |
||||
|
|
||||
|
ENTITY_EXPORT(DefaultTenantProfileConfiguration::getTenantEntityExportRateLimit), |
||||
|
ENTITY_IMPORT(DefaultTenantProfileConfiguration::getTenantEntityImportRateLimit), |
||||
|
NOTIFICATION_REQUESTS(DefaultTenantProfileConfiguration::getTenantNotificationRequestsRateLimit), |
||||
|
NOTIFICATION_REQUESTS_PER_RULE(DefaultTenantProfileConfiguration::getTenantNotificationRequestsPerRuleRateLimit), |
||||
|
REST_REQUESTS((profileConfiguration, level) -> ((EntityId) level).getEntityType() == EntityType.TENANT ? |
||||
|
profileConfiguration.getTenantServerRestLimitsConfiguration() : |
||||
|
profileConfiguration.getCustomerServerRestLimitsConfiguration()), |
||||
|
WS_UPDATES_PER_SESSION(DefaultTenantProfileConfiguration::getWsUpdatesPerSessionRateLimit), |
||||
|
CASSANDRA_QUERIES(DefaultTenantProfileConfiguration::getCassandraQueryTenantRateLimitsConfiguration), |
||||
|
PASSWORD_RESET(true), |
||||
|
TWO_FA_VERIFICATION_CODE_SEND(true), |
||||
|
TWO_FA_VERIFICATION_CODE_CHECK(true); |
||||
|
|
||||
|
private final BiFunction<DefaultTenantProfileConfiguration, Object, String> configExtractor; |
||||
|
@Getter |
||||
|
private final boolean refillRateLimitIntervally; |
||||
|
|
||||
|
LimitedApi(Function<DefaultTenantProfileConfiguration, String> configExtractor) { |
||||
|
this((profileConfiguration, level) -> configExtractor.apply(profileConfiguration)); |
||||
|
} |
||||
|
|
||||
|
LimitedApi(BiFunction<DefaultTenantProfileConfiguration, Object, String> configExtractor) { |
||||
|
this.configExtractor = configExtractor; |
||||
|
this.refillRateLimitIntervally = false; |
||||
|
} |
||||
|
|
||||
|
LimitedApi(boolean refillRateLimitIntervally) { |
||||
|
this.configExtractor = null; |
||||
|
this.refillRateLimitIntervally = refillRateLimitIntervally; |
||||
|
} |
||||
|
|
||||
|
public String getLimitConfig(DefaultTenantProfileConfiguration profileConfiguration, Object level) { |
||||
|
if (configExtractor != null) { |
||||
|
return configExtractor.apply(profileConfiguration, level); |
||||
|
} else { |
||||
|
throw new IllegalArgumentException("No tenant profile config for " + name() + " rate limits"); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
} |
||||
Loading…
Reference in new issue