Feature/openapi rate limit function

CHANGES.md

@@ -20,6 +20,7 @@ Apollo 2.4.0
 * [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)
 * [Fix: ensure clusters order in envClusters open api](https://github.com/apolloconfig/apollo/pull/5277)
+* [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-common/src/main/java/com/ctrip/framework/apollo/common/exception/BadRequestException.java

@@ -37,6 +37,10 @@ public static BadRequestException orgIdIsBlank() {
     return new BadRequestException("orgId can not be blank");
   }
 
+  public static BadRequestException rateLimitIsInvalid() {
+    return new BadRequestException("rate limit must be greater than 1");
+  }
+
   public static BadRequestException itemAlreadyExists(String itemKey) {
     return new BadRequestException("item already exists for itemKey:%s", itemKey);
   }

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

@@ -18,6 +18,7 @@
 
 import com.ctrip.framework.apollo.common.entity.BaseEntity;
 
+import javax.validation.constraints.PositiveOrZero;
 import org.hibernate.annotations.SQLDelete;
 import org.hibernate.annotations.Where;
 
@@ -41,6 +42,10 @@ public class ConsumerToken extends BaseEntity {
   @Column(name = "`Token`", nullable = false)
   private String token;
 
+  @PositiveOrZero
+  @Column(name = "`RateLimit`", nullable = false)
+  private Integer rateLimit;
+
   @Column(name = "`Expires`", nullable = false)
   private Date expires;
 
@@ -60,6 +65,14 @@ public void setToken(String token) {
     this.token = token;
   }
 
+  public Integer getRateLimit() {
+    return rateLimit;
+  }
+
+  public void setRateLimit(Integer rateLimit) {
+    this.rateLimit = rateLimit;
+  }
+
   public Date getExpires() {
     return expires;
   }
@@ -71,6 +84,7 @@ public void setExpires(Date expires) {
   @Override
   public String toString() {
     return toStringHelper().add("consumerId", consumerId).add("token", token)
+        .add("rateLimit", rateLimit)
         .add("expires", expires).toString();
   }
 }

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

@@ -16,11 +16,16 @@
  */
 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.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import javax.servlet.Filter;
 import javax.servlet.FilterChain;
 import javax.servlet.FilterConfig;
@@ -29,18 +34,35 @@
 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(song_s@ctrip.com)
  */
 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;
+
+  private static final int WARMUP_MILLIS = 1000; // ms
+  private static final int RATE_LIMITER_CACHE_MAX_SIZE = 50000;
+
+  private static final int TOO_MANY_REQUESTS = 429;
+
+  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) {
+  public ConsumerAuthenticationFilter(ConsumerAuthUtil consumerAuthUtil, ConsumerAuditUtil consumerAuditUtil, PortalConfig portalConfig) {
     this.consumerAuthUtil = consumerAuthUtil;
     this.consumerAuditUtil = consumerAuditUtil;
+    this.portalConfig = portalConfig;
   }
 
   @Override
@@ -55,14 +77,30 @@ 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) {
       response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Unauthorized");
       return;
     }
 
+    Integer rateLimit = consumerToken.getRateLimit();
+    if (null != rateLimit && rateLimit > 0) {
+      try {
+        ImmutablePair<Long, RateLimiter> rateLimiterPair = getOrCreateRateLimiterPair(token, rateLimit);
+        long warmupToMillis = rateLimiterPair.getLeft() + WARMUP_MILLIS;
+        if (System.currentTimeMillis() > warmupToMillis && !rateLimiterPair.getRight().tryAcquire()) {
+          response.sendError(TOO_MANY_REQUESTS, "Too Many 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) {
+    try{
+      return LIMITER.get(key, () ->
+          ImmutablePair.of(System.currentTimeMillis(), RateLimiter.create(limitCount)));
+    } catch (ExecutionException e) {
+      throw new RuntimeException("Failed to create rate limiter", e);
+    }
+  }
+
 }

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

@@ -120,10 +120,11 @@ public Consumer createConsumer(Consumer consumer) {
     return consumerRepository.save(consumer);
   }
 
-  public ConsumerToken generateAndSaveConsumerToken(Consumer consumer, Date expires) {
+  public ConsumerToken generateAndSaveConsumerToken(Consumer consumer, Integer rateLimit, Date expires) {
     Preconditions.checkArgument(consumer != null, "Consumer can not be null");
+    Preconditions.checkArgument(rateLimit != null && rateLimit >= 0, "Rate limit must be non-negative");
 
-    ConsumerToken consumerToken = generateConsumerToken(consumer, expires);
+    ConsumerToken consumerToken = generateConsumerToken(consumer, rateLimit, expires);
     consumerToken.setId(0);
 
     return consumerTokenRepository.save(consumerToken);
@@ -138,12 +139,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,17 +315,21 @@ public void createConsumerAudits(Iterable<ConsumerAudit> consumerAudits) {
   @Transactional
   public ConsumerToken createConsumerToken(ConsumerToken entity) {
     entity.setId(0); //for protection
-
     return consumerTokenRepository.save(entity);
   }
 
-  private ConsumerToken generateConsumerToken(Consumer consumer, Date expires) {
+  private ConsumerToken generateConsumerToken(Consumer consumer, Integer rateLimit, Date expires) {
     long consumerId = consumer.getId();
     String createdBy = userInfoHolder.getUser().getUserId();
     Date createdTime = new Date();
 
+    if (rateLimit == null || rateLimit < 0) {
+      rateLimit = 0;
+    }
+
     ConsumerToken consumerToken = new ConsumerToken();
     consumerToken.setConsumerId(consumerId);
+    consumerToken.setRateLimit(rateLimit);
     consumerToken.setExpires(expires);
     consumerToken.setDataChangeCreatedBy(createdBy);
     consumerToken.setDataChangeCreatedTime(createdTime);
@@ -350,7 +358,7 @@ String generateToken(String consumerAppId, Date generationTime, String consumerT
         (generationTime), consumerTokenSalt), Charsets.UTF_8).toString();
   }
 
-    ConsumerRole createConsumerRole(Long consumerId, Long roleId, String operator) {
+  ConsumerRole createConsumerRole(Long consumerId, Long roleId, String operator) {
     ConsumerRole consumerRole = new ConsumerRole();
 
     consumerRole.setConsumerId(consumerId);
@@ -389,7 +397,7 @@ private Set<String> findAppIdsByRoleIds(List<Long> roleIds) {
     return appIds;
   }
 
-  List<Consumer> findAllConsumer(Pageable page){
+  List<Consumer> findAllConsumer(Pageable page) {
     return this.consumerRepository.findAll(page).getContent();
   }
 
@@ -414,7 +422,7 @@ public List<ConsumerInfo> findConsumerInfoList(Pageable page) {
   }
 
   @Transactional
-  public void deleteConsumer(String appId){
+  public void deleteConsumer(String appId) {
     Consumer consumer = consumerRepository.findByAppId(appId);
     if (consumer == null) {
       throw new BadRequestException("ConsumerApp not exist");

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/controller/ConsumerController.java

@@ -81,13 +81,21 @@ public ConsumerInfo create(
       throw BadRequestException.orgIdIsBlank();
     }
 
+    if (requestVO.isRateLimitEnabled()) {
+      if (requestVO.getRateLimit() <= 0) {
+        throw BadRequestException.rateLimitIsInvalid();
+      }
+    } else {
+      requestVO.setRateLimit(0);
+    }
+
     Consumer createdConsumer = consumerService.createConsumer(convertToConsumer(requestVO));
 
     if (Objects.isNull(expires)) {
       expires = DEFAULT_EXPIRES;
     }
 
-    ConsumerToken consumerToken = consumerService.generateAndSaveConsumerToken(createdConsumer, expires);
+    ConsumerToken consumerToken = consumerService.generateAndSaveConsumerToken(createdConsumer, requestVO.getRateLimit(), expires);
     if (requestVO.isAllowCreateApplication()) {
       consumerService.assignCreateApplicationRoleToConsumer(consumerToken.getToken());
     }
@@ -127,7 +135,7 @@ public List<ConsumerRole> assignNamespaceRoleToConsumer(
     if (StringUtils.isEmpty(namespaceName)) {
       throw new BadRequestException("Params(NamespaceName) can not be empty.");
     }
-    if (null != envs){
+    if (null != envs) {
       String[] envArray = envs.split(",");
       List<String> envList = Lists.newArrayList();
       // validate env parameter
@@ -156,7 +164,7 @@ public List<ConsumerRole> assignNamespaceRoleToConsumer(
 
   @GetMapping("/consumers")
   @PreAuthorize(value = "@permissionValidator.isSuperAdmin()")
-  public List<ConsumerInfo> getConsumerList(Pageable page){
+  public List<ConsumerInfo> getConsumerList(Pageable page) {
     return consumerService.findConsumerInfoList(page);
   }
 

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/entity/vo/consumer/ConsumerCreateRequestVO.java

@@ -26,6 +26,8 @@ public class ConsumerCreateRequestVO {
   private String orgId;
   private String orgName;
   private String ownerName;
+  private boolean rateLimitEnabled;
+  private int rateLimit;
 
   public String getAppId() {
     return appId;
@@ -75,4 +77,20 @@ public void setOwnerName(String ownerName) {
     this.ownerName = ownerName;
   }
 
+  public boolean isRateLimitEnabled() {
+    return rateLimitEnabled;
+  }
+
+  public void setRateLimitEnabled(boolean rateLimitEnabled) {
+    this.rateLimitEnabled = rateLimitEnabled;
+  }
+
+  public int getRateLimit() {
+    return rateLimit;
+  }
+
+  public void setRateLimit(int rateLimit) {
+    this.rateLimit = rateLimit;
+  }
+
 }

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/main/resources/static/i18n/en.json

@@ -654,6 +654,11 @@
   "Open.Manage.Consumer.AllowCreateApplicationTips": "（Allow third-party applications to create apps and grant them app administrator privileges.",
   "Open.Manage.Consumer.AllowCreateApplication.No": "no",
   "Open.Manage.Consumer.AllowCreateApplication.Yes": "yes",
+  "Open.Manage.Consumer.RateLimit.Enabled": "Whether to enable current limit",
+  "Open.Manage.Consumer.RateLimit.Enabled.Tips": "(After enabling this feature, when third-party applications publish configurations on Apollo, their traffic will be controlled according to the configured QPS limit)",
+  "Open.Manage.Consumer.RateLimitValue": "Current limiting QPS",
+  "Open.Manage.Consumer.RateLimitValueTips": "(Unit: times/second, for example: 100 means that the configuration is published at most 100 times per second)",
+  "Open.Manage.Consumer.RateLimitValue.Error": "The minimum current limiting QPS is 1",
   "Namespace.Role.Title": "Permission Management",
   "Namespace.Role.GrantModifyTo": "Permission to edit",
   "Namespace.Role.GrantModifyTo2": "(Can edit the configuration)",

apollo-portal/src/main/resources/static/i18n/zh-CN.json

@@ -654,6 +654,11 @@
   "Open.Manage.Consumer.AllowCreateApplicationTips": "（允许第三方应用创建app，并且对创建出的app，拥有应用管理员的权限）",
   "Open.Manage.Consumer.AllowCreateApplication.No": "否",
   "Open.Manage.Consumer.AllowCreateApplication.Yes": "是",
+  "Open.Manage.Consumer.RateLimit.Enabled": "是否启用限流",
+  "Open.Manage.Consumer.RateLimit.Enabled.Tips": "（开启后，第三方应用在 Apollo 上发布配置时，会根据配置的 QPS 限制，控制其流量）",
+  "Open.Manage.Consumer.RateLimitValue": "限流QPS",
+  "Open.Manage.Consumer.RateLimitValueTips": "（单位:次/秒，例如: 100 表示每秒最多发布 100 次配置）",
+  "Open.Manage.Consumer.RateLimitValue.Error": "限流QPS最小为1",
   "Namespace.Role.Title": "权限管理",
   "Namespace.Role.GrantModifyTo": "修改权",
   "Namespace.Role.GrantModifyTo2": "（可以修改配置）",