Added rate limits for websocket updates and REST API

+/**
+ * Copyright © 2016-2018 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.config;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.stereotype.Component;
+import org.springframework.web.filter.GenericFilterBean;
+import org.thingsboard.server.common.data.EntityType;
+import org.thingsboard.server.common.data.id.CustomerId;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.msg.tools.TbRateLimits;
+import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
+import org.thingsboard.server.exception.ThingsboardErrorResponseHandler;
+import org.thingsboard.server.service.security.model.SecurityUser;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+@Component
+public class RateLimitProcessingFilter extends GenericFilterBean {
+
+    @Value("${server.rest.limits.tenant.enabled:false}")
+    private boolean perTenantLimitsEnabled;
+    @Value("${server.rest.limits.tenant.configuration:}")
+    private String perTenantLimitsConfiguration;
+    @Value("${server.rest.limits.customer.enabled:false}")
+    private boolean perCustomerLimitsEnabled;
+    @Value("${server.rest.limits.customer.configuration:}")
+    private String perCustomerLimitsConfiguration;
+
+    @Autowired
+    private ThingsboardErrorResponseHandler errorResponseHandler;
+
+    private ConcurrentMap<TenantId, TbRateLimits> perTenantLimits = new ConcurrentHashMap<>();
+    private ConcurrentMap<CustomerId, TbRateLimits> perCustomerLimits = new ConcurrentHashMap<>();
+
+    @Override
+    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
+        SecurityUser user = getCurrentUser();
+        if (user != null && !user.isSystemAdmin()) {
+            if (perTenantLimitsEnabled) {
+                TbRateLimits rateLimits = perTenantLimits.computeIfAbsent(user.getTenantId(), id -> new TbRateLimits(perTenantLimitsConfiguration));
+                if (!rateLimits.tryConsume()) {
+                    errorResponseHandler.handle(new TbRateLimitsException(EntityType.TENANT), (HttpServletResponse) response);
+                    return;
+                }
+            }
+            if (perCustomerLimitsEnabled && user.isCustomerUser()) {
+                TbRateLimits rateLimits = perCustomerLimits.computeIfAbsent(user.getCustomerId(), id -> new TbRateLimits(perCustomerLimitsConfiguration));
+                if (!rateLimits.tryConsume()) {
+                    errorResponseHandler.handle(new TbRateLimitsException(EntityType.CUSTOMER), (HttpServletResponse) response);
+                    return;
+                }
+            }
+        }
+        chain.doFilter(request, response);
+    }
+
+    protected SecurityUser getCurrentUser() {
+        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
+        if (authentication != null && authentication.getPrincipal() instanceof SecurityUser) {
+            return (SecurityUser) authentication.getPrincipal();
+        } else {
+            return null;
+        }
+    }
+
+}

@@ -91,6 +91,8 @@ public class ThingsboardSecurityConfiguration extends WebSecurityConfigurerAdapt
 
     @Autowired private ObjectMapper objectMapper;
 
+    @Autowired private RateLimitProcessingFilter rateLimitProcessingFilter;
+
     @Bean
     protected RestLoginProcessingFilter buildRestLoginProcessingFilter() throws Exception {
         RestLoginProcessingFilter filter = new RestLoginProcessingFilter(FORM_BASED_LOGIN_ENTRY_POINT, successHandler, failureHandler, objectMapper);
@@ -186,7 +188,8 @@ public class ThingsboardSecurityConfiguration extends WebSecurityConfigurerAdapt
                 .addFilterBefore(buildRestPublicLoginProcessingFilter(), UsernamePasswordAuthenticationFilter.class)
                 .addFilterBefore(buildJwtTokenAuthenticationProcessingFilter(), UsernamePasswordAuthenticationFilter.class)
                 .addFilterBefore(buildRefreshTokenProcessingFilter(), UsernamePasswordAuthenticationFilter.class)
-                .addFilterBefore(buildWsJwtTokenAuthenticationProcessingFilter(), UsernamePasswordAuthenticationFilter.class);
+                .addFilterBefore(buildWsJwtTokenAuthenticationProcessingFilter(), UsernamePasswordAuthenticationFilter.class)
+                .addFilterAfter(rateLimitProcessingFilter, UsernamePasswordAuthenticationFilter.class);
     }
 
 

@@ -52,6 +52,7 @@ import org.thingsboard.server.common.msg.TbMsgDataType;
 import org.thingsboard.server.common.msg.TbMsgMetaData;
 import org.thingsboard.server.common.msg.cluster.SendToClusterMsg;
 import org.thingsboard.server.common.msg.system.ServiceToRuleEngineMsg;
+import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
 import org.thingsboard.server.dao.alarm.AlarmService;
 import org.thingsboard.server.dao.asset.AssetService;
 import org.thingsboard.server.dao.attributes.AttributesService;

@@ -19,15 +19,17 @@ import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.BeanCreationNotAllowedException;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
-import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Service;
+import org.springframework.util.StringUtils;
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.handler.TextWebSocketHandler;
+import org.thingsboard.server.common.data.exception.ThingsboardErrorCode;
 import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.id.UserId;
+import org.thingsboard.server.common.msg.tools.TbRateLimits;
 import org.thingsboard.server.config.WebSocketConfiguration;
 import org.thingsboard.server.service.security.model.SecurityUser;
 import org.thingsboard.server.service.security.model.UserPrincipal;
@@ -63,11 +65,17 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
     @Value("${server.ws.limits.max_sessions_per_public_user:0}")
     private int maxSessionsPerPublicUser;
 
+    @Value("${server.ws.limits.max_updates_per_session:}")
+    private String perSessionUpdatesConfiguration;
+
+    private ConcurrentMap<String, TelemetryWebSocketSessionRef> blacklistedSessions = new ConcurrentHashMap<>();
+    private ConcurrentMap<String, TbRateLimits> perSessionUpdateLimits = new ConcurrentHashMap<>();
+
     private ConcurrentMap<TenantId, Set<String>> tenantSessionsMap = new ConcurrentHashMap<>();
     private ConcurrentMap<CustomerId, Set<String>> customerSessionsMap = new ConcurrentHashMap<>();
     private ConcurrentMap<UserId, Set<String>> regularUserSessionsMap = new ConcurrentHashMap<>();
     private ConcurrentMap<UserId, Set<String>> publicUserSessionsMap = new ConcurrentHashMap<>();
-    
+
     @Override
     public void handleTextMessage(WebSocketSession session, TextMessage message) {
         try {
@@ -168,13 +176,29 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
     }
 
     @Override
-    public void send(TelemetryWebSocketSessionRef sessionRef, String msg) throws IOException {
+    public void send(TelemetryWebSocketSessionRef sessionRef, int subscriptionId, String msg) throws IOException {
         String externalId = sessionRef.getSessionId();
         log.debug("[{}] Processing {}", externalId, msg);
         String internalId = externalSessionMap.get(externalId);
         if (internalId != null) {
             SessionMetaData sessionMd = internalSessionMap.get(internalId);
             if (sessionMd != null) {
+                if (!StringUtils.isEmpty(perSessionUpdatesConfiguration)) {
+                    TbRateLimits rateLimits = perSessionUpdateLimits.computeIfAbsent(sessionRef.getSessionId(), sid -> new TbRateLimits(perSessionUpdatesConfiguration));
+                    if (!rateLimits.tryConsume()) {
+                        if (blacklistedSessions.putIfAbsent(externalId, sessionRef) == null) {
+                            log.info("[{}][{}][{}] Failed to process session update. Max session updates limit reached"
+                                    , sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), externalId);
+                            synchronized (sessionMd) {
+                                sessionMd.session.sendMessage(new TextMessage("{\"subscriptionId\":" + subscriptionId + ", \"errorCode\":" + ThingsboardErrorCode.TOO_MANY_UPDATES.getErrorCode() + ", \"errorMsg\":\"Too many updates!\"}"));
+                            }
+                        }
+                        return;
+                    } else {
+                        log.debug("[{}][{}][{}] Session is no longer blacklisted.", sessionRef.getSecurityCtx().getTenantId(), sessionRef.getSecurityCtx().getId(), externalId);
+                        blacklistedSessions.remove(externalId);
+                    }
+                }
                 synchronized (sessionMd) {
                     sessionMd.session.sendMessage(new TextMessage(msg));
                 }
@@ -186,12 +210,6 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
         }
     }
 
-
-    @Override
-    public void close(TelemetryWebSocketSessionRef sessionRef) throws IOException {
-        close(sessionRef, CloseStatus.NORMAL);
-    }
-
     @Override
     public void close(TelemetryWebSocketSessionRef sessionRef, CloseStatus reason) throws IOException {
         String externalId = sessionRef.getSessionId();
@@ -271,6 +289,8 @@ public class TbWebSocketHandler extends TextWebSocketHandler implements Telemetr
 
     private void cleanupLimits(WebSocketSession session, TelemetryWebSocketSessionRef sessionRef) {
         String sessionId = session.getId();
+        perSessionUpdateLimits.remove(sessionRef.getSessionId());
+        blacklistedSessions.remove(sessionRef.getSessionId());
         if (maxSessionsPerTenant > 0) {
             Set<String> tenantSessions = tenantSessionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet());
             synchronized (tenantSessions) {

@@ -25,8 +25,10 @@ import org.springframework.security.authentication.BadCredentialsException;
 import org.springframework.security.core.AuthenticationException;
 import org.springframework.security.web.access.AccessDeniedHandler;
 import org.springframework.stereotype.Component;
+import org.thingsboard.server.common.data.EntityType;
 import org.thingsboard.server.common.data.exception.ThingsboardErrorCode;
 import org.thingsboard.server.common.data.exception.ThingsboardException;
+import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
 import org.thingsboard.server.service.security.exception.AuthMethodNotSupportedException;
 import org.thingsboard.server.service.security.exception.JwtExpiredTokenException;
 
@@ -34,6 +36,7 @@ import javax.servlet.ServletException;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
+
 @Component
 @Slf4j
 public class ThingsboardErrorResponseHandler implements AccessDeniedHandler {
@@ -62,6 +65,8 @@ public class ThingsboardErrorResponseHandler implements AccessDeniedHandler {
 
                 if (exception instanceof ThingsboardException) {
                     handleThingsboardException((ThingsboardException) exception, response);
+                } else if (exception instanceof TbRateLimitsException) {
+                    handleRateLimitException(response, (TbRateLimitsException) exception);
                 } else if (exception instanceof AccessDeniedException) {
                     handleAccessDeniedException(response);
                 } else if (exception instanceof AuthenticationException) {
@@ -77,6 +82,7 @@ public class ThingsboardErrorResponseHandler implements AccessDeniedHandler {
         }
     }
 
+
     private void handleThingsboardException(ThingsboardException thingsboardException, HttpServletResponse response) throws IOException {
 
         ThingsboardErrorCode errorCode = thingsboardException.getErrorCode();
@@ -110,6 +116,15 @@ public class ThingsboardErrorResponseHandler implements AccessDeniedHandler {
         mapper.writeValue(response.getWriter(), ThingsboardErrorResponse.of(thingsboardException.getMessage(), errorCode, status));
     }
 
+    private void handleRateLimitException(HttpServletResponse response, TbRateLimitsException exception) throws IOException {
+        response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
+        String message = "Too many requests for current " + exception.getEntityType().name().toLowerCase() + "!";
+        mapper.writeValue(response.getWriter(),
+                ThingsboardErrorResponse.of(message,
+                        ThingsboardErrorCode.TOO_MANY_REQUESTS, HttpStatus.TOO_MANY_REQUESTS));
+    }
+
+
     private void handleAccessDeniedException(HttpServletResponse response) throws IOException {
         response.setStatus(HttpStatus.FORBIDDEN.value());
         mapper.writeValue(response.getWriter(),

@@ -582,7 +582,7 @@ public class DefaultTelemetryWebSocketService implements TelemetryWebSocketServi
 
     private void sendWsMsg(TelemetryWebSocketSessionRef sessionRef, SubscriptionUpdate update) {
         try {
-            msgEndpoint.send(sessionRef, jsonMapper.writeValueAsString(update));
+            msgEndpoint.send(sessionRef, update.getSubscriptionId(), jsonMapper.writeValueAsString(update));
         } catch (JsonProcessingException e) {
             log.warn("[{}] Failed to encode reply: {}", sessionRef.getSessionId(), update, e);
         } catch (IOException e) {

@@ -24,9 +24,7 @@ import java.io.IOException;
  */
 public interface TelemetryWebSocketMsgEndpoint {
 
-    void send(TelemetryWebSocketSessionRef sessionRef, String msg) throws IOException;
-
-    void close(TelemetryWebSocketSessionRef sessionRef) throws IOException;
+    void send(TelemetryWebSocketSessionRef sessionRef, int subscriptionId, String msg) throws IOException;
 
     void close(TelemetryWebSocketSessionRef sessionRef, CloseStatus withReason) throws IOException;
 }

@@ -43,6 +43,15 @@ server:
       max_subscriptions_per_customer: "${TB_SERVER_WS_TENANT_RATE_LIMITS_MAX_SUBSCRIPTIONS_PER_CUSTOMER:0}"
       max_subscriptions_per_regular_user: "${TB_SERVER_WS_TENANT_RATE_LIMITS_MAX_SUBSCRIPTIONS_PER_REGULAR_USER:0}"
       max_subscriptions_per_public_user: "${TB_SERVER_WS_TENANT_RATE_LIMITS_MAX_SUBSCRIPTIONS_PER_PUBLIC_USER:0}"
+      max_updates_per_session: "${TB_SERVER_WS_TENANT_RATE_LIMITS_MAX_UPDATES_PER_SESSION:300:1,3000:60}"
+  rest:
+    limits:
+      tenant:
+        enabled: "${TB_SERVER_REST_LIMITS_TENANT_ENABLED:false}"
+        configuration: "${TB_SERVER_REST_LIMITS_TENANT_CONFIGURATION:100:1,2000:60}"
+      customer:
+        enabled: "${TB_SERVER_REST_LIMITS_CUSTOMER_ENABLED:false}"
+        configuration: "${TB_SERVER_REST_LIMITS_CUSTOMER_CONFIGURATION:50:1,1000:60}"
 
 # Zookeeper connection parameters. Used for service discovery.
 zk:

@@ -25,7 +25,9 @@ public enum ThingsboardErrorCode {
     PERMISSION_DENIED(20),
     INVALID_ARGUMENTS(30),
     BAD_REQUEST_PARAMS(31),
-    ITEM_NOT_FOUND(32);
+    ITEM_NOT_FOUND(32),
+    TOO_MANY_REQUESTS(33),
+    TOO_MANY_UPDATES(34);
 
     private int errorCode;
 

@@ -13,17 +13,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.thingsboard.server.common.transport.service;
+package org.thingsboard.server.common.msg.tools;
 
+import lombok.Getter;
 import org.thingsboard.server.common.data.EntityType;
 
 /**
  * Created by ashvayka on 22.10.18.
  */
-public class TbRateLimitsException extends Exception {
+public class TbRateLimitsException extends RuntimeException {
+    @Getter
     private final EntityType entityType;
 
-    TbRateLimitsException(EntityType entityType) {
+    public TbRateLimitsException(EntityType entityType) {
         this.entityType = entityType;
     }
 }

@@ -21,6 +21,7 @@ import org.thingsboard.server.common.data.EntityType;
 import org.thingsboard.server.common.data.id.DeviceId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.msg.tools.TbRateLimits;
+import org.thingsboard.server.common.msg.tools.TbRateLimitsException;
 import org.thingsboard.server.common.transport.SessionMsgListener;
 import org.thingsboard.server.common.transport.TransportService;
 import org.thingsboard.server.common.transport.TransportServiceCallback;

@@ -17,7 +17,6 @@ package org.thingsboard.server.dao.nosql;
 
 import com.datastax.driver.core.ResultSet;
 import com.datastax.driver.core.ResultSetFuture;
-import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.SettableFuture;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -26,7 +25,6 @@ import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Component;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.dao.entity.EntityService;
-import org.thingsboard.server.dao.tenant.TenantService;
 import org.thingsboard.server.dao.util.AbstractBufferedRateExecutor;
 import org.thingsboard.server.dao.util.AsyncTaskContext;
 import org.thingsboard.server.dao.util.NoSqlAnyDao;
@@ -34,7 +32,6 @@ import org.thingsboard.server.dao.util.NoSqlAnyDao;
 import javax.annotation.PreDestroy;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.ExecutionException;
 
 /**
  * Created by ashvayka on 24.10.18.

@@ -20,12 +20,10 @@ import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import lombok.extern.slf4j.Slf4j;
-import org.thingsboard.server.common.data.EntityType;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.msg.tools.TbRateLimits;
 
 import javax.annotation.Nullable;
-import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
@@ -37,7 +35,6 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Created by ashvayka on 24.10.18.