Implement token bucket rate limiting with IPv6 /64 normalization
- Add TokenBucket with nodelay semantics (nginx limit_req burst nodelay) - Per-IP rate limiting: IPv4 /32, IPv6 /64 prefix normalization - DashMap for concurrent access, ArcSwap for lock-free config reads - Background eviction task for stale entry cleanup - 429 response with plain text body, RATE_LIMIT log prefix - Config reload adopts new rate/burst on next request without clearing state - Unit tests for bucket algorithm and IPv6 normalization - Integration tests for 429 responses and per-IP independence
This commit is contained in:
@@ -1,2 +1,185 @@
|
||||
#[allow(dead_code)]
|
||||
pub struct TokenBucket;
|
||||
use std::net::{IpAddr, Ipv6Addr};
|
||||
use std::time::Instant;
|
||||
|
||||
pub struct TokenBucket {
|
||||
pub tokens: f64,
|
||||
pub last_refill: Instant,
|
||||
pub rate: f64,
|
||||
pub max: u32,
|
||||
pub last_access: Instant,
|
||||
}
|
||||
|
||||
impl TokenBucket {
|
||||
pub fn new(rate: f64, max: u32) -> Self {
|
||||
Self {
|
||||
tokens: max as f64,
|
||||
last_refill: Instant::now(),
|
||||
rate,
|
||||
max,
|
||||
last_access: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn try_consume(&mut self, rate: f64, max: u32) -> bool {
|
||||
self.refill(rate, max);
|
||||
self.last_access = Instant::now();
|
||||
|
||||
if self.tokens >= 1.0 {
|
||||
self.tokens -= 1.0;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn refill(&mut self, rate: f64, max: u32) {
|
||||
let now = Instant::now();
|
||||
let elapsed = now.duration_since(self.last_refill).as_millis() as f64;
|
||||
let tokens_to_add = (elapsed * rate) / 1000.0;
|
||||
self.tokens = (self.tokens + tokens_to_add).min(max as f64);
|
||||
self.last_refill = now;
|
||||
|
||||
if max < self.max {
|
||||
self.tokens = self.tokens.min(max as f64);
|
||||
}
|
||||
self.max = max;
|
||||
self.rate = rate;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_ip(ip: IpAddr) -> IpAddr {
|
||||
match ip {
|
||||
IpAddr::V4(_) => ip,
|
||||
IpAddr::V6(v6) => {
|
||||
let segments = v6.segments();
|
||||
let normalized = Ipv6Addr::new(
|
||||
segments[0],
|
||||
segments[1],
|
||||
segments[2],
|
||||
segments[3],
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
);
|
||||
IpAddr::V6(normalized)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::net::{Ipv4Addr, Ipv6Addr};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
#[test]
|
||||
fn token_bucket_starts_full() {
|
||||
let bucket = TokenBucket::new(10.0, 20);
|
||||
assert_eq!(bucket.tokens, 20.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_consume_within_burst() {
|
||||
let mut bucket = TokenBucket::new(10.0, 20);
|
||||
for _ in 0..20 {
|
||||
assert!(bucket.try_consume(10.0, 20));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_reject_when_empty() {
|
||||
let mut bucket = TokenBucket::new(10.0, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(bucket.try_consume(10.0, 5));
|
||||
}
|
||||
assert!(!bucket.try_consume(10.0, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_refills_over_time() {
|
||||
let mut bucket = TokenBucket::new(1000.0, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(bucket.try_consume(1000.0, 5));
|
||||
}
|
||||
assert!(!bucket.try_consume(1000.0, 5));
|
||||
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
assert!(bucket.try_consume(1000.0, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_caps_at_max() {
|
||||
let mut bucket = TokenBucket::new(10.0, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(bucket.try_consume(10.0, 5));
|
||||
}
|
||||
thread::sleep(Duration::from_millis(500));
|
||||
bucket.try_consume(10.0, 5);
|
||||
assert!(bucket.tokens <= 5.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_config_reload_caps_tokens() {
|
||||
let mut bucket = TokenBucket::new(10.0, 20);
|
||||
assert!(bucket.try_consume(10.0, 20));
|
||||
|
||||
let new_max = 10;
|
||||
bucket.refill(10.0, new_max);
|
||||
assert!(bucket.tokens <= new_max as f64 + 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_bucket_config_reload_adopts_new_rate() {
|
||||
let mut bucket = TokenBucket::new(10.0, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(bucket.try_consume(10.0, 5));
|
||||
}
|
||||
assert!(!bucket.try_consume(10.0, 5));
|
||||
|
||||
thread::sleep(Duration::from_millis(200));
|
||||
assert!(bucket.try_consume(50.0, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ipv4_unchanged() {
|
||||
let ip = IpAddr::from(Ipv4Addr::new(192, 168, 1, 1));
|
||||
assert_eq!(normalize_ip(ip), ip);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ipv6_truncates_to_64_prefix() {
|
||||
let ip = IpAddr::from(Ipv6Addr::new(
|
||||
0x2001, 0x0db8, 0x85a3, 0x0001, 0x1234, 0x5678, 0x9abc, 0xdef0,
|
||||
));
|
||||
let normalized = normalize_ip(ip);
|
||||
let expected = IpAddr::from(Ipv6Addr::new(0x2001, 0x0db8, 0x85a3, 0x0001, 0, 0, 0, 0));
|
||||
assert_eq!(normalized, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ipv6_already_64_prefix() {
|
||||
let ip = IpAddr::from(Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 0));
|
||||
let normalized = normalize_ip(ip);
|
||||
assert_eq!(normalized, ip);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_ipv6_link_local() {
|
||||
let ip = IpAddr::from(Ipv6Addr::new(
|
||||
0xfe80, 0, 0, 0, 0x1234, 0x5678, 0x9abc, 0xdef0,
|
||||
));
|
||||
let normalized = normalize_ip(ip);
|
||||
let expected = IpAddr::from(Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 0));
|
||||
assert_eq!(normalized, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_loopback_ipv6() {
|
||||
let ip = IpAddr::from(Ipv6Addr::LOCALHOST);
|
||||
let normalized = normalize_ip(ip);
|
||||
let expected = IpAddr::from(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0));
|
||||
assert_eq!(normalized, expected);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1 +1,208 @@
|
||||
pub mod bucket;
|
||||
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use axum::extract::Request;
|
||||
use axum::http::StatusCode;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::IntoResponse;
|
||||
use dashmap::DashMap;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::DynamicConfig;
|
||||
use bucket::{normalize_ip, TokenBucket};
|
||||
|
||||
pub struct RateLimiter {
|
||||
buckets: DashMap<IpAddr, TokenBucket>,
|
||||
config: Arc<ArcSwap<DynamicConfig>>,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
pub fn new(config: Arc<ArcSwap<DynamicConfig>>) -> Self {
|
||||
Self {
|
||||
buckets: DashMap::new(),
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn check_and_consume(&self, ip: IpAddr) -> bool {
|
||||
let normalized = normalize_ip(ip);
|
||||
let config = self.config.load();
|
||||
let rate = config.rate_limit.requests_per_second as f64;
|
||||
let burst = config.rate_limit.burst;
|
||||
|
||||
let entry = self.buckets.entry(normalized);
|
||||
match entry {
|
||||
dashmap::mapref::entry::Entry::Occupied(mut occupied) => {
|
||||
occupied.get_mut().try_consume(rate, burst)
|
||||
}
|
||||
dashmap::mapref::entry::Entry::Vacant(vacant) => {
|
||||
let mut bucket = TokenBucket::new(rate, burst);
|
||||
let result = bucket.try_consume(rate, burst);
|
||||
vacant.insert(bucket);
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn evict_stale(&self, max_age: Duration) {
|
||||
let cutoff = std::time::Instant::now() - max_age;
|
||||
self.buckets.retain(|_, bucket| bucket.last_access > cutoff);
|
||||
}
|
||||
|
||||
pub fn contains_ip(&self, ip: IpAddr) -> bool {
|
||||
self.buckets.contains_key(&normalize_ip(ip))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn rate_limit_middleware(
|
||||
axum::extract::State(limiter): axum::extract::State<Arc<RateLimiter>>,
|
||||
req: Request,
|
||||
next: Next,
|
||||
) -> axum::response::Response {
|
||||
let client_ip = req
|
||||
.headers()
|
||||
.get("x-forwarded-for")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.split(',').next())
|
||||
.and_then(|v| v.trim().parse::<IpAddr>().ok())
|
||||
.or_else(|| {
|
||||
req.extensions()
|
||||
.get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
|
||||
.map(|ci| ci.ip())
|
||||
});
|
||||
|
||||
let Some(ip) = client_ip else {
|
||||
return next.run(req).await;
|
||||
};
|
||||
|
||||
let host = req
|
||||
.headers()
|
||||
.get("host")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("-");
|
||||
|
||||
let path = req.uri().path();
|
||||
|
||||
if !limiter.check_and_consume(ip) {
|
||||
warn!(
|
||||
"RATE_LIMIT client_ip={} host={} path={} status=429",
|
||||
ip, host, path
|
||||
);
|
||||
return (StatusCode::TOO_MANY_REQUESTS, "Too Many Requests").into_response();
|
||||
}
|
||||
|
||||
next.run(req).await
|
||||
}
|
||||
|
||||
pub fn start_eviction_task(
|
||||
limiter: Arc<RateLimiter>,
|
||||
interval: Duration,
|
||||
max_age: Duration,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
let mut interval_timer = tokio::time::interval(interval);
|
||||
loop {
|
||||
interval_timer.tick().await;
|
||||
limiter.evict_stale(max_age);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::test_fixtures;
|
||||
use std::net::Ipv6Addr;
|
||||
use std::time::Duration;
|
||||
|
||||
fn make_limiter(rps: u32, burst: u32) -> Arc<RateLimiter> {
|
||||
let mut config = test_fixtures::test_dynamic_config();
|
||||
config.rate_limit = crate::config::RateLimitConfig {
|
||||
requests_per_second: rps,
|
||||
burst,
|
||||
};
|
||||
let config_arc = Arc::new(ArcSwap::from_pointee(config));
|
||||
Arc::new(RateLimiter::new(config_arc))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_and_consume_allows_within_burst() {
|
||||
let limiter = make_limiter(10, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_and_consume_rejects_above_burst() {
|
||||
let limiter = make_limiter(10, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
}
|
||||
assert!(!limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_and_consume_per_ip_independent() {
|
||||
let limiter = make_limiter(10, 5);
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
}
|
||||
assert!(!limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 2])));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_and_consume_ipv6_normalized_to_64() {
|
||||
let limiter = make_limiter(10, 5);
|
||||
let ip1 = IpAddr::from(Ipv6Addr::new(
|
||||
0x2001, 0x0db8, 0x85a3, 0x0001, 0x1111, 0x2222, 0x3333, 0x4444,
|
||||
));
|
||||
let ip2 = IpAddr::from(Ipv6Addr::new(
|
||||
0x2001, 0x0db8, 0x85a3, 0x0001, 0x5555, 0x6666, 0x7777, 0x8888,
|
||||
));
|
||||
for _ in 0..5 {
|
||||
assert!(limiter.check_and_consume(ip1));
|
||||
}
|
||||
assert!(!limiter.check_and_consume(ip2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evict_removes_stale_entries() {
|
||||
let limiter = make_limiter(10, 20);
|
||||
let ip1 = IpAddr::from([192, 168, 1, 1]);
|
||||
let ip2 = IpAddr::from([192, 168, 1, 2]);
|
||||
|
||||
limiter.check_and_consume(ip1);
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
limiter.check_and_consume(ip2);
|
||||
|
||||
limiter.evict_stale(Duration::from_millis(25));
|
||||
assert!(!limiter.buckets.contains_key(&normalize_ip(ip1)));
|
||||
assert!(limiter.buckets.contains_key(&normalize_ip(ip2)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_reload_adopted_on_next_request() {
|
||||
let config = test_fixtures::test_dynamic_config();
|
||||
let config_arc = Arc::new(ArcSwap::from_pointee(config));
|
||||
let limiter = Arc::new(RateLimiter::new(config_arc.clone()));
|
||||
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
|
||||
let mut new_config = test_fixtures::test_dynamic_config();
|
||||
new_config.rate_limit = crate::config::RateLimitConfig {
|
||||
requests_per_second: 10,
|
||||
burst: 2,
|
||||
};
|
||||
config_arc.store(Arc::new(new_config));
|
||||
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
assert!(limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
assert!(!limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user