v0.25.1
This commit is contained in:
1 parent
87e571c170
commit
fa7abc9e53
308 files changed
+49746
-50237
No files matched your search
@@ -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,
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -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;
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user