From 4464e3d1419261a81d92490c283f7ef9a70542e6 Mon Sep 17 00:00:00 2001
From: Wenchao Gong <hanyedxing2012@163.com>
Date: Sun, 17 Oct 2021 16:53:37 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=86=E9=92=9F=E3=80=81=E5=B0=8F=E6=97=B6?=
 =?UTF-8?q?=E3=80=81=E5=A4=A9=E9=99=90=E6=B5=81=E8=BF=94=E5=9B=9E=E5=A4=B4?=
 =?UTF-8?q?=E4=BF=A1=E6=81=AF=E8=87=AA=E5=AE=9A=E4=B9=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../gateway/filter/CustomerGlobalFilter.java  |  1 -
 .../ratelimit/CustomerRedisRateLimiter.java   | 23 +++++++++++++------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java
index 7de9df5b6..b5b8b4fe2 100644
--- a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java
+++ b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/filter/CustomerGlobalFilter.java
@@ -233,7 +233,6 @@ public class CustomerGlobalFilter implements GlobalFilter, Ordered {
                     return rateLimiterAfterSupplier.get();
                 }
                 return rateLimiter(exchange, rule, route, index + 1, rateLimiterAfterSupplier);
-
             }
             setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
             return exchange.getResponse().setComplete();
diff --git a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java
index 6d3b298cb..5c2acaea7 100644
--- a/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java
+++ b/ruoyi-extend/ruoyi-isc-gateway/src/main/java/com/ruoyi/gateway/ratelimit/CustomerRedisRateLimiter.java
@@ -1,6 +1,7 @@
 package com.ruoyi.gateway.ratelimit;
 
 import cn.hutool.core.date.LocalDateTimeUtil;
+import cn.hutool.core.util.StrUtil;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
 import org.springframework.cloud.gateway.support.ConfigurationService;
@@ -14,9 +15,7 @@ import java.time.Instant;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
 import java.time.temporal.ChronoUnit;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -67,7 +66,7 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
             default:
                 return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
         }
-        return isAllowed(routeConfig, keys, time);
+        return isAllowed(routeConfig, keys, time, timeUnit);
     }
 
     public Mono<Response> isAllowed(Config routeConfig, String id) {
@@ -124,7 +123,7 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
         return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
     }
 
-    public Mono<Response> isAllowed(Config routeConfig, List<String> keys, int time) {
+    public Mono<Response> isAllowed(Config routeConfig, List<String> keys, int time, TimeUnit timeUnit) {
         if (!this.initialized.get()) {
             throw new IllegalStateException("RedisRateLimiter is not initialized");
         }
@@ -151,7 +150,7 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
                 return Mono.just(-1L);
             }).map(tokensLeft -> {
                 boolean allowed = tokensLeft < replenishRate;
-                Response response = new Response(allowed, getHeaders(routeConfig, tokensLeft));
+                Response response = new Response(allowed, getHeaders(routeConfig, tokensLeft, timeUnit));
                 if (log.isDebugEnabled()) {
                     log.debug("response: " + response);
                 }
@@ -166,7 +165,17 @@ public class CustomerRedisRateLimiter extends RedisRateLimiter {
              */
             log.error("Error determining if user allowed from redis", e);
         }
-        return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
+        return Mono.just(new Response(true, getHeaders(routeConfig, -1L, timeUnit)));
+    }
+
+    public Map<String, String> getHeaders(Config config, Long tokensLeft, TimeUnit timeUnit) {
+        Map<String, String> headers = new HashMap<>();
+        if (isIncludeHeaders()) {
+            String suffix = StrUtil.upperFirst(timeUnit.name().toLowerCase());
+            headers.put(this.getRemainingHeader() + '-' + suffix, tokensLeft.toString());
+            headers.put(this.getReplenishRateHeader() + '-' + suffix, String.valueOf(config.getReplenishRate()));
+        }
+        return headers;
     }
 
     static List<String> getKeys(String id) {