use std::convert::Infallible;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Empty, Full};
use hyper::body::{Bytes, Incoming};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Request, Response};
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo};
use tokio::net::TcpListener;
use waf_core::{
Acquired, BucketParams, Config, Decision, FailMode,
Phase, RateLimitAction, RateLimitConfig, RateLimitKey, RequestContext, StateStore, WafMode, WafModule,
};
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use waf_proxy::Proxy;
struct BoomModule;
impl WafModule for BoomModule {
fn id(&self) -> &str {
"boom"
}
fn phase(&self) -> Phase {
Phase::RequestLine
}
fn init(&mut self, _: &Config) {}
fn inspect(&self, ctx: &RequestContext) -> Decision {
if ctx.normalized.headers.iter().any(|(k, _)| k == "x-boom") {
panic!("boom: simulated module defect");
}
Decision::Allow
}
}
struct KeepModule;
impl WafModule for KeepModule {
fn id(&self) -> &str {
"keep"
}
fn phase(&self) -> Phase {
Phase::Headers
}
fn init(&mut self, _: &Config) {}
fn structural(&self) -> bool {
true
}
fn inspect(&self, ctx: &RequestContext) -> Decision {
if ctx.normalized.headers.iter().any(|(k, _)| k == "x-keep") {
Decision::Block { rule_id: "keep".into(), reason: "x-keep present".into() }
} else {
Decision::Allow
}
}
}
type TestBody = BoxBody<Bytes, hyper::Error>;
fn bytes_body(data: impl Into<Bytes>) -> TestBody {
Full::new(data.into())
.map_err(|never| match never {})
.boxed()
}
fn empty_body() -> TestBody {
Empty::new().map_err(|never| match never {}).boxed()
}
fn test_client() -> Client<HttpConnector, TestBody> {
Client::builder(TokioExecutor::new()).build(HttpConnector::new())
}
async fn start_echo_backend() -> std::net::SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
tokio::spawn(async move {
let io = TokioIo::new(stream);
http1::Builder::new()
.serve_connection(
io,
service_fn(|req: Request<Incoming>| async move {
let pq = req
.uri()
.path_and_query()
.map(|pq| pq.as_str())
.unwrap_or("/")
.to_string();
let body = bytes_body(format!("ok:{pq}"));
Ok::<Response<TestBody>, Infallible>(Response::new(body))
}),
)
.await
.ok();
});
}
});
addr
}
fn make_config(backend_addr: std::net::SocketAddr) -> Config {
let mut c = Config::default();
c.proxy.listen = "127.0.0.1:0".parse().unwrap();
c.proxy.backend = format!("http://{backend_addr}");
c
}
fn make_config_rate_limited(backend_addr: std::net::SocketAddr, requests: u32) -> Config {
let mut cfg = make_config(backend_addr);
cfg.waf.mode = WafMode::Blocking;
cfg.rate_limit = RateLimitConfig {
enabled: true,
key: RateLimitKey::ClientIp,
requests,
window_seconds: 60,
burst: Some(requests),
action: RateLimitAction::Block,
score: 5,
max_tracked_keys: 1000,
};
cfg
}
struct DenyAllStore {
calls: Arc<AtomicUsize>,
}
impl StateStore for DenyAllStore {
fn try_acquire(&self, _key: &str, _cost: f64, _params: BucketParams) -> Acquired {
self.calls.fetch_add(1, Ordering::Relaxed);
Acquired { allowed: false, tokens_remaining: 0.0 }
}
}
#[tokio::test]
async fn injected_state_store_governs_rate_limit_decision() {
let backend = start_echo_backend().await;
let cfg = make_config_rate_limited(backend, 100);
let calls = Arc::new(AtomicUsize::new(0));
let store = Arc::new(DenyAllStore { calls: calls.clone() });
let proxy = Proxy::builder(&cfg).state_store(store).build().await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
let resp = client
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}/x"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), 429, "injected store must govern the rate-limit decision");
assert!(calls.load(Ordering::Relaxed) >= 1, "injected store must be consulted on the datapath");
}
#[tokio::test]
async fn rate_limit_returns_429_after_budget_exhausted() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config_rate_limited(backend, 1)).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
let send = |path: &str| {
client.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}{path}"))
.body(empty_body())
.unwrap(),
)
};
let first = send("/a").await.unwrap();
assert_eq!(first.status(), 200);
let second = send("/b").await.unwrap();
assert_eq!(second.status(), 429);
assert!(second.headers().contains_key("retry-after"), "429 must carry Retry-After");
}
#[tokio::test]
async fn distinct_clients_behind_same_lb_get_separate_buckets() {
let backend = start_echo_backend().await;
let mut cfg = make_config_rate_limited(backend, 1); cfg.network.trusted_proxies = vec!["127.0.0.1".to_string()];
cfg.network.client_ip_header = "x-forwarded-for".to_string();
cfg.network.trusted_hops = 1;
let proxy = Proxy::bind(&cfg).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
let send = |xff: &str| {
client.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}/x"))
.header("x-forwarded-for", xff)
.body(empty_body())
.unwrap(),
)
};
assert_eq!(send("1.2.3.4").await.unwrap().status(), 200);
assert_eq!(send("5.6.7.8").await.unwrap().status(), 200);
assert_eq!(send("1.2.3.4").await.unwrap().status(), 429);
}
#[tokio::test]
async fn passthrough_forwards_get_request() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let resp = test_client()
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}/hello"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), 200);
let body = resp.collect().await.unwrap().to_bytes();
assert_eq!(body.as_ref(), b"ok:/hello");
}
#[tokio::test]
async fn passthrough_forwards_path_and_query() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let resp = test_client()
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}/api/v1?foo=bar&x=1"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), 200);
let body = resp.collect().await.unwrap().to_bytes();
assert_eq!(body.as_ref(), b"ok:/api/v1?foo=bar&x=1");
}
#[tokio::test]
async fn passthrough_returns_502_when_backend_down() {
let dead_addr = {
let l = TcpListener::bind("127.0.0.1:0").await.unwrap();
let a = l.local_addr().unwrap();
drop(l);
a
};
let proxy = Proxy::bind(&make_config(dead_addr)).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let resp = test_client()
.request(
Request::builder()
.method("GET")
.uri(format!("http://{proxy_addr}/test"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), 502);
}
#[tokio::test]
async fn upstream_down_fail_open_returns_503() {
let dead_addr = {
let l = TcpListener::bind("127.0.0.1:0").await.unwrap();
let a = l.local_addr().unwrap();
drop(l);
a
};
let mut cfg = make_config(dead_addr);
cfg.resilience.on_upstream_error = FailMode::FailOpen;
let proxy = Proxy::bind(&cfg).await.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let resp = test_client()
.request(
Request::builder()
.uri(format!("http://{proxy_addr}/x"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), 503);
}
#[tokio::test]
async fn panic_in_module_is_isolated_other_clients_unaffected() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind_with_modules(&make_config(backend), vec![Box::new(BoomModule)])
.await
.unwrap();
let proxy_addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
let boom = client.request(
Request::builder()
.uri(format!("http://{proxy_addr}/boom"))
.header("x-boom", "1")
.body(empty_body())
.unwrap(),
);
let normal = client.request(
Request::builder()
.uri(format!("http://{proxy_addr}/ok"))
.body(empty_body())
.unwrap(),
);
let (boom_resp, normal_resp) = tokio::join!(boom, normal);
assert_eq!(boom_resp.unwrap().status(), 200);
assert_eq!(normal_resp.unwrap().status(), 200);
}
struct CountingBoom {
hits: Arc<AtomicUsize>,
}
impl WafModule for CountingBoom {
fn id(&self) -> &str {
"counting_boom"
}
fn phase(&self) -> Phase {
Phase::RequestLine
}
fn init(&mut self, _: &Config) {}
fn inspect(&self, ctx: &RequestContext) -> Decision {
if ctx.normalized.headers.iter().any(|(k, _)| k == "x-boom") {
self.hits.fetch_add(1, Ordering::SeqCst);
panic!("boom: simulated module defect");
}
Decision::Allow
}
}
fn make_config_resilience(backend: std::net::SocketAddr, policy: FailMode) -> Config {
let mut cfg = make_config(backend);
cfg.waf.mode = WafMode::Blocking;
cfg.waf.paranoia_level = 3; cfg.resilience.on_internal_error = policy;
cfg
}
async fn get_with_boom(
client: &Client<HttpConnector, TestBody>,
addr: std::net::SocketAddr,
query: &str,
) -> u16 {
client
.request(
Request::builder()
.uri(format!("http://{addr}/?{query}"))
.header("x-boom", "1")
.body(empty_body())
.unwrap(),
)
.await
.unwrap()
.status()
.as_u16()
}
const WEAK_CANDIDATE_QUERY: &str = "include=ftp%3A%2F%2Fhost.example%2Fresource";
const SQLI_QUERY: &str = "q=1%20UNION%20SELECT%20a%2Cb%20FROM%20users--";
#[tokio::test]
async fn fail_open_isolates_panic_and_keeps_surviving_protection() {
let backend = start_echo_backend().await;
let hits = Arc::new(AtomicUsize::new(0));
let cfg = make_config_resilience(backend, FailMode::FailOpen);
let proxy = Proxy::bind_with_modules(&cfg, vec![Box::new(CountingBoom { hits: Arc::clone(&hits) })])
.await
.unwrap();
let addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(
get_with_boom(&client, addr, WEAK_CANDIDATE_QUERY).await,
200,
"fail_open skips the panicking module; a sub-threshold request is served"
);
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"the module must actually have panicked (prefilter did not skip inspection)"
);
assert_eq!(
get_with_boom(&client, addr, SQLI_QUERY).await,
403,
"surviving modules still block a real attack after a peer module panics"
);
assert_eq!(
hits.load(Ordering::SeqCst),
2,
"the panicking module ran on the attack request too"
);
}
#[tokio::test]
async fn fail_closed_turns_panic_into_a_deny() {
let backend = start_echo_backend().await;
let hits = Arc::new(AtomicUsize::new(0));
let cfg = make_config_resilience(backend, FailMode::FailClosed);
let proxy = Proxy::bind_with_modules(&cfg, vec![Box::new(CountingBoom { hits: Arc::clone(&hits) })])
.await
.unwrap();
let addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(
get_with_boom(&client, addr, WEAK_CANDIDATE_QUERY).await,
403,
"fail_closed turns an internal panic into a deny"
);
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"the module must actually have panicked"
);
}
#[tokio::test]
async fn upstream_down_waf_still_inspects_and_blocks_attacks() {
let dead_addr = {
let l = TcpListener::bind("127.0.0.1:0").await.unwrap();
let a = l.local_addr().unwrap();
drop(l);
a
};
let mut cfg = make_config(dead_addr);
cfg.waf.mode = WafMode::Blocking;
let proxy = Proxy::bind(&cfg).await.unwrap();
let addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(
sqli_status(&client, addr).await,
403,
"the WAF must still inspect and block attacks even when the upstream is down"
);
let benign = client
.request(
Request::builder()
.uri(format!("http://{addr}/healthz"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap();
assert_eq!(
benign.status(),
502,
"benign traffic passes inspection and reaches the dead upstream → 502"
);
}
#[tokio::test]
async fn corrupt_reload_keeps_protection_no_unprotected_window() {
let backend = start_echo_backend().await;
let mut cfg = make_config(backend);
cfg.waf.mode = WafMode::Blocking;
cfg.waf.paranoia_level = 3;
let proxy = Proxy::bind(&cfg).await.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(sqli_status(&client, addr).await, 403);
let corrupt = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\nparanoia_level = 0\n"
);
let path = write_cfg("corrupt-pl0", &corrupt);
assert!(
reloader.reload_from(&path).is_err(),
"an invalid config (paranoia_level = 0) must be rejected"
);
assert_eq!(
sqli_status(&client, addr).await,
403,
"protection must never drop on a rejected reload (no unprotected window)"
);
std::fs::remove_file(&path).ok();
}
async fn raw_status_line(addr: std::net::SocketAddr, raw: &str) -> String {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let mut s = tokio::net::TcpStream::connect(addr).await.unwrap();
s.write_all(raw.as_bytes()).await.unwrap();
let mut buf = Vec::new();
s.read_to_end(&mut buf).await.unwrap();
String::from_utf8_lossy(&buf).lines().next().unwrap_or("").to_string()
}
#[tokio::test]
async fn smuggling_te_list_rejected_through_proxy() {
let backend = start_echo_backend().await;
let mut cfg = make_config(backend);
cfg.waf.mode = WafMode::Blocking;
let proxy = Proxy::bind(&cfg).await.unwrap();
let addr = proxy.local_addr().unwrap();
tokio::spawn(proxy.run());
let raw = "POST / HTTP/1.1\r\nHost: x\r\nTransfer-Encoding: gzip, chunked\r\nConnection: close\r\n\r\n0\r\n\r\n";
let status = raw_status_line(addr, raw).await;
assert!(status.contains("400"), "expected 400, got: {status:?}");
let raw_ok = "GET /ok HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n";
let status_ok = raw_status_line(addr, raw_ok).await;
assert!(status_ok.contains("200"), "expected 200, got: {status_ok:?}");
}
fn write_cfg(tag: &str, contents: &str) -> PathBuf {
let nanos = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
let mut p = std::env::temp_dir();
p.push(format!("waf-reload-{tag}-{nanos}.toml"));
std::fs::write(&p, contents).unwrap();
p
}
async fn sqli_status(client: &Client<HttpConnector, TestBody>, addr: std::net::SocketAddr) -> u16 {
client
.request(
Request::builder()
.uri(format!("http://{addr}/?q=1%20UNION%20SELECT%20a%2Cb%20FROM%20users--"))
.body(empty_body())
.unwrap(),
)
.await
.unwrap()
.status()
.as_u16()
}
async fn keep_status(client: &Client<HttpConnector, TestBody>, addr: std::net::SocketAddr) -> u16 {
client
.request(
Request::builder()
.uri(format!("http://{addr}/"))
.header("x-keep", "1")
.body(empty_body())
.unwrap(),
)
.await
.unwrap()
.status()
.as_u16()
}
#[tokio::test]
async fn factory_modules_survive_reload() {
let backend = start_echo_backend().await;
let mut cfg = make_config(backend);
cfg.waf.mode = WafMode::Blocking;
let proxy = Proxy::builder(&cfg)
.module_factory(|| Ok(vec![Box::new(KeepModule) as Box<dyn WafModule>]))
.build()
.await
.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(keep_status(&client, addr).await, 403, "factory module must block at boot");
let reload_cfg = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\n"
);
let path = write_cfg("factory-survive", &reload_cfg);
reloader.reload_from(&path).expect("valid reload should succeed");
assert_eq!(
keep_status(&client, addr).await,
403,
"injected module must survive the reload (core 0.3 module_factory)"
);
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn factory_error_on_reload_aborts_and_keeps_modules() {
let backend = start_echo_backend().await;
let mut cfg = make_config(backend);
cfg.waf.mode = WafMode::Blocking;
let calls = Arc::new(AtomicUsize::new(0));
let calls_f = Arc::clone(&calls);
let proxy = Proxy::builder(&cfg)
.module_factory(move || {
if calls_f.fetch_add(1, Ordering::SeqCst) == 0 {
Ok(vec![Box::new(KeepModule) as Box<dyn WafModule>])
} else {
Err("simulated: schema became invalid on disk".into())
}
})
.build()
.await
.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(keep_status(&client, addr).await, 403, "factory module must block at boot");
let reload_cfg = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\n"
);
let path = write_cfg("factory-err", &reload_cfg);
assert!(
reloader.reload_from(&path).is_err(),
"a factory error must abort the reload even when the config itself is valid"
);
assert_eq!(
keep_status(&client, addr).await,
403,
"protection must never drop on a factory-failed reload (modules preserved)"
);
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn reload_valid_activates_new_rules() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
assert_eq!(sqli_status(&client, addr).await, 200);
let cfg = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\nblock_threshold = 5\n"
);
let path = write_cfg("valid", &cfg);
reloader.reload_from(&path).expect("valid reload should succeed");
assert_eq!(sqli_status(&client, addr).await, 403);
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn reload_invalid_keeps_old_config() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap(); let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
let bad = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\n[network]\ntrusted_hops = 99\n"
);
let path = write_cfg("invalid", &bad);
assert!(reloader.reload_from(&path).is_err(), "invalid config must be rejected");
assert_eq!(sqli_status(&client, addr).await, 200);
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn reload_under_concurrent_load_has_no_race() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let cfg = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"detection-only\"\n"
);
let path = write_cfg("load", &cfg);
let client = test_client();
let mut handles = Vec::new();
for _ in 0..40 {
let c = client.clone();
handles.push(tokio::spawn(async move {
c.request(
Request::builder()
.uri(format!("http://{addr}/x"))
.body(empty_body())
.unwrap(),
)
.await
}));
}
for _ in 0..10 {
reloader.reload_from(&path).expect("reload should succeed");
}
for h in handles {
let resp = h.await.unwrap().unwrap();
assert_eq!(resp.status(), 200);
}
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn reload_listen_change_is_ignored_old_kept() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config(backend)).await.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
let cfg = format!(
"[proxy]\nlisten = \"127.0.0.1:9\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"detection-only\"\n"
);
let path = write_cfg("listen", &cfg);
reloader.reload_from(&path).expect("reload should still succeed");
let resp = client
.request(Request::builder().uri(format!("http://{addr}/x")).body(empty_body()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), 200);
std::fs::remove_file(&path).ok();
}
#[tokio::test]
async fn reload_does_not_reset_rate_limit_buckets() {
let backend = start_echo_backend().await;
let proxy = Proxy::bind(&make_config_rate_limited(backend, 1)).await.unwrap();
let addr = proxy.local_addr().unwrap();
let reloader = proxy.reloader();
tokio::spawn(proxy.run());
let client = test_client();
let send = |path: &str| {
client.request(
Request::builder().uri(format!("http://{addr}{path}")).body(empty_body()).unwrap(),
)
};
assert_eq!(send("/a").await.unwrap().status(), 200);
assert_eq!(send("/b").await.unwrap().status(), 429);
let cfg = format!(
"[proxy]\nlisten = \"{addr}\"\nbackend = \"http://{backend}\"\n[waf]\nmode = \"blocking\"\n[rate_limit]\nenabled = true\nrequests = 1\nwindow_seconds = 60\nburst = 1\n"
);
let path = write_cfg("rl", &cfg);
reloader.reload_from(&path).expect("reload should succeed");
assert_eq!(send("/c").await.unwrap().status(), 429, "bucket must survive reload");
std::fs::remove_file(&path).ok();
}