use axum::http::{header, HeaderMap, Method};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Refusal {
pub status: u16,
pub code: &'static str,
pub message: String,
}
#[derive(Debug, Clone)]
pub struct Policy {
hosts: Vec<String>,
origins: Vec<String>,
}
impl Policy {
pub fn new(port: u16, extra_origins: &[String]) -> Self {
let hosts = vec![
format!("127.0.0.1:{port}"),
format!("localhost:{port}"),
format!("[::1]:{port}"),
];
let mut origins: Vec<String> = hosts.iter().map(|host| format!("http://{host}")).collect();
origins.extend(
extra_origins
.iter()
.map(|origin| origin.trim_end_matches('/').to_ascii_lowercase()),
);
Self { hosts, origins }
}
pub fn origin(&self) -> &str {
&self.origins[0]
}
}
fn refuse(status: u16, code: &'static str, message: impl Into<String>) -> Refusal {
Refusal {
status,
code,
message: message.into(),
}
}
fn header_str(headers: &HeaderMap, name: impl header::AsHeaderName) -> Option<&str> {
headers.get(name).and_then(|value| value.to_str().ok())
}
pub fn check(
policy: &Policy,
method: &Method,
path: &str,
headers: &HeaderMap,
) -> Result<(), Refusal> {
let host = header_str(headers, header::HOST)
.unwrap_or_default()
.to_ascii_lowercase();
if !policy.hosts.contains(&host) {
return Err(refuse(
421,
"dev-ui-host-refused",
format!("The dev UI only answers as {}", policy.hosts[0]),
));
}
if let Some(origin) = headers.get(header::ORIGIN) {
let origin = origin.to_str().unwrap_or("null").to_ascii_lowercase();
if !policy.origins.contains(&origin) {
return Err(refuse(
403,
"dev-ui-origin-refused",
format!("Requests from {origin} may not use the a4 dev UI server"),
));
}
}
let api = path.starts_with("/dev/");
if api {
if let Some(site) = header_str(headers, "sec-fetch-site") {
if matches!(site, "cross-site" | "same-site") {
return Err(refuse(
403,
"dev-ui-cross-site-refused",
"Cross-site requests may not use the a4 dev UI server",
));
}
}
}
let writes = matches!(
*method,
Method::POST | Method::PUT | Method::PATCH | Method::DELETE
);
if writes {
let json = header_str(headers, header::CONTENT_TYPE)
.and_then(|value| value.split(';').next())
.is_some_and(|mime| mime.trim().eq_ignore_ascii_case("application/json"));
if !json {
return Err(refuse(
415,
"dev-ui-json-required",
"Writes take Content-Type: application/json",
));
}
}
if *method == Method::OPTIONS {
return Err(refuse(
405,
"dev-ui-preflight-refused",
"The a4 dev UI server answers no CORS preflight",
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use axum::http::HeaderValue;
fn headers(pairs: &[(&'static str, &str)]) -> HeaderMap {
let mut map = HeaderMap::new();
for (name, value) in pairs {
map.insert(*name, HeaderValue::from_str(value).unwrap());
}
map
}
fn code(result: Result<(), Refusal>) -> Option<&'static str> {
result.err().map(|refusal| refusal.code)
}
#[test]
fn allows_the_ui_its_own_origin_and_local_tools() {
let policy = Policy::new(50106, &[]);
let ui = headers(&[
("host", "127.0.0.1:50106"),
("origin", "http://127.0.0.1:50106"),
("sec-fetch-site", "same-origin"),
("content-type", "application/json"),
]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/steps/run", &ui)),
None
);
assert_eq!(
code(check(&policy, &Method::GET, "/dev/v1/events", &ui)),
None
);
let tool = headers(&[("host", "localhost:50106")]);
assert_eq!(
code(check(&policy, &Method::GET, "/dev/v1/events", &tool)),
None
);
let typed = headers(&[("host", "127.0.0.1:50106"), ("sec-fetch-site", "none")]);
assert_eq!(code(check(&policy, &Method::GET, "/", &typed)), None);
}
#[test]
fn refuses_dns_rebinding() {
let policy = Policy::new(50106, &[]);
for host in ["evil.example:50106", "127.0.0.1:50107", "127.0.0.1", ""] {
let request = headers(&[("host", host)]);
assert_eq!(
code(check(&policy, &Method::GET, "/dev/v1/protocol", &request)),
Some("dev-ui-host-refused"),
"{host}"
);
}
let request = headers(&[("host", "evil.example:50106")]);
assert_eq!(
code(check(&policy, &Method::GET, "/", &request)),
Some("dev-ui-host-refused")
);
}
#[test]
fn refuses_other_origins_and_cross_site_requests() {
let policy = Policy::new(50106, &[]);
for origin in ["https://evil.example", "http://localhost:3000", "null"] {
let request = headers(&[
("host", "127.0.0.1:50106"),
("origin", origin),
("content-type", "application/json"),
]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/steps/run", &request)),
Some("dev-ui-origin-refused"),
"{origin}"
);
}
for site in ["cross-site", "same-site"] {
let request = headers(&[("host", "127.0.0.1:50106"), ("sec-fetch-site", site)]);
assert_eq!(
code(check(&policy, &Method::GET, "/dev/v1/events", &request)),
Some("dev-ui-cross-site-refused")
);
assert_eq!(
code(check(&policy, &Method::GET, "/assets/app.js", &request)),
None
);
}
}
#[test]
fn writes_need_json_and_preflights_are_refused() {
let policy = Policy::new(50106, &[]);
let form = headers(&[
("host", "127.0.0.1:50106"),
("content-type", "application/x-www-form-urlencoded"),
]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/time-travel", &form)),
Some("dev-ui-json-required")
);
let bare = headers(&[("host", "127.0.0.1:50106")]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/time-travel", &bare)),
Some("dev-ui-json-required")
);
let json = headers(&[
("host", "127.0.0.1:50106"),
("content-type", "application/json; charset=utf-8"),
]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/time-travel", &json)),
None
);
assert_eq!(
code(check(
&policy,
&Method::OPTIONS,
"/dev/v1/time-travel",
&bare
)),
Some("dev-ui-preflight-refused")
);
}
#[test]
fn allows_the_vite_dev_server_origin() {
let policy = Policy::new(50106, &["http://localhost:5173/".to_string()]);
let request = headers(&[
("host", "127.0.0.1:50106"),
("origin", "http://localhost:5173"),
("content-type", "application/json"),
]);
assert_eq!(
code(check(&policy, &Method::POST, "/dev/v1/steps/run", &request)),
None
);
assert_eq!(policy.origin(), "http://127.0.0.1:50106");
}
}