Merge pull request #10230 from AndriiLandiak/edge-uplink-messages-rate-limit

Add rate limits for edge uplink message

application/src/main/java/org/thingsboard/server/service/edge/EdgeContextComponent.java

@@ -19,6 +19,7 @@ import lombok.Data;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Lazy;
 import org.springframework.stereotype.Component;
+import org.thingsboard.server.cache.limits.RateLimitService;
 import org.thingsboard.server.cluster.TbClusterService;
 import org.thingsboard.server.common.msg.notification.NotificationRuleProcessor;
 import org.thingsboard.server.dao.asset.AssetProfileService;
@@ -150,6 +151,9 @@ public class EdgeContextComponent {
     @Autowired
     private ResourceService resourceService;
 
+    @Autowired
+    private RateLimitService rateLimitService;
+
     @Autowired
     private NotificationRuleProcessor notificationRuleProcessor;
 

application/src/main/java/org/thingsboard/server/service/edge/rpc/EdgeGrpcSession.java

@@ -35,6 +35,7 @@ import org.thingsboard.server.common.data.kv.AttributeKvEntry;
 import org.thingsboard.server.common.data.kv.BaseAttributeKvEntry;
 import org.thingsboard.server.common.data.kv.LongDataEntry;
 import org.thingsboard.server.common.data.kv.StringDataEntry;
+import org.thingsboard.server.common.data.limit.LimitedApi;
 import org.thingsboard.server.common.data.notification.rule.trigger.EdgeCommunicationFailureTrigger;
 import org.thingsboard.server.common.data.page.PageData;
 import org.thingsboard.server.common.data.page.PageLink;
@@ -109,6 +110,7 @@ public final class EdgeGrpcSession implements Closeable {
 
     private static final String QUEUE_START_TS_ATTR_KEY = "queueStartTs";
     private static final String QUEUE_START_SEQ_ID_ATTR_KEY = "queueStartSeqId";
+    private static final String RATE_LIMIT_REACHED = "Rate limit reached";
 
     private final UUID sessionId;
     private final BiConsumer<EdgeId, EdgeGrpcSession> sessionOpenListener;
@@ -264,36 +266,51 @@ public final class EdgeGrpcSession implements Closeable {
     }
 
     private void onUplinkMsg(UplinkMsg uplinkMsg) {
+        if (isRateLimitViolated(uplinkMsg)) {
+            return;
+        }
         ListenableFuture<List<Void>> future = processUplinkMsg(uplinkMsg);
         Futures.addCallback(future, new FutureCallback<>() {
             @Override
             public void onSuccess(@Nullable List<Void> result) {
-                UplinkResponseMsg uplinkResponseMsg = UplinkResponseMsg.newBuilder()
+                sendResponseMessage(uplinkMsg.getUplinkMsgId(), true, null);
-                        .setUplinkMsgId(uplinkMsg.getUplinkMsgId())
-                        .setSuccess(true).build();
-                sendDownlinkMsg(ResponseMsg.newBuilder()
-                        .setUplinkResponseMsg(uplinkResponseMsg)
-                        .build());
             }
 
             @Override
             public void onFailure(Throwable t) {
                 String errorMsg = EdgeUtils.createErrorMsgFromRootCauseAndStackTrace(t);
-                UplinkResponseMsg uplinkResponseMsg = UplinkResponseMsg.newBuilder()
+                sendResponseMessage(uplinkMsg.getUplinkMsgId(), false, errorMsg);
-                        .setUplinkMsgId(uplinkMsg.getUplinkMsgId())
-                        .setSuccess(false).setErrorMsg(errorMsg).build();
-                sendDownlinkMsg(ResponseMsg.newBuilder()
-                        .setUplinkResponseMsg(uplinkResponseMsg)
-                        .build());
             }
         }, ctx.getGrpcCallbackExecutorService());
     }
 
+    private boolean isRateLimitViolated(UplinkMsg uplinkMsg) {
+        if (!ctx.getRateLimitService().checkRateLimit(LimitedApi.EDGE_UPLINK_MESSAGES, tenantId) ||
+                !ctx.getRateLimitService().checkRateLimit(LimitedApi.EDGE_UPLINK_MESSAGES_PER_EDGE, tenantId, edge.getId())) {
+            String errorMsg = String.format("Failed to process uplink message. %s", RATE_LIMIT_REACHED);
+            sendResponseMessage(uplinkMsg.getUplinkMsgId(), false, errorMsg);
+            return true;
+        }
+        return false;
+    }
+
+    private void sendResponseMessage(int uplinkMsgId, boolean success, String errorMsg) {
+        UplinkResponseMsg.Builder responseBuilder = UplinkResponseMsg.newBuilder()
+                .setUplinkMsgId(uplinkMsgId)
+                .setSuccess(success);
+        if (errorMsg != null) {
+            responseBuilder.setErrorMsg(errorMsg);
+        }
+        sendDownlinkMsg(ResponseMsg.newBuilder()
+                .setUplinkResponseMsg(responseBuilder.build())
+                .build());
+    }
+
     private void onDownlinkResponse(DownlinkResponseMsg msg) {
         try {
             if (msg.getSuccess()) {
                 sessionState.getPendingMsgsMap().remove(msg.getDownlinkMsgId());
-                log.debug("[{}][{}] Msg has been processed successfully!Msg Id: [{}], Msg: {}", this.tenantId, edge.getRoutingKey(), msg.getDownlinkMsgId(), msg);
+                log.debug("[{}][{}] Msg has been processed successfully! Msg Id: [{}], Msg: {}", this.tenantId, edge.getRoutingKey(), msg.getDownlinkMsgId(), msg);
             } else {
                 log.error("[{}][{}] Msg processing failed! Msg Id: [{}], Error msg: {}", this.tenantId, edge.getRoutingKey(), msg.getDownlinkMsgId(), msg.getErrorMsg());
             }

application/src/test/java/org/thingsboard/server/service/limits/RateLimitServiceTest.java

@@ -71,6 +71,8 @@ public class RateLimitServiceTest {
         profileConfiguration.setCassandraQueryTenantRateLimitsConfiguration(rateLimit);
         profileConfiguration.setEdgeEventRateLimits(rateLimit);
         profileConfiguration.setEdgeEventRateLimitsPerEdge(rateLimit);
+        profileConfiguration.setEdgeUplinkMessagesRateLimits(rateLimit);
+        profileConfiguration.setEdgeUplinkMessagesRateLimitsPerEdge(rateLimit);
         updateTenantProfileConfiguration(profileConfiguration);
 
         for (LimitedApi limitedApi : List.of(
@@ -80,7 +82,9 @@ public class RateLimitServiceTest {
                 LimitedApi.REST_REQUESTS_PER_CUSTOMER,
                 LimitedApi.CASSANDRA_QUERIES,
                 LimitedApi.EDGE_EVENTS,
-                LimitedApi.EDGE_EVENTS_PER_EDGE
+                LimitedApi.EDGE_EVENTS_PER_EDGE,
+                LimitedApi.EDGE_UPLINK_MESSAGES,
+                LimitedApi.EDGE_UPLINK_MESSAGES_PER_EDGE
         )) {
             testRateLimits(limitedApi, max, tenantId);
         }

common/data/src/main/java/org/thingsboard/server/common/data/limit/LimitedApi.java

@@ -33,6 +33,8 @@ public enum LimitedApi {
     CASSANDRA_QUERIES(DefaultTenantProfileConfiguration::getCassandraQueryTenantRateLimitsConfiguration, "Cassandra queries", true),
     EDGE_EVENTS(DefaultTenantProfileConfiguration::getEdgeEventRateLimits, "Edge events", true),
     EDGE_EVENTS_PER_EDGE(DefaultTenantProfileConfiguration::getEdgeEventRateLimitsPerEdge, "Edge events per edge", false),
+    EDGE_UPLINK_MESSAGES(DefaultTenantProfileConfiguration::getEdgeUplinkMessagesRateLimits, "Edge uplink messages", true),
+    EDGE_UPLINK_MESSAGES_PER_EDGE(DefaultTenantProfileConfiguration::getEdgeUplinkMessagesRateLimitsPerEdge, "Edge uplink messages per edge", false),
     PASSWORD_RESET(false, true),
     TWO_FA_VERIFICATION_CODE_SEND(false, true),
     TWO_FA_VERIFICATION_CODE_CHECK(false, true),

common/data/src/main/java/org/thingsboard/server/common/data/tenant/profile/DefaultTenantProfileConfiguration.java

@@ -83,6 +83,8 @@ public class DefaultTenantProfileConfiguration implements TenantProfileConfigura
 
     private String edgeEventRateLimits;
     private String edgeEventRateLimitsPerEdge;
+    private String edgeUplinkMessagesRateLimits;
+    private String edgeUplinkMessagesRateLimitsPerEdge;
 
     private int defaultStorageTtlDays;
     private int alarmsTtlDays;

ui-ngx/src/app/modules/home/components/profile/tenant/default-tenant-profile-configuration.component.html

@@ -571,6 +571,14 @@
                           [type]="rateLimitsType.EDGE_EVENTS_PER_EDGE_RATE_LIMIT">
           </tb-rate-limits>
         </div>
+        <div fxFlex fxLayout="row" fxLayout.xs="column" fxLayoutGap.gt-xs="16px">
+          <tb-rate-limits fxFlex="50" formControlName="edgeUplinkMessagesRateLimits"
+                          [type]="rateLimitsType.EDGE_UPLINK_MESSAGES_RATE_LIMIT">
+          </tb-rate-limits>
+          <tb-rate-limits fxFlex="50" formControlName="edgeUplinkMessagesRateLimitsPerEdge"
+                          [type]="rateLimitsType.EDGE_UPLINK_MESSAGES_PER_EDGE_RATE_LIMIT">
+          </tb-rate-limits>
+        </div>
       </ng-template>
     </mat-expansion-panel>
   </fieldset>

ui-ngx/src/app/modules/home/components/profile/tenant/default-tenant-profile-configuration.component.ts

@@ -108,7 +108,9 @@ export class DefaultTenantProfileConfigurationComponent implements ControlValueA
       wsUpdatesPerSessionRateLimit: [null, []],
       cassandraQueryTenantRateLimitsConfiguration: [null, []],
       edgeEventRateLimits: [null, []],
-      edgeEventRateLimitsPerEdge: [null, []]
+      edgeEventRateLimitsPerEdge: [null, []],
+      edgeUplinkMessagesRateLimits: [null, []],
+      edgeUplinkMessagesRateLimitsPerEdge: [null, []]
     });
 
     this.defaultTenantProfileConfigurationFormGroup.get('smsEnabled').valueChanges.pipe(

ui-ngx/src/app/modules/home/components/profile/tenant/rate-limits/rate-limits.models.ts

@@ -37,7 +37,9 @@ export enum RateLimitsType {
   TENANT_NOTIFICATION_REQUEST_RATE_LIMIT = 'TENANT_NOTIFICATION_REQUEST_RATE_LIMIT',
   TENANT_NOTIFICATION_REQUESTS_PER_RULE_RATE_LIMIT = 'TENANT_NOTIFICATION_REQUESTS_PER_RULE_RATE_LIMIT',
   EDGE_EVENTS_RATE_LIMIT = 'EDGE_EVENTS_RATE_LIMIT',
-  EDGE_EVENTS_PER_EDGE_RATE_LIMIT = 'EDGE_EVENTS_PER_EDGE_RATE_LIMIT'
+  EDGE_EVENTS_PER_EDGE_RATE_LIMIT = 'EDGE_EVENTS_PER_EDGE_RATE_LIMIT',
+  EDGE_UPLINK_MESSAGES_RATE_LIMIT = 'EDGE_UPLINK_MESSAGES_RATE_LIMIT',
+  EDGE_UPLINK_MESSAGES_PER_EDGE_RATE_LIMIT = 'EDGE_UPLINK_MESSAGES_PER_EDGE_RATE_LIMIT'
 }
 
 export const rateLimitsLabelTranslationMap = new Map<RateLimitsType, string>(
@@ -58,6 +60,8 @@ export const rateLimitsLabelTranslationMap = new Map<RateLimitsType, string>(
     [RateLimitsType.TENANT_NOTIFICATION_REQUESTS_PER_RULE_RATE_LIMIT, 'tenant-profile.tenant-notification-requests-per-rule-rate-limit'],
     [RateLimitsType.EDGE_EVENTS_RATE_LIMIT, 'tenant-profile.rate-limits.edge-events-rate-limit'],
     [RateLimitsType.EDGE_EVENTS_PER_EDGE_RATE_LIMIT, 'tenant-profile.rate-limits.edge-events-per-edge-rate-limit'],
+    [RateLimitsType.EDGE_UPLINK_MESSAGES_RATE_LIMIT, 'tenant-profile.rate-limits.edge-uplink-messages-rate-limit'],
+    [RateLimitsType.EDGE_UPLINK_MESSAGES_PER_EDGE_RATE_LIMIT, 'tenant-profile.rate-limits.edge-uplink-messages-per-edge-rate-limit'],
   ]
 );
 
@@ -79,6 +83,8 @@ export const rateLimitsDialogTitleTranslationMap = new Map<RateLimitsType, strin
     [RateLimitsType.TENANT_NOTIFICATION_REQUESTS_PER_RULE_RATE_LIMIT, 'tenant-profile.rate-limits.edit-tenant-notification-requests-per-rule-rate-limit-title'],
     [RateLimitsType.EDGE_EVENTS_RATE_LIMIT, 'tenant-profile.rate-limits.edit-edge-events-rate-limit'],
     [RateLimitsType.EDGE_EVENTS_PER_EDGE_RATE_LIMIT, 'tenant-profile.rate-limits.edit-edge-events-per-edge-rate-limit'],
+    [RateLimitsType.EDGE_UPLINK_MESSAGES_RATE_LIMIT, 'tenant-profile.rate-limits.edit-edge-uplink-messages-rate-limit'],
+    [RateLimitsType.EDGE_UPLINK_MESSAGES_PER_EDGE_RATE_LIMIT, 'tenant-profile.rate-limits.edit-edge-uplink-messages-per-edge-rate-limit'],
   ]
 );
 

ui-ngx/src/app/shared/models/limited-api.models.ts

@@ -26,7 +26,9 @@ export enum LimitedApi {
   TRANSPORT_MESSAGES_PER_TENANT = 'TRANSPORT_MESSAGES_PER_TENANT',
   TRANSPORT_MESSAGES_PER_DEVICE = 'TRANSPORT_MESSAGES_PER_DEVICE',
   EDGE_EVENTS = 'EDGE_EVENTS',
-  EDGE_EVENTS_PER_EDGE = 'EDGE_EVENTS_PER_EDGE'
+  EDGE_EVENTS_PER_EDGE = 'EDGE_EVENTS_PER_EDGE',
+  EDGE_UPLINK_MESSAGES = 'EDGE_UPLINK_MESSAGES',
+  EDGE_UPLINK_MESSAGES_PER_EDGE = 'EDGE_UPLINK_MESSAGES_PER_EDGE'
 }
 
 export const LimitedApiTranslationMap = new Map<LimitedApi, string>(
@@ -43,5 +45,7 @@ export const LimitedApiTranslationMap = new Map<LimitedApi, string>(
     [LimitedApi.TRANSPORT_MESSAGES_PER_DEVICE, 'api-limit.transport-messages-per-device'],
     [LimitedApi.EDGE_EVENTS, 'api-limit.edge-events'],
     [LimitedApi.EDGE_EVENTS_PER_EDGE, 'api-limit.edge-events-per-edge'],
+    [LimitedApi.EDGE_UPLINK_MESSAGES, 'api-limit.edge-uplink-messages'],
+    [LimitedApi.EDGE_UPLINK_MESSAGES_PER_EDGE, 'api-limit.edge-uplink-messages-per-edge']
   ]
 );

ui-ngx/src/assets/help/en_US/notification/rate_limits.md

@@ -11,7 +11,7 @@ Available template parameters:
 
 * `api` - rate-limited API label; one of: 'REST API requests', 'REST API requests per customer', 'transport messages', 
   'transport messages per device', 'Cassandra queries', 'WS updates per session', 'notification requests', 'notification requests per rule',
-  'entity version creation', 'entity version load', 'Edge events', 'Edge events per edge';
+  'entity version creation', 'entity version load', 'Edge events', 'Edge events per edge', 'Edge uplink messages', 'Edge uplink messages per edge';
 * `limitLevelEntityType` - entity type of the limit level entity, e.g. 'Tenant', 'Device', 'Notification rule', 'Customer', etc.;
 * `limitLevelEntityId` - id of the limit level entity;
 * `limitLevelEntityName` - name of the limit level entity;

ui-ngx/src/assets/locale/locale.constant-en_US.json

@@ -861,7 +861,9 @@
         "transport-messages-per-device": "Transport messages per device",
         "ws-updates-per-session": "WS updates per session",
         "edge-events": "Edge events",
-        "edge-events-per-edge": "Edge events per edge"
+        "edge-events-per-edge": "Edge events per edge",
+        "edge-uplink-messages": "Edge uplink messages",
+        "edge-uplink-messages-per-edge": "Edge uplink messages per edge"
     },
     "audit-log": {
         "audit": "Audit",
@@ -4210,6 +4212,10 @@
             "edit-edge-events-per-edge-rate-limit": "Edit edge events per edge rate limits",
             "edge-events-rate-limit": "Edge events",
             "edge-events-per-edge-rate-limit": "Edge events per edge",
+            "edit-edge-uplink-messages-rate-limit": "Edit edge uplink messages rate limits",
+            "edit-edge-uplink-messages-per-edge-rate-limit": "Edit edge uplink messages per edge rate limits",
+            "edge-uplink-messages-rate-limit": "Edge uplink messages",
+            "edge-uplink-messages-per-edge-rate-limit": "Edge uplink messages per edge",
             "messages-per": "messages per",
             "not-set": "Not set",
             "number-of-messages": "Number of messages",