This commit is contained in:
Mateusz Gruszczyński committed 2026-10-04 11:15:58 +02:00
1 parent 87e571c170
commit fa7abc9e53
308 files changed
+49746 -50237

No files matched your search

+418
View File
@@ -0,0 +1,418 @@
use super::body::Exchange;
use crate::{
runtime::App,
server::{ConnectionLease, ServerShutdown},
};
use anyhow::{ensure, Result};
use axum::{body::Body, extract::Request};
use base64::{engine::general_purpose::STANDARD, Engine};
use bytes::{Buf, Bytes};
use h3::error::Code;
use http::{HeaderMap, Method, Version};
use http_body::{Body as HttpBody, Frame, SizeHint};
use http_body_util::BodyExt;
use hyper::upgrade::{OnUpgrade, Upgraded};
use hyper_util::rt::TokioIo;
use rand::{rngs::OsRng, RngCore};
use sha1::{Digest, Sha1};
use serde::Serialize;
use std::{
collections::BTreeMap,
future::Future,
io,
pin::Pin,
sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc, Mutex, Weak,
},
task::{Context, Poll},
time::{Duration, Instant},
};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio_util::sync::CancellationToken;
fn has_token(headers: &HeaderMap, name: &str, token: &str) -> bool {
headers
.get_all(name)
.iter()
.filter_map(|v| v.to_str().ok())
.flat_map(|v| v.split(','))
.any(|v| v.trim().eq_ignore_ascii_case(token))
}
fn single<'a>(headers: &'a HeaderMap, name: &str) -> Result<&'a str> {
ensure!(
headers.get_all(name).iter().count() == 1,
"WebSocket header must occur once"
);
Ok(headers
.get(name)
.ok_or_else(|| anyhow::anyhow!("missing WebSocket header"))?
.to_str()?)
}
pub struct Handshake {
pub accept: String,
upstream_key: Option<String>,
protocols: Vec<String>,
bootstrap: Bootstrap,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Bootstrap {
Http1,
Http2,
Http3,
}
impl Handshake {
pub fn h1(&self) -> bool {
self.bootstrap == Bootstrap::Http1
}
pub fn h2(&self) -> bool {
self.bootstrap == Bootstrap::Http2
}
pub fn h3(&self) -> bool {
self.bootstrap == Bootstrap::Http3
}
pub fn upstream_key(&self) -> Option<&str> {
self.upstream_key.as_deref()
}
}
fn protocols(headers: &HeaderMap) -> Result<Vec<String>> {
let protocols = headers
.get_all("sec-websocket-protocol")
.iter()
.map(|v| v.to_str())
.collect::<std::result::Result<Vec<_>, _>>()?
.into_iter()
.flat_map(|v| v.split(','))
.map(|v| v.trim().to_owned())
.collect::<Vec<_>>();
ensure!(
protocols.len() <= 32 && protocols.iter().all(|v| !v.is_empty()),
"invalid WebSocket protocols"
);
Ok(protocols)
}
fn accept_for_key(key: &str) -> String {
let mut hasher = Sha1::new();
hasher.update(key.as_bytes());
hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
STANDARD.encode(hasher.finalize())
}
pub fn validate(request: &Request) -> Result<Handshake> {
ensure!(
request.method() == Method::GET && request.version() == Version::HTTP_11,
"WebSocket requires HTTP/1.1 GET"
);
ensure!(
single(request.headers(), "upgrade")?.eq_ignore_ascii_case("websocket")
&& has_token(request.headers(), "connection", "upgrade"),
"invalid WebSocket upgrade"
);
ensure!(
single(request.headers(), "sec-websocket-version")? == "13",
"WebSocket version must be 13"
);
let key = single(request.headers(), "sec-websocket-key")?;
ensure!(
key.len() == 24 && STANDARD.decode(key)?.len() == 16,
"invalid WebSocket key"
);
ensure!(
!request.headers().contains_key("transfer-encoding"),
"WebSocket handshake cannot contain a transfer body"
);
ensure!(
request.headers().get_all("content-length").iter().count() <= 1
&& request
.headers()
.get("content-length")
.is_none_or(|v| v == "0"),
"WebSocket handshake cannot contain a body"
);
for value in request.headers().get_all("connection").iter() {
for field in value.to_str()?.split(',') {
ensure!(
!field
.trim()
.to_ascii_lowercase()
.starts_with("sec-websocket-"),
"Connection must not nominate WebSocket handshake fields"
);
}
}
Ok(Handshake {
accept: accept_for_key(key),
upstream_key: None,
protocols: protocols(request.headers())?,
bootstrap: Bootstrap::Http1,
})
}
/// Validate the RFC 8441 bootstrap used by RFC 9220 on HTTP/3. The
/// Sec-WebSocket-Key/Accept proof belongs only to the HTTP/1.1 hop, so a fresh
/// key is synthesized later for the backend Upgrade request.
pub fn validate_h3(request: &Request) -> Result<Handshake> {
ensure!(
request.method() == Method::CONNECT && request.version() == Version::HTTP_3,
"HTTP/3 WebSocket requires Extended CONNECT"
);
ensure!(
request
.extensions()
.get::<h3::ext::Protocol>()
.is_some_and(|protocol| protocol.as_str() == "websocket"),
"Extended CONNECT protocol is not websocket"
);
for name in [
"connection",
"upgrade",
"sec-websocket-key",
"sec-websocket-accept",
"transfer-encoding",
"content-length",
] {
ensure!(
!request.headers().contains_key(name),
"HTTP/1.1-only WebSocket field on HTTP/3"
);
}
ensure!(
single(request.headers(), "sec-websocket-version")? == "13",
"WebSocket version must be 13"
);
let mut nonce = [0u8; 16];
OsRng.fill_bytes(&mut nonce);
let key = STANDARD.encode(nonce);
Ok(Handshake {
accept: accept_for_key(&key),
upstream_key: Some(key),
protocols: protocols(request.headers())?,
bootstrap: Bootstrap::Http3,
})
}
/// Validate RFC 8441 Extended CONNECT on HTTP/2. The downstream hop does not
/// use Sec-WebSocket-Key/Accept; a fresh RFC 6455 key is generated only when
/// RustProxy has to bridge the stream to an HTTP/1.1 backend.
pub fn validate_h2(request: &Request) -> Result<Handshake> {
ensure!(
request.method() == Method::CONNECT && request.version() == Version::HTTP_2,
"HTTP/2 WebSocket requires Extended CONNECT"
);
ensure!(
request
.extensions()
.get::<hyper::ext::Protocol>()
.is_some_and(|protocol| protocol.as_str() == "websocket"),
"Extended CONNECT protocol is not websocket"
);
for name in [
"connection",
"upgrade",
"sec-websocket-key",
"sec-websocket-accept",
"transfer-encoding",
"content-length",
] {
ensure!(
!request.headers().contains_key(name),
"HTTP/1.1-only WebSocket field on HTTP/2"
);
}
ensure!(
single(request.headers(), "sec-websocket-version")? == "13",
"WebSocket version must be 13"
);
let mut nonce = [0u8; 16];
OsRng.fill_bytes(&mut nonce);
let key = STANDARD.encode(nonce);
Ok(Handshake {
accept: accept_for_key(&key),
upstream_key: Some(key),
protocols: protocols(request.headers())?,
bootstrap: Bootstrap::Http2,
})
}
pub fn h2_protocol(request: &Request) -> Option<&str> {
request
.extensions()
.get::<hyper::ext::Protocol>()
.map(|protocol| protocol.as_str())
}
pub fn h3_protocol(request: &Request) -> Option<&str> {
request
.extensions()
.get::<h3::ext::Protocol>()
.map(|protocol| protocol.as_str())
}
pub fn validate_response(
headers: &HeaderMap,
expected: &Handshake,
legacy_support: bool,
) -> Result<()> {
if legacy_support {
// Some older embedded servers return a valid 101 + Sec-WebSocket-Accept
// while omitting Connection: Upgrade, duplicating Upgrade, or attaching
// zero/meaningless HTTP body-framing headers. The connection is dedicated
// to the upgraded tunnel and those framing headers are stripped before the
// downstream response, so keep the cryptographic accept proof strict while
// tolerating the legacy response syntax.
ensure!(
has_token(headers, "upgrade", "websocket"),
"invalid legacy upstream upgrade"
);
if headers.contains_key("content-length") {
for value in headers.get_all("content-length").iter() {
ensure!(
value.to_str()?.trim() == "0",
"legacy 101 response declared a non-zero body"
);
}
}
} else {
ensure!(
single(headers, "upgrade")?.eq_ignore_ascii_case("websocket")
&& has_token(headers, "connection", "upgrade"),
"invalid upstream upgrade"
);
ensure!(
!headers.contains_key("transfer-encoding")
&& !headers.contains_key("content-length"),
"101 response must not include body framing"
);
for value in headers.get_all("connection").iter() {
for field in value.to_str()?.split(',') {
ensure!(
!field
.trim()
.to_ascii_lowercase()
.starts_with("sec-websocket-"),
"upstream nominated a protected handshake field"
);
}
}
}
ensure!(
single(headers, "sec-websocket-accept")? == expected.accept,
"upstream WebSocket accept mismatch"
);
if headers.contains_key("sec-websocket-protocol") {
let selected = single(headers, "sec-websocket-protocol")?;
ensure!(
expected.protocols.iter().any(|v| v == selected),
"upstream chose an unoffered subprotocol"
);
}
Ok(())
}
pub fn validate_h2_response(headers: &HeaderMap, expected: &Handshake) -> Result<()> {
for name in [
"connection",
"upgrade",
"sec-websocket-key",
"sec-websocket-accept",
"transfer-encoding",
"content-length",
] {
ensure!(
!headers.contains_key(name),
"HTTP/1.1-only WebSocket field on HTTP/2 response"
);
}
if headers.contains_key("sec-websocket-protocol") {
let selected = single(headers, "sec-websocket-protocol")?;
ensure!(
expected.protocols.iter().any(|value| value == selected),
"upstream chose an unoffered subprotocol"
);
}
Ok(())
}
struct MeteredIo<T> {
inner: T,
exchange: Arc<Exchange>,
session: Arc<TunnelSession>,
client_to_backend: bool,
started: Instant,
activity: Arc<AtomicU64>,
}
fn record_bytes(
exchange: &Exchange,
bytes: u64,
client_to_backend: bool,
activity: &AtomicU64,
started: Instant,
) {
if bytes == 0 {
return;
}
exchange.backend_bytes(bytes, client_to_backend);
activity.store(started.elapsed().as_millis() as u64, Ordering::Relaxed);
if client_to_backend {
exchange
.metrics
.websocket_bytes_in
.fetch_add(bytes, Ordering::Relaxed);
exchange.request_bytes.fetch_add(bytes, Ordering::Relaxed);
} else {
exchange
.metrics
.websocket_bytes_out
.fetch_add(bytes, Ordering::Relaxed);
exchange.response_bytes.fetch_add(bytes, Ordering::Relaxed);
}
}
impl<T: AsyncRead + Unpin> AsyncRead for MeteredIo<T> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
let before = buf.filled().len();
let result = Pin::new(&mut this.inner).poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &result {
let bytes = (buf.filled().len() - before) as u64;
record_bytes(
&this.exchange,
bytes,
this.client_to_backend,
&this.activity,
this.started,
);
this.session.record(bytes, this.client_to_backend);
}
result
}
}
impl<T: AsyncWrite + Unpin> AsyncWrite for MeteredIo<T> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bytes: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.get_mut().inner).poll_write(cx, bytes)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
}
}
#[derive(Debug, Clone, Copy)]
enum TunnelTransport {
H1,
H2,
H3,
}
+207
View File
@@ -0,0 +1,207 @@
mod tests {
use super::*;
use axum::body::Body;
fn request() -> Request {
Request::builder()
.method("GET")
.version(Version::HTTP_11)
.header("upgrade", "websocket")
.header("connection", "Upgrade")
.header("sec-websocket-version", "13")
.header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==")
.body(Body::empty())
.unwrap()
}
fn h3_request() -> Request {
let mut request = Request::builder()
.method(Method::CONNECT)
.version(Version::HTTP_3)
.uri("https://example.test/socket")
.header("sec-websocket-version", "13")
.body(Body::empty())
.unwrap();
request
.extensions_mut()
.insert("websocket".parse::<h3::ext::Protocol>().unwrap_or_else(|_| panic!("static websocket H3 protocol token must parse")));
request
}
fn h2_request() -> Request {
let mut request = Request::builder()
.method(Method::CONNECT)
.version(Version::HTTP_2)
.uri("https://example.test/socket")
.header("sec-websocket-version", "13")
.body(Body::empty())
.unwrap();
request
.extensions_mut()
.insert(hyper::ext::Protocol::from_static("websocket"));
request
}
#[test]
fn rfc_handshake_vector() {
let handshake = validate(&request()).unwrap();
assert_eq!(handshake.accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
}
#[test]
fn rfc9220_bootstrap_uses_extended_connect_without_h1_proof_headers() {
let request = h3_request();
let handshake = validate_h3(&request).unwrap();
assert!(handshake.h3());
let key = handshake.upstream_key().expect("backend nonce");
assert_eq!(STANDARD.decode(key).unwrap().len(), 16);
assert!(!request.headers().contains_key("connection"));
assert!(!request.headers().contains_key("upgrade"));
assert!(!request.headers().contains_key("sec-websocket-key"));
assert!(!request.headers().contains_key("sec-websocket-accept"));
let mut backend = HeaderMap::new();
backend.insert("connection", "Upgrade".parse().unwrap());
backend.insert("upgrade", "websocket".parse().unwrap());
backend.insert("sec-websocket-accept", accept_for_key(key).parse().unwrap());
assert!(validate_response(&backend, &handshake, false).is_ok());
}
#[test]
fn rfc8441_h2_bootstrap_uses_extended_connect_without_h1_proof_headers() {
let request = h2_request();
let handshake = validate_h2(&request).unwrap();
assert!(handshake.h2());
assert!(!handshake.h1());
let key = handshake.upstream_key().expect("HTTP/1 bridge nonce");
assert_eq!(STANDARD.decode(key).unwrap().len(), 16);
assert!(!request.headers().contains_key("connection"));
assert!(!request.headers().contains_key("upgrade"));
assert!(!request.headers().contains_key("sec-websocket-key"));
let mut h2_response = HeaderMap::new();
h2_response.insert("sec-websocket-protocol", "chat".parse().unwrap());
assert!(validate_h2_response(&h2_response, &handshake).is_err());
}
#[test]
fn rfc8441_h2_response_rejects_h1_framing_fields() {
let handshake = validate_h2(&h2_request()).unwrap();
for (name, value) in [("connection", "close"), ("upgrade", "websocket"), ("content-length", "0")] {
let mut headers = HeaderMap::new();
headers.insert(
http::header::HeaderName::from_bytes(name.as_bytes()).unwrap(),
value.parse().unwrap(),
);
assert!(validate_h2_response(&headers, &handshake).is_err(), "accepted forbidden response {name}");
}
}
#[test]
fn tunnel_transport_metrics_classify_frontend_protocol() {
assert!(matches!(TunnelTransport::from_protocol("h1-upgrade"), TunnelTransport::H1));
assert!(matches!(TunnelTransport::from_protocol("h2-rfc8441-h1-upstream"), TunnelTransport::H2));
assert!(matches!(TunnelTransport::from_protocol("h2-rfc8441-h2-upstream"), TunnelTransport::H2));
assert!(matches!(TunnelTransport::from_protocol("h3-rfc9220"), TunnelTransport::H3));
}
#[test]
fn rfc8441_h2_rejects_h1_only_fields_and_wrong_protocol() {
for (name, value) in [
("connection", "Upgrade"),
("upgrade", "websocket"),
("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ=="),
("sec-websocket-accept", "ignored"),
("content-length", "0"),
("transfer-encoding", "chunked"),
] {
let mut request = h2_request();
request.headers_mut().insert(
http::header::HeaderName::from_bytes(name.as_bytes()).unwrap(),
value.parse().unwrap(),
);
assert!(validate_h2(&request).is_err(), "accepted forbidden {name}");
}
let mut request = h2_request();
request
.extensions_mut()
.insert(hyper::ext::Protocol::from_static("connect-udp"));
assert!(validate_h2(&request).is_err());
}
#[test]
fn rfc9220_rejects_h1_only_fields_and_non_websocket_protocols() {
for (name, value) in [
("connection", "Upgrade"),
("upgrade", "websocket"),
("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ=="),
("sec-websocket-accept", "ignored"),
("content-length", "0"),
("transfer-encoding", "chunked"),
] {
let mut request = h3_request();
request.headers_mut().insert(
http::header::HeaderName::from_bytes(name.as_bytes()).unwrap(),
value.parse().unwrap(),
);
assert!(validate_h3(&request).is_err(), "accepted forbidden {name}");
}
let mut request = h3_request();
request
.extensions_mut()
.insert("connect-udp".parse::<h3::ext::Protocol>().unwrap_or_else(|_| panic!("static connect-udp H3 protocol token must parse")));
assert!(validate_h3(&request).is_err());
}
#[test]
fn malformed_or_body_bearing_upgrade_fails() {
let mut req = request();
req.headers_mut()
.insert("content-length", "1".parse().unwrap());
assert!(validate(&req).is_err());
let mut req = request();
req.headers_mut()
.insert("sec-websocket-key", "invalid".parse().unwrap());
assert!(validate(&req).is_err());
let mut req = request();
*req.version_mut() = Version::HTTP_2;
assert!(validate(&req).is_err());
}
#[test]
fn legacy_upstream_response_accepts_old_101_syntax_but_keeps_accept_proof_strict() {
let expected = validate(&request()).unwrap();
let mut headers = HeaderMap::new();
headers.append("upgrade", "websocket".parse().unwrap());
headers.append("upgrade", "WebSocket".parse().unwrap());
headers.insert("content-length", "0".parse().unwrap());
headers.insert("transfer-encoding", "chunked".parse().unwrap());
headers.insert("sec-websocket-accept", expected.accept.parse().unwrap());
assert!(validate_response(&headers, &expected, false).is_err());
assert!(validate_response(&headers, &expected, true).is_ok());
headers.insert("content-length", "1".parse().unwrap());
assert!(validate_response(&headers, &expected, true).is_err());
headers.insert("content-length", "0".parse().unwrap());
headers.insert("sec-websocket-accept", "wrong".parse().unwrap());
assert!(validate_response(&headers, &expected, true).is_err());
}
#[test]
fn normal_tcp_close_is_not_a_websocket_proxy_error() {
for kind in [
io::ErrorKind::UnexpectedEof,
io::ErrorKind::ConnectionReset,
io::ErrorKind::ConnectionAborted,
io::ErrorKind::BrokenPipe,
io::ErrorKind::NotConnected,
] {
assert!(benign_tunnel_close(&io::Error::new(kind, "closed")));
}
assert!(!benign_tunnel_close(&io::Error::new(
io::ErrorKind::TimedOut,
"idle timeout"
)));
}
#[test]
fn upstream_cannot_select_unoffered_protocol() {
let expected = validate(&request()).unwrap();
let mut headers = HeaderMap::new();
headers.insert("connection", "Upgrade".parse().unwrap());
headers.insert("upgrade", "websocket".parse().unwrap());
headers.insert("sec-websocket-accept", expected.accept.parse().unwrap());
assert!(validate_response(&headers, &expected, false).is_ok());
headers.insert("sec-websocket-protocol", "unoffered".parse().unwrap());
assert!(validate_response(&headers, &expected, false).is_err());
}
}
+425
View File
@@ -0,0 +1,425 @@
impl TunnelTransport {
fn from_protocol(protocol: &str) -> Self {
if protocol.starts_with("h2-") {
Self::H2
} else if protocol.starts_with("h3-") {
Self::H3
} else {
Self::H1
}
}
fn increment(self, metrics: &crate::metrics::Metrics) {
match self {
Self::H1 => {
metrics.websocket_h1_total.fetch_add(1, Ordering::Relaxed);
metrics.websocket_h1_active.fetch_add(1, Ordering::Relaxed);
}
Self::H2 => {
metrics.websocket_h2_total.fetch_add(1, Ordering::Relaxed);
metrics.websocket_h2_active.fetch_add(1, Ordering::Relaxed);
}
Self::H3 => {
metrics.websocket_h3_total.fetch_add(1, Ordering::Relaxed);
metrics.websocket_h3_active.fetch_add(1, Ordering::Relaxed);
}
}
}
fn decrement_active(self, metrics: &crate::metrics::Metrics) {
match self {
Self::H1 => {
metrics.websocket_h1_active.fetch_sub(1, Ordering::Relaxed);
}
Self::H2 => {
metrics.websocket_h2_active.fetch_sub(1, Ordering::Relaxed);
}
Self::H3 => {
metrics.websocket_h3_active.fetch_sub(1, Ordering::Relaxed);
}
}
}
}
struct ActiveTunnel {
exchange: Arc<Exchange>,
started: Instant,
transport: TunnelTransport,
established: AtomicBool,
}
impl ActiveTunnel {
fn establish(&self) {
if self
.established
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
self.exchange.websocket_phase("tunnel_started");
self.exchange
.metrics
.websocket_total
.fetch_add(1, Ordering::Relaxed);
self.exchange
.metrics
.websocket_active
.fetch_add(1, Ordering::Relaxed);
self.transport.increment(&self.exchange.metrics);
}
}
fn is_established(&self) -> bool {
self.established.load(Ordering::Acquire)
}
}
impl Drop for ActiveTunnel {
fn drop(&mut self) {
if !self.is_established() {
return;
}
self.exchange
.metrics
.websocket_active
.fetch_sub(1, Ordering::Relaxed);
self.transport.decrement_active(&self.exchange.metrics);
self.exchange
.metrics
.websocket_duration_us
.fetch_add(self.started.elapsed().as_micros() as u64, Ordering::Relaxed);
self.exchange
.metrics
.websocket_closed
.fetch_add(1, Ordering::Relaxed);
}
}
#[derive(Debug, Clone, Serialize)]
pub struct TunnelInfo {
pub id: u64,
pub client: String,
pub host: String,
pub upstream: String,
pub protocol: String,
pub duration_ms: u64,
pub rx_bytes: u64,
pub tx_bytes: u64,
}
#[derive(Default)]
pub struct TunnelRegistry {
entries: Mutex<BTreeMap<u64, Weak<TunnelSession>>>,
}
pub struct TunnelSession {
registry: Weak<TunnelRegistry>,
id: u64,
client: String,
host: String,
upstream: String,
protocol: String,
started: Instant,
rx_bytes: AtomicU64,
tx_bytes: AtomicU64,
activity: AtomicU64,
cancel: CancellationToken,
_active: ActiveTunnel,
_lease: Option<Arc<ConnectionLease>>,
}
impl TunnelRegistry {
fn register_inner(
self: &Arc<Self>,
exchange: Arc<Exchange>,
upstream: String,
protocol: &str,
stop: CancellationToken,
lease: Option<Arc<ConnectionLease>>,
established: bool,
) -> Arc<TunnelSession> {
let started = Instant::now();
let transport = TunnelTransport::from_protocol(protocol);
let client = {
exchange
.client_ip
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
.unwrap_or_else(|| "-".into())
};
let session = Arc::new(TunnelSession {
registry: Arc::downgrade(self),
id: exchange.id,
client,
host: exchange.host.clone(),
upstream,
protocol: protocol.to_owned(),
started,
rx_bytes: AtomicU64::new(0),
tx_bytes: AtomicU64::new(0),
activity: AtomicU64::new(0),
cancel: stop.child_token(),
_active: ActiveTunnel {
exchange,
started,
transport,
established: AtomicBool::new(false),
},
_lease: lease,
});
self.entries
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(session.id, Arc::downgrade(&session));
if established {
session.establish();
}
session
}
pub fn register(
self: &Arc<Self>,
exchange: Arc<Exchange>,
upstream: String,
protocol: &str,
stop: CancellationToken,
lease: Option<Arc<ConnectionLease>>,
) -> Arc<TunnelSession> {
self.register_inner(exchange, upstream, protocol, stop, lease, true)
}
pub fn register_pending(
self: &Arc<Self>,
exchange: Arc<Exchange>,
upstream: String,
protocol: &str,
stop: CancellationToken,
lease: Option<Arc<ConnectionLease>>,
) -> Arc<TunnelSession> {
self.register_inner(exchange, upstream, protocol, stop, lease, false)
}
pub fn list(&self) -> Vec<TunnelInfo> {
let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
let mut result = Vec::with_capacity(entries.len());
entries.retain(|_, weak| {
if let Some(session) = weak.upgrade() {
if session._active.is_established() {
result.push(session.info());
}
true
} else {
false
}
});
result
}
pub fn kill(&self, id: u64) -> Option<TunnelInfo> {
let session = self
.entries
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&id)
.and_then(Weak::upgrade)?;
let info = session.info();
session.cancel.cancel();
Some(info)
}
}
impl TunnelSession {
pub fn establish(&self) {
self._active.establish();
}
fn info(&self) -> TunnelInfo {
TunnelInfo {
id: self.id,
client: self.client.clone(),
host: self.host.clone(),
upstream: self.upstream.clone(),
protocol: self.protocol.clone(),
duration_ms: self.started.elapsed().as_millis().min(u64::MAX as u128) as u64,
rx_bytes: self.rx_bytes.load(Ordering::Relaxed),
tx_bytes: self.tx_bytes.load(Ordering::Relaxed),
}
}
fn record(&self, bytes: u64, client_to_backend: bool) {
if bytes == 0 {
return;
}
if client_to_backend {
self.rx_bytes.fetch_add(bytes, Ordering::Relaxed);
} else {
self.tx_bytes.fetch_add(bytes, Ordering::Relaxed);
}
self.activity.store(
self.started.elapsed().as_millis().min(u64::MAX as u128) as u64,
Ordering::Relaxed,
);
}
pub fn cancellation(&self) -> CancellationToken {
self.cancel.clone()
}
fn inactive_for(&self) -> Duration {
self.started
.elapsed()
.saturating_sub(Duration::from_millis(self.activity.load(Ordering::Relaxed)))
}
}
impl Drop for TunnelSession {
fn drop(&mut self) {
if let Some(registry) = self.registry.upgrade() {
let mut entries = registry.entries.lock().unwrap_or_else(|e| e.into_inner());
if entries
.get(&self.id)
.and_then(Weak::upgrade)
.is_none_or(|session| std::ptr::eq(session.as_ref(), self))
{
entries.remove(&self.id);
}
}
}
}
fn record_session_bytes(
exchange: &Exchange,
session: &TunnelSession,
bytes: u64,
client_to_backend: bool,
) {
if bytes == 0 {
return;
}
exchange.backend_bytes(bytes, client_to_backend);
if client_to_backend {
exchange
.metrics
.websocket_bytes_in
.fetch_add(bytes, Ordering::Relaxed);
exchange.request_bytes.fetch_add(bytes, Ordering::Relaxed);
} else {
exchange
.metrics
.websocket_bytes_out
.fetch_add(bytes, Ordering::Relaxed);
exchange.response_bytes.fetch_add(bytes, Ordering::Relaxed);
}
session.record(bytes, client_to_backend);
}
struct H2MeteredBody<B> {
inner: Pin<Box<B>>,
session: Arc<TunnelSession>,
exchange: Arc<Exchange>,
client_to_backend: bool,
cancelled: Pin<Box<dyn Future<Output = ()> + Send>>,
}
impl<B> H2MeteredBody<B>
where
B: HttpBody<Data = Bytes> + Send + 'static,
{
fn new(
inner: B,
session: Arc<TunnelSession>,
exchange: Arc<Exchange>,
client_to_backend: bool,
) -> Self {
let cancelled = session.cancellation().cancelled_owned();
Self {
inner: Box::pin(inner),
session,
exchange,
client_to_backend,
cancelled: Box::pin(cancelled),
}
}
}
impl<B> HttpBody for H2MeteredBody<B>
where
B: HttpBody<Data = Bytes> + Send + 'static,
{
type Data = Bytes;
type Error = B::Error;
fn poll_frame(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
if self.cancelled.as_mut().poll(cx).is_ready() {
return Poll::Ready(None);
}
match self.inner.as_mut().poll_frame(cx) {
Poll::Ready(Some(Ok(frame))) => {
if let Some(data) = frame.data_ref() {
record_session_bytes(
&self.exchange,
&self.session,
data.len() as u64,
self.client_to_backend,
);
}
Poll::Ready(Some(Ok(frame)))
}
other => other,
}
}
fn is_end_stream(&self) -> bool {
self.inner.as_ref().get_ref().is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.inner.as_ref().get_ref().size_hint()
}
}
pub fn wrap_h2_request_body(
body: Body,
session: Arc<TunnelSession>,
exchange: Arc<Exchange>,
) -> Body {
Body::new(H2MeteredBody::new(body, session, exchange, true))
}
pub fn wrap_h2_response_body(
body: hyper::body::Incoming,
session: Arc<TunnelSession>,
exchange: Arc<Exchange>,
) -> Body {
Body::new(H2MeteredBody::new(body, session, exchange, false))
}
pub fn spawn_h2_idle_watch(state: Arc<App>, session: &Arc<TunnelSession>) {
let weak = Arc::downgrade(session);
let idle = state.websocket_idle_timeout;
let tracker = state.tunnels.clone();
tracker.spawn(async move {
loop {
tokio::time::sleep(idle).await;
let Some(session) = weak.upgrade() else {
return;
};
if session.inactive_for() >= idle {
session._active.exchange.failed.store(true, Ordering::Relaxed);
session
._active
.exchange
.websocket_error("tunnel_error", "idle_timeout");
session
._active
.exchange
.metrics
.websocket_errors
.fetch_add(1, Ordering::Relaxed);
session.cancel.cancel();
return;
}
}
});
}
+382
View File
@@ -0,0 +1,382 @@
/// Bridge an HTTP/2 Extended CONNECT stream to an already upgraded HTTP/1.1
/// backend WebSocket. DATA bytes remain byte-transparent in both directions.
pub fn bridge_h2_to_h1(
state: Arc<App>,
mut downstream: Body,
upstream: Upgraded,
exchange: Arc<Exchange>,
upstream_label: String,
lease: Option<Arc<ConnectionLease>>,
stop: Option<ServerShutdown>,
) -> Body {
let stop = stop.map(|value| value.0).unwrap_or_else(|| state.shutdown.clone());
let session = state.tunnel_registry.register(
exchange.clone(),
upstream_label,
"h2-rfc8441-h1-upstream",
stop,
lease,
);
spawn_h2_idle_watch(state.clone(), &session);
let upstream = TokioIo::new(upstream);
let (upstream_read, mut upstream_write) = tokio::io::split(upstream);
let request_session = session.clone();
let request_exchange = exchange.clone();
let tracker = state.tunnels.clone();
tracker.spawn(async move {
let cancel = request_session.cancellation();
loop {
let frame = tokio::select! {
_ = cancel.cancelled() => return,
frame = downstream.frame() => frame,
};
match frame {
Some(Ok(frame)) => {
if let Ok(data) = frame.into_data() {
if !data.is_empty() {
if upstream_write.write_all(&data).await.is_err() {
return;
}
record_session_bytes(
&request_exchange,
&request_session,
data.len() as u64,
true,
);
}
}
}
Some(Err(_)) => return,
None => {
let _ = upstream_write.shutdown().await;
return;
}
}
}
});
let stream = futures_util::stream::unfold(
(upstream_read, session, exchange),
|(mut upstream_read, session, exchange)| async move {
let mut buffer = vec![0u8; 16 * 1024];
let cancel = session.cancellation();
let read = tokio::select! {
_ = cancel.cancelled() => return None,
read = upstream_read.read(&mut buffer) => read,
};
match read {
Ok(0) => None,
Ok(read) => {
buffer.truncate(read);
record_session_bytes(&exchange, &session, read as u64, false);
Some((Ok::<Bytes, io::Error>(Bytes::from(buffer)), (upstream_read, session, exchange)))
}
Err(error) => Some((Err(error), (upstream_read, session, exchange))),
}
},
);
Body::from_stream(stream)
}
fn benign_tunnel_close(error: &io::Error) -> bool {
matches!(
error.kind(),
io::ErrorKind::UnexpectedEof
| io::ErrorKind::ConnectionReset
| io::ErrorKind::ConnectionAborted
| io::ErrorKind::BrokenPipe
| io::ErrorKind::NotConnected
)
}
struct H3TunnelState {
upstream: Upgraded,
exchange: Arc<Exchange>,
upstream_label: String,
lease: Option<Arc<ConnectionLease>>,
stop: Option<ServerShutdown>,
}
/// Response extension consumed by the HTTP/3 transport after it has sent the
/// successful 200 response. The upstream connection is already upgraded before
/// this extension is created, so downstream success never races backend takeover.
#[derive(Clone)]
pub struct H3Tunnel(Arc<Mutex<Option<H3TunnelState>>>);
impl H3Tunnel {
pub fn new(
upstream: Upgraded,
exchange: Arc<Exchange>,
upstream_label: String,
lease: Option<Arc<ConnectionLease>>,
stop: Option<ServerShutdown>,
) -> Self {
Self(Arc::new(Mutex::new(Some(H3TunnelState {
upstream,
exchange,
upstream_label,
lease,
stop,
}))))
}
fn take(&self) -> Option<H3TunnelState> {
self.0.lock().unwrap_or_else(|e| e.into_inner()).take()
}
}
fn h3_io_error(error: impl std::fmt::Display) -> io::Error {
io::Error::other(error.to_string())
}
async fn h3_client_to_backend<W>(
receive: &mut h3::server::RequestStream<h3::quinn::RecvStream, Bytes>,
upstream: &mut W,
exchange: &Exchange,
activity: &AtomicU64,
started: Instant,
session: &TunnelSession,
) -> io::Result<()>
where
W: AsyncWrite + Unpin,
{
while let Some(mut data) = receive.recv_data().await.map_err(h3_io_error)? {
let bytes = data.copy_to_bytes(data.remaining());
if !bytes.is_empty() {
upstream.write_all(&bytes).await?;
record_bytes(exchange, bytes.len() as u64, true, activity, started);
session.record(bytes.len() as u64, true);
}
}
if receive
.recv_trailers()
.await
.map_err(h3_io_error)?
.is_some()
{
return Err(io::Error::other(
"WebSocket Extended CONNECT must not end with HTTP trailers",
));
}
upstream.shutdown().await
}
async fn backend_to_h3<R>(
upstream: &mut R,
send: &mut h3::server::RequestStream<h3::quinn::SendStream<Bytes>, Bytes>,
exchange: &Exchange,
activity: &AtomicU64,
started: Instant,
session: &TunnelSession,
) -> io::Result<()>
where
R: AsyncRead + Unpin,
{
let mut buffer = [0u8; 16 * 1024];
loop {
let read = upstream.read(&mut buffer).await?;
if read == 0 {
send.finish().await.map_err(h3_io_error)?;
return Ok(());
}
send.send_data(Bytes::copy_from_slice(&buffer[..read]))
.await
.map_err(h3_io_error)?;
record_bytes(exchange, read as u64, false, activity, started);
session.record(read as u64, false);
}
}
/// Bridge one RFC 9220 HTTP/3 stream to the already validated HTTP/1.1
/// WebSocket upgrade at the selected backend. DATA payload bytes are copied
/// unchanged; FIN remains a half-close and abnormal termination maps to
/// H3_REQUEST_CANCELLED.
pub async fn run_h3(
state: Arc<App>,
send: &mut h3::server::RequestStream<h3::quinn::SendStream<Bytes>, Bytes>,
receive: &mut h3::server::RequestStream<h3::quinn::RecvStream, Bytes>,
tunnel: H3Tunnel,
) -> io::Result<()> {
let Some(H3TunnelState {
upstream,
exchange,
upstream_label,
lease,
stop,
}) = tunnel.take()
else {
send.stop_stream(Code::H3_REQUEST_CANCELLED);
receive.stop_sending(Code::H3_REQUEST_CANCELLED);
return Err(io::Error::other("HTTP/3 WebSocket tunnel already consumed"));
};
let stop = stop.map(|v| v.0).unwrap_or_else(|| state.shutdown.clone());
let started = Instant::now();
let activity = Arc::new(AtomicU64::new(0));
let session = state.tunnel_registry.register(
exchange.clone(),
upstream_label,
"h3-rfc9220",
stop,
lease,
);
let stop = session.cancellation();
let upstream = TokioIo::new(upstream);
let (mut upstream_read, mut upstream_write) = tokio::io::split(upstream);
let result = {
let copying = async {
tokio::try_join!(
h3_client_to_backend(receive, &mut upstream_write, &exchange, &activity, started, &session),
backend_to_h3(&mut upstream_read, send, &exchange, &activity, started, &session)
)?;
Ok::<(), io::Error>(())
};
tokio::pin!(copying);
let mut remaining = state.websocket_idle_timeout;
loop {
tokio::select! {
result = &mut copying => break result,
_ = stop.cancelled() => break Err(io::Error::new(io::ErrorKind::ConnectionAborted, "WebSocket tunnel shutdown")),
_ = tokio::time::sleep(remaining) => {
let inactive = started.elapsed().saturating_sub(Duration::from_millis(activity.load(Ordering::Relaxed)));
if inactive >= state.websocket_idle_timeout {
break Err(io::Error::new(io::ErrorKind::TimedOut, "WebSocket idle timeout"));
}
remaining = state.websocket_idle_timeout - inactive;
}
}
}
};
if let Err(error) = &result {
send.stop_stream(Code::H3_REQUEST_CANCELLED);
receive.stop_sending(Code::H3_REQUEST_CANCELLED);
exchange.h3_reset("H3_REQUEST_CANCELLED");
if !benign_tunnel_close(error) {
exchange.failed.store(true, Ordering::Relaxed);
exchange.websocket_error("tunnel_error", "h3_bridge_error");
state
.metrics
.websocket_errors
.fetch_add(1, Ordering::Relaxed);
} else {
exchange.websocket_phase("tunnel_closed");
}
} else {
exchange.websocket_phase("tunnel_closed");
}
result
}
#[allow(clippy::too_many_arguments)]
pub fn spawn(
state: Arc<App>,
downstream: OnUpgrade,
upstream: OnUpgrade,
exchange: Arc<Exchange>,
upstream_label: String,
lease: Option<Arc<ConnectionLease>>,
stop: Option<ServerShutdown>,
) {
let tracker = state.tunnels.clone();
tracker.spawn(async move {
exchange.websocket_phase("upgrade_wait");
let stop = stop.map(|v| v.0).unwrap_or_else(|| state.shutdown.clone());
let upgraded = tokio::select! {
_ = stop.cancelled() => return,
result = tokio::time::timeout(Duration::from_secs(10), async {
tokio::try_join!(downstream, upstream)
}) => result,
};
let (downstream, upstream) = match upgraded {
Ok(Ok(streams)) => streams,
_ => {
exchange.failed.store(true, Ordering::Relaxed);
exchange.websocket_error("upgrade_failed", "on_upgrade");
state
.metrics
.websocket_errors
.fetch_add(1, Ordering::Relaxed);
return;
}
};
let started = Instant::now();
let activity = Arc::new(AtomicU64::new(0));
let session = state.tunnel_registry.register(
exchange.clone(),
upstream_label,
"h1-upgrade",
stop,
lease,
);
let stop = session.cancellation();
let mut downstream = MeteredIo {
inner: TokioIo::new(downstream),
exchange: exchange.clone(),
session: session.clone(),
client_to_backend: true,
started,
activity: activity.clone(),
};
let mut upstream = MeteredIo {
inner: TokioIo::new(upstream),
exchange: exchange.clone(),
session: session.clone(),
client_to_backend: false,
started,
activity: activity.clone(),
};
if let Err(error) = tunnel(
&mut downstream,
&mut upstream,
stop,
activity,
started,
state.websocket_idle_timeout,
)
.await
{
// Browsers and upstreams commonly close upgraded TCP streams without
// a WebSocket close frame. Reset/EOF/broken-pipe are normal tunnel
// termination signals, not malformed HTTP requests or proxy errors.
if !benign_tunnel_close(&error) {
exchange.failed.store(true, Ordering::Relaxed);
exchange.websocket_error("tunnel_error", "bridge_error");
state
.metrics
.websocket_errors
.fetch_add(1, Ordering::Relaxed);
} else {
exchange.websocket_phase("tunnel_closed");
}
} else {
exchange.websocket_phase("tunnel_closed");
}
});
}
async fn tunnel<A, B>(
a: &mut A,
b: &mut B,
stop: CancellationToken,
activity: Arc<AtomicU64>,
started: Instant,
idle: Duration,
) -> io::Result<()>
where
A: AsyncRead + AsyncWrite + Unpin,
B: AsyncRead + AsyncWrite + Unpin,
{
// Tokio uses fixed-size buffers in both directions. Writes apply backpressure.
let copying = tokio::io::copy_bidirectional(a, b);
tokio::pin!(copying);
let mut remaining = idle;
loop {
tokio::select! {
result = &mut copying => return result.map(|_| ()),
_ = stop.cancelled() => return Ok(()),
_ = tokio::time::sleep(remaining) => {
let inactive = started.elapsed().saturating_sub(Duration::from_millis(activity.load(Ordering::Relaxed)));
if inactive >= idle { return Err(io::Error::new(io::ErrorKind::TimedOut, "WebSocket idle timeout")); }
remaining = idle - inactive;
},
}
}
}