Menu
AkurAI-Build
publicLatest 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()
}
}