use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
use axum::response::{IntoResponse, Response};
pub fn enforce_ws_origin(headers: &HeaderMap) -> Option<Response> {
let host = host_from_headers(headers)?;
let origin = headers.get(header::ORIGIN)?;
if origin_matches_host(origin, host) {
return None;
}
tracing::warn!("ws origin rejected: origin={:?} host={}", origin, host);
Some((StatusCode::FORBIDDEN, "origin not allowed").into_response())
}
pub fn host_from_headers(headers: &HeaderMap) -> Option<&str> {
headers.get(header::HOST).and_then(|v| v.to_str().ok())
}
pub fn origin_host(origin: &HeaderValue) -> Option<&str> {
let o = origin.to_str().ok()?;
let authority = o.split_once("://").map(|(_, rest)| rest)?;
let host = authority.split(['/', '?', '#']).next().unwrap_or("");
let host = strip_port(host);
(!host.is_empty()).then_some(host)
}
pub fn origin_matches_host(origin: &HeaderValue, host: &str) -> bool {
match origin_host(origin) {
Some(origin_host) => origin_host.eq_ignore_ascii_case(strip_port(host)),
None => false,
}
}
pub fn strip_port(s: &str) -> &str {
let Some(rest) = s.strip_prefix('[') else {
return s.split(':').next().unwrap_or(s);
};
let Some((host, tail)) = rest.split_once(']') else { return s };
let port_ok = match tail.strip_prefix(':') {
Some(port) => !port.is_empty() && port.bytes().all(|b| b.is_ascii_digit()),
None => tail.is_empty(),
};
if port_ok { host } else { s }
}
#[cfg(test)]
mod tests {
use super::*;
fn header_map(entries: &[(&str, &str)]) -> HeaderMap {
let mut h = HeaderMap::new();
for (k, v) in entries {
let name: header::HeaderName = k.parse().unwrap();
h.insert(name, HeaderValue::from_str(v).unwrap());
}
h
}
#[test]
fn enforce_ws_origin_allows_same_origin() {
let h = header_map(&[("host", "127.0.0.1:9077"), ("origin", "http://127.0.0.1:9077")]);
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_allows_origin_port_differing_from_host() {
let h = header_map(&[("host", "omniterm.lan:443"), ("origin", "https://omniterm.lan")]);
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_allows_subdomain_host() {
let h = header_map(&[
("host", "3000.omniterm.lan:9777"),
("origin", "http://3000.omniterm.lan:9777"),
]);
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_rejects_cross_site() {
let h = header_map(&[("host", "127.0.0.1:9077"), ("origin", "https://evil.com")]);
let resp = enforce_ws_origin(&h).expect("cross-site must be rejected");
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[test]
fn enforce_ws_origin_rejects_malformed_origin() {
let h = header_map(&[("host", "127.0.0.1:9077"), ("origin", "127.0.0.1:9077")]);
assert!(enforce_ws_origin(&h).is_some());
}
#[test]
fn enforce_ws_origin_allows_missing_origin() {
let h = header_map(&[("host", "127.0.0.1:9077")]);
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_allows_missing_host() {
let h = header_map(&[("origin", "http://127.0.0.1:9077")]);
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_allows_empty_header_map() {
assert!(enforce_ws_origin(&HeaderMap::new()).is_none());
}
#[test]
fn origin_matches_subdomain_host() {
let o = HeaderValue::from_static("http://3000.omniterm.lan:9777");
assert!(origin_matches_host(&o, "3000.omniterm.lan:9777"));
let o2 = HeaderValue::from_static("https://3000.omniterm.lan");
assert!(origin_matches_host(&o2, "3000.omniterm.lan"));
}
#[test]
fn origin_rejects_cross_site() {
let o = HeaderValue::from_static("https://evil.com");
assert!(!origin_matches_host(&o, "3000.omniterm.lan"));
let o2 = HeaderValue::from_static("3000.omniterm.lan");
assert!(!origin_matches_host(&o2, "3000.omniterm.lan"));
}
#[test]
fn strip_port_keeps_ipv6_literal() {
assert_eq!(strip_port("3000.omniterm.lan:9777"), "3000.omniterm.lan");
assert_eq!(strip_port("192.168.5.216:9077"), "192.168.5.216");
assert_eq!(strip_port("3000.omniterm.lan"), "3000.omniterm.lan");
assert_eq!(strip_port("[::1]:8080"), "::1");
assert_eq!(strip_port("[::1]"), "::1");
}
#[test]
fn strip_port_rejects_garbage_after_bracket() {
assert_eq!(strip_port("[::1]@evil.com"), "[::1]@evil.com");
assert_eq!(strip_port("[::1]:80@evil.com"), "[::1]:80@evil.com");
assert_eq!(strip_port("[::1"), "[::1");
assert_eq!(strip_port("[::1]:8080"), "::1");
assert_eq!(strip_port("[::1]"), "::1");
}
#[test]
fn origin_matches_ipv6_literal_host() {
let o = HeaderValue::from_static("http://[::1]:9077");
assert!(origin_matches_host(&o, "[::1]:9077"));
}
#[test]
fn origin_is_case_insensitive_on_host() {
let o = HeaderValue::from_static("http://Omniterm.LAN");
assert!(origin_matches_host(&o, "omniterm.lan"));
assert!(origin_matches_host(&o, "OMNITERM.lan:9077"));
}
#[test]
fn origin_rejects_empty_host_segment() {
let o = HeaderValue::from_static("http://:9077");
assert!(!origin_matches_host(&o, "127.0.0.1:9077"));
}
#[test]
fn origin_stops_at_path_and_query() {
let o = HeaderValue::from_static("http://evil.com/x?y=1");
assert!(!origin_matches_host(&o, "omniterm.lan"));
}
#[test]
fn enforce_ws_origin_rejects_when_first_of_multiple_values_is_cross_site() {
let mut h = HeaderMap::new();
h.insert(header::HOST, HeaderValue::from_static("127.0.0.1:9077"));
h.append(header::ORIGIN, HeaderValue::from_static("http://evil.com"));
h.append(header::ORIGIN, HeaderValue::from_static("http://127.0.0.1:9077"));
assert!(enforce_ws_origin(&h).is_some());
}
#[test]
fn enforce_ws_origin_rejects_non_utf8_origin() {
let mut h = HeaderMap::new();
h.insert(header::HOST, HeaderValue::from_static("127.0.0.1:9077"));
h.append(header::ORIGIN, HeaderValue::from_bytes(b"\xff\xfe").unwrap());
assert!(enforce_ws_origin(&h).is_some());
}
#[test]
fn enforce_ws_origin_allows_when_host_header_absent() {
let mut h = HeaderMap::new();
h.append(header::ORIGIN, HeaderValue::from_static("http://127.0.0.1:9077"));
assert!(enforce_ws_origin(&h).is_none());
}
#[test]
fn enforce_ws_origin_allows_when_host_is_non_utf8() {
let mut h = HeaderMap::new();
h.insert(header::HOST, HeaderValue::from_bytes(b"\xff\xfe").unwrap());
h.insert(header::ORIGIN, HeaderValue::from_static("http://127.0.0.1:9077"));
assert!(enforce_ws_origin(&h).is_none());
}
}