use axum::{
body::{Body, Bytes},
extract::{Path, Request, State},
http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header},
response::{IntoResponse, Response},
};
use reqwest::Method;
use tracing::{debug, warn};
use crate::server::AppState;
static HOP_BY_HOP: &[&str] = &[
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailers",
"transfer-encoding",
"upgrade",
"host",
];
fn full_id(daemon_key: &str) -> Option<&'static str> {
match daemon_key {
"search" => Some("trusty-search"),
"memory" => Some("trusty-memory"),
"analyze" => Some("trusty-analyze"),
"review" => Some("trusty-review"),
_ => None,
}
}
fn is_local_upstream(url: &str) -> bool {
url.starts_with("http://127.")
|| url.starts_with("http://[::1]")
|| url.starts_with("http://localhost")
}
pub fn build_upstream_url(base_url: &str, subpath: &str, query: Option<&str>) -> String {
let base = base_url.trim_end_matches('/');
let path = subpath.trim_start_matches('/');
let url = if path.is_empty() {
format!("{base}/")
} else {
format!("{base}/{path}")
};
match query {
Some(q) if !q.is_empty() => format!("{url}?{q}"),
_ => url,
}
}
fn filter_headers(headers: &HeaderMap) -> HeaderMap {
let mut out = HeaderMap::new();
for (name, value) in headers {
if !HOP_BY_HOP.contains(&name.as_str()) {
out.append(name.clone(), value.clone());
}
}
out
}
fn error_response(status: StatusCode, body: &'static str) -> Response {
Response::builder()
.status(status)
.header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
.body(Body::from(body))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}
pub async fn proxy_handler(
State(state): State<AppState>,
Path((daemon_key, subpath)): Path<(String, String)>,
req: Request,
) -> Response {
let Some(full_daemon_id) = full_id(&daemon_key) else {
warn!("proxy: unknown daemon key '{daemon_key}'");
return error_response(StatusCode::BAD_REQUEST, "unknown daemon");
};
let base_url = {
let snap = state.poller_cache().snapshot().await;
match snap {
None => {
warn!("proxy: cache not yet populated for '{daemon_key}'");
return error_response(StatusCode::SERVICE_UNAVAILABLE, "cache not ready");
}
Some(s) => {
let map = s.url_map();
match map.get(full_daemon_id).cloned() {
Some(url) => url,
None => {
warn!("proxy: daemon '{daemon_key}' is not running");
return error_response(StatusCode::BAD_GATEWAY, "daemon not running");
}
}
}
}
};
if !is_local_upstream(&base_url) {
warn!("proxy: upstream '{base_url}' is not a local address — rejecting (SSRF guard)");
return error_response(StatusCode::BAD_GATEWAY, "upstream not local");
}
let (parts, body) = req.into_parts();
let query = parts.uri.query();
let upstream_url = build_upstream_url(&base_url, &subpath, query);
debug!("proxy: {daemon_key} → {upstream_url}");
let method = match Method::from_bytes(parts.method.as_str().as_bytes()) {
Ok(m) => m,
Err(_) => {
return error_response(StatusCode::BAD_REQUEST, "unsupported method");
}
};
let safe_headers = filter_headers(&parts.headers);
const BODY_LIMIT: usize = 64 * 1024 * 1024;
let body_bytes: Bytes = match axum::body::to_bytes(body, BODY_LIMIT).await {
Ok(b) => b,
Err(e) => {
warn!("proxy: failed to read request body: {e}");
return error_response(
StatusCode::PAYLOAD_TOO_LARGE,
"request body exceeds proxy limit of 64 MiB",
);
}
};
let client = state.http_client();
let upstream_req = client
.request(method, &upstream_url)
.headers(safe_headers)
.body(body_bytes);
let upstream_resp = match upstream_req.send().await {
Ok(r) => r,
Err(e) => {
warn!("proxy: upstream request failed for '{daemon_key}': {e}");
return error_response(StatusCode::BAD_GATEWAY, "upstream request failed");
}
};
let status = StatusCode::from_u16(upstream_resp.status().as_u16())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
let mut resp_builder = Response::builder().status(status);
for (name, value) in upstream_resp.headers() {
if !HOP_BY_HOP.contains(&name.as_str())
&& let Ok(n) = HeaderName::from_bytes(name.as_str().as_bytes())
&& let Ok(v) = HeaderValue::from_bytes(value.as_bytes())
{
resp_builder = resp_builder.header(n, v);
}
}
let resp_body = match upstream_resp.bytes().await {
Ok(b) => Body::from(b),
Err(e) => {
warn!("proxy: failed to read upstream body: {e}");
Body::from("upstream body error")
}
};
resp_builder
.body(resp_body)
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_build_upstream_url_simple_path() {
assert_eq!(
build_upstream_url("http://127.0.0.1:7878", "health", None),
"http://127.0.0.1:7878/health"
);
}
#[test]
fn test_build_upstream_url_with_query() {
assert_eq!(
build_upstream_url(
"http://127.0.0.1:7879",
"indexes/abc/complexity_hotspots",
Some("top_k=5")
),
"http://127.0.0.1:7879/indexes/abc/complexity_hotspots?top_k=5"
);
}
#[test]
fn test_build_upstream_url_empty_path() {
assert_eq!(
build_upstream_url("http://127.0.0.1:7070", "", None),
"http://127.0.0.1:7070/"
);
}
#[test]
fn test_build_upstream_url_base_trailing_slash() {
assert_eq!(
build_upstream_url("http://127.0.0.1:7878/", "health", None),
"http://127.0.0.1:7878/health"
);
}
#[test]
fn test_build_upstream_url_empty_query_omitted() {
assert_eq!(
build_upstream_url("http://127.0.0.1:7878", "health", Some("")),
"http://127.0.0.1:7878/health"
);
}
#[test]
fn test_is_local_upstream_accepted() {
assert!(is_local_upstream("http://127.0.0.1:7878"));
assert!(is_local_upstream("http://127.0.0.1:7878/health"));
assert!(is_local_upstream("http://127.1.2.3:9000"));
assert!(is_local_upstream("http://[::1]:8080"));
assert!(is_local_upstream("http://localhost:7070"));
assert!(is_local_upstream("http://localhost"));
}
#[test]
fn test_is_local_upstream_rejected() {
assert!(!is_local_upstream("http://192.168.1.1:7878"));
assert!(!is_local_upstream("http://10.0.0.1:7879"));
assert!(!is_local_upstream("http://evil.example.com/steal"));
assert!(!is_local_upstream("https://127.0.0.1:7878")); assert!(!is_local_upstream("http://0.0.0.0:7878"));
}
#[test]
fn test_daemon_key_mapping() {
assert_eq!(full_id("search"), Some("trusty-search"));
assert_eq!(full_id("memory"), Some("trusty-memory"));
assert_eq!(full_id("analyze"), Some("trusty-analyze"));
assert_eq!(full_id("review"), Some("trusty-review"));
assert_eq!(full_id("unknown"), None);
}
#[test]
fn test_filter_headers_strips_hop_by_hop() {
let mut h = HeaderMap::new();
h.insert("connection", HeaderValue::from_static("keep-alive"));
h.insert("x-custom", HeaderValue::from_static("hello"));
let filtered = filter_headers(&h);
assert!(!filtered.contains_key("connection"));
assert!(filtered.contains_key("x-custom"));
}
}