#![forbid(unsafe_code)]
use std::net::SocketAddr;
use std::time::Duration;
use arcature::prelude::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
const HANG_GUARD: Duration = Duration::from_secs(10);
struct RunningApp {
addr: SocketAddr,
shutdown: tokio::sync::oneshot::Sender<()>,
join: tokio::task::JoinHandle<()>,
}
impl RunningApp {
async fn start(app: Application<()>) -> RunningApp {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind ephemeral listener");
let addr = listener.local_addr().expect("read bound address");
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let join = tokio::spawn(async move {
app.serve_with_shutdown(listener, async {
let _ = shutdown_rx.await;
})
.await
.expect("application served without engine error");
});
RunningApp {
addr,
shutdown: shutdown_tx,
join,
}
}
fn addr(&self) -> SocketAddr {
self.addr
}
async fn stop(self) {
self.shutdown
.send(())
.expect("server task still alive to receive shutdown");
tokio::time::timeout(HANG_GUARD, self.join)
.await
.expect("server did not hang on shutdown")
.expect("server task did not panic");
}
}
async fn http_get(addr: SocketAddr, path: &str) -> String {
let mut stream = tokio::time::timeout(HANG_GUARD, TcpStream::connect(addr))
.await
.expect("connect did not hang")
.expect("connect succeeds");
let request = format!("GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n");
stream
.write_all(request.as_bytes())
.await
.expect("write request");
let mut buffer = Vec::new();
tokio::time::timeout(HANG_GUARD, stream.read_to_end(&mut buffer))
.await
.expect("read did not hang")
.expect("read succeeds");
String::from_utf8_lossy(&buffer).into_owned()
}
fn header<'a>(response: &'a str, name: &str) -> Option<&'a str> {
let lower = name.to_ascii_lowercase();
for line in response.split("\r\n") {
if let Some((k, v)) = line.split_once(": ")
&& k.to_ascii_lowercase() == lower
{
return Some(v.trim());
}
}
None
}
#[tokio::test]
async fn error_mapping_applies_to_handler_response() {
let app = Application::new()
.routes(Routes::new().route(
"/error",
get(|| async { StatusCode::INTERNAL_SERVER_ERROR }),
))
.error_mapping(|response| {
if response.status().is_server_error() {
let mut response = response;
response
.headers_mut()
.insert("x-mapped", "yes".parse().expect("valid header"));
response
} else {
response
}
})
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/error").await;
assert!(response.starts_with("HTTP/1.1 500"));
assert_eq!(header(&response, "x-mapped"), Some("yes"));
server.stop().await;
}
#[tokio::test]
async fn error_mapping_does_not_affect_success_responses() {
let app = Application::new()
.routes(Routes::new().route("/", get(|| async { "ok" })))
.error_mapping(|response| {
if response.status().is_server_error() {
let mut response = response;
response
.headers_mut()
.insert("x-mapped", "yes".parse().expect("valid header"));
response
} else {
response
}
})
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/").await;
assert!(response.starts_with("HTTP/1.1 200"));
assert_eq!(header(&response, "x-mapped"), None);
server.stop().await;
}
#[tokio::test]
async fn error_mapping_does_not_affect_proxy_short_circuit() {
let app = Application::new()
.routes(Routes::new().route("/", get(|| async { "ok" })))
.proxy(|_req| ProxyAction::ShortCircuit {
status: StatusCode::SERVICE_UNAVAILABLE,
response: None,
})
.error_mapping(|response| {
let mut response = response;
response
.headers_mut()
.insert("x-mapped", "yes".parse().expect("valid header"));
response
})
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/").await;
assert!(response.starts_with("HTTP/1.1 503"));
assert_eq!(header(&response, "x-mapped"), None);
server.stop().await;
}
#[tokio::test]
async fn no_error_mapping_passes_through() {
let app = Application::new()
.routes(Routes::new().route("/", get(|| async { "ok" })))
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/").await;
assert!(response.starts_with("HTTP/1.1 200"));
server.stop().await;
}
#[tokio::test]
async fn error_mapping_with_proxy_continue_works() {
let app = Application::new()
.routes(Routes::new().route(
"/error",
get(|| async { StatusCode::INTERNAL_SERVER_ERROR }),
))
.proxy(|_req| ProxyAction::continue_default())
.error_mapping(|response| {
if response.status().is_server_error() {
let mut response = response;
response
.headers_mut()
.insert("x-mapped", "yes".parse().expect("valid header"));
response
} else {
response
}
})
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/error").await;
assert!(response.starts_with("HTTP/1.1 500"));
assert_eq!(header(&response, "x-mapped"), Some("yes"));
server.stop().await;
}
#[tokio::test]
async fn error_mapping_sees_404_fallback() {
let app = Application::new()
.routes(Routes::new().route("/", get(|| async { "ok" })))
.error_mapping(|response| {
if response.status() == StatusCode::NOT_FOUND {
let mut response = response;
response
.headers_mut()
.insert("x-mapped-404", "yes".parse().expect("valid header"));
response
} else {
response
}
})
.build();
let server = RunningApp::start(app).await;
let response = http_get(server.addr(), "/nonexistent").await;
assert!(response.starts_with("HTTP/1.1 404"));
assert_eq!(header(&response, "x-mapped-404"), Some("yes"));
server.stop().await;
}