Feature/openapi rate limit function

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,9 +16,13 @@
  */
 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 javax.servlet.Filter;
@@ -29,18 +33,29 @@
 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 Cache<String, ImmutablePair<Long, RateLimiter>> LIMITER = CacheBuilder.newBuilder().build();
+  private static final int WARMUP_MILLIS = 1000; // ms
+
+  public ConsumerAuthenticationFilter(ConsumerAuthUtil consumerAuthUtil, ConsumerAuditUtil consumerAuditUtil, PortalConfig portalConfig) {
     this.consumerAuthUtil = consumerAuthUtil;
     this.consumerAuditUtil = consumerAuditUtil;
+    this.portalConfig = portalConfig;
   }
 
   @Override
@@ -55,14 +70,28 @@ 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 (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);
+      }
+    }
+
+    long consumerId = consumerToken.getConsumerId();
     consumerAuthUtil.storeConsumerId(request, consumerId);
     consumerAuditUtil.audit(request, consumerId);
 
@@ -73,4 +102,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/openapi/filter/ConsumerAuthenticationFilterTest.java

@@ -16,9 +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 java.io.IOException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import javax.servlet.ServletException;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -33,6 +40,9 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.atMost;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -48,6 +58,9 @@ public class ConsumerAuthenticationFilterTest {
   private ConsumerAuthUtil consumerAuthUtil;
   @Mock
   private ConsumerAuditUtil consumerAuditUtil;
+  @Mock
+  private PortalConfig portalConfig;
+
   @Mock
   private HttpServletRequest request;
   @Mock
@@ -57,7 +70,7 @@ public class ConsumerAuthenticationFilterTest {
 
   @Before
   public void setUp() throws Exception {
-    authenticationFilter = new ConsumerAuthenticationFilter(consumerAuthUtil, consumerAuditUtil);
+    authenticationFilter = new ConsumerAuthenticationFilter(consumerAuthUtil, consumerAuditUtil, portalConfig);
   }
 
   @Test
@@ -89,4 +102,101 @@ public void testAuthFailed() throws Exception {
     verify(consumerAuditUtil, never()).audit(eq(request), anyLong());
     verify(filterChain, never()).doFilter(request, response);
   }
+
+
+  @Test
+  public void testRateLimitSuccessfully() throws Exception {
+    String someToken = "someToken";
+    Long someConsumerId = 1L;
+    int qps = 5;
+    int durationInSeconds = 10;
+
+    setupRateLimitMocks(someToken, someConsumerId, qps);
+
+    Runnable task = () -> {
+      try {
+        authenticationFilter.doFilter(request, response, filterChain);
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      } catch (ServletException e) {
+        throw new RuntimeException(e);
+      }
+    };
+
+    executeWithQps(qps, task, durationInSeconds);
+
+    int total = qps * durationInSeconds;
+
+    verify(consumerAuthUtil, times(total)).storeConsumerId(request, someConsumerId);
+    verify(consumerAuditUtil, times(total)).audit(request, someConsumerId);
+    verify(filterChain, times(total)).doFilter(request, response);
+
+  }
+
+
+  @Test
+  public void testRateLimitPartFailure() throws Exception {
+    String someToken = "someToken";
+    Long someConsumerId = 1L;
+    int qps = 5;
+    int durationInSeconds = 10;
+
+    setupRateLimitMocks(someToken, someConsumerId, qps);
+
+    Runnable task = () -> {
+      try {
+        authenticationFilter.doFilter(request, response, filterChain);
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      } catch (ServletException e) {
+        throw new RuntimeException(e);
+      }
+    };
+
+    executeWithQps(qps + 1, task, durationInSeconds);
+
+    int leastTimes = qps * durationInSeconds;
+    int mostTimes = (qps + 1) * durationInSeconds;
+
+    verify(response, atLeastOnce()).sendError(eq(HttpServletResponse.SC_FORBIDDEN), anyString());
+
+    verify(consumerAuthUtil, atLeast(leastTimes)).storeConsumerId(request, someConsumerId);
+    verify(consumerAuthUtil, atMost(mostTimes)).storeConsumerId(request, someConsumerId);
+    verify(consumerAuditUtil, atLeast(leastTimes)).audit(request, someConsumerId);
+    verify(consumerAuditUtil, atMost(mostTimes)).audit(request, someConsumerId);
+    verify(filterChain, atLeast(leastTimes)).doFilter(request, response);
+    verify(filterChain, atMost(mostTimes)).doFilter(request, response);
+
+  }
+
+
+  private void setupRateLimitMocks(String someToken, Long someConsumerId, int qps) {
+    ConsumerToken someConsumerToken = new ConsumerToken();
+    someConsumerToken.setConsumerId(someConsumerId);
+    someConsumerToken.setLimitCount(qps);
+
+    when(request.getHeader(HttpHeaders.AUTHORIZATION)).thenReturn(someToken);
+    when(consumerAuthUtil.getConsumerId(someToken)).thenReturn(someConsumerId);
+    when(consumerAuthUtil.getConsumerToken(someToken)).thenReturn(someConsumerToken);
+    when(portalConfig.isOpenApiLimitEnabled()).thenReturn(true);
+  }
+
+
+  public static void executeWithQps(int qps, Runnable task, int durationInSeconds) {
+    ExecutorService executor = Executors.newFixedThreadPool(qps);
+    long totalTasks = qps * durationInSeconds;
+
+    for (int i = 0; i < totalTasks; i++) {
+      executor.submit(task);
+      try {
+        TimeUnit.MILLISECONDS.sleep(1000 / qps); // Control QPS
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+    }
+
+    executor.shutdown();
+  }
+
 }

scripts/sql/profiles/h2-default/apolloportaldb.sql

@@ -161,6 +161,7 @@ CREATE TABLE `ConsumerToken` (
   `Id` int(11) unsigned NOT NULL AUTO_INCREMENT COMMENT '自增Id',
   `ConsumerId` int(11) unsigned DEFAULT NULL COMMENT 'ConsumerId',
   `Token` varchar(128) NOT NULL DEFAULT '' COMMENT 'token',
+  `LimitCount` int NOT NULL DEFAULT 20 COMMENT '限流值',
   `Expires` datetime NOT NULL DEFAULT '2099-01-01 00:00:00' COMMENT 'token失效时间',
   `IsDeleted` boolean NOT NULL DEFAULT FALSE COMMENT '1: deleted, 0: normal',
   `DeletedAt` BIGINT(20) NOT NULL DEFAULT '0' COMMENT 'Delete timestamp based on milliseconds',