AkurAI Build
Menu

AkurAI-Build

public

Latest change 41f5eba4919e59f92e84857cf4b4314e159da9fb - tests: exercise native listener end to end by Olafur Bui

#![allow(dead_code)]

use std::{
    collections::BTreeMap,
    fmt::Display,
    io::{Read, Write},
    net::{TcpListener, TcpStream},
    path::Path,
    process::{Child, Command, Stdio},
    thread,
    time::{Duration, Instant},
};

fn result_or_panic<T, E: Display>(result: Result<T, E>, context: &str) -> T {
    result.unwrap_or_else(|error| panic!("{context}: {error}"))
}

fn option_or_panic<T>(value: Option<T>, context: &str) -> T {
    value.unwrap_or_else(|| panic!("{context}"))
}

pub struct ProcessServer {
    child: Option<Child>,
    address: String,
}

impl ProcessServer {
    pub fn start(mut command: Command, address: String) -> Self {
        command
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        let mut child = result_or_panic(command.spawn(), "spawn bunfork server");
        let deadline = Instant::now() + Duration::from_secs(10);
        loop {
            if TcpStream::connect(&address).is_ok() {
                break;
            }
            if let Some(status) = result_or_panic(child.try_wait(), "poll bunfork server") {
                let output = take_output(&mut child);
                panic!("bunfork exited before listening ({status}):\n{output}");
            }
            if Instant::now() >= deadline {
                let _ = child.kill();
                let _ = child.wait();
                let output = take_output(&mut child);
                panic!("bunfork did not listen on {address}:\n{output}");
            }
            thread::sleep(Duration::from_millis(20));
        }
        Self {
            child: Some(child),
            address,
        }
    }

    pub fn address(&self) -> &str {
        &self.address
    }

    pub fn request(&self, method: &str, target: &str) -> HttpResponse {
        self.request_with_body(method, target, &[], &[])
    }

    pub fn request_with_headers(
        &self,
        method: &str,
        target: &str,
        headers: &[(&str, &str)],
    ) -> HttpResponse {
        self.request_with_body(method, target, headers, &[])
    }

    pub fn request_with_body(
        &self,
        method: &str,
        target: &str,
        headers: &[(&str, &str)],
        body: &[u8],
    ) -> HttpResponse {
        let mut stream = result_or_panic(
            TcpStream::connect(&self.address),
            "connect to bunfork server",
        );
        result_or_panic(
            stream.set_read_timeout(Some(Duration::from_secs(5))),
            "set HTTP read timeout",
        );
        result_or_panic(
            stream.set_write_timeout(Some(Duration::from_secs(5))),
            "set HTTP write timeout",
        );
        let mut request = format!("{method} {target} HTTP/1.1\r\nHost: {}\r\n", self.address);
        for (name, value) in headers {
            request.push_str(name);
            request.push_str(": ");
            request.push_str(value);
            request.push_str("\r\n");
        }
        request.push_str(&format!(
            "Content-Length: {}\r\nConnection: close\r\n\r\n",
            body.len()
        ));
        result_or_panic(
            stream.write_all(request.as_bytes()),
            "write HTTP request head",
        );
        result_or_panic(stream.write_all(body), "write HTTP request body");
        let mut bytes = Vec::new();
        result_or_panic(stream.read_to_end(&mut bytes), "read HTTP response");
        HttpResponse::parse(&bytes)
    }

    pub fn shutdown_with_sigterm(&mut self) {
        let child = option_or_panic(self.child.as_mut(), "server process missing");
        let signal = result_or_panic(
            Command::new("kill")
                .args(["-TERM", &child.id().to_string()])
                .status(),
            "send SIGTERM",
        );
        assert!(signal.success(), "kill -TERM failed: {signal}");

        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            if let Some(status) = result_or_panic(child.try_wait(), "wait for SIGTERM shutdown") {
                let output = take_output(child);
                assert!(
                    status.success(),
                    "bunfork did not exit cleanly after SIGTERM ({status}):\n{output}"
                );
                self.child = None;
                return;
            }
            if Instant::now() >= deadline {
                let _ = child.kill();
                let _ = child.wait();
                let output = take_output(child);
                panic!("bunfork did not stop after SIGTERM:\n{output}");
            }
            thread::sleep(Duration::from_millis(20));
        }
    }
}

impl Drop for ProcessServer {
    fn drop(&mut self) {
        if let Some(child) = &mut self.child {
            let _ = child.kill();
            let _ = child.wait();
        }
    }
}

#[derive(Debug)]
pub struct HttpResponse {
    pub status: u16,
    pub headers: BTreeMap<String, String>,
    pub body: Vec<u8>,
}

impl HttpResponse {
    fn parse(bytes: &[u8]) -> Self {
        let separator = option_or_panic(
            bytes.windows(4).position(|window| window == b"\r\n\r\n"),
            "HTTP response omitted header terminator",
        );
        let head = result_or_panic(
            std::str::from_utf8(&bytes[..separator]),
            "HTTP headers are not UTF-8",
        );
        let mut lines = head.split("\r\n");
        let status = option_or_panic(
            lines
                .next()
                .and_then(|line| line.split_whitespace().nth(1))
                .and_then(|value| value.parse().ok()),
            "HTTP response omitted status",
        );
        let mut headers = BTreeMap::<String, String>::new();
        for line in lines {
            let (name, value) = option_or_panic(line.split_once(':'), "malformed HTTP header");
            headers
                .entry(name.trim().to_ascii_lowercase())
                .and_modify(|existing| {
                    existing.push_str(", ");
                    existing.push_str(value.trim());
                })
                .or_insert_with(|| value.trim().to_owned());
        }
        let raw_body = &bytes[separator + 4..];
        let body = if raw_body.is_empty() {
            Vec::new()
        } else if headers
            .get("transfer-encoding")
            .is_some_and(|value| value.eq_ignore_ascii_case("chunked"))
        {
            decode_chunked(raw_body)
        } else {
            raw_body.to_vec()
        };
        Self {
            status,
            headers,
            body,
        }
    }

    pub fn assert_status(&self, expected: u16) {
        assert_eq!(self.status, expected, "unexpected HTTP response: {self:?}");
    }

    pub fn assert_header_contains(&self, name: &str, expected: &str) {
        let value = self
            .headers
            .get(name)
            .unwrap_or_else(|| panic!("missing {name} header: {self:?}"));
        assert!(
            value
                .to_ascii_lowercase()
                .contains(&expected.to_ascii_lowercase()),
            "{name}={value:?} does not contain {expected:?}"
        );
    }

    pub fn assert_body_contains(&self, expected: &str) {
        let body = String::from_utf8_lossy(&self.body);
        assert!(
            body.contains(expected),
            "body does not contain {expected:?}: {body:?}"
        );
    }

    pub fn assert_body_excludes(&self, unexpected: &str) {
        let body = String::from_utf8_lossy(&self.body);
        assert!(
            !body.contains(unexpected),
            "body unexpectedly contains {unexpected:?}: {body:?}"
        );
    }

    pub fn assert_security_headers(&self) {
        self.assert_header_contains("x-content-type-options", "nosniff");
        self.assert_header_contains("x-frame-options", "deny");
        self.assert_header_contains("referrer-policy", "no-referrer");
        self.assert_header_contains("content-security-policy", "default-src 'self'");
        self.assert_header_contains("permissions-policy", "camera=()");
    }
}

fn decode_chunked(bytes: &[u8]) -> Vec<u8> {
    let mut decoded = Vec::new();
    let mut cursor = 0;
    loop {
        let line_end = option_or_panic(
            bytes[cursor..]
                .windows(2)
                .position(|window| window == b"\r\n")
                .map(|offset| cursor + offset),
            "invalid chunk size line",
        );
        let chunk_line = result_or_panic(
            std::str::from_utf8(&bytes[cursor..line_end]),
            "chunk size is not ASCII",
        );
        let size_text = option_or_panic(chunk_line.split(';').next(), "empty chunk size");
        let size = result_or_panic(usize::from_str_radix(size_text, 16), "invalid chunk size");
        cursor = line_end + 2;
        if size == 0 {
            break;
        }
        let end = option_or_panic(cursor.checked_add(size), "chunk size overflow");
        assert!(end + 2 <= bytes.len(), "truncated chunked response");
        decoded.extend_from_slice(&bytes[cursor..end]);
        assert_eq!(&bytes[end..end + 2], b"\r\n", "invalid chunk terminator");
        cursor = end + 2;
    }
    decoded
}

pub fn free_port() -> u16 {
    let listener = result_or_panic(TcpListener::bind(("127.0.0.1", 0)), "reserve loopback port");
    result_or_panic(listener.local_addr(), "read reserved loopback port").port()
}

fn take_output(child: &mut Child) -> String {
    let mut output = String::new();
    if let Some(mut stdout) = child.stdout.take() {
        let _ = stdout.read_to_string(&mut output);
    }
    if let Some(mut stderr) = child.stderr.take() {
        let mut value = String::new();
        let _ = stderr.read_to_string(&mut value);
        output.push_str(&value);
    }
    output
}

pub struct ScratchDir(tempfile::TempDir);

impl ScratchDir {
    pub fn new(label: &str) -> Self {
        let directory = result_or_panic(
            tempfile::Builder::new()
                .prefix(&format!("bunfork-{label}-"))
                .tempdir(),
            "create test scratch directory",
        );
        Self(directory)
    }

    pub fn path(&self) -> &Path {
        self.0.path()
    }
}