use std::time::Duration;
use axum::Router;
use axum::body::Body;
use axum::extract::{FromRequest, Request, State, WebSocketUpgrade};
use axum::http::{HeaderMap, HeaderName, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::any;
use futures_util::{SinkExt, StreamExt};
use reqwest::Identity;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use crate::state::AppState;
use super::config::ProxySurface;
#[derive(Clone)]
struct ProxyState {
client: reqwest::Client,
upstream: String,
ws_upstream: String,
}
pub fn router(_state: &AppState, surface: &ProxySurface) -> Router<AppState> {
let mut builder = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.timeout(Duration::from_secs(300));
if let (Ok(cert_pem), Ok(key_pem)) = (
pem_from_env("NAUTALID_PROXY_CLIENT_CERT_PEM", "NAUTALID_PROXY_CLIENT_CERT_PATH"),
pem_from_env("NAUTALID_PROXY_CLIENT_KEY_PEM", "NAUTALID_PROXY_CLIENT_KEY_PATH"),
) {
if let Ok(id) = Identity::from_pem(format!("{cert_pem}\n{key_pem}").as_bytes()) {
builder = builder.identity(id);
}
}
if let Ok(ca) = pem_from_env("NAUTALID_PROXY_CLIENT_CA_PEM", "NAUTALID_PROXY_CLIENT_CA_PATH") {
builder = builder.tls_built_in_root_certs(false);
if let Ok(cert) = reqwest::Certificate::from_pem(ca.as_bytes()) {
builder = builder.add_root_certificate(cert);
}
}
let client = builder.build().expect("proxy client");
let ws_scheme = if surface.upstream.starts_with("https://") {
"wss"
} else {
"ws"
};
let ws_host = surface
.upstream
.trim_start_matches("https://")
.trim_start_matches("http://");
let ws_upstream = format!("{ws_scheme}://{ws_host}");
let inner = Router::new()
.route("/{*path}", any(proxy_handler))
.fallback(any(proxy_handler))
.with_state(ProxyState {
client,
upstream: surface.upstream.clone(),
ws_upstream,
});
Router::new().nest(&surface.prefix, inner)
}
async fn proxy_handler(State(state): State<ProxyState>, req: Request) -> Response {
if is_websocket_upgrade(req.headers()) {
let path_and_query = req
.uri()
.path_and_query()
.map(|pq| pq.as_str())
.unwrap_or("/")
.to_string();
let target = format!("{}{}", state.ws_upstream, path_and_query);
return match WebSocketUpgrade::from_request(req, &()).await {
Ok(ws) => ws
.on_upgrade(move |socket| proxy_websocket(socket, target))
.into_response(),
Err(_) => StatusCode::BAD_REQUEST.into_response(),
};
}
match proxy_http(State(state), req).await {
Ok(res) => res,
Err(status) => status.into_response(),
}
}
async fn proxy_http(
State(state): State<ProxyState>,
req: Request,
) -> Result<Response, StatusCode> {
let path_and_query = req
.uri()
.path_and_query()
.map(|pq| pq.as_str())
.unwrap_or("/");
let target = format!("{}{}", state.upstream, path_and_query);
let method = req.method().clone();
let headers = copy_forward_headers(req.headers());
let (parts, body) = req.into_parts();
let stream = body.into_data_stream().map(|result| {
result.map_err(std::io::Error::other)
});
let req_body = reqwest::Body::wrap_stream(stream);
let _ = parts;
let upstream = state
.client
.request(method, &target)
.headers(headers)
.body(req_body)
.send()
.await
.map_err(|_| StatusCode::BAD_GATEWAY)?;
let status = upstream.status();
let headers = upstream.headers().clone();
let stream = upstream.bytes_stream();
let body = Body::from_stream(stream);
let mut response = Response::new(body);
*response.status_mut() = status;
for (name, value) in headers.iter() {
if name == reqwest::header::TRANSFER_ENCODING {
continue;
}
if let (Some(n), Ok(v)) = (http::HeaderName::from_bytes(name.as_str().as_bytes()).ok(), http::HeaderValue::from_bytes(value.as_bytes())) {
response.headers_mut().insert(n, v);
}
}
let _ = parts;
Ok(response)
}
async fn proxy_websocket(
mut client_ws: axum::extract::ws::WebSocket,
target: String,
) {
let Ok((upstream, _)) = connect_async(&target).await else {
return;
};
let (mut upstream_tx, mut upstream_rx) = upstream.split();
loop {
tokio::select! {
msg = client_ws.recv() => {
match msg {
Some(Ok(axum::extract::ws::Message::Text(t))) => {
if upstream_tx
.send(Message::Text(t.to_string().into()))
.await
.is_err()
{
break;
}
}
Some(Ok(axum::extract::ws::Message::Binary(b))) => {
if upstream_tx.send(Message::Binary(b)).await.is_err() { break; }
}
Some(Ok(axum::extract::ws::Message::Close(_))) | None => break,
Some(Ok(axum::extract::ws::Message::Ping(p))) => {
if upstream_tx.send(Message::Ping(p)).await.is_err() { break; }
}
Some(Ok(axum::extract::ws::Message::Pong(p))) => {
if upstream_tx.send(Message::Pong(p)).await.is_err() { break; }
}
Some(Err(_)) => break,
}
}
msg = upstream_rx.next() => {
match msg {
Some(Ok(Message::Text(t))) => {
if client_ws
.send(axum::extract::ws::Message::Text(t.to_string().into()))
.await
.is_err()
{
break;
}
}
Some(Ok(Message::Binary(b))) => {
if client_ws.send(axum::extract::ws::Message::Binary(b)).await.is_err() { break; }
}
Some(Ok(Message::Close(_))) | None => break,
Some(Ok(Message::Ping(p))) => {
if client_ws.send(axum::extract::ws::Message::Ping(p)).await.is_err() { break; }
}
Some(Ok(Message::Pong(p))) => {
if client_ws.send(axum::extract::ws::Message::Pong(p)).await.is_err() { break; }
}
Some(Ok(Message::Frame(_))) => {}
Some(Err(_)) => break,
}
}
}
}
}
fn is_websocket_upgrade(headers: &HeaderMap) -> bool {
headers
.get("upgrade")
.and_then(|v| v.to_str().ok())
.is_some_and(|u| u.eq_ignore_ascii_case("websocket"))
}
fn copy_forward_headers(headers: &HeaderMap) -> HeaderMap {
let mut out = HeaderMap::new();
for (name, value) in headers.iter() {
if name == HeaderName::from_static("host") || name == HeaderName::from_static("connection") {
continue;
}
out.insert(name.clone(), value.clone());
}
out
}
fn pem_from_env(inline: &str, path_var: &str) -> Result<String, ()> {
if let Ok(pem) = std::env::var(inline) {
if !pem.is_empty() {
return Ok(pem);
}
}
if let Ok(path) = std::env::var(path_var) {
if !path.is_empty() {
return std::fs::read_to_string(path).map_err(|_| ());
}
}
Err(())
}