urus/tests/integration.rs

//! End-to-end integration tests.
//!
//! Each test spins up the server in a background OS thread on an
//! ephemeral port, opens a regular `std::net::TcpStream` against it, and
//! exercises the wire. Single thread for the smarm runtime per test so
//! teardown happens cleanly when the test thread is dropped (the process
//! exits at the end of the test binary).

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::time::Duration;

use urus::{serve_with, Conn, Config, Next, Pipeline, Router};

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Reserve a free localhost port by binding briefly and reading back the
/// assigned port number. The port is released the instant we drop the
/// listener; there's a tiny race window before the urus server claims it
/// but it's well below the threshold for flakiness in practice.
fn free_port() -> u16 {
    let l = TcpListener::bind("127.0.0.1:0").unwrap();
    l.local_addr().unwrap().port()
}

fn spawn_server(pipeline: Pipeline) -> u16 {
    let port = free_port();
    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    std::thread::spawn(move || {
        let cfg = Config {
            listener_pool:     2,
            scheduler_threads: Some(2),
            ..Config::new(addr)
        };
        serve_with(cfg, pipeline).unwrap();
    });
    // Wait for the server to actually be listening.
    for _ in 0..50 {
        if TcpStream::connect(addr).is_ok() {
            return port;
        }
        std::thread::sleep(Duration::from_millis(50));
    }
    panic!("server didn't come up on {addr}");
}

fn send_request(port: u16, req: &[u8]) -> Vec<u8> {
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    s.write_all(req).unwrap();
    s.shutdown(std::net::Shutdown::Write).ok();
    let mut buf = Vec::new();
    s.read_to_end(&mut buf).unwrap();
    buf
}

fn http_status(resp: &[u8]) -> u16 {
    let s = std::str::from_utf8(&resp[..resp.len().min(64)]).unwrap();
    let parts: Vec<&str> = s.splitn(3, ' ').collect();
    parts[1].parse().unwrap()
}

fn http_body(resp: &[u8]) -> &[u8] {
    // Split on the first \r\n\r\n.
    for i in 0..resp.len().saturating_sub(3) {
        if &resp[i..i + 4] == b"\r\n\r\n" {
            return &resp[i + 4..];
        }
    }
    &[]
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[test]
fn hello_world() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| {
            c.put_status(200).put_body("hello urus")
        })
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"GET / HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 200);
    assert_eq!(http_body(&resp), b"hello urus");
}

#[test]
fn server_and_listeners_are_registered() {
    // `whereis` must run inside the runtime (the test itself is a foreign
    // OS thread with no runtime in its TLS), so probe from a handler —
    // connection actors live in the runtime by construction.
    let pipe = Pipeline::new().plug(
        Router::new().get("/whereis", |c: Conn, _n: Next| {
            let ok = smarm::whereis("urus.server").is_some()
                && smarm::whereis("urus.listener.0").is_some()
                && smarm::whereis("urus.listener.1").is_some(); // pool of 2
            c.put_status(200).put_body(if ok { "registered" } else { "missing" })
        })
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"GET /whereis HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 200);
    assert_eq!(http_body(&resp), b"registered");
}

#[test]
fn echo_body() {
    let pipe = Pipeline::new().plug(
        Router::new().post("/echo", |c: Conn, _n: Next| {
            let body = c.body.as_bytes().to_vec();
            c.put_status(200).put_body(body)
        })
    );
    let port = spawn_server(pipe);
    let resp = send_request(
        port,
        b"POST /echo HTTP/1.1\r\nHost: x\r\nContent-Length: 5\r\nConnection: close\r\n\r\nhello",
    );
    assert_eq!(http_status(&resp), 200);
    assert_eq!(http_body(&resp), b"hello");
}

#[test]
fn path_params() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/users/:id", |c: Conn, _n: Next| {
            let id = c.params.get("id").unwrap_or("?").to_string();
            c.put_status(200).put_body(format!("user={id}"))
        })
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"GET /users/42 HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 200);
    assert_eq!(http_body(&resp), b"user=42");
}

#[test]
fn method_not_allowed() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/only-get", |c: Conn, _n: Next| c.put_status(200))
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"POST /only-get HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 405);
}

#[test]
fn unknown_route_falls_to_404() {
    // Router falls through; connection actor's default emits 404.
    let pipe = Pipeline::new().plug(
        Router::new().get("/known", |c: Conn, _n: Next| c.put_status(200))
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"GET /missing HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 404);
}

#[test]
fn keep_alive_two_requests_on_one_connection() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| c.put_status(200).put_body("ok"))
    );
    let port = spawn_server(pipe);
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();

    // First request, no Connection header (HTTP/1.1 default = keep-alive).
    s.write_all(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n").unwrap();

    // Read the first response (Content-Length: 2 -> body "ok").
    let mut buf = vec![0u8; 1024];
    let mut total = 0;
    let mut first_end = None;
    while first_end.is_none() {
        let n = s.read(&mut buf[total..]).unwrap();
        assert!(n > 0, "server hung up early");
        total += n;
        // First response ends at header-terminator + 2 bytes of body.
        for i in 0..total.saturating_sub(3) {
            if &buf[i..i + 4] == b"\r\n\r\n" {
                let body_start = i + 4;
                if total >= body_start + 2 {
                    first_end = Some(body_start + 2);
                }
                break;
            }
        }
    }
    let first = &buf[..first_end.unwrap()];
    assert_eq!(http_status(first), 200);
    assert_eq!(http_body(first), b"ok");

    // Second request on the same socket — proves keep-alive works.
    s.write_all(b"GET / HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n").unwrap();
    let mut rest = Vec::new();
    s.read_to_end(&mut rest).unwrap();
    assert_eq!(http_status(&rest), 200);
    assert_eq!(http_body(&rest), b"ok");
}

#[test]
fn panicking_handler_yields_500() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/boom", |_c: Conn, _n: Next| -> Conn {
            panic!("handler explosion")
        })
    );
    let port = spawn_server(pipe);
    let resp = send_request(port, b"GET /boom HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 500);
}

#[test]
fn middleware_can_short_circuit() {
    // Middleware that requires `X-Auth: secret`; otherwise returns 401
    // without calling `next.run(conn)`.
    let auth = |c: Conn, n: Next| {
        if c.headers.get("x-auth").map(|s| s == "secret").unwrap_or(false) {
            n.run(c)
        } else {
            c.put_status(401).halt()
        }
    };
    let pipe = Pipeline::new()
        .plug(auth)
        .plug(Router::new().get("/secret", |c: Conn, _n: Next| c.put_status(200).put_body("ok")));

    let port = spawn_server(pipe);

    // Without auth → 401.
    let resp = send_request(port, b"GET /secret HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 401);

    // With auth → 200 + body.
    let resp = send_request(
        port,
        b"GET /secret HTTP/1.1\r\nHost: x\r\nX-Auth: secret\r\nConnection: close\r\n\r\n",
    );
    assert_eq!(http_status(&resp), 200);
    assert_eq!(http_body(&resp), b"ok");
}

// ---------------------------------------------------------------------------
// Supervision (v0.2 chunk 1)
// ---------------------------------------------------------------------------

/// A panicking listener must be restarted by the pool supervisor, and the
/// connection that was pending when it died must still be served.
///
/// Pool size is 1 on purpose: with a sibling listener the test would pass
/// even if restarts were broken (the sibling would pick up the accept).
/// With one listener, the request after the injected panic can only
/// succeed if a fresh listener came up on the same fd.
#[test]
fn panicking_listener_restarts() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/ping", |c: Conn, _n: Next| c.put_status(200).put_body("pong"))
    );

    let port = free_port();
    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    std::thread::spawn(move || {
        let cfg = Config {
            listener_pool:     1,
            scheduler_threads: Some(2),
            ..Config::new(addr)
        };
        serve_with(cfg, pipe).unwrap();
    });
    for _ in 0..50 {
        if TcpStream::connect(addr).is_ok() {
            break;
        }
        std::thread::sleep(Duration::from_millis(50));
    }

    // Sanity: server answers before the fault.
    let resp = send_request(port, b"GET /ping HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 200);

    // Arm the fault: the listener's next accept-loop iteration panics
    // *before* accepting, so our connection waits in the kernel backlog
    // until the restarted listener picks it up.
    urus::serve::INJECT_LISTENER_PANICS.store(1, std::sync::atomic::Ordering::Relaxed);

    let resp = send_request(port, b"GET /ping HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    assert_eq!(http_status(&resp), 200, "request pending across the panic was not served");
    assert_eq!(http_body(&resp), b"pong");
    assert_eq!(
        urus::serve::INJECT_LISTENER_PANICS.load(std::sync::atomic::Ordering::Relaxed),
        0,
        "fault was never consumed — listener didn't wake for the connection"
    );

    // The restarted listener keeps serving.
    for _ in 0..5 {
        let resp = send_request(port, b"GET /ping HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
        assert_eq!(http_status(&resp), 200);
    }
}

// ---------------------------------------------------------------------------
// Graceful shutdown (v0.2 chunk 2)
// ---------------------------------------------------------------------------

/// Boot a server with a shutdown handle. Returns (port, handle, done_rx)
/// where done_rx fires when serve_with_shutdown returns.
fn spawn_server_with_handle(
    pipeline: Pipeline,
    drain: Duration,
) -> (u16, urus::Handle, std::sync::mpsc::Receiver<()>) {
    let port = free_port();
    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    let (handle, signal) = urus::shutdown_handle();
    let (done_tx, done_rx) = std::sync::mpsc::channel();
    std::thread::spawn(move || {
        let cfg = Config {
            listener_pool:     2,
            scheduler_threads: Some(2),
            drain_timeout:     drain,
            ..Config::new(addr)
        };
        urus::serve_with_shutdown(cfg, pipeline, signal).unwrap();
        let _ = done_tx.send(());
    });
    for _ in 0..50 {
        if TcpStream::connect(addr).is_ok() {
            return (port, handle, done_rx);
        }
        std::thread::sleep(Duration::from_millis(50));
    }
    panic!("server didn't come up on {addr}");
}

/// Shutdown with nothing in flight returns promptly — well inside the
/// (deliberately huge) drain window, proving idle conns don't run the
/// clock out.
#[test]
fn shutdown_with_no_connections_returns() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| c.put_status(200))
    );
    let (_port, handle, done_rx) = spawn_server_with_handle(pipe, Duration::from_secs(30));
    handle.shutdown();
    done_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("serve_with_shutdown did not return after shutdown()");
}

/// An idle keep-alive connection is closed immediately on shutdown (the
/// drain deadline of 30s must NOT be what gates the return), and the
/// server refuses new connections afterwards.
#[test]
fn shutdown_closes_idle_keepalive_promptly() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| c.put_status(200).put_body("ok"))
    );
    let (port, handle, done_rx) = spawn_server_with_handle(pipe, Duration::from_secs(30));

    // Complete one request on a keep-alive connection and leave it open.
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    s.write_all(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n").unwrap();
    let mut buf = [0u8; 1024];
    let n = s.read(&mut buf).unwrap();
    assert!(n > 0 && buf.starts_with(b"HTTP/1.1 200"));

    let t0 = std::time::Instant::now();
    handle.shutdown();

    // The parked conn actor is stopped; its socket closes → we see EOF.
    let n = s.read(&mut buf).expect("read after shutdown");
    assert_eq!(n, 0, "expected EOF on the idle keep-alive connection");

    done_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("serve did not return");
    assert!(
        t0.elapsed() < Duration::from_secs(5),
        "shutdown waited for the drain deadline instead of stopping the idle conn"
    );

    // Listeners are gone: new connections are refused.
    assert!(
        TcpStream::connect(("127.0.0.1", port)).is_err(),
        "server still accepting after shutdown"
    );
}

/// A request in flight when shutdown is signalled completes with a real
/// response before its connection is closed.
#[test]
fn shutdown_drains_in_flight_request() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/slow", |c: Conn, _n: Next| {
            smarm::sleep(Duration::from_millis(400));
            c.put_status(200).put_body("made it")
        })
    );
    let (port, handle, done_rx) = spawn_server_with_handle(pipe, Duration::from_secs(10));

    // Fire the slow request, then shut down while it is in flight.
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    s.write_all(b"GET /slow HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n").unwrap();
    std::thread::sleep(Duration::from_millis(100)); // let the head land
    handle.shutdown();

    let mut resp = Vec::new();
    s.read_to_end(&mut resp).unwrap();
    assert_eq!(http_status(&resp), 200, "in-flight request was cut off by shutdown");
    assert_eq!(http_body(&resp), b"made it");

    done_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("serve did not return after drain");
}

/// A request still in flight at the drain deadline is force-stopped: the
/// connection dies without a response, but serve returns near the
/// deadline rather than hanging on the stuck handler's full duration.
#[test]
fn shutdown_force_stops_at_drain_deadline() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/stuck", |c: Conn, _n: Next| {
            smarm::sleep(Duration::from_secs(60));
            c.put_status(200)
        })
    );
    let (port, handle, done_rx) = spawn_server_with_handle(pipe, Duration::from_millis(300));

    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(10))).unwrap();
    s.write_all(b"GET /stuck HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n").unwrap();
    std::thread::sleep(Duration::from_millis(100));
    handle.shutdown();

    done_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("serve did not return after force-stop deadline");
}

// ---------------------------------------------------------------------------
// Timeouts (v0.2 chunk 3)
// ---------------------------------------------------------------------------

fn spawn_server_with_timeouts(
    pipeline: Pipeline,
    keep_alive: Duration,
    request: Duration,
) -> u16 {
    let port = free_port();
    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    std::thread::spawn(move || {
        let cfg = Config {
            listener_pool:      2,
            scheduler_threads:  Some(2),
            keep_alive_timeout: keep_alive,
            request_timeout:    request,
            ..Config::new(addr)
        };
        serve_with(cfg, pipeline).unwrap();
    });
    for _ in 0..50 {
        if TcpStream::connect(addr).is_ok() {
            return port;
        }
        std::thread::sleep(Duration::from_millis(50));
    }
    panic!("server didn't come up on {addr}");
}

/// Read from `s` until the response head is complete (double CRLF). Only
/// suitable for responses with an empty body.
fn read_response_head(s: &mut TcpStream) -> Vec<u8> {
    let mut resp = Vec::new();
    let mut byte = [0u8; 1];
    while !resp.ends_with(b"\r\n\r\n") {
        match s.read(&mut byte) {
            Ok(0) => break,
            Ok(_) => resp.push(byte[0]),
            Err(e) => panic!("read failed mid-response: {e}"),
        }
    }
    resp
}

/// An idle keep-alive connection is reaped once keep_alive_timeout passes
/// with no next request: the server closes the socket (clean EOF on our
/// side) well before the much larger request_timeout.
#[test]
fn idle_keepalive_reaped_at_keep_alive_timeout() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| c.put_status(200))
    );
    let port = spawn_server_with_timeouts(
        pipe,
        Duration::from_millis(300), // keep_alive_timeout under test
        Duration::from_secs(10),    // request_timeout out of the way
    );

    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    s.write_all(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n").unwrap(); // keep-alive
    let head = read_response_head(&mut s);
    assert_eq!(http_status(&head), 200);

    // Now go quiet. The server should close us at ~300ms; the 5s read
    // timeout on our side is the failure detector.
    let start = std::time::Instant::now();
    let mut byte = [0u8; 1];
    match s.read(&mut byte) {
        Ok(0) => {} // clean EOF: reaped
        Ok(n) => panic!("expected EOF, got {n} unexpected byte(s)"),
        Err(e) => panic!("expected EOF, read errored: {e}"),
    }
    assert!(
        start.elapsed() < Duration::from_secs(3),
        "reap took {:?}, expected ~300ms", start.elapsed()
    );
}

/// A slowloris client that sends a partial head and then stalls is killed
/// at request_timeout with a best-effort 408, even though the (large)
/// keep-alive budget hasn't expired.
#[test]
fn slowloris_partial_head_killed_at_request_timeout() {
    let pipe = Pipeline::new().plug(
        Router::new().get("/", |c: Conn, _n: Next| c.put_status(200))
    );
    let port = spawn_server_with_timeouts(
        pipe,
        Duration::from_secs(10),    // keep_alive_timeout out of the way
        Duration::from_millis(300), // request_timeout under test
    );

    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
    // Partial head: request clock starts on these bytes, never completes.
    s.write_all(b"GET / HTTP/1.1\r\nHost: x\r\n").unwrap();

    let start = std::time::Instant::now();
    let mut resp = Vec::new();
    s.read_to_end(&mut resp).expect("expected 408+EOF or EOF");
    assert!(
        start.elapsed() < Duration::from_secs(3),
        "kill took {:?}, expected ~300ms", start.elapsed()
    );
    // The 408 is best-effort (single non-parking write), but with our read
    // side live it should land.
    assert!(!resp.is_empty(), "expected a best-effort 408 before close");
    assert_eq!(http_status(&resp), 408);
}