Add Rate Limit provider

src/main/java/com/auth0/jwk/Bucket.java

@@ -0,0 +1,37 @@
+package com.auth0.jwk;
+
+/**
+ * Token Bucket interface.
+ */
+interface Bucket {
+
+    /**
+     * Calculates the wait time before one token will be available in the Bucket.
+     *
+     * @return the wait time in milliseconds in which one token will be available in the Bucket.
+     */
+    long willLeakIn();
+
+    /**
+     * Calculates the wait time before the given amount of tokens will be available in the Bucket.
+     *
+     * @param count the amount of tokens to check how much time to wait for.
+     * @return the wait time in milliseconds in which the given amount of tokens will be available in the Bucket.
+     */
+    long willLeakIn(long count);
+
+    /**
+     * Tries to consume one token from the Bucket.
+     *
+     * @return true if it could consume the token or false if the Bucket doesn't have tokens available now.
+     */
+    boolean consume();
+
+    /**
+     * Tries to consume the given amount of tokens from the Bucket.
+     *
+     * @param count the amount of tokens to try to consume.
+     * @return true if it could consume the given amount of tokens or false if the Bucket doesn't have that amount of tokens available now.
+     */
+    boolean consume(long count);
+}

src/main/java/com/auth0/jwk/BucketImpl.java

@@ -0,0 +1,99 @@
+package com.auth0.jwk;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Token Bucket implementation to guarantee availability of a fixed amount of tokens in a given time rate.
+ */
+class BucketImpl implements Bucket {
+
+    private final long size;
+    private final long rate;
+    private final TimeUnit rateUnit;
+    private AtomicLong available;
+    private AtomicLong lastTokenAddedAt;
+
+    BucketImpl(long size, long rate, TimeUnit rateUnit) {
+        assertPositiveValue(size, "Invalid bucket size.");
+        assertPositiveValue(rate, "Invalid bucket refill rate.");
+        this.size = size;
+        this.available = new AtomicLong(size);
+        this.lastTokenAddedAt = new AtomicLong(System.currentTimeMillis());
+        this.rate = rate;
+        this.rateUnit = rateUnit;
+
+        beginRefillAtRate();
+    }
+
+    private void beginRefillAtRate() {
+        ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
+        Runnable refillTask = new Runnable() {
+            public void run() {
+                try {
+                    rateUnit.sleep(rate);
+                } catch (InterruptedException e) {
+                    e.printStackTrace();
+                }
+                addToken();
+            }
+        };
+        executorService.scheduleAtFixedRate(refillTask, 0, rate, rateUnit);
+    }
+
+    private void addToken() {
+        if (available.get() < size) {
+            available.incrementAndGet();
+        }
+        lastTokenAddedAt.set(System.currentTimeMillis());
+    }
+
+    private void assertPositiveValue(long value, long maxValue, String exceptionMessage) {
+        if (value < 1 || value > maxValue) {
+            throw new IllegalArgumentException(exceptionMessage);
+        }
+    }
+
+    private void assertPositiveValue(Number value, String exceptionMessage) {
+        this.assertPositiveValue(value.intValue(), value.intValue(), exceptionMessage);
+    }
+
+    @Override
+    public long willLeakIn() {
+        return willLeakIn(1);
+    }
+
+    @Override
+    public long willLeakIn(long count) {
+        assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size));
+        long av = available.get();
+        if (av >= count) {
+            return 0;
+        }
+
+        long nextIn = rateUnit.toMillis(rate) - (System.currentTimeMillis() - lastTokenAddedAt.get());
+        final long remaining = count - av - 1;
+        if (remaining > 0) {
+            nextIn += rateUnit.toMillis(rate) * remaining;
+        }
+        return nextIn;
+    }
+
+    @Override
+    public boolean consume() {
+        return consume(1);
+    }
+
+    @Override
+    public boolean consume(long count) {
+        assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size));
+        if (count <= available.get()) {
+            available.addAndGet(-count);
+            return true;
+        }
+        System.out.println();
+        return false;
+    }
+}

src/main/java/com/auth0/jwk/GuavaCachedJwkProvider.java

@@ -1,5 +1,6 @@
 package com.auth0.jwk;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
 
@@ -52,4 +53,9 @@ public Jwk call() throws Exception {
             throw new SigningKeyNotFoundException("Failed to get key with kid " + keyId, e);
         }
     }
+
+    @VisibleForTesting
+    JwkProvider getBaseProvider() {
+        return provider;
+    }
 }

src/main/java/com/auth0/jwk/JwkException.java

@@ -3,6 +3,10 @@
 @SuppressWarnings("WeakerAccess")
 public class JwkException extends Exception {
 
+    public JwkException(String message) {
+        super(message);
+    }
+
     public JwkException(String message, Throwable cause) {
         super(message, cause);
     }

src/main/java/com/auth0/jwk/JwkProviderBuilder.java

@@ -13,15 +13,19 @@ public class JwkProviderBuilder {
     private long expiresIn;
     private long cacheSize;
     private boolean cached;
+    private BucketImpl bucket;
+    private boolean rateLimited;
 
     /**
-     * Creates a new builder
+     * Creates a new builder.
      */
     public JwkProviderBuilder() {
         this.cached = true;
         this.expiresIn = 10;
         this.expiresUnit = TimeUnit.HOURS;
         this.cacheSize = 5;
+        this.rateLimited = true;
+        this.bucket = new BucketImpl(10, 1, TimeUnit.MINUTES);
     }
 
     /**
@@ -36,7 +40,7 @@ public JwkProviderBuilder forDomain(String domain) {
     }
 
     /**
-     * Toggle the cache of Jwk
+     * Toggle the cache of Jwk. By default the provider will use cache.
      * @param cached if the provider should cache jwks
      * @return the builder
      */
@@ -60,6 +64,28 @@ public JwkProviderBuilder cached(long cacheSize, long expiresIn, TimeUnit unit)
         return this;
     }
 
+    /**
+     * Toggle the rate limit of Jwk. By default the Provider will use rate limit.
+     * @param rateLimited if the provider should rate limit jwks
+     * @return the builder
+     */
+    public JwkProviderBuilder rateLimited(boolean rateLimited) {
+        this.rateLimited = rateLimited;
+        return this;
+    }
+
+    /**
+     * Enable the cache specifying size and expire time.
+     * @param bucketSize max number of jwks to deliver in the given rate.
+     * @param refillRate amount of time to wait before a jwk can the jwk will be cached
+     * @param unit unit of time for the expire of jwk
+     * @return the builder
+     */
+    public JwkProviderBuilder rateLimited(long bucketSize, long refillRate, TimeUnit unit) {
+        bucket = new BucketImpl(bucketSize, refillRate, unit);
+        return this;
+    }
+
     /**
      * Creates a {@link JwkProvider}
      * @return a newly created {@link JwkProvider}
@@ -68,13 +94,14 @@ public JwkProvider build() {
         if (url == null) {
             throw new IllegalStateException("Cannot build provider without domain");
         }
-
-        final UrlJwkProvider urlProvider = new UrlJwkProvider(url);
-        if (!this.cached) {
-            return urlProvider;
+        JwkProvider urlProvider = new UrlJwkProvider(url);
+        if (this.rateLimited) {
+            urlProvider = new RateLimitedJwkProvider(urlProvider, bucket);
         }
-
-        return new GuavaCachedJwkProvider(urlProvider, cacheSize, expiresIn, expiresUnit);
+        if (this.cached) {
+            urlProvider = new GuavaCachedJwkProvider(urlProvider, cacheSize, expiresIn, expiresUnit);
+        }
+        return urlProvider;
     }
 
     private String normalizeDomain(String domain) {

src/main/java/com/auth0/jwk/RateLimitReachedException.java

@@ -0,0 +1,21 @@
+package com.auth0.jwk;
+
+@SuppressWarnings("WeakerAccess")
+public class RateLimitReachedException extends JwkException {
+    private final long availableInMs;
+
+    public RateLimitReachedException(long availableInMs) {
+        super(String.format("The Rate Limit has been reached! Please wait %d milliseconds.", availableInMs));
+        this.availableInMs = availableInMs;
+    }
+
+    /**
+     * Returns the delay in which the jwk request can be retried.
+     *
+     * @return the time to wait in milliseconds before retrying the request.
+     */
+    public long getAvailableIn() {
+        return availableInMs;
+    }
+
+}

src/main/java/com/auth0/jwk/RateLimitedJwkProvider.java

@@ -0,0 +1,37 @@
+package com.auth0.jwk;
+
+import com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Jwk provider that limits the amount of Jwks to deliver in a given rate.
+ */
+@SuppressWarnings("WeakerAccess")
+public class RateLimitedJwkProvider implements JwkProvider {
+
+    private final JwkProvider provider;
+    private final Bucket bucket;
+
+    /**
+     * Creates a new provider that will check the given Bucket if a jwks can be provided now.
+     *
+     * @param bucket   bucket to limit the amount of jwk requested in a given amount of time.
+     * @param provider provider to use to request jwk when the bucket allows it.
+     */
+    public RateLimitedJwkProvider(JwkProvider provider, Bucket bucket) {
+        this.provider = provider;
+        this.bucket = bucket;
+    }
+
+    @Override
+    public Jwk get(final String keyId) throws JwkException {
+        if (!bucket.consume()) {
+            throw new RateLimitReachedException(bucket.willLeakIn());
+        }
+        return provider.get(keyId);
+    }
+
+    @VisibleForTesting
+    JwkProvider getBaseProvider() {
+        return provider;
+    }
+}

src/main/java/com/auth0/jwk/UrlJwkProvider.java

@@ -37,13 +37,17 @@ public UrlJwkProvider(URL url) {
      * @param domain domain where to look for the jwks.json file
      */
     public UrlJwkProvider(String domain) {
+        this(urlForDomain(domain));
+    }
+
+    private static URL urlForDomain(String domain) {
         if (Strings.isNullOrEmpty(domain)) {
             throw new IllegalArgumentException("A domain is required");
         }
 
         try {
             final URL url = new URL(domain);
-            this.url = new URL(url, "/.well-known/jwks.json");
+            return new URL(url, "/.well-known/jwks.json");
         } catch (MalformedURLException e) {
             throw new IllegalArgumentException("Invalid jwks uri", e);
         }