From 413b518e218945cf75eb1986bd6e35b427cb42f7 Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Tue, 15 Nov 2016 19:18:41 -0300 Subject: [PATCH 1/6] refactor UrlJwkProvider constructors --- src/main/java/com/auth0/jwk/UrlJwkProvider.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/auth0/jwk/UrlJwkProvider.java b/src/main/java/com/auth0/jwk/UrlJwkProvider.java index f141381..5cccb65 100644 --- a/src/main/java/com/auth0/jwk/UrlJwkProvider.java +++ b/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); } From cc4a1d9a16850ff5de44e84387dd3b0a0ae26810 Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Tue, 15 Nov 2016 19:19:05 -0300 Subject: [PATCH 2/6] add token-bucket implementation --- src/main/java/com/auth0/jwk/Bucket.java | 11 ++ src/main/java/com/auth0/jwk/BucketImpl.java | 108 ++++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 src/main/java/com/auth0/jwk/Bucket.java create mode 100644 src/main/java/com/auth0/jwk/BucketImpl.java diff --git a/src/main/java/com/auth0/jwk/Bucket.java b/src/main/java/com/auth0/jwk/Bucket.java new file mode 100644 index 0000000..d40881f --- /dev/null +++ b/src/main/java/com/auth0/jwk/Bucket.java @@ -0,0 +1,11 @@ +package com.auth0.jwk; + +interface Bucket { + long willLeakIn(); + + long willLeakIn(int count); + + boolean consume(); + + boolean consume(int count); +} diff --git a/src/main/java/com/auth0/jwk/BucketImpl.java b/src/main/java/com/auth0/jwk/BucketImpl.java new file mode 100644 index 0000000..62cd142 --- /dev/null +++ b/src/main/java/com/auth0/jwk/BucketImpl.java @@ -0,0 +1,108 @@ +package com.auth0.jwk; + +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Leaky Bucket implementation to make guarantee availability of a fixed amount of tokens in a given time rate. + */ +class BucketImpl implements Bucket { + + private static final int MAX_BUCKET_SIZE = Integer.MAX_VALUE - 5; + + private final int size; + private final long rate; + private final TimeUnit rateUnit; + private final long beginTime = System.currentTimeMillis(); + private AtomicInteger available; + private AtomicLong lastTokenAddedAt; + + public BucketImpl(int size, long rate, TimeUnit rateUnit) { + assertPositiveValue(size, MAX_BUCKET_SIZE, "Invalid bucket size."); + this.size = size; + this.available = new AtomicInteger(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, rate, rate, rateUnit); + } + + private void addToken() { + if (available.get() < size) { + available.incrementAndGet(); + log(String.format("Added 1 token.. Current state: %d/%d", available.get(), size)); + } + lastTokenAddedAt.set(System.currentTimeMillis()); + } + + private void assertPositiveValue(int value, int maxValue, String exceptionMessage) { + if (value < 1 || value > maxValue) { + throw new IllegalArgumentException(exceptionMessage); + } + } + + private void log(String message) { + long msDiff = System.currentTimeMillis() - beginTime; + System.out.println(String.format("%-8d - %s", msDiff, message)); + } + + @Override + public long willLeakIn() { + return willLeakIn(1); + } + + @Override + public long willLeakIn(int count) { + assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size)); + int av = available.get(); + if (av >= count) { + log(String.format("%d Tokens are available already.", count)); + return 0; + } + + long nextIn = rateUnit.toMillis(rate) - System.currentTimeMillis() - lastTokenAddedAt.get(); + final int remaining = count - av - 1; + if (remaining > 0) { + nextIn += rateUnit.toMillis(rate) * remaining; + } + log(String.format("Can't consume %d. Actual state is: %d/%d. Retry in %d ms.", count, available.get(), size, nextIn)); + return nextIn; + } + + @Override + public boolean consume() { + return consume(1); + } + + @Override + public boolean consume(int 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); + log(String.format("Consumed %d tokens.. Current state: %d/%d", count, available.get(), size)); + return true; + } + log(String.format("Not enough tokens available to consume %d", count)); + System.out.println(); + return false; + } +} From a039d7c26771d0dd06d8f12cab8bcfa08f51fc5b Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Tue, 15 Nov 2016 19:19:41 -0300 Subject: [PATCH 3/6] allow to decorate provider with rate limit feature --- src/main/java/com/auth0/jwk/Bucket.java | 24 +++++++ src/main/java/com/auth0/jwk/BucketImpl.java | 2 +- .../com/auth0/jwk/GuavaCachedJwkProvider.java | 6 ++ src/main/java/com/auth0/jwk/JwkException.java | 4 ++ .../com/auth0/jwk/JwkProviderBuilder.java | 43 +++++++++--- .../auth0/jwk/RateLimitReachedException.java | 21 ++++++ .../com/auth0/jwk/RateLimitedJwkProvider.java | 37 +++++++++++ .../auth0/jwk/GuavaCachedJwkProviderTest.java | 12 ++-- .../com/auth0/jwk/JwkProviderBuilderTest.java | 65 ++++++++++++++++++- .../jwk/RateLimitReachedExceptionTest.java | 18 +++++ .../auth0/jwk/RateLimitedJwkProviderTest.java | 61 +++++++++++++++++ 11 files changed, 275 insertions(+), 18 deletions(-) create mode 100644 src/main/java/com/auth0/jwk/RateLimitReachedException.java create mode 100644 src/main/java/com/auth0/jwk/RateLimitedJwkProvider.java create mode 100644 src/test/java/com/auth0/jwk/RateLimitReachedExceptionTest.java create mode 100644 src/test/java/com/auth0/jwk/RateLimitedJwkProviderTest.java diff --git a/src/main/java/com/auth0/jwk/Bucket.java b/src/main/java/com/auth0/jwk/Bucket.java index d40881f..873c269 100644 --- a/src/main/java/com/auth0/jwk/Bucket.java +++ b/src/main/java/com/auth0/jwk/Bucket.java @@ -1,11 +1,35 @@ 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. + * + * @return the wait time in milliseconds in which the given amount of tokens will be available in the Bucket. + */ long willLeakIn(int 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. + * + * @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(int count); } diff --git a/src/main/java/com/auth0/jwk/BucketImpl.java b/src/main/java/com/auth0/jwk/BucketImpl.java index 62cd142..6748afb 100644 --- a/src/main/java/com/auth0/jwk/BucketImpl.java +++ b/src/main/java/com/auth0/jwk/BucketImpl.java @@ -7,7 +7,7 @@ import java.util.concurrent.atomic.AtomicLong; /** - * Leaky Bucket implementation to make guarantee availability of a fixed amount of tokens in a given time rate. + * Token Bucket implementation to guarantee availability of a fixed amount of tokens in a given time rate. */ class BucketImpl implements Bucket { diff --git a/src/main/java/com/auth0/jwk/GuavaCachedJwkProvider.java b/src/main/java/com/auth0/jwk/GuavaCachedJwkProvider.java index 8af5c06..c988a20 100644 --- a/src/main/java/com/auth0/jwk/GuavaCachedJwkProvider.java +++ b/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; + } } diff --git a/src/main/java/com/auth0/jwk/JwkException.java b/src/main/java/com/auth0/jwk/JwkException.java index 2e48c15..a330e29 100644 --- a/src/main/java/com/auth0/jwk/JwkException.java +++ b/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); } diff --git a/src/main/java/com/auth0/jwk/JwkProviderBuilder.java b/src/main/java/com/auth0/jwk/JwkProviderBuilder.java index 59a8c44..f3074fc 100644 --- a/src/main/java/com/auth0/jwk/JwkProviderBuilder.java +++ b/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(int 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) { diff --git a/src/main/java/com/auth0/jwk/RateLimitReachedException.java b/src/main/java/com/auth0/jwk/RateLimitReachedException.java new file mode 100644 index 0000000..455039f --- /dev/null +++ b/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; + } + +} diff --git a/src/main/java/com/auth0/jwk/RateLimitedJwkProvider.java b/src/main/java/com/auth0/jwk/RateLimitedJwkProvider.java new file mode 100644 index 0000000..cc6d9ee --- /dev/null +++ b/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; + } +} diff --git a/src/test/java/com/auth0/jwk/GuavaCachedJwkProviderTest.java b/src/test/java/com/auth0/jwk/GuavaCachedJwkProviderTest.java index ed6ad14..cc1a516 100644 --- a/src/test/java/com/auth0/jwk/GuavaCachedJwkProviderTest.java +++ b/src/test/java/com/auth0/jwk/GuavaCachedJwkProviderTest.java @@ -1,9 +1,5 @@ package com.auth0.jwk; -import com.auth0.jwk.GuavaCachedJwkProvider; -import com.auth0.jwk.Jwk; -import com.auth0.jwk.JwkProvider; -import com.auth0.jwk.SigningKeyNotFoundException; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -13,8 +9,7 @@ import org.mockito.runners.MockitoJUnitRunner; import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.hasItem; -import static org.junit.Assert.*; +import static org.junit.Assert.assertThat; import static org.mockito.Matchers.anyString; import static org.mockito.Matchers.eq; import static org.mockito.Mockito.*; @@ -60,4 +55,9 @@ public void shouldUseCachedValue() throws Exception { assertThat(provider.get(KID), equalTo(jwk)); verify(fallback, only()).get(KID); } + + @Test + public void shouldGetBaseProvider() throws Exception { + assertThat(provider.getBaseProvider(), equalTo(fallback)); + } } \ No newline at end of file diff --git a/src/test/java/com/auth0/jwk/JwkProviderBuilderTest.java b/src/test/java/com/auth0/jwk/JwkProviderBuilderTest.java index 56c175d..cdaeab1 100644 --- a/src/test/java/com/auth0/jwk/JwkProviderBuilderTest.java +++ b/src/test/java/com/auth0/jwk/JwkProviderBuilderTest.java @@ -6,8 +6,9 @@ import java.util.concurrent.TimeUnit; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; -import static org.junit.Assert.*; +import static org.junit.Assert.assertThat; public class JwkProviderBuilderTest { @@ -40,12 +41,70 @@ public void shouldFailWhenOnlySpecifyingCache() throws Exception { @Test public void shouldCreateCachedProvider() throws Exception { - assertThat(new JwkProviderBuilder().cached(true).forDomain("samples.auth0.com").build(), notNullValue()); + JwkProvider provider = new JwkProviderBuilder().rateLimited(false).cached(true).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(GuavaCachedJwkProvider.class)); + assertThat(((GuavaCachedJwkProvider) provider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); } @Test public void shouldCreateCachedProviderWithCustomValues() throws Exception { - assertThat(new JwkProviderBuilder().cached(10, 24, TimeUnit.HOURS).forDomain("samples.auth0.com").build(), notNullValue()); + JwkProvider provider = new JwkProviderBuilder().rateLimited(false).cached(10, 24, TimeUnit.HOURS).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(GuavaCachedJwkProvider.class)); + assertThat(((GuavaCachedJwkProvider) provider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); } + @Test + public void shouldFailWhenOnlySpecifyingRateLimit() throws Exception { + expectedException.expect(IllegalStateException.class); + expectedException.expectMessage("Cannot build provider without domain"); + new JwkProviderBuilder().rateLimited(false).build(); + } + + @Test + public void shouldCreateRateLimitedProvider() throws Exception { + JwkProvider provider = new JwkProviderBuilder().cached(false).rateLimited(true).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(RateLimitedJwkProvider.class)); + assertThat(((RateLimitedJwkProvider) provider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); + } + + @Test + public void shouldCreateRateLimitedProviderWithCustomValues() throws Exception { + JwkProvider provider = new JwkProviderBuilder().cached(false).rateLimited(10, 24, TimeUnit.HOURS).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(RateLimitedJwkProvider.class)); + assertThat(((RateLimitedJwkProvider) provider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); + } + + @Test + public void shouldCreateCachedAndRateLimitedProvider() throws Exception { + JwkProvider provider = new JwkProviderBuilder().cached(true).rateLimited(true).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(GuavaCachedJwkProvider.class)); + JwkProvider baseProvider = ((GuavaCachedJwkProvider) provider).getBaseProvider(); + assertThat(baseProvider, instanceOf(RateLimitedJwkProvider.class)); + assertThat(((RateLimitedJwkProvider) baseProvider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); + } + + @Test + public void shouldCreateCachedAndRateLimitedProviderWithCustomValues() throws Exception { + JwkProvider provider = new JwkProviderBuilder().cached(10, 24, TimeUnit.HOURS).rateLimited(10, 24, TimeUnit.HOURS).forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(GuavaCachedJwkProvider.class)); + JwkProvider baseProvider = ((GuavaCachedJwkProvider) provider).getBaseProvider(); + assertThat(baseProvider, instanceOf(RateLimitedJwkProvider.class)); + assertThat(((RateLimitedJwkProvider) baseProvider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); + } + + @Test + public void shouldCreateCachedAndRateLimitedProviderByDefault() throws Exception { + JwkProvider provider = new JwkProviderBuilder().forDomain("samples.auth0.com").build(); + assertThat(provider, notNullValue()); + assertThat(provider, instanceOf(GuavaCachedJwkProvider.class)); + JwkProvider baseProvider = ((GuavaCachedJwkProvider) provider).getBaseProvider(); + assertThat(baseProvider, instanceOf(RateLimitedJwkProvider.class)); + assertThat(((RateLimitedJwkProvider) baseProvider).getBaseProvider(), instanceOf(UrlJwkProvider.class)); + } } \ No newline at end of file diff --git a/src/test/java/com/auth0/jwk/RateLimitReachedExceptionTest.java b/src/test/java/com/auth0/jwk/RateLimitReachedExceptionTest.java new file mode 100644 index 0000000..eef1952 --- /dev/null +++ b/src/test/java/com/auth0/jwk/RateLimitReachedExceptionTest.java @@ -0,0 +1,18 @@ +package com.auth0.jwk; + +import org.junit.Test; + +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.notNullValue; +import static org.junit.Assert.assertThat; + +public class RateLimitReachedExceptionTest { + @Test + public void shouldGetAvailableIn() throws Exception { + RateLimitReachedException exception = new RateLimitReachedException(123456789); + assertThat(exception, notNullValue()); + assertThat(exception.getMessage(), equalTo("The Rate Limit has been reached! Please wait 123456789 milliseconds.")); + assertThat(exception.getAvailableIn(), equalTo(123456789L)); + } + +} \ No newline at end of file diff --git a/src/test/java/com/auth0/jwk/RateLimitedJwkProviderTest.java b/src/test/java/com/auth0/jwk/RateLimitedJwkProviderTest.java new file mode 100644 index 0000000..d9631a9 --- /dev/null +++ b/src/test/java/com/auth0/jwk/RateLimitedJwkProviderTest.java @@ -0,0 +1,61 @@ +package com.auth0.jwk; + +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.runners.MockitoJUnitRunner; + +import static org.hamcrest.Matchers.equalTo; +import static org.junit.Assert.assertThat; +import static org.mockito.Matchers.eq; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@RunWith(MockitoJUnitRunner.class) +public class RateLimitedJwkProviderTest { + + private static final String KID = "KID"; + private RateLimitedJwkProvider provider; + + @Mock + private JwkProvider fallback; + + @Mock + private Jwk jwk; + + @Mock + private Bucket bucket; + + @Rule + public ExpectedException expectedException = ExpectedException.none(); + + @Before + public void setUp() throws Exception { + provider = new RateLimitedJwkProvider(fallback, bucket); + } + + @Test + public void shouldFailToGetWhenBucketIsEmpty() throws Exception { + when(bucket.consume()).thenReturn(false); + expectedException.expect(RateLimitReachedException.class); + when(fallback.get(eq(KID))).thenReturn(jwk); + provider.get(KID); + } + + @Test + public void shouldGetWhenBucketHasTokensAvailable() throws Exception { + when(bucket.consume()).thenReturn(true); + when(fallback.get(eq(KID))).thenReturn(jwk); + assertThat(provider.get(KID), equalTo(jwk)); + verify(fallback).get(eq(KID)); + } + + @Test + public void shouldGetBaseProvider() throws Exception { + assertThat(provider.getBaseProvider(), equalTo(fallback)); + } + +} \ No newline at end of file From 77c6cab12bbd0a4db95bb8aae2bf50c6683571c2 Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Wed, 16 Nov 2016 12:07:13 -0300 Subject: [PATCH 4/6] add bucket tests --- src/main/java/com/auth0/jwk/BucketImpl.java | 11 ++- .../java/com/auth0/jwk/BucketImplTest.java | 95 +++++++++++++++++++ 2 files changed, 103 insertions(+), 3 deletions(-) create mode 100644 src/test/java/com/auth0/jwk/BucketImplTest.java diff --git a/src/main/java/com/auth0/jwk/BucketImpl.java b/src/main/java/com/auth0/jwk/BucketImpl.java index 6748afb..10a7ebf 100644 --- a/src/main/java/com/auth0/jwk/BucketImpl.java +++ b/src/main/java/com/auth0/jwk/BucketImpl.java @@ -20,8 +20,9 @@ class BucketImpl implements Bucket { private AtomicInteger available; private AtomicLong lastTokenAddedAt; - public BucketImpl(int size, long rate, TimeUnit rateUnit) { + BucketImpl(int size, long rate, TimeUnit rateUnit) { assertPositiveValue(size, MAX_BUCKET_SIZE, "Invalid bucket size."); + assertPositiveValue(rate, "Invalid bucket refill rate."); this.size = size; this.available = new AtomicInteger(size); this.lastTokenAddedAt = new AtomicLong(System.currentTimeMillis()); @@ -43,7 +44,7 @@ public void run() { addToken(); } }; - executorService.scheduleAtFixedRate(refillTask, rate, rate, rateUnit); + executorService.scheduleAtFixedRate(refillTask, 0, rate, rateUnit); } private void addToken() { @@ -60,6 +61,10 @@ private void assertPositiveValue(int value, int maxValue, String exceptionMessag } } + private void assertPositiveValue(Number value, String exceptionMessage) { + this.assertPositiveValue(value.intValue(), value.intValue(), exceptionMessage); + } + private void log(String message) { long msDiff = System.currentTimeMillis() - beginTime; System.out.println(String.format("%-8d - %s", msDiff, message)); @@ -79,7 +84,7 @@ public long willLeakIn(int count) { return 0; } - long nextIn = rateUnit.toMillis(rate) - System.currentTimeMillis() - lastTokenAddedAt.get(); + long nextIn = rateUnit.toMillis(rate) - (System.currentTimeMillis() - lastTokenAddedAt.get()); final int remaining = count - av - 1; if (remaining > 0) { nextIn += rateUnit.toMillis(rate) * remaining; diff --git a/src/test/java/com/auth0/jwk/BucketImplTest.java b/src/test/java/com/auth0/jwk/BucketImplTest.java new file mode 100644 index 0000000..7a1b6e0 --- /dev/null +++ b/src/test/java/com/auth0/jwk/BucketImplTest.java @@ -0,0 +1,95 @@ +package com.auth0.jwk; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +import java.util.concurrent.TimeUnit; + +import static org.hamcrest.Matchers.*; +import static org.junit.Assert.assertThat; + +public class BucketImplTest { + @Rule + public ExpectedException exception = ExpectedException.none(); + + private static final long RATE = 500; + private static final int SIZE = 5; + + @Test + public void shouldThrowOnCreateWithNegativeSize() throws Exception { + exception.expect(IllegalArgumentException.class); + exception.expectMessage("Invalid bucket size."); + new BucketImpl(-1, 10, TimeUnit.SECONDS); + } + + @Test + public void shouldThrowOnCreateWithMoreThanMaxSize() throws Exception { + exception.expect(IllegalArgumentException.class); + exception.expectMessage("Invalid bucket size."); + new BucketImpl(Integer.MAX_VALUE, 10, TimeUnit.SECONDS); + } + + @Test + public void shouldThrowOnCreateWithNegativeRate() throws Exception { + exception.expect(IllegalArgumentException.class); + exception.expectMessage("Invalid bucket refill rate."); + new BucketImpl(10, -1, TimeUnit.SECONDS); + } + + @Test + public void shouldThrowWhenLeakingMoreThanBucketSize() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.SECONDS); + exception.expect(IllegalArgumentException.class); + exception.expectMessage(String.format("Cannot consume %d tokens when the BucketImpl size is %d!", SIZE + 1, SIZE)); + bucket.willLeakIn(SIZE + 1); + } + + @Test + public void shouldThrowWhenConsumingMoreThanBucketSize() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.SECONDS); + exception.expect(IllegalArgumentException.class); + exception.expectMessage(String.format("Cannot consume %d tokens when the BucketImpl size is %d!", SIZE + 1, SIZE)); + bucket.consume(SIZE + 1); + } + + @Test + public void shouldCreateFullBucket() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); + assertThat(bucket, notNullValue()); + assertThat(bucket.willLeakIn(SIZE), equalTo(0L)); + assertThat(bucket.willLeakIn(), equalTo(0L)); + } + + @Test + public void shouldAddOneTokenPerRate() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); + assertThat(bucket, notNullValue()); + assertThat(bucket.consume(SIZE), equalTo(true)); + + assertThat(bucket.willLeakIn(), lessThan(RATE)); + Thread.sleep(RATE); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.willLeakIn(), lessThan(RATE)); + } + + @Test + public void shouldConsumeAllBucket() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); + assertThat(bucket, notNullValue()); + assertThat(bucket.consume(SIZE), equalTo(true)); + assertThat(bucket.consume(), equalTo(false)); + } + + @Test + public void shouldConsume() throws Exception { + Bucket bucket = new BucketImpl(SIZE, RATE, TimeUnit.MILLISECONDS); + assertThat(bucket, notNullValue()); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.consume(), equalTo(true)); + assertThat(bucket.consume(), equalTo(false)); + } +} \ No newline at end of file From 9ff974a5706facb4136b1542161bbd26a9cf06f5 Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Wed, 16 Nov 2016 12:22:18 -0300 Subject: [PATCH 5/6] change bucket size units to Long --- src/main/java/com/auth0/jwk/Bucket.java | 6 +++-- src/main/java/com/auth0/jwk/BucketImpl.java | 23 ++++++++----------- .../com/auth0/jwk/JwkProviderBuilder.java | 2 +- .../java/com/auth0/jwk/BucketImplTest.java | 11 ++------- 4 files changed, 17 insertions(+), 25 deletions(-) diff --git a/src/main/java/com/auth0/jwk/Bucket.java b/src/main/java/com/auth0/jwk/Bucket.java index 873c269..d601dc7 100644 --- a/src/main/java/com/auth0/jwk/Bucket.java +++ b/src/main/java/com/auth0/jwk/Bucket.java @@ -15,9 +15,10 @@ interface Bucket { /** * 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(int count); + long willLeakIn(long count); /** * Tries to consume one token from the Bucket. @@ -29,7 +30,8 @@ interface Bucket { /** * 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(int count); + boolean consume(long count); } diff --git a/src/main/java/com/auth0/jwk/BucketImpl.java b/src/main/java/com/auth0/jwk/BucketImpl.java index 10a7ebf..3e1d21a 100644 --- a/src/main/java/com/auth0/jwk/BucketImpl.java +++ b/src/main/java/com/auth0/jwk/BucketImpl.java @@ -3,7 +3,6 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; /** @@ -11,20 +10,18 @@ */ class BucketImpl implements Bucket { - private static final int MAX_BUCKET_SIZE = Integer.MAX_VALUE - 5; - - private final int size; + private final long size; private final long rate; private final TimeUnit rateUnit; private final long beginTime = System.currentTimeMillis(); - private AtomicInteger available; + private AtomicLong available; private AtomicLong lastTokenAddedAt; - BucketImpl(int size, long rate, TimeUnit rateUnit) { - assertPositiveValue(size, MAX_BUCKET_SIZE, "Invalid bucket size."); + BucketImpl(long size, long rate, TimeUnit rateUnit) { + assertPositiveValue(size, "Invalid bucket size."); assertPositiveValue(rate, "Invalid bucket refill rate."); this.size = size; - this.available = new AtomicInteger(size); + this.available = new AtomicLong(size); this.lastTokenAddedAt = new AtomicLong(System.currentTimeMillis()); this.rate = rate; this.rateUnit = rateUnit; @@ -55,7 +52,7 @@ private void addToken() { lastTokenAddedAt.set(System.currentTimeMillis()); } - private void assertPositiveValue(int value, int maxValue, String exceptionMessage) { + private void assertPositiveValue(long value, long maxValue, String exceptionMessage) { if (value < 1 || value > maxValue) { throw new IllegalArgumentException(exceptionMessage); } @@ -76,16 +73,16 @@ public long willLeakIn() { } @Override - public long willLeakIn(int count) { + public long willLeakIn(long count) { assertPositiveValue(count, size, String.format("Cannot consume %d tokens when the BucketImpl size is %d!", count, size)); - int av = available.get(); + long av = available.get(); if (av >= count) { log(String.format("%d Tokens are available already.", count)); return 0; } long nextIn = rateUnit.toMillis(rate) - (System.currentTimeMillis() - lastTokenAddedAt.get()); - final int remaining = count - av - 1; + final long remaining = count - av - 1; if (remaining > 0) { nextIn += rateUnit.toMillis(rate) * remaining; } @@ -99,7 +96,7 @@ public boolean consume() { } @Override - public boolean consume(int count) { + 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); diff --git a/src/main/java/com/auth0/jwk/JwkProviderBuilder.java b/src/main/java/com/auth0/jwk/JwkProviderBuilder.java index f3074fc..26f5249 100644 --- a/src/main/java/com/auth0/jwk/JwkProviderBuilder.java +++ b/src/main/java/com/auth0/jwk/JwkProviderBuilder.java @@ -81,7 +81,7 @@ public JwkProviderBuilder rateLimited(boolean rateLimited) { * @param unit unit of time for the expire of jwk * @return the builder */ - public JwkProviderBuilder rateLimited(int bucketSize, long refillRate, TimeUnit unit) { + public JwkProviderBuilder rateLimited(long bucketSize, long refillRate, TimeUnit unit) { bucket = new BucketImpl(bucketSize, refillRate, unit); return this; } diff --git a/src/test/java/com/auth0/jwk/BucketImplTest.java b/src/test/java/com/auth0/jwk/BucketImplTest.java index 7a1b6e0..7f77d4a 100644 --- a/src/test/java/com/auth0/jwk/BucketImplTest.java +++ b/src/test/java/com/auth0/jwk/BucketImplTest.java @@ -13,8 +13,8 @@ public class BucketImplTest { @Rule public ExpectedException exception = ExpectedException.none(); - private static final long RATE = 500; - private static final int SIZE = 5; + private static final long RATE = 500L; + private static final long SIZE = 5L; @Test public void shouldThrowOnCreateWithNegativeSize() throws Exception { @@ -23,13 +23,6 @@ public void shouldThrowOnCreateWithNegativeSize() throws Exception { new BucketImpl(-1, 10, TimeUnit.SECONDS); } - @Test - public void shouldThrowOnCreateWithMoreThanMaxSize() throws Exception { - exception.expect(IllegalArgumentException.class); - exception.expectMessage("Invalid bucket size."); - new BucketImpl(Integer.MAX_VALUE, 10, TimeUnit.SECONDS); - } - @Test public void shouldThrowOnCreateWithNegativeRate() throws Exception { exception.expect(IllegalArgumentException.class); From 4300878acdf656503b61cc1bcf3aa5e39d8411a4 Mon Sep 17 00:00:00 2001 From: Luciano Balmaceda Date: Wed, 16 Nov 2016 15:42:03 -0300 Subject: [PATCH 6/6] remove bucket logs --- src/main/java/com/auth0/jwk/BucketImpl.java | 11 ----------- src/test/java/com/auth0/jwk/BucketImplTest.java | 4 ++-- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/auth0/jwk/BucketImpl.java b/src/main/java/com/auth0/jwk/BucketImpl.java index 3e1d21a..248d6d5 100644 --- a/src/main/java/com/auth0/jwk/BucketImpl.java +++ b/src/main/java/com/auth0/jwk/BucketImpl.java @@ -13,7 +13,6 @@ class BucketImpl implements Bucket { private final long size; private final long rate; private final TimeUnit rateUnit; - private final long beginTime = System.currentTimeMillis(); private AtomicLong available; private AtomicLong lastTokenAddedAt; @@ -47,7 +46,6 @@ public void run() { private void addToken() { if (available.get() < size) { available.incrementAndGet(); - log(String.format("Added 1 token.. Current state: %d/%d", available.get(), size)); } lastTokenAddedAt.set(System.currentTimeMillis()); } @@ -62,11 +60,6 @@ private void assertPositiveValue(Number value, String exceptionMessage) { this.assertPositiveValue(value.intValue(), value.intValue(), exceptionMessage); } - private void log(String message) { - long msDiff = System.currentTimeMillis() - beginTime; - System.out.println(String.format("%-8d - %s", msDiff, message)); - } - @Override public long willLeakIn() { return willLeakIn(1); @@ -77,7 +70,6 @@ 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) { - log(String.format("%d Tokens are available already.", count)); return 0; } @@ -86,7 +78,6 @@ public long willLeakIn(long count) { if (remaining > 0) { nextIn += rateUnit.toMillis(rate) * remaining; } - log(String.format("Can't consume %d. Actual state is: %d/%d. Retry in %d ms.", count, available.get(), size, nextIn)); return nextIn; } @@ -100,10 +91,8 @@ 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); - log(String.format("Consumed %d tokens.. Current state: %d/%d", count, available.get(), size)); return true; } - log(String.format("Not enough tokens available to consume %d", count)); System.out.println(); return false; } diff --git a/src/test/java/com/auth0/jwk/BucketImplTest.java b/src/test/java/com/auth0/jwk/BucketImplTest.java index 7f77d4a..152ecf0 100644 --- a/src/test/java/com/auth0/jwk/BucketImplTest.java +++ b/src/test/java/com/auth0/jwk/BucketImplTest.java @@ -60,10 +60,10 @@ public void shouldAddOneTokenPerRate() throws Exception { assertThat(bucket, notNullValue()); assertThat(bucket.consume(SIZE), equalTo(true)); - assertThat(bucket.willLeakIn(), lessThan(RATE)); + assertThat(bucket.willLeakIn(), lessThanOrEqualTo(RATE)); Thread.sleep(RATE); assertThat(bucket.consume(), equalTo(true)); - assertThat(bucket.willLeakIn(), lessThan(RATE)); + assertThat(bucket.willLeakIn(), lessThanOrEqualTo(RATE)); } @Test