added new X-RateLimit-* headers to the rate-limiting policy. these he…

…ader names are customizable
apiman · Feb 3, 2015 · c46f938 · c46f938
1 parent 74c59ac
commit c46f938
Show file tree

Hide file tree

Showing 11 changed files with 272 additions and 14 deletions.
diff --git a/.../engine/core/src/main/java/io/apiman/gateway/engine/components/IRateLimiterComponent.java b/.../engine/core/src/main/java/io/apiman/gateway/engine/components/IRateLimiterComponent.java
@@ -17,6 +17,7 @@
 
 import io.apiman.gateway.engine.IComponent;
 import io.apiman.gateway.engine.async.IAsyncResultHandler;
+import io.apiman.gateway.engine.components.rate.RateLimitResponse;
 import io.apiman.gateway.engine.rates.RateBucketPeriod;
 
 /**
@@ -37,6 +38,6 @@ public interface IRateLimiterComponent extends IComponent {
      * @param bucketId
      * @param limit
      */
-    void accept(String bucketId, RateBucketPeriod period, int limit, IAsyncResultHandler<Boolean> handler);
+    void accept(String bucketId, RateBucketPeriod period, int limit, IAsyncResultHandler<RateLimitResponse> handler);
 
 }
diff --git a/...engine/core/src/main/java/io/apiman/gateway/engine/components/rate/RateLimitResponse.java b/...engine/core/src/main/java/io/apiman/gateway/engine/components/rate/RateLimitResponse.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright 2015 JBoss Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.apiman.gateway.engine.components.rate;
+
+/**
+ * A simple bean that is returned when using the rate limiter 
+ * component.
+ *
+ * @author eric.wittmann@redhat.com
+ */
+public class RateLimitResponse {
+
+    private boolean accepted;
+    private int remaining;
+    private long reset;
+
+    /**
+     * Constructor.
+     */
+    public RateLimitResponse() {
+        // TODO Auto-generated constructor stub
+    }
+
+    /**
+     * @return the accepted
+     */
+    public boolean isAccepted() {
+        return accepted;
+    }
+
+    /**
+     * @param accepted the accepted to set
+     */
+    public void setAccepted(boolean accepted) {
+        this.accepted = accepted;
+    }
+
+    /**
+     * @return the remaining
+     */
+    public int getRemaining() {
+        return remaining;
+    }
+
+    /**
+     * @param remaining the remaining to set
+     */
+    public void setRemaining(int remaining) {
+        this.remaining = remaining;
+    }
+
+    /**
+     * @return the reset
+     */
+    public long getReset() {
+        return reset;
+    }
+
+    /**
+     * @param reset the reset to set
+     */
+    public void setReset(long reset) {
+        this.reset = reset;
+    }
+
+}
diff --git a/...engine/core/src/main/java/io/apiman/gateway/engine/impl/InMemoryRateLimiterComponent.java b/...engine/core/src/main/java/io/apiman/gateway/engine/impl/InMemoryRateLimiterComponent.java
@@ -18,6 +18,7 @@
 import io.apiman.gateway.engine.async.AsyncResultImpl;
 import io.apiman.gateway.engine.async.IAsyncResultHandler;
 import io.apiman.gateway.engine.components.IRateLimiterComponent;
+import io.apiman.gateway.engine.components.rate.RateLimitResponse;
 import io.apiman.gateway.engine.rates.RateBucketPeriod;
 import io.apiman.gateway.engine.rates.RateLimiterBucket;
 
@@ -44,7 +45,7 @@ public InMemoryRateLimiterComponent() {
      * @see io.apiman.gateway.engine.components.IRateLimiterComponent#accept(java.lang.String, io.apiman.gateway.engine.rates.RateBucketPeriod, int, io.apiman.gateway.engine.async.IAsyncResultHandler)
      */
     @Override
-    public void accept(String bucketId, RateBucketPeriod period, int limit, IAsyncResultHandler<Boolean> handler) {
+    public void accept(String bucketId, RateBucketPeriod period, int limit, IAsyncResultHandler<RateLimitResponse> handler) {
         RateLimiterBucket bucket = null;
         synchronized (buckets) {
             bucket = buckets.get(bucketId);
@@ -53,13 +54,19 @@ public void accept(String bucketId, RateBucketPeriod period, int limit, IAsyncRe
                 buckets.put(bucketId, bucket);
             }
             bucket.resetIfNecessary(period);
+
+            RateLimitResponse response = new RateLimitResponse();
             if (bucket.count >= limit) {
-                handler.handle(AsyncResultImpl.<Boolean>create(Boolean.FALSE));
+                response.setAccepted(false);
             } else {
                 bucket.count++;
                 bucket.last = System.currentTimeMillis();
-                handler.handle(AsyncResultImpl.<Boolean>create(Boolean.TRUE));
+                response.setAccepted(true);
             }
+            int reset = (int) (bucket.getResetMillis(period) / 1000L);
+            response.setReset(reset);
+            response.setRemaining(limit - bucket.count);
+            handler.handle(AsyncResultImpl.<RateLimitResponse>create(response));
         }
     }
 

diff --git a/gateway/engine/core/src/main/java/io/apiman/gateway/engine/rates/RateLimiterBucket.java b/gateway/engine/core/src/main/java/io/apiman/gateway/engine/rates/RateLimiterBucket.java
@@ -48,15 +48,35 @@ public void resetIfNecessary(RateBucketPeriod period) {
             count = 0;
         }
     }
+
+    /**
+     * Returns the number of millis until the period resets.
+     * @param period
+     */
+    public long getResetMillis(RateBucketPeriod period) {
+        long now = System.currentTimeMillis();
+        long periodBoundary = getPeriodBoundary(now, period);
+        return periodBoundary - now;
+    }
 
     /**
      * Gets the period boundary for the period bounding the last
      * request.
      * @param period
      */
     private long getLastPeriodBoundary(RateBucketPeriod period) {
+        return getPeriodBoundary(last, period);
+    }
+
+    /**
+     * Gets the boundary timestamp for the given rate bucket period.  In other words,
+     * returns the timestamp associated with when the rate period will reset.
+     * @param timestamp
+     * @param period
+     */
+    private static long getPeriodBoundary(long timestamp, RateBucketPeriod period) {
         Calendar lastCal = Calendar.getInstance();
-        lastCal.setTimeInMillis(last);
+        lastCal.setTimeInMillis(timestamp);
         switch (period) {
         case Second:
             lastCal.set(Calendar.MILLISECOND, 0);

diff --git a/...gine/ispn/src/main/java/io/apiman/gateway/engine/ispn/InfinispanRateLimiterComponent.java b/...gine/ispn/src/main/java/io/apiman/gateway/engine/ispn/InfinispanRateLimiterComponent.java
@@ -18,6 +18,7 @@
 import io.apiman.gateway.engine.async.AsyncResultImpl;
 import io.apiman.gateway.engine.async.IAsyncResultHandler;
 import io.apiman.gateway.engine.components.IRateLimiterComponent;
+import io.apiman.gateway.engine.components.rate.RateLimitResponse;
 import io.apiman.gateway.engine.rates.RateBucketPeriod;
 import io.apiman.gateway.engine.rates.RateLimiterBucket;
 
@@ -93,7 +94,7 @@ private Cache<Object, Object> getCache() {
      */
     @Override
     public void accept(String bucketId, RateBucketPeriod period, int limit,
-            IAsyncResultHandler<Boolean> handler) {
+            IAsyncResultHandler<RateLimitResponse> handler) {
         RateLimiterBucket bucket = null;
         synchronized (mutex) {
             bucket = (RateLimiterBucket) getCache().get(bucketId);
@@ -102,13 +103,19 @@ public void accept(String bucketId, RateBucketPeriod period, int limit,
                 getCache().put(bucketId, bucket);
             }
             bucket.resetIfNecessary(period);
+
+            RateLimitResponse response = new RateLimitResponse();
             if (bucket.count >= limit) {
-                handler.handle(AsyncResultImpl.<Boolean>create(Boolean.FALSE));
+                response.setAccepted(false);
             } else {
                 bucket.count++;
                 bucket.last = System.currentTimeMillis();
-                handler.handle(AsyncResultImpl.<Boolean>create(Boolean.TRUE));
+                response.setAccepted(true);
             }
+            int reset = (int) (bucket.getResetMillis(period) / 1000L);
+            response.setReset(reset);
+            response.setRemaining(limit - bucket.count);
+            handler.handle(AsyncResultImpl.<RateLimitResponse>create(response));
             getCache().put(bucketId, bucket);
         }
     }

diff --git a/...y/engine/policies/src/main/java/io/apiman/gateway/engine/policies/RateLimitingPolicy.java b/...y/engine/policies/src/main/java/io/apiman/gateway/engine/policies/RateLimitingPolicy.java
@@ -20,8 +20,10 @@
 import io.apiman.gateway.engine.beans.PolicyFailure;
 import io.apiman.gateway.engine.beans.PolicyFailureType;
 import io.apiman.gateway.engine.beans.ServiceRequest;
+import io.apiman.gateway.engine.beans.ServiceResponse;
 import io.apiman.gateway.engine.components.IPolicyFailureFactoryComponent;
 import io.apiman.gateway.engine.components.IRateLimiterComponent;
+import io.apiman.gateway.engine.components.rate.RateLimitResponse;
 import io.apiman.gateway.engine.policies.config.RateLimitingConfig;
 import io.apiman.gateway.engine.policies.config.rates.RateLimitingGranularity;
 import io.apiman.gateway.engine.policies.config.rates.RateLimitingPeriod;
@@ -30,6 +32,9 @@
 import io.apiman.gateway.engine.policy.IPolicyContext;
 import io.apiman.gateway.engine.rates.RateBucketPeriod;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Policy that enforces rate limits.
  *
@@ -40,6 +45,10 @@ public class RateLimitingPolicy extends AbstractMappedPolicy<RateLimitingConfig>
     private static final String NO_USER_AVAILABLE = new String();
     private static final String NO_APPLICATION_AVAILABLE = new String();
 
+    private static final String DEFAULT_LIMIT_HEADER = "X-RateLimit-Limit"; //$NON-NLS-1$
+    private static final String DEFAULT_REMAINING_HEADER = "X-RateLimit-Remaining"; //$NON-NLS-1$
+    private static final String DEFAULT_RESET_HEADER = "X-RateLimit-Reset"; //$NON-NLS-1$
+
     /**
      * Constructor.
      */
@@ -61,7 +70,7 @@ protected Class<RateLimitingConfig> getConfigurationClass() {
     protected void doApply(final ServiceRequest request, final IPolicyContext context, final RateLimitingConfig config,
             final IPolicyChain<ServiceRequest> chain) {
         String bucketId = createBucketId(request, config);
-        RateBucketPeriod period = getPeriod(config);
+        final RateBucketPeriod period = getPeriod(config);
 
         if (bucketId == NO_USER_AVAILABLE) {
             IPolicyFailureFactoryComponent failureFactory = context.getComponent(IPolicyFailureFactoryComponent.class);
@@ -77,25 +86,58 @@ protected void doApply(final ServiceRequest request, final IPolicyContext contex
         }
 
         IRateLimiterComponent rateLimiter = context.getComponent(IRateLimiterComponent.class);
-        rateLimiter.accept(bucketId, period, config.getLimit(), new IAsyncResultHandler<Boolean>() {
+        rateLimiter.accept(bucketId, period, config.getLimit(), new IAsyncResultHandler<RateLimitResponse>() {
             @Override
-            public void handle(IAsyncResult<Boolean> result) {
+            public void handle(IAsyncResult<RateLimitResponse> result) {
                 if (result.isError()) {
                     chain.throwError(result.getError());
                 } else {
-                    boolean accepted = result.getResult();
-                    if (accepted) {
+                    RateLimitResponse rtr = result.getResult();
+
+                    Map<String, String> responseHeaders = new HashMap<>();
+                    String limitHeader = config.getHeaderLimit();
+                    if (limitHeader == null) {
+                        limitHeader = DEFAULT_LIMIT_HEADER;
+                    }
+                    String remainingHeader = config.getHeaderRemaining();
+                    if (remainingHeader == null) {
+                        remainingHeader = DEFAULT_REMAINING_HEADER;
+                    }
+                    String resetHeader = config.getHeaderReset();
+                    if (resetHeader == null) {
+                        resetHeader = DEFAULT_RESET_HEADER;
+                    }
+                    responseHeaders.put(limitHeader, String.valueOf(config.getLimit()));
+                    responseHeaders.put(remainingHeader, String.valueOf(rtr.getRemaining()));
+                    responseHeaders.put(resetHeader, String.valueOf(rtr.getReset()));
+
+                    if (rtr.isAccepted()) {
+                        context.setAttribute("rate-limit-response-headers", responseHeaders); //$NON-NLS-1$
                         chain.doApply(request);
                     } else {
                         IPolicyFailureFactoryComponent failureFactory = context.getComponent(IPolicyFailureFactoryComponent.class);
                         PolicyFailure failure = failureFactory.createFailure(PolicyFailureType.Other, PolicyFailureCodes.RATE_LIMIT_EXCEEDED, Messages.i18n.format("RateLimitingPolicy.RateExceeded")); //$NON-NLS-1$
+                        failure.getHeaders().putAll(responseHeaders);
                         failure.setResponseCode(429);
                         chain.doFailure(failure);
                     }
                 }
             }
         });
     }
+
+    /**
+     * @see io.apiman.gateway.engine.policies.AbstractMappedPolicy#doApply(io.apiman.gateway.engine.beans.ServiceResponse, io.apiman.gateway.engine.policy.IPolicyContext, java.lang.Object, io.apiman.gateway.engine.policy.IPolicyChain)
+     */
+    @Override
+    protected void doApply(ServiceResponse response, IPolicyContext context, RateLimitingConfig config,
+            IPolicyChain<ServiceResponse> chain) {
+        Map<String, String> headers = context.getAttribute("rate-limit-response-headers", null); //$NON-NLS-1$
+        if (headers != null) {
+            response.getHeaders().putAll(headers);
+        }
+        super.doApply(response, context, config, chain);
+    }
 
     /**
      * Creates the ID of the rate bucket to use.  The ID is composed differently 

diff --git a/...e/policies/src/main/java/io/apiman/gateway/engine/policies/config/RateLimitingConfig.java b/...e/policies/src/main/java/io/apiman/gateway/engine/policies/config/RateLimitingConfig.java
@@ -32,6 +32,9 @@ public class RateLimitingConfig {
     private RateLimitingGranularity granularity;
     private RateLimitingPeriod period;
     private String userHeader;
+    private String headerRemaining;
+    private String headerLimit;
+    private String headerReset;
 
     /**
      * Constructor.
@@ -95,4 +98,46 @@ public void setLimit(int limit) {
         this.limit = limit;
     }
 
+    /**
+     * @return the headerRemaining
+     */
+    public String getHeaderRemaining() {
+        return headerRemaining;
+    }
+
+    /**
+     * @param headerRemaining the headerRemaining to set
+     */
+    public void setHeaderRemaining(String headerRemaining) {
+        this.headerRemaining = headerRemaining;
+    }
+
+    /**
+     * @return the headerLimit
+     */
+    public String getHeaderLimit() {
+        return headerLimit;
+    }
+
+    /**
+     * @param headerLimit the headerLimit to set
+     */
+    public void setHeaderLimit(String headerLimit) {
+        this.headerLimit = headerLimit;
+    }
+
+    /**
+     * @return the headerReset
+     */
+    public String getHeaderReset() {
+        return headerReset;
+    }
+
+    /**
+     * @param headerReset the headerReset to set
+     */
+    public void setHeaderReset(String headerReset) {
+        this.headerReset = headerReset;
+    }
+
 }