Skip to content

Commit e5ca35a

Browse files
jknackjknack
committed
Rate limit handler
1 parent 6fb133e commit e5ca35a

File tree

7 files changed

+334
-4
lines changed

7 files changed

+334
-4
lines changed

‎docs/asciidoc/handlers.adoc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ include::handlers/cors.adoc[]
66

77
include::handlers/head.adoc[]
88

9+
include::handlers/rate-limit.adoc[]
10+
911
include::handlers/ssl.adoc[]
1012

1113
include::handlers/trace.adoc[]

‎docs/asciidoc/handlers/rate-limit.adoc‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
=== RateLimitHandler
2+
3+
Rate limit handler using https://github.com/vladimir-bukhtoyarov/bucket4j[Bucket4j].
4+
5+
Add the dependency to your project:
6+
7+
[dependency, artifactId="bucket4j-core"]
8+
.
9+
10+
.10 requests per minute
11+
[source, java, role="primary"]
12+
----
13+
{
14+
Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
15+
Bucket bucket = Bucket4j.builder().addLimit(limit).build(); <1>
16+
17+
before(new RateLimitHandler(bucket)); <2>
18+
}
19+
----
20+
21+
.Kotlin
22+
[source, kotlin, role="secondary"]
23+
----
24+
{
25+
val limit = Bandwidth.simple(10, Duration.ofMinutes(1))
26+
val bucket = Bucket4j.builder().addLimit(limit).build() <1>
27+
28+
before(RateLimitHandler(bucket)) <2>
29+
}
30+
----
31+
32+
<1> Creates a bucket
33+
<2> Install the RateLimitHandler
34+
35+
.10 requests per minute per ip/remote address
36+
[source, java, role="primary"]
37+
----
38+
{
39+
before(new RateLimitHandler(remoteAddress -> {
40+
Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
41+
return Bucket4j.builder().addLimit(limit).build();
42+
}));
43+
}
44+
----
45+
46+
.Kotlin
47+
[source, kotlin, role="secondary"]
48+
----
49+
{
50+
before(RateLimitHandler {remoteAddress ->
51+
val limit = Bandwidth.simple(10, Duration.ofMinutes(1))
52+
Bucket4j.builder().addLimit(limit).build()
53+
})
54+
}
55+
----
56+
57+
.10 requests per minute per header value
58+
[source, java, role="primary"]
59+
----
60+
{
61+
before(new RateLimitHandler(key -> {
62+
Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
63+
return Bucket4j.builder().addLimit(limit).build();
64+
}, "ApiKey"));
65+
}
66+
----
67+
68+
.Kotlin
69+
[source, kotlin, role="secondary"]
70+
----
71+
{
72+
before(RateLimitHandler {key ->
73+
val limit = Bandwidth.simple(10, Duration.ofMinutes(1))
74+
Bucket4j.builder().addLimit(limit).build()
75+
}, "ApiKey")
76+
}
77+
----
78+
79+
.10 requests per minute
80+
[source, java, role="primary"]
81+
----
82+
{
83+
Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
84+
Bucket bucket = Bucket4j.builder().addLimit(limit).build(); <1>
85+
86+
before(new RateLimitHandler(bucket)); <2>
87+
}
88+
----
89+
90+
.Kotlin
91+
[source, kotlin, role="secondary"]
92+
----
93+
{
94+
val limit = Bandwidth.simple(10, Duration.ofMinutes(1))
95+
val bucket = Bucket4j.builder().addLimit(limit).build() <1>
96+
97+
before(RateLimitHandler(bucket)) <2>
98+
}
99+
----
100+
101+
<1> Creates a bucket
102+
<2> Install the RateLimitHandler
103+
104+
.Rate limit in a cluster
105+
[source, java, role="primary"]
106+
----
107+
{
108+
ProxyManager<String> buckets = ...;
109+
before(RateLimitHandler.cluster(key -> {
110+
return buckets.getProxy(key, () -> {
111+
return ...;
112+
});
113+
}));
114+
}
115+
----
116+
117+
.Kotlin
118+
[source, kotlin, role="secondary"]
119+
----
120+
{
121+
val buckets = ...;
122+
before(RateLimitHandler.cluster {key ->
123+
buckets.getProxy(key) {
124+
....
125+
}
126+
})
127+
}
128+
----
129+
130+
For using it inside a cluster you need to configure one of the bucket4j options for https://github.com/vladimir-bukhtoyarov/bucket4j#supported-back-ends[clustering].

‎docs/asciidoc/routing.adoc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1139,7 +1139,7 @@ class App: Kooby({
11391139
[TIP]
11401140
====
11411141
Composition is a great option for modularization. You can easily develop/test/deploy each
1142-
application indendepently and compose them all in another application.
1142+
application independently and compose them all in another application.
11431143
11441144
We do provide <<mvc-api, MVC API>> as another alternative for modularization.
11451145
====

‎jooby/pom.xml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,14 @@
213213
<optional>true</optional>
214214
</dependency>
215215

216+
<!-- bucket4j -->
217+
<dependency>
218+
<groupId>com.github.vladimir-bukhtoyarov</groupId>
219+
<artifactId>bucket4j-core</artifactId>
220+
<optional>true</optional>
221+
</dependency>
222+
223+
216224
<!-- Test dependencies -->
217225
<dependency>
218226
<groupId>org.junit.jupiter</groupId>

‎jooby/src/main/java/io/jooby/RateLimitHandler.java‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
/**
2+
* Jooby https://jooby.io
3+
* Apache License Version 2.0 https://jooby.io/LICENSE.txt
4+
* Copyright 2014 Edgar Espina
5+
*/
6+
package io.jooby;
7+
8+
import io.github.bucket4j.Bucket;
9+
import io.github.bucket4j.ConsumptionProbe;
10+
11+
import javax.annotation.Nonnull;
12+
import java.util.Map;
13+
import java.util.concurrent.ConcurrentHashMap;
14+
import java.util.function.Function;
15+
16+
import static java.util.concurrent.TimeUnit.NANOSECONDS;
17+
18+
/**
19+
* Rate limit handler using https://github.com/vladimir-bukhtoyarov/bucket4j.
20+
*
21+
* NOTE: bucket4j must be included as part of your project dependencies (classpath).
22+
*
23+
* Example 1: 10 requests per minute
24+
* <pre>{@code
25+
* {
26+
* Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
27+
* Bucket bucket = Bucket4j.builder().addLimit(limit).build();
28+
*
29+
* before(new RateLimitHandler(bucket));
30+
* }
31+
* }</pre>
32+
*
33+
* Example 2: 10 requests per minute per IP address
34+
* <pre>{@code
35+
* {
36+
* before(new RateLimitHandler(remoteAddress -> {
37+
* Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
38+
* return Bucket4j.builder().addLimit(limit).build();
39+
* }));
40+
* }
41+
* }</pre>
42+
*
43+
* Example 3: 10 requests per minute using an <code>ApiKey</code> header.
44+
* <pre>{@code
45+
* {
46+
* before(new RateLimitHandler(key -> {
47+
* Bandwidth limit = Bandwidth.simple(10, Duration.ofMinutes(1));
48+
* return Bucket4j.builder().addLimit(limit).build();
49+
* }, "ApiKey"));
50+
* }
51+
* }</pre>
52+
*
53+
* Example 4: Rate limit in a cluster
54+
* <pre>{@code
55+
* {
56+
* // Get one of the proxy manager from bucket4j
57+
* ProxyManager<String> buckets = ...;
58+
* before(RateLimitHandler.cluster(key -> {
59+
* buckets.getProxy(key, () -> {
60+
* return Bucket4j.configurationBuilder()
61+
* .addLimit(Bandwidth.classic(100, Refill.intervally(100, Duration.ofMinutes(1))))
62+
* .build();
63+
* });
64+
* }));
65+
* }
66+
* }</pre>
67+
*
68+
* @author edgar
69+
* @since 2.5.2
70+
*/
71+
public class RateLimitHandler implements Route.Before {
72+
73+
private final Function<Context, Bucket> factory;
74+
75+
/**
76+
* Rate limit per IP/Remote Address.
77+
*
78+
* @param bucketFactory Bucket factory.
79+
*/
80+
public RateLimitHandler(@Nonnull SneakyThrows.Function<String, Bucket> bucketFactory) {
81+
this(bucketFactory, Context::getRemoteAddress);
82+
}
83+
84+
/**
85+
* Rate limit per header key.
86+
*
87+
* @param bucketFactory Bucket factory.
88+
* @param headerName Header to use as key.
89+
*/
90+
public RateLimitHandler(@Nonnull SneakyThrows.Function<String, Bucket> bucketFactory,
91+
@Nonnull String headerName) {
92+
this(bucketFactory, ctx -> ctx.header(headerName).value());
93+
}
94+
95+
/**
96+
* Rate limiter with a custom key provider.
97+
*
98+
* @param bucketFactory Bucket factory.
99+
* @param classifier Key provider.
100+
*/
101+
public RateLimitHandler(@Nonnull SneakyThrows.Function<String, Bucket> bucketFactory,
102+
@Nonnull SneakyThrows.Function<Context, String> classifier) {
103+
this(byKey(bucketFactory, classifier));
104+
}
105+
106+
/**
107+
* Rate limiter with a shared/global bucket.
108+
*
109+
* @param bucket Bucket to use.
110+
*/
111+
public RateLimitHandler(@Nonnull Bucket bucket) {
112+
this((Function<Context, Bucket>) ctx -> bucket);
113+
}
114+
115+
private RateLimitHandler(Function<Context, Bucket> factory) {
116+
this.factory = factory;
117+
}
118+
119+
/**
120+
* Rate limiter per IP/Remote address using a cluster.
121+
*
122+
* @param proxyManager Cluster bucket configuration.
123+
* @return Rate limiter.
124+
*/
125+
public static @Nonnull RateLimitHandler cluster(
126+
@Nonnull SneakyThrows.Function<String, Bucket> proxyManager) {
127+
return cluster(proxyManager, Context::getRemoteAddress);
128+
}
129+
130+
/**
131+
* Rate limiter per header key using a cluster.
132+
*
133+
* @param proxyManager Cluster bucket configuration.
134+
* @param headerName Header to use as key.
135+
* @return Rate limiter.
136+
*/
137+
public static @Nonnull RateLimitHandler cluster(
138+
@Nonnull SneakyThrows.Function<String, Bucket> proxyManager, @Nonnull String headerName) {
139+
return cluster(proxyManager, ctx -> ctx.header(headerName).value());
140+
}
141+
142+
/**
143+
* Rate limiter per key using a cluster.
144+
*
145+
* @param proxyManager Cluster bucket configuration.
146+
* @param classifier Key provider.
147+
* @return Rate limiter.
148+
*/
149+
public static RateLimitHandler cluster(
150+
@Nonnull SneakyThrows.Function<String, Bucket> proxyManager,
151+
@Nonnull SneakyThrows.Function<Context, String> classifier) {
152+
return new RateLimitHandler(
153+
(Function<Context, Bucket>) ctx -> proxyManager.apply(classifier.apply(ctx)));
154+
}
155+
156+
@Override public void apply(@Nonnull Context ctx) throws Exception {
157+
Bucket bucket = factory.apply(ctx);
158+
// tryConsume returns false immediately if no tokens available with the bucket
159+
ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
160+
if (probe.isConsumed()) {
161+
ctx.setResponseHeader("X-Rate-Limit-Remaining", probe.getRemainingTokens());
162+
} else {
163+
ctx.setResponseHeader("X-Rate-Limit-Retry-After-Milliseconds",
164+
NANOSECONDS.toMillis(probe.getNanosToWaitForRefill()));
165+
ctx.send(StatusCode.TOO_MANY_REQUESTS);
166+
}
167+
}
168+
169+
private static Function<Context, Bucket> byKey(
170+
SneakyThrows.Function<String, Bucket> bucketFactory,
171+
SneakyThrows.Function<Context, String> classifier) {
172+
Map<String, Bucket> buckets = new ConcurrentHashMap<>();
173+
return ctx -> buckets.computeIfAbsent(classifier.apply(ctx), bucketFactory);
174+
}
175+
}

‎modules/jooby-bom/pom.xml‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,13 @@
1616
<!-- THIS FILE IS AUTO GENERATED. DON'T EDIT -->
1717

1818
<properties>
19-
<jooby.version>2.5.1</jooby.version>
19+
<jooby.version>2.5.2-SNAPSHOT</jooby.version>
2020
<HikariCP.version>3.4.1</HikariCP.version>
2121
<archetype-packaging.version>3.1.2</archetype-packaging.version>
2222
<asm.version>7.2</asm.version>
2323
<auto-service.version>1.0-rc6</auto-service.version>
2424
<boringssl.version>2.0.27.Final</boringssl.version>
25+
<bucket4j-core.version>4.7.0</bucket4j-core.version>
2526
<checkstyle.version>8.28</checkstyle.version>
2627
<commons-io.version>2.6</commons-io.version>
2728
<compile-testing.version>0.18</compile-testing.version>
@@ -48,8 +49,8 @@
4849
<jdbi.version>3.12.0</jdbi.version>
4950
<jetty.version>9.4.25.v20191220</jetty.version>
5051
<jfiglet.version>0.0.8</jfiglet.version>
51-
<jooby-maven-plugin.version>2.5.1</jooby-maven-plugin.version>
52-
<jooby.version>2.5.1</jooby.version>
52+
<jooby-maven-plugin.version>2.5.2-SNAPSHOT</jooby-maven-plugin.version>
53+
<jooby.version>2.5.2-SNAPSHOT</jooby.version>
5354
<json.version>20190722</json.version>
5455
<jsonwebtoken.version>0.10.7</jsonwebtoken.version>
5556
<jsr305.version>3.0.2</jsr305.version>
@@ -511,6 +512,11 @@
511512
<artifactId>guice</artifactId>
512513
<version>${guice.version}</version>
513514
</dependency>
515+
<dependency>
516+
<groupId>com.github.vladimir-bukhtoyarov</groupId>
517+
<artifactId>bucket4j-core</artifactId>
518+
<version>${bucket4j-core.version}</version>
519+
</dependency>
514520
<dependency>
515521
<groupId>org.jetbrains.kotlin</groupId>
516522
<artifactId>kotlin-stdlib</artifactId>

0 commit comments

Comments
 (0)