Feature/openapi rate limit function

CHANGES.md

@@ -19,6 +19,7 @@ Apollo 2.4.0
 * [Feature: openapi query namespace support not fill item](https://github.com/apolloconfig/apollo/pull/5249)
 * [Refactor: align database ClusterName and NamespaceName fields lengths](https://github.com/apolloconfig/apollo/pull/5263)
 * [Feature: Added the value length limit function for AppId-level configuration items](https://github.com/apolloconfig/apollo/pull/5264)
+* [Feature: Added current limiting function to ConsumerToken](https://github.com/apolloconfig/apollo/pull/5267)
 
 ------------------
 All issues and pull requests are [here](https://github.com/apolloconfig/apollo/milestone/15?closed=1)

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/entity/ConsumerToken.java

@@ -41,6 +41,9 @@ public class ConsumerToken extends BaseEntity {
   @Column(name = "`Token`", nullable = false)
   private String token;
 
+  @Column(name = "LimitCount")
+  private Integer limitCount;
+
   @Column(name = "`Expires`", nullable = false)
   private Date expires;
 
@@ -60,6 +63,14 @@ public void setToken(String token) {
     this.token = token;
   }
 
+  public Integer getLimitCount() {
+    return limitCount;
+  }
+
+  public void setLimitCount(Integer limitCount) {
+    this.limitCount = limitCount;
+  }
+
   public Date getExpires() {
     return expires;
   }
@@ -71,6 +82,7 @@ public void setExpires(Date expires) {
   @Override
   public String toString() {
     return toStringHelper().add("consumerId", consumerId).add("token", token)
+        .add("limitCount", limitCount)
         .add("expires", expires).toString();
   }
 }

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/filter/ConsumerAuthenticationFilter.java

@@ -16,11 +16,15 @@
  */
 package com.ctrip.framework.apollo.openapi.filter;
 
+import com.ctrip.framework.apollo.openapi.entity.ConsumerToken;
 import com.ctrip.framework.apollo.openapi.util.ConsumerAuditUtil;
 import com.ctrip.framework.apollo.openapi.util.ConsumerAuthUtil;
-
+import com.ctrip.framework.apollo.portal.component.config.PortalConfig;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.RateLimiter;
 import java.io.IOException;
-
+import java.util.concurrent.TimeUnit;
 import javax.servlet.Filter;
 import javax.servlet.FilterChain;
 import javax.servlet.FilterConfig;
@@ -29,18 +33,33 @@
 import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
+import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpHeaders;
 
 /**
  * @author Jason Song([email protected])
  */
 public class ConsumerAuthenticationFilter implements Filter {
+
+  private static final Logger logger = LoggerFactory.getLogger(ConsumerAuthenticationFilter.class);
+
   private final ConsumerAuthUtil consumerAuthUtil;
   private final ConsumerAuditUtil consumerAuditUtil;
+  private final PortalConfig portalConfig;
 
-  public ConsumerAuthenticationFilter(ConsumerAuthUtil consumerAuthUtil, ConsumerAuditUtil consumerAuditUtil) {
+  private static final int WARMUP_MILLIS = 1000; // ms
+  private static final int RATE_LIMITER_CACHE_MAX_SIZE = 50000;
+
+  private static final Cache<String, ImmutablePair<Long, RateLimiter>> LIMITER = CacheBuilder.newBuilder()
+      .expireAfterWrite(1, TimeUnit.DAYS)
+      .maximumSize(RATE_LIMITER_CACHE_MAX_SIZE).build();
+
+  public ConsumerAuthenticationFilter(ConsumerAuthUtil consumerAuthUtil, ConsumerAuditUtil consumerAuditUtil, PortalConfig portalConfig) {
     this.consumerAuthUtil = consumerAuthUtil;
     this.consumerAuditUtil = consumerAuditUtil;
+    this.portalConfig = portalConfig;
   }
 
   @Override
@@ -55,14 +74,33 @@ public void doFilter(ServletRequest req, ServletResponse resp, FilterChain chain
     HttpServletResponse response = (HttpServletResponse) resp;
 
     String token = request.getHeader(HttpHeaders.AUTHORIZATION);
+    ConsumerToken consumerToken = consumerAuthUtil.getConsumerToken(token);
 
-    Long consumerId = consumerAuthUtil.getConsumerId(token);
-
-    if (consumerId == null) {
+    if (null == consumerToken || consumerToken.getConsumerId() <= 0) {
       response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Unauthorized");
       return;
     }
 
+    Integer limitCount = consumerToken.getLimitCount();
+    if (limitCount == null) {
+      limitCount = 0;
+    }
+    if (portalConfig.isOpenApiLimitEnabled() && limitCount > 0) {
+      try {
+        ImmutablePair<Long, RateLimiter> rateLimiterPair = getOrCreateRateLimiterPair(token, limitCount);
+        long warmupToMillis = rateLimiterPair.getLeft() + WARMUP_MILLIS;
+        if (System.currentTimeMillis() > warmupToMillis && !rateLimiterPair.getRight().tryAcquire()) {
+          response.sendError(HttpServletResponse.SC_FORBIDDEN, "Too many call requests, the flow is limited");
+          return;
+        }
+      } catch (Exception e) {
+        logger.error("ConsumerAuthenticationFilter ratelimit error", e);
+        response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, "Rate limiting failed");
+        return;
+      }
+    }
+
+    long consumerId = consumerToken.getConsumerId();
     consumerAuthUtil.storeConsumerId(request, consumerId);
     consumerAuditUtil.audit(request, consumerId);
 
@@ -73,4 +111,14 @@ public void doFilter(ServletRequest req, ServletResponse resp, FilterChain chain
   public void destroy() {
     //nothing
   }
+
+  private ImmutablePair<Long, RateLimiter> getOrCreateRateLimiterPair(String key, Integer limitCount) {
+    ImmutablePair<Long, RateLimiter> rateLimiterPair = LIMITER.getIfPresent(key);
+    if (rateLimiterPair == null) {
+      rateLimiterPair = ImmutablePair.of(System.currentTimeMillis(), RateLimiter.create(limitCount));
+      LIMITER.put(key, rateLimiterPair);
+    }
+    return rateLimiterPair;
+  }
+
 }

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/service/ConsumerService.java

@@ -138,12 +138,15 @@ public ConsumerToken getConsumerTokenByAppId(String appId) {
     return consumerTokenRepository.findByConsumerId(consumer.getId());
   }
 
-  public Long getConsumerIdByToken(String token) {
+  public ConsumerToken getConsumerTokenByToken(String token) {
     if (Strings.isNullOrEmpty(token)) {
       return null;
     }
-    ConsumerToken consumerToken = consumerTokenRepository.findTopByTokenAndExpiresAfter(token,
-                                                                                        new Date());
+    return consumerTokenRepository.findTopByTokenAndExpiresAfter(token, new Date());
+  }
+
+  public Long getConsumerIdByToken(String token) {
+    ConsumerToken consumerToken = getConsumerTokenByToken(token);
     return consumerToken == null ? null : consumerToken.getConsumerId();
   }
 
@@ -311,7 +314,9 @@ public void createConsumerAudits(Iterable<ConsumerAudit> consumerAudits) {
   @Transactional
   public ConsumerToken createConsumerToken(ConsumerToken entity) {
     entity.setId(0); //for protection
-
+    if (entity.getLimitCount() <= 0) {
+      entity.setLimitCount(portalConfig.openApiLimitCount());
+    }
     return consumerTokenRepository.save(entity);
   }
 
@@ -322,6 +327,7 @@ private ConsumerToken generateConsumerToken(Consumer consumer, Date expires) {
 
     ConsumerToken consumerToken = new ConsumerToken();
     consumerToken.setConsumerId(consumerId);
+    consumerToken.setLimitCount(portalConfig.openApiLimitCount());
     consumerToken.setExpires(expires);
     consumerToken.setDataChangeCreatedBy(createdBy);
     consumerToken.setDataChangeCreatedTime(createdTime);

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/util/ConsumerAuthUtil.java

@@ -16,6 +16,7 @@
  */
 package com.ctrip.framework.apollo.openapi.util;
 
+import com.ctrip.framework.apollo.openapi.entity.ConsumerToken;
 import com.ctrip.framework.apollo.openapi.service.ConsumerService;
 import org.springframework.stereotype.Service;
 
@@ -37,6 +38,10 @@ public Long getConsumerId(String token) {
     return consumerService.getConsumerIdByToken(token);
   }
 
+  public ConsumerToken getConsumerToken(String token) {
+    return consumerService.getConsumerTokenByToken(token);
+  }
+
   public void storeConsumerId(HttpServletRequest request, Long consumerId) {
     request.setAttribute(CONSUMER_ID, consumerId);
   }

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/component/config/PortalConfig.java

@@ -242,6 +242,14 @@ public String consumerTokenSalt() {
     return getValue("consumer.token.salt", "apollo-portal");
   }
 
+  public int openApiLimitCount() {
+    return getIntProperty("open.api.limit.count", 20);
+  }
+
+  public boolean isOpenApiLimitEnabled() {
+    return getBooleanProperty("open.api.limit.enabled", false);
+  }
+
   public boolean isEmailEnabled() {
     return getBooleanProperty("email.enabled", false);
   }

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/spi/configuration/AuthFilterConfiguration.java

@@ -19,6 +19,7 @@
 import com.ctrip.framework.apollo.openapi.filter.ConsumerAuthenticationFilter;
 import com.ctrip.framework.apollo.openapi.util.ConsumerAuditUtil;
 import com.ctrip.framework.apollo.openapi.util.ConsumerAuthUtil;
+import com.ctrip.framework.apollo.portal.component.config.PortalConfig;
 import org.springframework.boot.web.servlet.FilterRegistrationBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -28,12 +29,13 @@ public class AuthFilterConfiguration {
 
   @Bean
   public FilterRegistrationBean<ConsumerAuthenticationFilter> openApiAuthenticationFilter(
-          ConsumerAuthUtil consumerAuthUtil,
-          ConsumerAuditUtil consumerAuditUtil) {
+      ConsumerAuthUtil consumerAuthUtil,
+      ConsumerAuditUtil consumerAuditUtil,
+      PortalConfig portalConfig) {
 
     FilterRegistrationBean<ConsumerAuthenticationFilter> openApiFilter = new FilterRegistrationBean<>();
 
-    openApiFilter.setFilter(new ConsumerAuthenticationFilter(consumerAuthUtil, consumerAuditUtil));
+    openApiFilter.setFilter(new ConsumerAuthenticationFilter(consumerAuthUtil, consumerAuditUtil, portalConfig));
     openApiFilter.addUrlPatterns("/openapi/*");
 
     return openApiFilter;

apollo-portal/src/test/java/com/ctrip/framework/apollo/SkipAuthorizationConfiguration.java

@@ -17,6 +17,7 @@
 package com.ctrip.framework.apollo;
 
 import com.ctrip.framework.apollo.openapi.auth.ConsumerPermissionValidator;
+import com.ctrip.framework.apollo.openapi.entity.ConsumerToken;
 import com.ctrip.framework.apollo.openapi.util.ConsumerAuthUtil;
 import com.ctrip.framework.apollo.portal.component.PermissionValidator;
 import org.springframework.context.annotation.Bean;
@@ -50,6 +51,11 @@ public ConsumerPermissionValidator consumerPermissionValidator() {
   public ConsumerAuthUtil consumerAuthUtil() {
     final ConsumerAuthUtil mock = mock(ConsumerAuthUtil.class);
     when(mock.getConsumerId(any())).thenReturn(1L);
+
+    ConsumerToken someConsumerToken = new ConsumerToken();
+    someConsumerToken.setConsumerId(1L);
+    someConsumerToken.setLimitCount(20);
+    when(mock.getConsumerToken(any())).thenReturn(someConsumerToken);
     return mock;
   }