use std::sync::Arc;
use std::thread;
use std::time::Duration;
use waf_core::testkit::Request;
use waf_core::{Decision, RequestContext, WafModule};
use waf_wasm::{WasmModule, WasmOptions};
const BLOCK_WAT: &str = r#"
(module
(import "env" "proxy_get_header_map_pairs" (func $get_headers (param i32 i32 i32) (result i32)))
(import "env" "proxy_send_local_response" (func $send (param i32 i32 i32 i32 i32 i32 i32 i32) (result i32)))
(memory (export "memory") 1)
(global $bump (mut i32) (i32.const 1024))
(func (export "proxy_on_memory_allocate") (param $n i32) (result i32)
(local $p i32) (local $need i32)
(local.set $p (global.get $bump))
(local.set $need (i32.add (local.get $p) (local.get $n)))
(global.set $bump (local.get $need))
(block $d (loop $g
(br_if $d (i32.le_u (local.get $need) (i32.mul (memory.size) (i32.const 65536))))
(if (i32.eq (memory.grow (i32.const 1)) (i32.const -1)) (then (return (i32.const -1))))
(br $g)))
(local.get $p))
(func (export "proxy_on_context_create") (param i32 i32))
(func (export "proxy_on_request_headers") (param i32 i32 i32) (result i32)
(local $ptr i32) (local $size i32) (local $i i32) (local $base i32)
(drop (call $get_headers (i32.const 0) (i32.const 0) (i32.const 4)))
(local.set $ptr (i32.load (i32.const 0)))
(local.set $size (i32.load (i32.const 4)))
(local.set $i (i32.const 0))
(block $found
(block $notfound
(loop $scan
(br_if $notfound (i32.gt_s (local.get $i) (i32.sub (local.get $size) (i32.const 3))))
(local.set $base (i32.add (local.get $ptr) (local.get $i)))
(if (i32.and
(i32.and
(i32.eq (i32.load8_u (local.get $base)) (i32.const 120)) ;; x
(i32.eq (i32.load8_u (i32.add (local.get $base) (i32.const 1))) (i32.const 45))) ;; -
(i32.eq (i32.load8_u (i32.add (local.get $base) (i32.const 2))) (i32.const 98))) ;; b
(then (br $found)))
(local.set $i (i32.add (local.get $i) (i32.const 1)))
(br $scan)))
(return (i32.const 0))) ;; Continue
(drop (call $send (i32.const 403) (i32.const 0) (i32.const 0)
(i32.const 0) (i32.const 0) (i32.const 0) (i32.const 0) (i32.const 0)))
(i32.const 1)) ;; Pause
(func (export "proxy_on_request_body") (param i32 i32 i32) (result i32) (i32.const 0))
(func (export "proxy_on_done") (param i32) (result i32) (i32.const 1))
(func (export "proxy_on_delete") (param i32)))
"#;
const SPIN_WAT: &str = r#"
(module
(memory (export "memory") 1)
(global $bump (mut i32) (i32.const 1024))
(func (export "proxy_on_memory_allocate") (param i32) (result i32) (global.get $bump))
(func (export "proxy_on_context_create") (param i32 i32))
(func (export "proxy_on_request_headers") (param i32 i32 i32) (result i32)
(loop $l (br $l)) (i32.const 0))
(func (export "proxy_on_done") (param i32) (result i32) (i32.const 1))
(func (export "proxy_on_delete") (param i32)))
"#;
const MEMBOMB_WAT: &str = r#"
(module
(memory (export "memory") 1)
(global $bump (mut i32) (i32.const 1024))
(func (export "proxy_on_memory_allocate") (param i32) (result i32) (global.get $bump))
(func (export "proxy_on_context_create") (param i32 i32))
(func (export "proxy_on_request_headers") (param i32 i32 i32) (result i32)
(block $denied
(loop $g
(br_if $denied (i32.eq (memory.grow (i32.const 1)) (i32.const -1)))
(br $g)))
(i32.store (i32.const 0x7fffff00) (i32.const 1)) ;; OOB once growth is capped -> trap
(i32.const 0))
(func (export "proxy_on_done") (param i32) (result i32) (i32.const 1))
(func (export "proxy_on_delete") (param i32)))
"#;
fn wasm(wat: &str) -> Vec<u8> {
wat::parse_str(wat).expect("parse WAT")
}
fn module(wat: &str, opts: WasmOptions) -> WasmModule {
let (m, _report) = WasmModule::from_bytes("test", &wasm(wat), b"", &opts).expect("build");
m
}
fn block_module(opts: WasmOptions) -> WasmModule {
module(BLOCK_WAT, opts)
}
fn req_with_header(name: &str, value: &str) -> RequestContext {
Request::new().header(name, value).build()
}
#[test]
fn blocks_on_marker_header() {
let m = block_module(WasmOptions::default());
let d = m.inspect(&req_with_header("x-block", "1"));
match d {
Decision::Block { rule_id, .. } => assert_eq!(rule_id, "wasm-test"),
other => panic!("expected Block, got {other:?}"),
}
}
#[test]
fn allows_benign_request() {
let m = block_module(WasmOptions::default());
assert_eq!(m.inspect(&req_with_header("accept", "text/html")), Decision::Allow);
}
#[test]
fn structural_and_phase_contract() {
let m = block_module(WasmOptions::default());
assert!(m.structural(), "wasm module must run on every request");
assert_eq!(m.id(), "wasm:test");
}
#[test]
fn pseudo_headers_reach_the_guest_header_map() {
let m = block_module(WasmOptions::default());
let ctx = Request::new().path("/x-blocked/page").header("accept", "ok").build();
assert!(matches!(m.inspect(&ctx), Decision::Block { .. }), ":path must reach the header map");
}
#[test]
fn no_disposition_leakage_sequential_same_instance() {
let opts = WasmOptions { pool_size: 1, ..WasmOptions::default() };
let m = block_module(opts);
let first = m.inspect(&req_with_header("x-block", "1"));
assert!(matches!(first, Decision::Block { .. }), "first must block");
let second = m.inspect(&req_with_header("accept", "text/html"));
assert_eq!(second, Decision::Allow, "captured disposition leaked across requests");
}
#[test]
fn no_leakage_under_concurrency() {
let opts = WasmOptions { pool_size: 2, ..WasmOptions::default() };
let m = Arc::new(block_module(opts));
let mut handles = Vec::new();
for t in 0..32u32 {
let m = Arc::clone(&m);
handles.push(thread::spawn(move || {
if t % 2 == 0 {
assert!(matches!(m.inspect(&req_with_header("x-block", "1")), Decision::Block { .. }));
} else {
assert_eq!(m.inspect(&req_with_header("accept", "ok")), Decision::Allow);
}
}));
}
for h in handles {
h.join().expect("a worker panicked -> Store sharing bug");
}
}
#[test]
fn fuel_exhaustion_fails_closed() {
let opts = WasmOptions { fuel_per_request: 100_000, ..WasmOptions::default() };
let m = module(SPIN_WAT, opts);
match m.inspect(&req_with_header("accept", "ok")) {
Decision::Reject { status, .. } => assert_eq!(status, 500),
other => panic!("expected Reject 500, got {other:?}"),
}
}
#[test]
fn memory_cap_fails_closed() {
let opts = WasmOptions {
max_memory_bytes: 64 * 1024, ..WasmOptions::default()
};
let m = module(MEMBOMB_WAT, opts);
match m.inspect(&req_with_header("accept", "ok")) {
Decision::Reject { status, .. } => assert_eq!(status, 500),
other => panic!("expected Reject 500, got {other:?}"),
}
}
#[test]
fn checkout_timeout_is_short() {
let opts = WasmOptions {
pool_size: 1,
checkout_timeout: Duration::from_millis(10),
..WasmOptions::default()
};
let m = block_module(opts);
assert!(matches!(m.inspect(&req_with_header("x-block", "1")), Decision::Block { .. }));
}