SpringCloudGateway自定义限流过滤器

一 背景

• spring cloud gateway自带到限流过滤器无法满足我们的需求。
• 1 RequestRateLimiterGatewayFilterFactory不支持指定uri
• 2 RequestRateLimiterGatewayFilterFactory不会具体到route
• 3 RequestRateLimiterGatewayFilterFactory不支持指定速率单位
• 4 RequestRateLimiterGatewayFilterFactory不支持指定响应body

  二 第一步-部分重写RedisRateLimiter

  import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
  import reactor.core.publisher.Mono;
  
  public interface CustomRateLimiter<C> extends RateLimiter<C> {
      Mono<RateLimiter.Response> isAllowed(CustomRequestRateLimiterGatewayFilterFactory.Config customConfig, String id);
  }

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.BeansException;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import org.springframework.context.ApplicationContext;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.util.DigestUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * redis限流扩展 - 解决一个路由配置多个Filter的问题
 *
 * @author mafgwo
 * @since 2020/02/06
 */
@ConfigurationProperties("spring.cloud.gateway.redis-rate-limiter")
public class CustomRedisRateLimiter extends RedisRateLimiter implements CustomRateLimiter<RedisRateLimiter.Config> {

    private Log log = LogFactory.getLog(getClass());

    private ReactiveStringRedisTemplate redisTemplate;

    private RedisScript<List<Long>> script;

    private AtomicBoolean initialized = new AtomicBoolean(false);

    private Integer defaultReplenishRate;
    private Integer defaultBurstCapacity;
    private Integer defaultPerSeconds;

    public CustomRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity, int defaultPerSeconds) {
        super(defaultReplenishRate, defaultBurstCapacity);
        this.defaultReplenishRate = defaultReplenishRate;
        this.defaultBurstCapacity = defaultBurstCapacity;
        this.defaultPerSeconds = defaultPerSeconds;
    }


    /**
     * Used when setting default configuration in constructor.
     *
     * @param context the ApplicationContext object to be used by this object
     * @throws BeansException if thrown by application context methods
     */
    @Override
    @SuppressWarnings("unchecked")
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        if (initialized.compareAndSet(false, true)) {
            if (this.redisTemplate == null) {
                this.redisTemplate = context.getBean(ReactiveStringRedisTemplate.class);
            }
            this.script = context.getBean(REDIS_SCRIPT_NAME, RedisScript.class);
            if (context.getBeanNamesForType(ConfigurationService.class).length > 0) {
                setConfigurationService(context.getBean(ConfigurationService.class));
            }
        }
    }

    static List<String> getKeys(String routeId, String id, Integer perSeconds) {
        // use `{}` around keys to use Redis Key hash tags
        // this allows for using redis cluster

        // Make a unique key per user.
        String idMd5 = DigestUtils.md5DigestAsHex((id + "_" + perSeconds).getBytes());
        String prefix = "cgi:" + routeId + ":{" + idMd5;

        // You need two Redis keys for Token Bucket.
        String tokenKey = prefix + "}:tokens";
        String timestampKey = prefix + "}:timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    public Map<String, String> getHeaders(Integer replenishRate, Integer burstCapacity, Long tokensLeft) {
        return super.getHeaders(new Config().setReplenishRate(replenishRate).setBurstCapacity(burstCapacity), tokensLeft);
    }

    @Override
    public Mono<Response> isAllowed(CustomRequestRateLimiterGatewayFilterFactory.Config customConfig, String id) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("CustomRedisRateLimiter is not initialized");
        }
        Integer rate = customConfig.getReplenishRate() != null ? customConfig.getReplenishRate() : defaultReplenishRate;
        Integer capacity = customConfig.getBurstCapacity() != null ? customConfig.getBurstCapacity() : defaultBurstCapacity;
        Integer perSeconds = customConfig.getPerSeconds() != null ? customConfig.getPerSeconds() : defaultPerSeconds;
        try {
            // 不同的填充耗时
            List<String> keys = getKeys(customConfig.getRouteId(), id, perSeconds);
            // The arguments to the LUA script. time() returns unixtime in seconds.
            List<String> scriptArgs = Arrays.asList(rate + "",
                capacity + "", Instant.now().getEpochSecond() + "", "1", perSeconds + "");
            // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
            Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys,
                scriptArgs);
            // .log("redisratelimiter", Level.FINER);
            return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
                .reduce(new ArrayList<Long>(), (longs, l) -> {
                    longs.addAll(l);
                    return longs;
                }).map(results -> {
                    boolean allowed = results.get(0) == 1L;
                    Long tokensLeft = results.get(1);

                    Response response = new Response(allowed,
                        getHeaders(rate, capacity, tokensLeft));

                    if (log.isDebugEnabled()) {
                        log.debug("response: " + response);
                    }
                    return response;
                });
        } catch (Exception e) {
            /*
             * We don't want a hard dependency on Redis to allow traffic. Make sure to set
             * an alert so you know if this is happening too much. Stripe's observed
             * failure rate is 0.01%.
             */
            log.error("Error determining if user allowed from redis", e);
        }
        return Mono.just(new Response(true, getHeaders(rate, capacity, -1L)));
    }
}

@Bean
public CustomRateLimiter rateLimiter() {
    CustomRedisRateLimiter rateLimiter = new CustomRedisRateLimiter(properties.getDefaultReplenishRate(), properties.getDefaultBurstCapacity(), properties.getDefaultPerSeconds());
    rateLimiter.setIncludeHeaders(false);
    return rateLimiter;
}

三 第二步-部分重写lua脚本，支持指定速率单位per_seconds

META-INF/scripts/request_rate_limiter.lua

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
local per_seconds = tonumber(ARGV[5])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*per_seconds*2)

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

local add_delta = math.floor((now-last_refreshed)/per_seconds)
local delta = math.max(0, add_delta)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

redis.call("setex", tokens_key, ttl, new_tokens)

if delta > 0 then
  redis.call("setex", timestamp_key, ttl, now)
else
  redis.call("setex", timestamp_key, ttl, last_refreshed)
end

return { allowed_num, new_tokens }

四 第三步-重写RequestRateLimiterGatewayFilterFactory

/**
 * @author mafgwo
 * @since 2020/02/04
 */
public class RateLImiterConstant {

    public static final String RATE_LIMITER_FALLBACK_URI = "/ratelimiter-fallback";

    public static final String RATE_LIMITER_FALLBACK_PATH = "forward:" + RATE_LIMITER_FALLBACK_URI + "?httpStatus=%s&code=%s&message=%s";

}

import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;

import java.net.URI;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.containsEncodedParts;

/**
 * @author mafgwo
 * @since 2020/02/06
 */
public class FailResponseUtil {


    public static ServerWebExchange processResponse(ServerWebExchange exchange, HttpStatus statusCode, String errorCode, String errorMsg) {
        URI uri = exchange.getRequest().getURI();
        boolean encoded = containsEncodedParts(uri);

        String path = String.format(RateLImiterConstant.RATE_LIMITER_FALLBACK_PATH, statusCode.name(),
            errorCode == null ? "" : errorCode,
            errorMsg == null ? "" : errorMsg);

        URI requestUrl = UriComponentsBuilder.fromUri(uri).host(null).port(null)
            .uri(URI.create(path)).scheme(null).build(encoded).toUri();
        exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl);

        ServerHttpRequest request = exchange.getRequest().mutate()
            .uri(requestUrl).build();
        return exchange.mutate().request(request).build();
    }
}

import lombok.Data;
import org.springframework.util.StringUtils;

/**
 * 根据业务需要定义响应body
 * @author mafgwo
 * @since 2020/02/04
 **/
@Data
public class FallbackResponse {
    public static final String DEFAULT_ERROR_CODE = "CGI0000001";
    public static final String DEFAULT_ERROR_MESSAGE = "服务繁忙，请稍后再试";

    public static final FallbackResponse DEFAULT = new FallbackResponse();

    private String code;
    private String message;

    public FallbackResponse() {
        this(null, null);
    }

    public FallbackResponse(String code, String message) {
        this.code = code;
        this.message = message;
        if (StringUtils.isEmpty(code)) {
            this.code = DEFAULT_ERROR_CODE;
        }

        if (StringUtils.isEmpty(message)) {
            this.message = DEFAULT_ERROR_MESSAGE;
        }
    }
}

import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;

/**
 * @author mafgwo
 * @since 2020/02/04
 **/
@RestController
public class FallbackController {

    @GetMapping(value = RateLImiterConstant.RATE_LIMITER_FALLBACK_URI)
    public Mono<ResponseEntity<FallbackResponse>> fallback(@RequestParam HttpStatus httpStatus,
                                                           @RequestParam(required = false) String code,
                                                           @RequestParam(required = false) String message) {
        FallbackResponse fallbackResponse = new FallbackResponse(code, message);
        ResponseEntity<FallbackResponse> body = ResponseEntity.status(httpStatus).body(fallbackResponse);
        return Mono.just(body);
    }
}

package com.xx.mall.gateway.ratelimiter;

import lombok.Data;
import lombok.experimental.Accessors;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HasRouteId;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.util.CollectionUtils;
import org.springframework.web.reactive.DispatcherHandler;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;

/**
 * 从RequestRateLimiterGatewayFilterFactory扩展
 *
 * @author mafgwo
 * @since 2020/02/03
 */
@ConfigurationProperties("spring.cloud.gateway.filter.request-rate-limiter")
public class CustomRequestRateLimiterGatewayFilterFactory extends
    AbstractGatewayFilterFactory<CustomRequestRateLimiterGatewayFilterFactory.Config> {

    private final ObjectProvider<DispatcherHandler> dispatcherHandlerProvider;

    // do not use this dispatcherHandler directly, use getDispatcherHandler() instead.
    private volatile DispatcherHandler dispatcherHandler;

    private static final String EMPTY_KEY = "____EMPTY_KEY__";

    private final CustomRateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    /**
     * Switch to deny requests if the Key Resolver returns an empty key, defaults to true.
     */
    private boolean denyEmptyKey = true;

    /**
     * HttpStatus to return when denyEmptyKey is true, defaults to FORBIDDEN.
     */
    private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();


    public CustomRequestRateLimiterGatewayFilterFactory(CustomRateLimiter defaultRateLimiter,
                                                    KeyResolver defaultKeyResolver,
                                                    ObjectProvider<DispatcherHandler> dispatcherHandlerProvider) {
        super(Config.class);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
        this.dispatcherHandlerProvider = dispatcherHandlerProvider;
    }

    private DispatcherHandler getDispatcherHandler() {
        if (dispatcherHandler == null) {
            dispatcherHandler = dispatcherHandlerProvider.getIfAvailable();
        }

        return dispatcherHandler;
    }

    public KeyResolver getDefaultKeyResolver() {
        return defaultKeyResolver;
    }

    public RateLimiter getDefaultRateLimiter() {
        return defaultRateLimiter;
    }

    public boolean isDenyEmptyKey() {
        return denyEmptyKey;
    }

    public void setDenyEmptyKey(boolean denyEmptyKey) {
        this.denyEmptyKey = denyEmptyKey;
    }

    public String getEmptyKeyStatusCode() {
        return emptyKeyStatusCode;
    }

    public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
        this.emptyKeyStatusCode = emptyKeyStatusCode;
    }

    @SuppressWarnings("unchecked")
    @Override
    public GatewayFilter apply(Config config) {
        KeyResolver resolver = getOrDefault(config.keyResolver, defaultKeyResolver);
        CustomRateLimiter<Object> limiter = getOrDefault(config.rateLimiter,
            defaultRateLimiter);
        boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder
            .parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));

        return (exchange, chain) -> {
            // allow path exists and allow path not contain current path
            if (!CollectionUtils.isEmpty(config.getLimitPaths()) && config.getLimitPaths().stream().noneMatch(item -> exchange.getRequest().getPath().value().matches(item))) {
                return chain.filter(exchange);
            }
            return resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY)
                .flatMap(key -> {
                    if (EMPTY_KEY.equals(key)) {
                        if (denyEmpty) {
                            return this.processResponse(exchange, emptyKeyStatus.getHttpStatus(), config);
                        }
                        return chain.filter(exchange);
                    }
                    String routeId = config.getRouteId();
                    if (routeId == null) {
                        Route route = exchange
                            .getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                        config.setRouteId(route.getId());
                    }
                    return limiter.isAllowed(config, key).flatMap(response -> {

                        for (Map.Entry<String, String> header : response.getHeaders()
                            .entrySet()) {
                            exchange.getResponse().getHeaders().add(header.getKey(),
                                header.getValue());
                        }

                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        }
                        return this.processResponse(exchange, config.getStatusCode(), config);
                    });
                });
        };
    }

    private Mono<Void> processResponse(ServerWebExchange exchange, HttpStatus statusCode, Config config) {
        ServerWebExchange mutated = FailResponseUtil.processResponse(exchange, statusCode, config.getErrorCode(), config.getErrorMsg());
        return getDispatcherHandler().handle(mutated);
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }

    @Accessors(chain = true)
    @Data
    public static class Config implements HasRouteId {

        private KeyResolver keyResolver;

        private List<String> limitPaths;

        private CustomRateLimiter rateLimiter;

        private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;

        private String errorCode;

        private String errorMsg;

        private Boolean denyEmptyKey;

        private String emptyKeyStatus;

        private String routeId;

        private Integer replenishRate;
        private Integer burstCapacity;
        private Integer perSeconds;

        @Override
        public void setRouteId(String routeId) {
            this.routeId = routeId;
        }

        @Override
        public String getRouteId() {
            return this.routeId;
        }
    }
}

@Bean
public CustomerRequestRateLimiterGatewayFilterFactory customerRequestRateLimiterGatewayFilterFactory(CustomerRateLimiter rateLimiter, KeyResolver principalKeyResolver, ObjectProvider<DispatcherHandler> dispatcherHandlerProvider) {
    return new CustomerRequestRateLimiterGatewayFilterFactory(rateLimiter, principalKeyResolver, dispatcherHandlerProvider);
}

五 route网关配置示例

- id: products
  uri: http://xxxx.xxxx.com
  predicates:
  - Path=/v2/products/**
  filters:
  - name: CustomerRequestRateLimiter
    args:
      key-resolver: '#{@apiTokenKeyResolver}'
      replenishRate: 1
      burstCapacity: 3
      error-code: TEST0001002
      error-msg: 单URI单个用户请求过于频繁，请稍后重试...
  - name: CustomerRequestRateLimiter
    args:
      key-resolver: '#{@principalKeyResolver}'
      limit-paths: /v2/products.*
      replenishRate: 1
      burstCapacity: 1
      per-seconds: 10
      error-code: TEST0001003
      error-msg: 请求过于频繁10，请稍后重试...
  - name: CustomerRequestRateLimiter
    args:
      key-resolver: '#{@principalKeyResolver}'
      limit-paths: /v2/aaaa.*
      replenishRate: 5
      burstCapacity: 5
      per-seconds: 10
      error-code: TEST0001004
      error-msg: 请求过于频繁15，请稍后重试...