Custom hide or show header in response (#150)

#### Changes proposed in this pull request: - custom hide or show header in the response
marcosbarbero · Feb 26, 2019 · 9104dcc · 9104dcc
1 parent bb038ef
commit 9104dcc
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/README.adoc b/README.adoc
@@ -148,6 +148,7 @@ zuul:
     enabled: true 
     repository: REDIS
     behind-proxy: true
+    add-response-headers: true
     default-policy-list: #optional - will apply unless specific policy exists
       - limit: 10 #optional - request number limit per refresh interval window
         quota: 1000 #optional - request time limit per refresh interval window (in seconds)
@@ -216,6 +217,7 @@ Property namespace: __zuul.ratelimit__
 
 |enabled             |true/false                   |false
 |behind-proxy        |true/false                   |false
+|add-response-headers|true/false                   |true
 |key-prefix          |String                       |${spring.application.name:rate-limit-application}
 |repository          |CONSUL, REDIS, JPA, BUCKET4J_JCACHE, BUCKET4J_HAZELCAST, BUCKET4J_INFINISPAN, BUCKET4J_IGNITE| -
 |default-policy-list |List of link:./spring-cloud-zuul-ratelimit-core/src/main/java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/config/properties/RateLimitProperties.java#L82[Policy]| -

diff --git a/...rcosbarbero/cloud/autoconfigure/zuul/ratelimit/config/properties/RateLimitProperties.java b/...rcosbarbero/cloud/autoconfigure/zuul/ratelimit/config/properties/RateLimitProperties.java
@@ -62,6 +62,7 @@ public class RateLimitProperties {
     private Map<String, List<Policy>> policyList = Maps.newHashMap();
     private boolean behindProxy;
     private boolean enabled;
+    private boolean addResponseHeaders = true;
     @NotNull
     @Value("${spring.application.name:rate-limit-application}")
     private String keyPrefix;

diff --git a/...java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/RateLimitPreFilter.java b/...java/com/marcosbarbero/cloud/autoconfigure/zuul/ratelimit/filters/RateLimitPreFilter.java
@@ -25,6 +25,7 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.PRE_TYPE;
 
+import com.google.common.collect.Maps;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.Rate;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimitKeyGenerator;
 import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimitUtils;
@@ -39,6 +40,8 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.web.util.UrlPathHelper;
 
+import java.util.Map;
+
 /**
  * @author Marcos Barbero
  * @author Michal Šváb
@@ -77,27 +80,35 @@ public Object run() {
         final Route route = route(request);
 
         policy(route, request).forEach(policy -> {
+            Map<String, String> responseHeaders = Maps.newHashMap();
+
             final String key = rateLimitKeyGenerator.key(request, route, policy);
             final Rate rate = rateLimiter.consume(policy, key, null);
             final String httpHeaderKey = key.replaceAll("[^A-Za-z0-9-.]", "_").replaceAll("__", "_");
 
             final Long limit = policy.getLimit();
             final Long remaining = rate.getRemaining();
             if (limit != null) {
-                response.setHeader(HEADER_LIMIT + httpHeaderKey, String.valueOf(limit));
-                response.setHeader(HEADER_REMAINING + httpHeaderKey, String.valueOf(Math.max(remaining, 0)));
+                responseHeaders.put(HEADER_LIMIT + httpHeaderKey, String.valueOf(limit));
+                responseHeaders.put(HEADER_REMAINING + httpHeaderKey, String.valueOf(Math.max(remaining, 0)));
             }
 
             final Long quota = policy.getQuota();
             final Long remainingQuota = rate.getRemainingQuota();
             if (quota != null) {
                 request.setAttribute(REQUEST_START_TIME, System.currentTimeMillis());
-                response.setHeader(HEADER_QUOTA + httpHeaderKey, String.valueOf(quota));
-                response.setHeader(HEADER_REMAINING_QUOTA + httpHeaderKey,
+                responseHeaders.put(HEADER_QUOTA + httpHeaderKey, String.valueOf(quota));
+                responseHeaders.put(HEADER_REMAINING_QUOTA + httpHeaderKey,
                     String.valueOf(MILLISECONDS.toSeconds(Math.max(remainingQuota, 0))));
             }
 
-            response.setHeader(HEADER_RESET + httpHeaderKey, String.valueOf(rate.getReset()));
+            responseHeaders.put(HEADER_RESET + httpHeaderKey, String.valueOf(rate.getReset()));
+
+            if (properties.isAddResponseHeaders()) {
+                for (Map.Entry<String, String> headersEntry : responseHeaders.entrySet()) {
+                    response.setHeader(headersEntry.getKey(), headersEntry.getValue());
+                }
+            }
 
             if ((limit != null && remaining < 0) || (quota != null && remainingQuota < 0)) {
                 ctx.setResponseStatusCode(HttpStatus.TOO_MANY_REQUESTS.value());