v0.25.1
This commit is contained in:
Mateusz Gruszczyński committed
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,
|
||||
}
|
||||
Reference in new issue
Block a user