Feature/openapi rate limit function

CHANGES.md

@@ -23,5 +23,6 @@ Apollo 2.4.0
 * [Fix: bump xstream from 1.4.20 to 1.4.21 to fix CVE-2024-47072](https://github.com/apolloconfig/apollo/pull/5280)
 * [Feature: highlight diffs for properties](https://github.com/apolloconfig/apollo/pull/5282)
 
+* [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 @@
     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 @@
   @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 @@
     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 @@
   @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,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.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,15 +33,30 @@
 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 static final int WARMUP_MILLIS = 1000; // ms
+  private static final int RATE_LIMITER_CACHE_MAX_SIZE = 20000;
+
+  private static final int TOO_MANY_REQUESTS = 429;
+
+  private static final Cache<String, ImmutablePair<Long, RateLimiter>> LIMITER = CacheBuilder.newBuilder()
+      .expireAfterAccess(1, TimeUnit.HOURS)
+      .maximumSize(RATE_LIMITER_CACHE_MAX_SIZE).build();
+
   public ConsumerAuthenticationFilter(ConsumerAuthUtil consumerAuthUtil, ConsumerAuditUtil consumerAuditUtil) {
     this.consumerAuthUtil = consumerAuthUtil;
     this.consumerAuditUtil = consumerAuditUtil;
@@ -55,14 +74,30 @@
     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(consumerToken.getToken(), 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 +108,14 @@
   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/repository/ConsumerTokenRepository.java

@@ -18,6 +18,7 @@
 
 import com.ctrip.framework.apollo.openapi.entity.ConsumerToken;
 
+import java.util.List;
 import org.springframework.data.repository.PagingAndSortingRepository;
 
 import java.util.Date;
@@ -35,4 +36,7 @@ public interface ConsumerTokenRepository extends PagingAndSortingRepository<Cons
   ConsumerToken findTopByTokenAndExpiresAfter(String token, Date validDate);
 
   ConsumerToken findByConsumerId(Long consumerId);
+
+  List<ConsumerToken> findByConsumerIdIn(List<Long> consumerIds);
+
 }

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

@@ -44,6 +44,7 @@
 import com.google.common.hash.Hashing;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
 import org.apache.commons.lang3.time.FastDateFormat;
 import org.springframework.data.domain.Pageable;
@@ -56,6 +57,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import org.springframework.util.CollectionUtils;
 
 /**
  * @author Jason Song([email protected])
@@ -120,10 +122,10 @@
     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");
 
-    ConsumerToken consumerToken = generateConsumerToken(consumer, expires);
+    ConsumerToken consumerToken = generateConsumerToken(consumer, rateLimit, expires);
     consumerToken.setId(0);
 
     return consumerTokenRepository.save(consumerToken);
@@ -138,12 +140,15 @@
     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();
   }
 
@@ -195,7 +200,8 @@
   private ConsumerInfo convert(
       Consumer consumer,
       String token,
-      boolean allowCreateApplication
+      boolean allowCreateApplication,
+      Integer rateLimit
   ) {
     ConsumerInfo consumerInfo = new ConsumerInfo();
     consumerInfo.setConsumerId(consumer.getId());
@@ -205,6 +211,7 @@
     consumerInfo.setOwnerEmail(consumer.getOwnerEmail());
     consumerInfo.setOrgId(consumer.getOrgId());
     consumerInfo.setOrgName(consumer.getOrgName());
+    consumerInfo.setRateLimit(rateLimit);
 
     consumerInfo.setToken(token);
     consumerInfo.setAllowCreateApplication(allowCreateApplication);
@@ -220,13 +227,21 @@
     if (consumer == null) {
       return null;
     }
-    return convert(consumer, consumerToken.getToken(), isAllowCreateApplication(consumer.getId()));
+    return convert(consumer, consumerToken.getToken(), isAllowCreateApplication(consumer.getId()), getRateLimit(consumer.getId()));
   }
 
   private boolean isAllowCreateApplication(Long consumerId) {
     return isAllowCreateApplication(Collections.singletonList(consumerId)).get(0);
   }
 
+  private Integer getRateLimit(Long consumerId) {
+    List<Integer> list = getRateLimit(Collections.singletonList(consumerId));
+    if (CollectionUtils.isEmpty(list)) {
+      return 0;
+    }
+    return list.get(0);
+  }
+
   private List<Boolean> isAllowCreateApplication(List<Long> consumerIdList) {
     Role createAppRole = getCreateAppRole();
     if (createAppRole == null) {
@@ -249,6 +264,19 @@
     return list;
   }
 
+  private List<Integer> getRateLimit(List<Long> consumerIds) {
+    List<ConsumerToken> consumerTokens = consumerTokenRepository.findByConsumerIdIn(consumerIds);
+    Map<Long, Integer> consumerRateLimits = consumerTokens.stream()
+        .collect(Collectors.toMap(
+            ConsumerToken::getConsumerId,
+            consumerToken -> consumerToken.getRateLimit() != null ? consumerToken.getRateLimit() : 0
+        ));
+
+    return consumerIds.stream()
+        .map(id -> consumerRateLimits.getOrDefault(id, 0))
+        .collect(Collectors.toList());
+  }
+
   private Role getCreateAppRole() {
     return rolePermissionService.findRoleByRoleName(CREATE_APPLICATION_ROLE_NAME);
   }
@@ -311,17 +339,21 @@
   @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 +382,7 @@
         (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 +421,7 @@
     return appIds;
   }
 
-  List<Consumer> findAllConsumer(Pageable page){
+  List<Consumer> findAllConsumer(Pageable page) {
     return this.consumerRepository.findAll(page).getContent();
   }
 
@@ -398,14 +430,15 @@
     List<Long> consumerIdList = consumerList.stream()
         .map(Consumer::getId).collect(Collectors.toList());
     List<Boolean> allowCreateApplicationList = isAllowCreateApplication(consumerIdList);
+    List<Integer> rateLimitList = getRateLimit(consumerIdList);
 
     List<ConsumerInfo> consumerInfoList = new ArrayList<>(consumerList.size());
 
     for (int i = 0; i < consumerList.size(); i++) {
       Consumer consumer = consumerList.get(i);
       // without token
       ConsumerInfo consumerInfo = convert(
-          consumer, null, allowCreateApplicationList.get(i)
+          consumer, null, allowCreateApplicationList.get(i), rateLimitList.get(i)
       );
       consumerInfoList.add(consumerInfo);
     }
@@ -414,7 +447,7 @@
   }
 
   @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);
   }