feat(rate-limiter): add ConnectInfo-based tests for rate limiter (ADR-025)
This commit is contained in:
@@ -206,4 +206,108 @@ mod tests {
|
|||||||
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])));
|
assert!(!limiter.check_and_consume(IpAddr::from([192, 168, 1, 1])));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn middleware_uses_connect_info_without_xff_header() {
|
||||||
|
let limiter = make_limiter(10, 5);
|
||||||
|
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route("/", axum::routing::get(|| async { "ok" }))
|
||||||
|
.layer(axum::middleware::from_fn_with_state(
|
||||||
|
limiter,
|
||||||
|
rate_limit_middleware,
|
||||||
|
))
|
||||||
|
.into_make_service_with_connect_info::<std::net::SocketAddr>();
|
||||||
|
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
|
||||||
|
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
for _ in 0..5 {
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
}
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn middleware_rejects_without_connect_info() {
|
||||||
|
let limiter = make_limiter(10, 20);
|
||||||
|
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route("/", axum::routing::get(|| async { "ok" }))
|
||||||
|
.layer(axum::middleware::from_fn_with_state(
|
||||||
|
limiter,
|
||||||
|
rate_limit_middleware,
|
||||||
|
));
|
||||||
|
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
|
||||||
|
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn middleware_ignores_xff_header_same_bucket() {
|
||||||
|
let limiter = make_limiter(10, 2);
|
||||||
|
|
||||||
|
let app = axum::Router::new()
|
||||||
|
.route("/", axum::routing::get(|| async { "ok" }))
|
||||||
|
.layer(axum::middleware::from_fn_with_state(
|
||||||
|
limiter,
|
||||||
|
rate_limit_middleware,
|
||||||
|
))
|
||||||
|
.into_make_service_with_connect_info::<std::net::SocketAddr>();
|
||||||
|
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
|
||||||
|
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.1")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.2")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.3")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -228,6 +228,81 @@ async fn test_rate_limit_per_ip_independent() {
|
|||||||
assert_eq!(resp2.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
assert_eq!(resp2.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_rate_limit_without_connect_info_rejected_with_429() {
|
||||||
|
let mut config = reverse_proxy::config::test_fixtures::test_dynamic_config();
|
||||||
|
config.rate_limit = reverse_proxy::config::RateLimitConfig {
|
||||||
|
requests_per_second: 10,
|
||||||
|
burst: 20,
|
||||||
|
};
|
||||||
|
let config_arc = Arc::new(ArcSwap::from_pointee(config));
|
||||||
|
let limiter = Arc::new(reverse_proxy::rate_limit::RateLimiter::new(config_arc));
|
||||||
|
|
||||||
|
let app = Router::new().route("/", get(|| async { "ok" })).layer(
|
||||||
|
axum::middleware::from_fn_with_state(
|
||||||
|
limiter,
|
||||||
|
reverse_proxy::rate_limit::rate_limit_middleware,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
|
||||||
|
tokio::spawn(async { axum::serve(listener, app).await.unwrap() });
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
let body = resp.text().await.unwrap();
|
||||||
|
assert_eq!(body, "Too Many Requests");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_rate_limit_xff_header_ignored_same_bucket() {
|
||||||
|
let mut config = reverse_proxy::config::test_fixtures::test_dynamic_config();
|
||||||
|
config.rate_limit = reverse_proxy::config::RateLimitConfig {
|
||||||
|
requests_per_second: 10,
|
||||||
|
burst: 2,
|
||||||
|
};
|
||||||
|
let config_arc = Arc::new(ArcSwap::from_pointee(config));
|
||||||
|
let limiter = Arc::new(reverse_proxy::rate_limit::RateLimiter::new(config_arc));
|
||||||
|
|
||||||
|
let app = make_rate_limit_app(limiter);
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let addr = listener.local_addr().unwrap();
|
||||||
|
|
||||||
|
tokio::spawn(async { axum::serve(listener, app).await.unwrap() });
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.1")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.2")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::OK);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.get(format!("http://127.0.0.1:{}/", addr.port()))
|
||||||
|
.header("X-Forwarded-For", "10.0.0.3")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(resp.status(), reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_rate_limit_eviction_task() {
|
async fn test_rate_limit_eviction_task() {
|
||||||
let mut config = reverse_proxy::config::test_fixtures::test_dynamic_config();
|
let mut config = reverse_proxy::config::test_fixtures::test_dynamic_config();
|
||||||
|
|||||||
Reference in New Issue
Block a user