mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(rust): websocket upgrade + bidirectional pump (phase-1)
When the catch-all sees an Upgrade: websocket request, hand it to the ws module: axum upgrades the client side, tokio-tungstenite connects to the upstream (rewriting http->ws / https->wss while preserving path + query), and two pumps shovel messages until either side closes. Forwarded headers exclude what tungstenite manages (Host, Upgrade, Connection, Sec-*) but preserve Authorization, Sec-WebSocket-Protocol, etc. Supports text, binary, ping, pong, and close frames in both directions.
This commit is contained in:
chopratejas
parent
0493ebf1fe
commit
dd6ebef8cf
1 changed files with 252 additions and 0 deletions
252
crates/headroom-proxy/src/websocket.rs
Normal file
252
crates/headroom-proxy/src/websocket.rs
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
//! WebSocket reverse-proxy handler.
|
||||
//!
|
||||
//! Accepts a client upgrade via axum, opens a tungstenite connection to the
|
||||
//! upstream (rewriting scheme http->ws / https->wss), and bidirectionally
|
||||
//! pumps messages until either side closes.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::extract::ws::{CloseFrame, Message as AxMsg, WebSocket, WebSocketUpgrade};
|
||||
use axum::http::{HeaderName, HeaderValue, Request, Response, StatusCode};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::protocol::CloseFrame as TgCloseFrame;
|
||||
use tokio_tungstenite::tungstenite::Message as TgMsg;
|
||||
|
||||
use crate::headers::build_forward_request_headers;
|
||||
use crate::proxy::AppState;
|
||||
|
||||
/// Entry point invoked from the catch-all when an upgrade is detected.
|
||||
pub async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
state: AppState,
|
||||
client_addr: SocketAddr,
|
||||
req: Request<Body>,
|
||||
) -> Response<Body> {
|
||||
let request_id = req
|
||||
.headers()
|
||||
.get("x-request-id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
// Build the upstream WS URL.
|
||||
let upstream_url = match build_upstream_ws_url(&state.config.upstream, req.uri()) {
|
||||
Ok(u) => u,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "failed to build upstream ws url");
|
||||
return (StatusCode::BAD_GATEWAY, e).into_response_body();
|
||||
}
|
||||
};
|
||||
|
||||
// Build forwarded headers (drop hop-by-hop EXCEPT Upgrade/Connection — but
|
||||
// tungstenite generates its own; we forward only the user-meaningful ones
|
||||
// such as Authorization, Sec-WebSocket-Protocol, etc.).
|
||||
let forwarded_host = req
|
||||
.headers()
|
||||
.get(http::header::HOST)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let forward_headers = build_forward_request_headers(
|
||||
req.headers(),
|
||||
client_addr.ip(),
|
||||
"http",
|
||||
forwarded_host.as_deref(),
|
||||
&request_id,
|
||||
);
|
||||
// Sec-WebSocket-Protocol must be propagated for subprotocol negotiation.
|
||||
let subprotocols: Option<String> = req
|
||||
.headers()
|
||||
.get("sec-websocket-protocol")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
ws.on_upgrade(move |client_ws| async move {
|
||||
if let Err(e) = run_ws_pump(
|
||||
client_ws,
|
||||
upstream_url,
|
||||
forward_headers,
|
||||
subprotocols,
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "websocket pump ended with error");
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
trait IntoResponseBody {
|
||||
fn into_response_body(self) -> Response<Body>;
|
||||
}
|
||||
impl IntoResponseBody for (StatusCode, String) {
|
||||
fn into_response_body(self) -> Response<Body> {
|
||||
Response::builder()
|
||||
.status(self.0)
|
||||
.body(Body::from(self.1))
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_upstream_ws_url(base: &url::Url, req_uri: &http::Uri) -> Result<url::Url, String> {
|
||||
let mut joined = base.clone();
|
||||
let new_scheme = match joined.scheme() {
|
||||
"http" => "ws".to_string(),
|
||||
"https" => "wss".to_string(),
|
||||
s @ ("ws" | "wss") => s.to_string(),
|
||||
other => return Err(format!("unsupported upstream scheme: {other}")),
|
||||
};
|
||||
joined
|
||||
.set_scheme(&new_scheme)
|
||||
.map_err(|()| "failed to set ws scheme".to_string())?;
|
||||
let path = req_uri.path();
|
||||
let query = req_uri.query();
|
||||
let base_path = joined.path().trim_end_matches('/').to_string();
|
||||
let combined = if path.is_empty() || path == "/" {
|
||||
if base_path.is_empty() {
|
||||
"/".to_string()
|
||||
} else {
|
||||
base_path
|
||||
}
|
||||
} else if base_path.is_empty() {
|
||||
path.to_string()
|
||||
} else {
|
||||
format!("{base_path}{path}")
|
||||
};
|
||||
joined.set_path(&combined);
|
||||
joined.set_query(query);
|
||||
Ok(joined)
|
||||
}
|
||||
|
||||
async fn run_ws_pump(
|
||||
client_ws: WebSocket,
|
||||
upstream_url: url::Url,
|
||||
forward_headers: http::HeaderMap,
|
||||
subprotocols: Option<String>,
|
||||
request_id: String,
|
||||
) -> Result<(), String> {
|
||||
// Build the upstream handshake request manually so we can inject headers.
|
||||
let mut req = upstream_url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|e| format!("ws into_client_request: {e}"))?;
|
||||
{
|
||||
let h = req.headers_mut();
|
||||
// Tungstenite will set Host, Upgrade, Connection, Sec-WebSocket-Key,
|
||||
// Sec-WebSocket-Version itself. We add user-meaningful pass-throughs.
|
||||
for (name, value) in forward_headers.iter() {
|
||||
// Skip headers tungstenite manages.
|
||||
let n = name.as_str().to_ascii_lowercase();
|
||||
if matches!(
|
||||
n.as_str(),
|
||||
"host"
|
||||
| "upgrade"
|
||||
| "connection"
|
||||
| "sec-websocket-key"
|
||||
| "sec-websocket-version"
|
||||
| "sec-websocket-extensions"
|
||||
| "content-length"
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
h.append(name, value.clone());
|
||||
}
|
||||
if let Some(sp) = subprotocols {
|
||||
if let Ok(v) = HeaderValue::from_str(&sp) {
|
||||
h.insert(HeaderName::from_static("sec-websocket-protocol"), v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (upstream_ws, _resp) = tokio_tungstenite::connect_async(req)
|
||||
.await
|
||||
.map_err(|e| format!("upstream ws connect: {e}"))?;
|
||||
|
||||
let (mut upstream_sink, mut upstream_stream) = upstream_ws.split();
|
||||
let (mut client_sink, mut client_stream) = client_ws.split();
|
||||
|
||||
// Pump client -> upstream.
|
||||
let c2u = async {
|
||||
while let Some(msg) = client_stream.next().await {
|
||||
let m = match msg {
|
||||
Ok(m) => m,
|
||||
Err(_) => break,
|
||||
};
|
||||
let tg = match ax_to_tg(m) {
|
||||
Some(tg) => tg,
|
||||
None => continue,
|
||||
};
|
||||
let close = matches!(tg, TgMsg::Close(_));
|
||||
if upstream_sink.send(tg).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if close {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = upstream_sink.close().await;
|
||||
};
|
||||
|
||||
// Pump upstream -> client.
|
||||
let u2c = async {
|
||||
while let Some(msg) = upstream_stream.next().await {
|
||||
let m = match msg {
|
||||
Ok(m) => m,
|
||||
Err(_) => break,
|
||||
};
|
||||
let ax = match tg_to_ax(m) {
|
||||
Some(ax) => ax,
|
||||
None => continue,
|
||||
};
|
||||
let close = matches!(ax, AxMsg::Close(_));
|
||||
if client_sink.send(ax).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if close {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = client_sink.close().await;
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
request_id = %request_id,
|
||||
protocol = "ws",
|
||||
upstream = %upstream_url,
|
||||
"ws session opened"
|
||||
);
|
||||
tokio::join!(c2u, u2c);
|
||||
tracing::info!(request_id = %request_id, protocol = "ws", "ws session closed");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ax_to_tg(m: AxMsg) -> Option<TgMsg> {
|
||||
Some(match m {
|
||||
AxMsg::Text(t) => TgMsg::Text(t.to_string()),
|
||||
AxMsg::Binary(b) => TgMsg::Binary(b.to_vec()),
|
||||
AxMsg::Ping(p) => TgMsg::Ping(p.to_vec()),
|
||||
AxMsg::Pong(p) => TgMsg::Pong(p.to_vec()),
|
||||
AxMsg::Close(Some(cf)) => TgMsg::Close(Some(TgCloseFrame {
|
||||
code: tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::from(cf.code),
|
||||
reason: cf.reason.to_string().into(),
|
||||
})),
|
||||
AxMsg::Close(None) => TgMsg::Close(None),
|
||||
})
|
||||
}
|
||||
|
||||
fn tg_to_ax(m: TgMsg) -> Option<AxMsg> {
|
||||
Some(match m {
|
||||
TgMsg::Text(t) => AxMsg::Text(t.as_str().to_string()),
|
||||
TgMsg::Binary(b) => AxMsg::Binary(b.to_vec()),
|
||||
TgMsg::Ping(p) => AxMsg::Ping(p.to_vec()),
|
||||
TgMsg::Pong(p) => AxMsg::Pong(p.to_vec()),
|
||||
TgMsg::Close(Some(cf)) => AxMsg::Close(Some(CloseFrame {
|
||||
code: cf.code.into(),
|
||||
reason: cf.reason.to_string().into(),
|
||||
})),
|
||||
TgMsg::Close(None) => AxMsg::Close(None),
|
||||
TgMsg::Frame(_) => return None,
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue