Skip to content

Commit 4b3033e

Browse files
connyayconnyay
authored
feat: add Rate Limiter binding support (#845)
- Add `Env::rate_limiter()` and `RouteContext::rate_limiter()` methods for accessing rate limiter bindings - Export `RateLimitOutcome` for public API usage - Implement test endpoints for basic, custom key, bulk, and reset scenarios
1 parent 5f5d27f commit 4b3033e

File tree

9 files changed

+286
-3
lines changed

9 files changed

+286
-3
lines changed

‎test/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ mod kv;
2828
mod put_raw;
2929
mod queue;
3030
mod r2;
31+
mod rate_limit;
3132
mod request;
3233
mod router;
3334
mod secret_store;

‎test/src/rate_limit.rs‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
use super::SomeSharedData;
2+
use std::collections::HashMap;
3+
use worker::{js_sys, Env, Request, Response, Result};
4+
5+
#[worker::send]
6+
pub async fn handle_rate_limit_check(
7+
_req: Request,
8+
env: Env,
9+
_data: SomeSharedData,
10+
) -> Result<Response> {
11+
let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;
12+
13+
// Use a fixed key for testing
14+
let outcome = rate_limiter.limit("test-key".to_string()).await?;
15+
16+
Response::from_json(&serde_json::json!({
17+
"success": outcome.success,
18+
}))
19+
}
20+
21+
#[worker::send]
22+
pub async fn handle_rate_limit_with_key(
23+
req: Request,
24+
env: Env,
25+
_data: SomeSharedData,
26+
) -> Result<Response> {
27+
let uri = req.url()?;
28+
let segments = uri.path_segments().unwrap().collect::<Vec<_>>();
29+
let key = segments.get(2).unwrap_or(&"default-key");
30+
31+
let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;
32+
let outcome = rate_limiter.limit(key.to_string()).await?;
33+
34+
Response::from_json(&serde_json::json!({
35+
"success": outcome.success,
36+
"key": key,
37+
}))
38+
}
39+
40+
#[worker::send]
41+
pub async fn handle_rate_limit_bulk_test(
42+
_req: Request,
43+
env: Env,
44+
_data: SomeSharedData,
45+
) -> Result<Response> {
46+
let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;
47+
48+
// Test multiple requests to verify rate limiting behavior
49+
let mut results = Vec::new();
50+
for i in 0..15 {
51+
let key = format!("bulk-test-{}", i % 3); // Use 3 different keys
52+
let outcome = rate_limiter.limit(key.clone()).await?;
53+
results.push(serde_json::json!({
54+
"index": i,
55+
"key": key,
56+
"success": outcome.success,
57+
}));
58+
}
59+
60+
Response::from_json(&serde_json::json!({
61+
"results": results,
62+
}))
63+
}
64+
65+
#[worker::send]
66+
pub async fn handle_rate_limit_reset(
67+
_req: Request,
68+
env: Env,
69+
_data: SomeSharedData,
70+
) -> Result<Response> {
71+
let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;
72+
73+
// Use a unique key to avoid interference with other tests
74+
let key = format!("reset-test-{}", js_sys::Date::now());
75+
76+
// Make multiple requests with the same key
77+
let mut outcomes = HashMap::new();
78+
for i in 0..12 {
79+
let outcome = rate_limiter.limit(key.clone()).await?;
80+
outcomes.insert(format!("request_{}", i + 1), outcome.success);
81+
}
82+
83+
Response::from_json(&outcomes)
84+
}

‎test/src/router.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::{
22
alarm, analytics_engine, assets, auto_response, cache, container, counter, d1, durable, fetch,
3-
form, js_snippets, kv, put_raw, queue, r2, request, secret_store, service, socket, sql_counter,
4-
sql_iterator, user, ws, SomeSharedData, GLOBAL_STATE,
3+
form, js_snippets, kv, put_raw, queue, r2, rate_limit, request, secret_store, service, socket,
4+
sql_counter, sql_iterator, user, ws, SomeSharedData, GLOBAL_STATE,
55
};
66
#[cfg(feature = "http")]
77
use std::convert::TryInto;
@@ -227,6 +227,10 @@ macro_rules! add_routes (
227227
add_route!($obj, get, sync, "/test-panic", handle_test_panic);
228228
add_route!($obj, post, "/container/echo", container::handle_container);
229229
add_route!($obj, get, "/container/ws", container::handle_container);
230+
add_route!($obj, get, "/rate-limit/check", rate_limit::handle_rate_limit_check);
231+
add_route!($obj, get, format_route!("/rate-limit/key/{}", "key"), rate_limit::handle_rate_limit_with_key);
232+
add_route!($obj, get, "/rate-limit/bulk-test", rate_limit::handle_rate_limit_bulk_test);
233+
add_route!($obj, get, "/rate-limit/reset", rate_limit::handle_rate_limit_reset);
230234
});
231235

232236
#[cfg(feature = "http")]

‎test/tests/mf.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,14 @@ const mf_instance = new Miniflare({
112112
HTTP_ANALYTICS: {
113113
scriptName: "mini-analytics-engine" // mock out analytics engine binding to the "mini-analytics-engine" worker
114114
}
115+
},
116+
ratelimits: {
117+
TEST_RATE_LIMITER: {
118+
simple: {
119+
limit: 10,
120+
period: 60,
121+
}
122+
}
115123
}
116124
},
117125
{

‎test/tests/rate_limit.spec.ts‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
import { describe, test, expect } from "vitest";
2+
import { mf, mfUrl } from "./mf";
3+
4+
describe("rate limit", () => {
5+
test("basic rate limit check", async () => {
6+
const resp = await mf.dispatchFetch(`${mfUrl}rate-limit/check`);
7+
expect(resp.status).toBe(200);
8+
const data = await resp.json() as { success: boolean };
9+
expect(data).toHaveProperty("success");
10+
expect(data.success).toBe(true);
11+
});
12+
13+
test("rate limit with custom key", async () => {
14+
const key = "test-key-123";
15+
const resp = await mf.dispatchFetch(`${mfUrl}rate-limit/key/${key}`);
16+
expect(resp.status).toBe(200);
17+
const data = await resp.json() as { success: boolean; key: string };
18+
expect(data).toHaveProperty("success");
19+
expect(data).toHaveProperty("key");
20+
expect(data.key).toBe(key);
21+
expect(data.success).toBe(true);
22+
});
23+
24+
test("different keys have independent limits", async () => {
25+
// Test that different keys have separate rate limits
26+
const key1 = "user-1";
27+
const key2 = "user-2";
28+
29+
const resp1 = await mf.dispatchFetch(`${mfUrl}rate-limit/key/${key1}`);
30+
const resp2 = await mf.dispatchFetch(`${mfUrl}rate-limit/key/${key2}`);
31+
32+
expect(resp1.status).toBe(200);
33+
expect(resp2.status).toBe(200);
34+
35+
const data1 = await resp1.json() as { success: boolean; key: string };
36+
const data2 = await resp2.json() as { success: boolean; key: string };
37+
38+
expect(data1.success).toBe(true);
39+
expect(data2.success).toBe(true);
40+
expect(data1.key).toBe(key1);
41+
expect(data2.key).toBe(key2);
42+
});
43+
44+
test("bulk rate limit test", async () => {
45+
const resp = await mf.dispatchFetch(`${mfUrl}rate-limit/bulk-test`);
46+
expect(resp.status).toBe(200);
47+
const data = await resp.json() as { results: Array<{ index: number; key: string; success: boolean }> };
48+
expect(data).toHaveProperty("results");
49+
expect(Array.isArray(data.results)).toBe(true);
50+
expect(data.results.length).toBe(15);
51+
52+
// Check that results have the expected structure
53+
data.results.forEach((result, index: number) => {
54+
expect(result).toHaveProperty("index");
55+
expect(result).toHaveProperty("key");
56+
expect(result).toHaveProperty("success");
57+
expect(result.index).toBe(index);
58+
expect(typeof result.success).toBe("boolean");
59+
});
60+
61+
// We're using 3 different keys (bulk-test-0, bulk-test-1, bulk-test-2)
62+
// with a limit of 10 per 60 seconds. Each key is used 5 times (15 requests / 3 keys).
63+
// All requests should succeed since each key stays under the limit of 10.
64+
65+
// Group results by key
66+
const resultsByKey: Record<string, Array<{ index: number; key: string; success: boolean }>> = {};
67+
data.results.forEach((result) => {
68+
if (!resultsByKey[result.key]) {
69+
resultsByKey[result.key] = [];
70+
}
71+
resultsByKey[result.key].push(result);
72+
});
73+
74+
// Should have exactly 3 keys
75+
expect(Object.keys(resultsByKey).length).toBe(3);
76+
77+
// Each key should have 5 requests, all successful (under limit of 10)
78+
Object.entries(resultsByKey).forEach(([key, results]) => {
79+
expect(results.length).toBe(5);
80+
results.forEach((result) => {
81+
expect(result.success).toBe(true);
82+
});
83+
});
84+
});
85+
86+
test("rate limit reset with unique keys", async () => {
87+
const resp = await mf.dispatchFetch(`${mfUrl}rate-limit/reset`);
88+
expect(resp.status).toBe(200);
89+
const data = await resp.json() as Record<string, boolean>;
90+
91+
// Should have 12 request results
92+
expect(Object.keys(data).length).toBe(12);
93+
94+
// Check that we have the expected keys
95+
for (let i = 1; i <= 12; i++) {
96+
expect(data).toHaveProperty(`request_${i}`);
97+
expect(typeof data[`request_${i}`]).toBe("boolean");
98+
}
99+
100+
// With a limit of 10 per 60 seconds, the first 10 requests MUST succeed
101+
// and requests 11 and 12 MUST fail
102+
for (let i = 1; i <= 10; i++) {
103+
expect(data[`request_${i}`]).toBe(true);
104+
}
105+
106+
// Requests 11 and 12 must be rate limited
107+
expect(data["request_11"]).toBe(false);
108+
expect(data["request_12"]).toBe(false);
109+
});
110+
111+
test("multiple rapid requests with same key", async () => {
112+
// Generate a unique key for this test
113+
const testKey = `rapid-test-${Date.now()}`;
114+
115+
// Make multiple rapid requests with the same key
116+
const promises = [];
117+
for (let i = 0; i < 5; i++) {
118+
promises.push(mf.dispatchFetch(`${mfUrl}rate-limit/key/${testKey}`));
119+
}
120+
121+
const responses = await Promise.all(promises);
122+
123+
// All responses should be successful (200 status)
124+
responses.forEach(resp => {
125+
expect(resp.status).toBe(200);
126+
});
127+
128+
// Parse the responses
129+
const results = await Promise.all(responses.map(r => r.json())) as Array<{ success: boolean; key: string }>;
130+
131+
// All should have the same key
132+
results.forEach(data => {
133+
expect(data.key).toBe(testKey);
134+
expect(data).toHaveProperty("success");
135+
});
136+
137+
// With limit of 10, all 5 requests should succeed
138+
results.forEach((data) => {
139+
expect(data.success).toBe(true);
140+
});
141+
});
142+
143+
test("sequential requests enforce rate limit", async () => {
144+
// Generate a unique key for this test to avoid interference
145+
const testKey = `sequential-test-${Date.now()}`;
146+
147+
// Make 15 sequential requests with the same key
148+
// With a limit of 10 per 60 seconds, first 10 should succeed, rest should fail
149+
const results: Array<{ success: boolean; key: string }> = [];
150+
for (let i = 0; i < 15; i++) {
151+
const resp = await mf.dispatchFetch(`${mfUrl}rate-limit/key/${testKey}`);
152+
expect(resp.status).toBe(200);
153+
const data = await resp.json() as { success: boolean; key: string };
154+
results.push(data);
155+
}
156+
157+
// Verify first 10 requests succeed
158+
for (let i = 0; i < 10; i++) {
159+
expect(results[i].success).toBe(true);
160+
expect(results[i].key).toBe(testKey);
161+
}
162+
163+
// Verify requests 11-15 are rate limited
164+
for (let i = 10; i < 15; i++) {
165+
expect(results[i].success).toBe(false);
166+
expect(results[i].key).toBe(testKey);
167+
}
168+
});
169+
});

‎test/wrangler.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,3 +84,8 @@ secret_name = "secret-name"
8484
class_name = "EchoContainer"
8585
image = "./container-echo/Dockerfile"
8686
max_instances = 1
87+
88+
[[ratelimits]]
89+
name = "TEST_RATE_LIMITER"
90+
namespace_id = "1"
91+
simple = { limit = 10, period = 60 }

‎worker/src/env.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::analytics_engine::AnalyticsEngineDataset;
44
#[cfg(feature = "d1")]
55
use crate::d1::D1Database;
66
use crate::kv::KvStore;
7+
use crate::rate_limit::RateLimiter;
78
use crate::Ai;
89
#[cfg(feature = "queue")]
910
use crate::Queue;
@@ -122,6 +123,11 @@ impl Env {
122123
pub fn secret_store(&self, binding: &str) -> Result<SecretStore> {
123124
self.get_binding(binding)
124125
}
126+
127+
/// Access a Rate Limiter by the binding name configured in your wrangler.toml file.
128+
pub fn rate_limiter(&self, binding: &str) -> Result<RateLimiter> {
129+
self.get_binding(binding)
130+
}
125131
}
126132

127133
pub trait EnvBinding: Sized + JsCast {

‎worker/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ pub use crate::kv::{KvError, KvStore};
187187
#[cfg(feature = "queue")]
188188
pub use crate::queue::*;
189189
pub use crate::r2::*;
190-
pub use crate::rate_limit::RateLimiter;
190+
pub use crate::rate_limit::{RateLimitOutcome, RateLimiter};
191191
pub use crate::request::{FromRequest, Request};
192192
pub use crate::request_init::*;
193193
pub use crate::response::{EncodeBody, IntoResponse, Response, ResponseBody, ResponseBuilder};

‎worker/src/router.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use crate::{
77
durable::ObjectNamespace,
88
env::{Env, Secret, Var},
99
http::Method,
10+
rate_limit::RateLimiter,
1011
request::Request,
1112
response::Response,
1213
Bucket, Fetcher, KvStore, Result,
@@ -118,6 +119,11 @@ impl<D> RouteContext<D> {
118119
pub fn d1(&self, binding: &str) -> Result<crate::D1Database> {
119120
self.env.d1(binding)
120121
}
122+
123+
/// Access a Rate Limiter by the binding name configured in your wrangler.toml file.
124+
pub fn rate_limiter(&self, binding: &str) -> Result<RateLimiter> {
125+
self.env.rate_limiter(binding)
126+
}
121127
}
122128

123129
impl Router<'_, ()> {

0 commit comments

Comments
 (0)