From 415c99147e46150aa8252fc1f3939055806a78b6 Mon Sep 17 00:00:00 2001 From: YevhenBondarenko Date: Wed, 22 May 2024 12:02:13 +0200 Subject: [PATCH] added check rate limits for CreateDeviceFromGateway --- .../mqtt/MqttGatewayRateLimitsTest.java | 3 ++- .../DefaultTransportRateLimitService.java | 2 +- .../service/DefaultTransportService.java | 24 +++++++++++++++---- 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/application/src/test/java/org/thingsboard/server/transport/mqtt/MqttGatewayRateLimitsTest.java b/application/src/test/java/org/thingsboard/server/transport/mqtt/MqttGatewayRateLimitsTest.java index 794d93a44a..a1986ca2c2 100644 --- a/application/src/test/java/org/thingsboard/server/transport/mqtt/MqttGatewayRateLimitsTest.java +++ b/application/src/test/java/org/thingsboard/server/transport/mqtt/MqttGatewayRateLimitsTest.java @@ -81,7 +81,8 @@ public class MqttGatewayRateLimitsTest extends AbstractControllerTest { @Test public void transportGatewayMsgRateLimitTest() throws Exception { - transportGatewayRateLimitTest(profileConfiguration -> profileConfiguration.setTransportGatewayMsgRateLimit("1:600")); + // Device A 2 msgs success ('create device', 'to device actor'), Device B 'create device success' , 'to device actor' - limited + transportGatewayRateLimitTest(profileConfiguration -> profileConfiguration.setTransportGatewayMsgRateLimit("3:600")); } @Test diff --git a/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/limits/DefaultTransportRateLimitService.java b/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/limits/DefaultTransportRateLimitService.java index f0a5d73f11..975dd21fab 100644 --- a/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/limits/DefaultTransportRateLimitService.java +++ b/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/limits/DefaultTransportRateLimitService.java @@ -84,7 +84,7 @@ public class DefaultTransportRateLimitService implements TransportRateLimitServi return TbPair.of(EntityType.DEVICE, true); } - if (!checkEntityRateLimit(dataPoints, getDeviceRateLimits(tenantId, deviceId))) { + if (deviceId != null && !checkEntityRateLimit(dataPoints, getDeviceRateLimits(tenantId, deviceId))) { return TbPair.of(EntityType.DEVICE, false); } return null; diff --git a/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/service/DefaultTransportService.java b/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/service/DefaultTransportService.java index 9312370c34..598a5aa7a6 100644 --- a/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/service/DefaultTransportService.java +++ b/common/transport/transport-api/src/main/java/org/thingsboard/server/common/transport/service/DefaultTransportService.java @@ -452,8 +452,13 @@ public class DefaultTransportService extends TransportActivityManager implements @Override public void process(TenantId tenantId, TransportProtos.GetOrCreateDeviceFromGatewayRequestMsg requestMsg, TransportServiceCallback callback) { - TbProtoQueueMsg protoMsg = new TbProtoQueueMsg<>(UUID.randomUUID(), TransportApiRequestMsg.newBuilder().setGetOrCreateDeviceRequestMsg(requestMsg).build()); log.trace("Processing msg: {}", requestMsg); + DeviceId gatewayid = new DeviceId(new UUID(requestMsg.getGatewayIdMSB(), requestMsg.getGatewayIdLSB())); + if (!checkLimits(tenantId, gatewayid, null, requestMsg.getDeviceName(), requestMsg, callback, 0)) { + return; + } + + TbProtoQueueMsg protoMsg = new TbProtoQueueMsg<>(UUID.randomUUID(), TransportApiRequestMsg.newBuilder().setGetOrCreateDeviceRequestMsg(requestMsg).build()); var key = new EntityLimitKey(tenantId, StringUtils.truncate(requestMsg.getDeviceName(), 256)); if (entityLimitsCache.get(key)) { transportCallbackExecutor.submit(() -> callback.onError(new RuntimeException(DataConstants.MAXIMUM_NUMBER_OF_DEVICES_REACHED))); @@ -857,6 +862,14 @@ public class DefaultTransportService extends TransportActivityManager implements gatewayId = new DeviceId(new UUID(sessionInfo.getGatewayIdMSB(), sessionInfo.getGatewayIdLSB())); } + return checkLimits(tenantId, gatewayId, deviceId, sessionInfo.getDeviceName(), msg, callback, dataPoints); + } + + private boolean checkLimits(TenantId tenantId, DeviceId gatewayId, DeviceId deviceId, String deviceName, Object msg, TransportServiceCallback callback, int dataPoints) { + if (log.isTraceEnabled()) { + log.trace("[{}][{}] Processing msg: {}", tenantId, deviceName, msg); + } + var rateLimitedPair = rateLimitService.checkLimits(tenantId, gatewayId, deviceId, dataPoints); if (rateLimitedPair == null) { return true; @@ -868,13 +881,16 @@ public class DefaultTransportService extends TransportActivityManager implements if (rateLimitedEntityType == EntityType.DEVICE || rateLimitedEntityType == EntityType.TENANT) { LimitedApi limitedApi = - rateLimitedEntityType == EntityType.TENANT ? LimitedApi.TRANSPORT_MESSAGES_PER_TENANT : + rateLimitedEntityType == EntityType.TENANT ? LimitedApi.TRANSPORT_MESSAGES_PER_TENANT : rateLimitedPair.getSecond() ? LimitedApi.TRANSPORT_MESSAGES_PER_GATEWAY : LimitedApi.TRANSPORT_MESSAGES_PER_DEVICE; + + EntityId limitLevel = rateLimitedEntityType == EntityType.DEVICE ? deviceId == null ? gatewayId : deviceId : tenantId; + notificationRuleProcessor.process(RateLimitsTrigger.builder() .tenantId(tenantId) .api(limitedApi) - .limitLevel(rateLimitedEntityType == EntityType.DEVICE ? deviceId : tenantId) - .limitLevelEntityName(rateLimitedEntityType == EntityType.DEVICE ? sessionInfo.getDeviceName() : null) + .limitLevel(limitLevel) + .limitLevelEntityName(rateLimitedEntityType == EntityType.DEVICE ? deviceName : null) .build()); } return false;