waf-wasm 0.4.0

Proxy-Wasm runtime (on wasmi) for Light WAF: load .wasm filters as detection modules, with instance pooling and fuel/memory DoS caps.
Documentation
// SPDX-License-Identifier: Apache-2.0
//! Driver-level tests for [`waf_wasm::WasmModule`] using hand-written Proxy-Wasm guests
//! (WAT, compiled inline). These migrate the B3-0 probe's load-bearing cases as PERMANENT
//! tests: per-request isolation (sequential + concurrent on a pooled instance) and the two
//! DoS ceilings (fuel / memory) failing closed to `Reject{500}`.

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};

/// A Proxy-Wasm guest that blocks (403) when the request headers contain `x-block`.
/// Exercises the real host path: `get_header_map_pairs` → reentrant alloc into guest
/// memory → scan → `send_local_response`.
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)))
"#;

/// A guest that spins forever in `on_request_headers` — exhausts the per-request fuel.
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)))
"#;

/// A guest that, once memory growth is denied by the cap, stores out of bounds → trap.
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");
}

/// Pseudo-headers (`:path` etc.) are injected into the header map: the block guest scans
/// for "x-b", so a request whose PATH contains it blocks even with NO `x-block` header —
/// proving `:path` reaches the guest's header map (Envoy convention, real filters rely on it).
#[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");
}

/// Paletto #1 (sequential): a pool of ONE instance, reused. A blocking request followed by
/// a benign one must NOT leak the captured disposition — the second is `Allow`.
#[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");
}

/// Paletto #1 (concurrent): many threads hammer a small pool with mixed requests; each
/// decision must be correct and no thread may panic (which a shared `Store` would cause).
#[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");
    }
}

/// DoS (fuel): a runaway guest is bounded by the per-request fuel budget and fails closed.
#[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:?}"),
    }
}

/// DoS (memory): hitting the memory cap traps and fails closed.
#[test]
fn memory_cap_fails_closed() {
    let opts = WasmOptions {
        max_memory_bytes: 64 * 1024, // 1 page == the guest's initial size: growth always denied
        ..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:?}"),
    }
}

/// Pool exhaustion: with a tiny timeout and a guest that holds its instance via a slow
/// body, a burst of concurrent callers that cannot get an instance fails closed rather
/// than blocking forever. Here we only assert the timeout path is reachable and clean.
#[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);
    // A normal request still succeeds (instance returns to the pool well within timeout).
    assert!(matches!(m.inspect(&req_with_header("x-block", "1")), Decision::Block { .. }));
}