RustProxy/src/proxy/websocket/handshake.rs

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::SysRng, TryRng};
use serde::Serialize;
use sha1::{Digest, Sha1};
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];
    SysRng.try_fill_bytes(&mut nonce).expect("secure operating system random source unavailable");
    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];
    SysRng.try_fill_bytes(&mut nonce).expect("secure operating system random source unavailable");
    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"),
            "101 response must not use Transfer-Encoding"
        );
        ensure!(
            headers.get_all("content-length").iter().count() <= 1
                && headers
                    .get("content-length")
                    .is_none_or(|value| value.to_str().is_ok_and(|value| value.trim() == "0")),
            "101 response Content-Length must be 0 when present"
        );
        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,
}