From 4e618139f4201791a2884ce7153e060b01fb34e8 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Fri, 28 Aug 2026 16:59:01 -0700 Subject: [PATCH 01/11] Add standalone hostname-filtered CONNECT proxy Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- .github/workflows/ci.yml | 6 + Cargo.lock | 15 + litebox_egress_proxy/Cargo.toml | 41 ++- litebox_egress_proxy/src/config.rs | 153 +++++++++ litebox_egress_proxy/src/headers.rs | 99 ++++++ litebox_egress_proxy/src/lib.rs | 117 ++++++- litebox_egress_proxy/src/limits.rs | 37 +++ litebox_egress_proxy/src/listener.rs | 229 ++++++++++++++ litebox_egress_proxy/src/main.rs | 49 +++ litebox_egress_proxy/src/proxy.rs | 341 ++++++++++++++++++++ litebox_egress_proxy/src/request_head.rs | 108 +++++++ litebox_egress_proxy/src/stream.rs | 264 ++++++++++++++++ litebox_egress_proxy/src/upstream.rs | 58 ++++ litebox_egress_proxy/tests/loopback.rs | 380 +++++++++++++++++++++++ 14 files changed, 1888 insertions(+), 9 deletions(-) create mode 100644 litebox_egress_proxy/src/config.rs create mode 100644 litebox_egress_proxy/src/headers.rs create mode 100644 litebox_egress_proxy/src/limits.rs create mode 100644 litebox_egress_proxy/src/listener.rs create mode 100644 litebox_egress_proxy/src/main.rs create mode 100644 litebox_egress_proxy/src/proxy.rs create mode 100644 litebox_egress_proxy/src/request_head.rs create mode 100644 litebox_egress_proxy/src/stream.rs create mode 100644 litebox_egress_proxy/src/upstream.rs create mode 100644 litebox_egress_proxy/tests/loopback.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 357bafe4c..d2aaf1c53 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -91,6 +91,7 @@ jobs: AARCH64_CRATES: >- -p litebox -p litebox_common_linux + -p litebox_egress_proxy -p litebox_syscall_rewriter -p litebox_packager -p litebox_platform_linux_userland @@ -199,6 +200,7 @@ jobs: WINDOWS_CRATES: >- -p litebox_common_linux -p litebox_common_windows + -p litebox_egress_proxy -p litebox_syscall_rewriter -p litebox_packager -p litebox_broker_protocol @@ -290,6 +292,9 @@ jobs: # - `litebox_broker_userland` is allowed to have `std` access, # since it is the hosted userland broker executable. # + # - `litebox_egress_proxy` is allowed to have `std` access, since it + # is the hosted userland CONNECT proxy executable. + # # - `litebox_platform_lvbs` has a custom target (`no_std`), so it does # not work with the current no_std checker. # @@ -351,6 +356,7 @@ jobs: -not -path './litebox_broker_transport_linux_userland/Cargo.toml' \ -not -path './litebox_broker_transport_windows_userland/Cargo.toml' \ -not -path './litebox_broker_userland/Cargo.toml' \ + -not -path './litebox_egress_proxy/Cargo.toml' \ -not -path './litebox_platform_linux_userland/Cargo.toml' \ -not -path './litebox_platform_windows_userland/Cargo.toml' \ -not -path './litebox_runner_linux_on_windows_userland/Cargo.toml' \ diff --git a/Cargo.lock b/Cargo.lock index 99f8d0fe4..1f23a61c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -990,6 +990,12 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "hyper" version = "1.8.1" @@ -1003,6 +1009,7 @@ dependencies = [ "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "pin-utils", @@ -1534,8 +1541,16 @@ dependencies = [ name = "litebox_egress_proxy" version = "0.1.0" dependencies = [ + "bytes", + "clap", "hashbrown", + "http-body-util", + "httparse", + "hyper", + "hyper-util", + "libc", "thiserror", + "tokio", ] [[package]] diff --git a/litebox_egress_proxy/Cargo.toml b/litebox_egress_proxy/Cargo.toml index b1630eb63..2ded1634f 100644 --- a/litebox_egress_proxy/Cargo.toml +++ b/litebox_egress_proxy/Cargo.toml @@ -4,8 +4,47 @@ version = "0.1.0" edition = "2024" [dependencies] +bytes = { version = "1.10", default-features = false, features = ["std"] } +clap = { version = "4.5", default-features = false, features = [ + "derive", + "error-context", + "help", + "std", + "usage", +] } hashbrown = "0.15.2" -thiserror = { version = "2.0", default-features = false } +http-body-util = { version = "0.1.3", default-features = false } +httparse = { version = "1.10.1", default-features = false } +hyper = { version = "1.8", default-features = false, features = [ + "http1", + "server", +] } +hyper-util = { version = "0.1.20", default-features = false, features = [ + "tokio", +] } +thiserror = { version = "2.0", default-features = false, features = ["std"] } +tokio = { version = "1.50", default-features = false, features = [ + "io-util", + "net", + "rt", + "sync", + "time", +] } + +[dev-dependencies] +tokio = { version = "1.50", default-features = false, features = [ + "io-util", + "macros", + "net", + "rt", + "rt-multi-thread", + "sync", + "test-util", + "time", +] } + +[target.'cfg(target_os = "linux")'.dependencies] +libc = { version = "0.2", default-features = false } [lints] workspace = true diff --git a/litebox_egress_proxy/src/config.rs b/litebox_egress_proxy/src/config.rs new file mode 100644 index 000000000..aa70f8fdf --- /dev/null +++ b/litebox_egress_proxy/src/config.rs @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Executable configuration. + +use std::net::{Ipv4Addr, SocketAddrV4}; + +use clap::{ArgGroup, Parser}; +use thiserror::Error; + +use crate::listener::ListenerSource; +use crate::policy::{HostPolicy, PolicyError}; + +/// The standalone egress proxy for LiteBox sandboxes. +#[derive(Debug, Parser)] +#[command( + name = "litebox_egress_proxy", + about = "Hostname-filtering CONNECT egress proxy", + group(ArgGroup::new("listener").required(true).args(["listen", "listener_fd"])) +)] +pub struct Cli { + /// Loopback address to bind, for example `127.0.0.1:0`. + #[arg(long, value_name = "IPV4:PORT")] + listen: Option, + + /// Inherited, already-bound loopback listener descriptor. + #[arg(long, value_name = "FD", conflicts_with = "listen")] + listener_fd: Option, + + /// Allowed hostname and destination ports, repeatable. + #[arg(long = "allow-host", value_name = "HOST:PORT[-PORT]")] + allow_host: Vec, +} + +/// Reason the arguments were rejected. +#[derive(Debug, Error)] +pub enum ConfigError { + /// `--listen` was not a socket address. + #[error("--listen must be an IPv4 address and port, for example 127.0.0.1:0")] + ListenAddress, + /// `--listen` was not canonical IPv4 loopback. + #[error("--listen must use the canonical loopback address 127.0.0.1")] + ListenNotLoopback, + /// An `--allow-host` rule was invalid. + #[error("invalid --allow-host rule: {0}")] + Policy(#[from] PolicyError), +} + +/// The validated configuration of one proxy process. +#[derive(Clone, Debug)] +pub struct ProxyConfig { + /// Where the listener comes from. + pub listener: ListenerSource, + /// The immutable hostname policy. + pub policy: HostPolicy, +} + +impl Cli { + /// Converts parsed arguments into a validated configuration. + pub fn into_config(self) -> Result { + let listener = match (self.listen, self.listener_fd) { + (Some(address), _) => ListenerSource::Bind(parse_listen_address(&address)?), + (None, Some(descriptor)) => ListenerSource::Inherited(descriptor), + (None, None) => return Err(ConfigError::ListenAddress), + }; + let policy = HostPolicy::from_rules(&self.allow_host)?; + + Ok(ProxyConfig { listener, policy }) + } +} + +fn parse_listen_address(raw: &str) -> Result { + let address: SocketAddrV4 = raw.parse().map_err(|_| ConfigError::ListenAddress)?; + if *address.ip() != Ipv4Addr::LOCALHOST { + return Err(ConfigError::ListenNotLoopback); + } + Ok(address) +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::policy::Hostname; + + fn parse(arguments: &[&str]) -> Result { + let mut all = vec!["litebox_egress_proxy"]; + all.extend_from_slice(arguments); + Cli::try_parse_from(all).unwrap().into_config() + } + + #[test] + fn parses_a_standalone_configuration() { + let config = parse(&[ + "--listen", + "127.0.0.1:0", + "--allow-host", + "Example.COM:443", + "--allow-host", + "example.com:8000-8100", + ]) + .unwrap(); + + assert_eq!( + config.listener, + ListenerSource::Bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)) + ); + + let host = Hostname::parse("example.com").unwrap(); + assert!(config.policy.allows(&host, 443)); + assert!(config.policy.allows(&host, 8100)); + assert!(!config.policy.allows(&host, 80)); + } + + #[test] + fn parses_an_inherited_listener_configuration() { + let config = parse(&["--listener-fd", "7"]).unwrap(); + let host = Hostname::parse("example.com").unwrap(); + assert_eq!(config.listener, ListenerSource::Inherited(7)); + assert!(!config.policy.allows(&host, 443)); + } + + #[test] + fn listener_modes_are_mutually_exclusive_and_required() { + assert!( + Cli::try_parse_from([ + "litebox_egress_proxy", + "--listen", + "127.0.0.1:0", + "--listener-fd", + "3", + ]) + .is_err() + ); + assert!(Cli::try_parse_from(["litebox_egress_proxy"]).is_err()); + } + + #[test] + fn rejects_non_loopback_listen_addresses() { + assert!(matches!( + parse(&["--listen", "0.0.0.0:8080"]), + Err(ConfigError::ListenNotLoopback) + )); + assert!(matches!( + parse(&["--listen", "127.0.0.2:8080"]), + Err(ConfigError::ListenNotLoopback) + )); + assert!(matches!( + parse(&["--listen", "localhost:8080"]), + Err(ConfigError::ListenAddress) + )); + } +} diff --git a/litebox_egress_proxy/src/headers.rs b/litebox_egress_proxy/src/headers.rs new file mode 100644 index 000000000..f7f4db4f5 --- /dev/null +++ b/litebox_egress_proxy/src/headers.rs @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! CONNECT request framing validation. + +use hyper::header; +use hyper::header::HeaderMap; +use thiserror::Error; + +/// Reason a CONNECT request was rejected as framing-ambiguous. +#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)] +pub enum FramingError { + /// More than one `Host` header was present. + #[error("message carries more than one Host header")] + DuplicateHost, + /// A CONNECT request carried body framing. + #[error("CONNECT request carries a body")] + BodyOnConnect, +} + +/// Validates request framing before `hyper` normalizes the raw header list. +pub(crate) fn validate_raw_request_framing( + method: &str, + _version: u8, + headers: &[httparse::Header<'_>], +) -> Result<(), FramingError> { + if raw_header_count(headers, b"host") > 1 { + return Err(FramingError::DuplicateHost); + } + if method == "CONNECT" + && (raw_header_count(headers, b"transfer-encoding") != 0 + || raw_header_count(headers, b"content-length") != 0) + { + return Err(FramingError::BodyOnConnect); + } + Ok(()) +} + +/// Validates that a parsed CONNECT request is unambiguously bodyless. +pub fn validate_connect_framing(headers: &HeaderMap) -> Result<(), FramingError> { + if headers.contains_key(header::TRANSFER_ENCODING) + || headers.contains_key(header::CONTENT_LENGTH) + { + return Err(FramingError::BodyOnConnect); + } + if headers.get_all(header::HOST).iter().count() > 1 { + return Err(FramingError::DuplicateHost); + } + Ok(()) +} + +fn raw_header_count(headers: &[httparse::Header<'_>], name: &[u8]) -> usize { + headers + .iter() + .filter(|header| header.name.as_bytes().eq_ignore_ascii_case(name)) + .count() +} + +#[cfg(test)] +mod tests { + use super::*; + + use hyper::header::{HeaderName, HeaderValue}; + + fn headers(pairs: &[(&str, &str)]) -> HeaderMap { + let mut map = HeaderMap::new(); + for (name, value) in pairs { + map.append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_str(value).unwrap(), + ); + } + map + } + + #[test] + fn connect_requests_must_be_bodyless() { + assert!(validate_connect_framing(&headers(&[("host", "example.com:443")])).is_ok()); + assert_eq!( + validate_connect_framing(&headers(&[("content-length", "0")])), + Err(FramingError::BodyOnConnect) + ); + assert_eq!( + validate_connect_framing(&headers(&[("transfer-encoding", "chunked")])), + Err(FramingError::BodyOnConnect) + ); + } + + #[test] + fn duplicate_host_is_rejected() { + assert_eq!( + validate_connect_framing(&headers(&[ + ("host", "example.com:443"), + ("host", "example.com:443"), + ])), + Err(FramingError::DuplicateHost) + ); + } +} diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index 77af2d3ed..b1740aad2 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -1,15 +1,116 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT license. -//! Exact hostname and destination-port policy primitives for the LiteBox -//! egress proxy. +//! A standalone, hostname-filtering CONNECT egress proxy for LiteBox. //! -//! Policy rules and request authorities share one canonical hostname type, so -//! authorization is an exact match after normalization. - -#![no_std] - -extern crate alloc; +//! Authorized hostnames are resolved on demand through the trusted host +//! resolver, then connected using the returned numeric addresses. CONNECT +//! tunnels relay bytes without inspecting or terminating TLS. pub mod authority; +pub mod config; +pub mod headers; +pub mod limits; +pub mod listener; pub mod policy; +pub mod proxy; +pub mod stream; +pub mod upstream; + +mod request_head; + +use std::io::{self, Write}; +use std::net::{SocketAddr, SocketAddrV4}; +use std::sync::Arc; + +use thiserror::Error; +use tokio::net::TcpListener; + +use crate::config::ProxyConfig; +use crate::listener::ListenerError; +use crate::proxy::ProxyState; +use crate::upstream::TcpUpstreamConnector; + +/// Reason the proxy could not start. +#[derive(Debug, Error)] +pub enum StartupError { + /// The listener could not be acquired or validated. + #[error(transparent)] + Listener(#[from] ListenerError), + /// An I/O operation failed during startup or while serving. + #[error(transparent)] + Io(#[from] io::Error), +} + +/// A proxy that has completed startup and is ready to serve. +pub struct StartedProxy { + listener: TcpListener, + state: Arc, + local_address: SocketAddrV4, +} + +impl StartedProxy { + /// Returns the loopback address the proxy listens on. + pub fn local_address(&self) -> SocketAddrV4 { + self.local_address + } + + /// Serves client connections until the listener fails. + pub async fn serve(self) -> io::Result<()> { + proxy::serve(self.listener, self.state).await + } +} + +/// Acquires the listener and prepares the shared state. +pub fn start(config: &ProxyConfig) -> Result { + let listener = listener::acquire(config.listener)?; + let listener = TcpListener::from_std(listener)?; + let SocketAddr::V4(local_address) = listener.local_addr()? else { + return Err(StartupError::Listener(ListenerError::NotLoopback( + listener.local_addr()?, + ))); + }; + + let state = Arc::new(ProxyState::new( + config.policy.clone(), + Arc::new(TcpUpstreamConnector), + )); + + Ok(StartedProxy { + listener, + state, + local_address, + }) +} + +/// Writes the single readiness line. +pub fn write_readiness(writer: &mut impl Write, address: SocketAddrV4) -> io::Result<()> { + writeln!(writer, "READY {address}")?; + writer.flush() +} + +/// Runs the proxy: startup, readiness announcement, then serving. +pub async fn run(config: &ProxyConfig) -> Result<(), StartupError> { + let started = start(config)?; + + let mut stdout = io::stdout().lock(); + write_readiness(&mut stdout, started.local_address())?; + drop(stdout); + + started.serve().await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::net::Ipv4Addr; + + #[test] + fn readiness_line_is_exactly_one_line() { + let mut output = Vec::new(); + write_readiness(&mut output, SocketAddrV4::new(Ipv4Addr::LOCALHOST, 34567)).unwrap(); + assert_eq!(output, b"READY 127.0.0.1:34567\n"); + } +} diff --git a/litebox_egress_proxy/src/limits.rs b/litebox_egress_proxy/src/limits.rs new file mode 100644 index 000000000..028a1d7d8 --- /dev/null +++ b/litebox_egress_proxy/src/limits.rs @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Fixed resource limits for the egress proxy. +//! +//! None of these limits is caller-configurable: the proxy is a trusted +//! component whose behaviour must be identical for every sandbox. + +use core::time::Duration; + +/// Maximum number of client connections served concurrently. +/// +/// Additional connections stay in the listener backlog until a slot frees up. +pub const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; + +/// Maximum number of bytes buffered for a client request head. +pub const MAX_REQUEST_HEADER_BYTES: usize = 16 * 1024; + +/// Maximum number of individual header fields parsed per message. +pub const MAX_HEADER_FIELDS: usize = 100; + +/// Total timeout for hostname resolution and connection attempts. +pub const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +/// Idle timeout applied to HTTP bodies and CONNECT tunnels. +/// +/// A stream that makes no read or write progress for this long is torn down. +pub const IDLE_TIMEOUT: Duration = Duration::from_secs(60); + +/// Maximum time a client may take to send a complete request head. +pub const REQUEST_HEADER_READ_TIMEOUT: Duration = Duration::from_secs(30); + +/// Maximum time spent draining client input after a non-upgraded response. +pub const CLIENT_CLOSE_DRAIN_TIMEOUT: Duration = Duration::from_secs(1); + +/// Maximum client input discarded while closing a non-upgraded connection. +pub const MAX_CLIENT_CLOSE_DRAIN_BYTES: usize = 64 * 1024; diff --git a/litebox_egress_proxy/src/listener.rs b/litebox_egress_proxy/src/listener.rs new file mode 100644 index 000000000..ab7994592 --- /dev/null +++ b/litebox_egress_proxy/src/listener.rs @@ -0,0 +1,229 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Listener acquisition for the standalone and broker modes. +//! +//! Both modes end with the same invariant: the proxy only ever serves a bound, +//! listening IPv4 loopback TCP socket. The standalone mode binds it itself; the +//! broker mode adopts a listener that a launcher bound and inherited to this +//! process. + +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener}; + +use thiserror::Error; + +/// Where the proxy's listener comes from. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ListenerSource { + /// Bind a fresh loopback listener. Port zero requests an ephemeral port. + Bind(SocketAddrV4), + /// Adopt an inherited, already-bound listener by descriptor number. + Inherited(i32), +} + +/// Reason a listener could not be acquired. +#[derive(Debug, Error)] +pub enum ListenerError { + /// Binding the requested address failed. + #[error("failed to bind {address}: {source}")] + Bind { + /// The address that could not be bound. + address: SocketAddrV4, + /// The underlying failure. + source: std::io::Error, + }, + /// Reading the listener's local address failed. + #[error("failed to read the listener address: {0}")] + LocalAddress(#[source] std::io::Error), + /// The listener was not bound to canonical IPv4 loopback. + #[error("listener is bound to {0}, which is not 127.0.0.1")] + NotLoopback(SocketAddr), + /// The listener was bound to port zero, which an inherited listener never + /// is once it has been bound. + #[error("inherited listener is not bound to a concrete port")] + UnboundPort, + /// The descriptor number was negative. + #[error("inherited listener descriptor is not a valid descriptor number")] + InvalidDescriptor, + /// The descriptor did not refer to an open file. + #[error("inherited listener descriptor is not open")] + DescriptorNotOpen, + /// Inspecting the socket failed. + #[error("failed to inspect the inherited listener: {0}")] + Inspect(#[source] std::io::Error), + /// The descriptor was not an IPv4 stream socket in the listening state. + #[error("inherited descriptor is not a listening IPv4 TCP socket")] + NotAnIpv4Listener, + /// The platform has no inherited-listener contract. + #[error("--listener-fd is only supported on Linux")] + InheritanceUnsupported, + /// Configuring the listener for asynchronous use failed. + #[error("failed to configure the listener: {0}")] + Configure(#[source] std::io::Error), +} + +/// Acquires the listener described by `source`. +/// +/// The returned listener is non-blocking and validated to be bound to +/// canonical IPv4 loopback. +pub fn acquire(source: ListenerSource) -> Result { + let listener = match source { + ListenerSource::Bind(address) => { + TcpListener::bind(address).map_err(|source| ListenerError::Bind { address, source })? + } + ListenerSource::Inherited(descriptor) => adopt_inherited(descriptor)?, + }; + + let local = listener.local_addr().map_err(ListenerError::LocalAddress)?; + let SocketAddr::V4(local) = local else { + return Err(ListenerError::NotLoopback(local)); + }; + if *local.ip() != Ipv4Addr::LOCALHOST { + return Err(ListenerError::NotLoopback(SocketAddr::V4(local))); + } + if local.port() == 0 { + return Err(ListenerError::UnboundPort); + } + + listener + .set_nonblocking(true) + .map_err(ListenerError::Configure)?; + Ok(listener) +} + +/// Adopts an inherited descriptor after validating that it really is a bound, +/// listening IPv4 TCP socket. +#[cfg(target_os = "linux")] +fn adopt_inherited(descriptor: i32) -> Result { + use std::os::fd::FromRawFd; + + if descriptor < 0 { + return Err(ListenerError::InvalidDescriptor); + } + + // SAFETY: `fcntl(F_GETFD)` only reads the descriptor flags of `descriptor`. + // It neither takes ownership nor mutates process state, and it reports an + // invalid descriptor as `-1` instead of causing undefined behaviour. + let flags = unsafe { libc::fcntl(descriptor, libc::F_GETFD) }; + if flags < 0 { + return Err(ListenerError::DescriptorNotOpen); + } + + if socket_option(descriptor, libc::SO_DOMAIN)? != libc::AF_INET + || socket_option(descriptor, libc::SO_TYPE)? != libc::SOCK_STREAM + || socket_option(descriptor, libc::SO_ACCEPTCONN)? != 1 + { + return Err(ListenerError::NotAnIpv4Listener); + } + + // SAFETY: the checks above established that `descriptor` is an open, + // listening IPv4 stream socket. The launcher contract for `--listener-fd` + // transfers ownership of that descriptor to this process, and nothing else + // in this process holds or closes it, so wrapping it in a `TcpListener` + // gives a single unique owner. + Ok(unsafe { TcpListener::from_raw_fd(descriptor) }) +} + +/// Reads a `SOL_SOCKET` integer option. +#[cfg(target_os = "linux")] +fn socket_option(descriptor: i32, option: libc::c_int) -> Result { + let mut value: libc::c_int = 0; + let mut length = libc::socklen_t::try_from(size_of::()) + .expect("the size of a C int fits in socklen_t"); + + // SAFETY: `value` and `length` are valid, correctly sized and aligned + // locals that outlive the call. `getsockopt` writes at most `length` bytes + // into `value` and updates `length` accordingly, and reports failure as + // `-1` rather than writing out of bounds. + let result = unsafe { + libc::getsockopt( + descriptor, + libc::SOL_SOCKET, + option, + std::ptr::from_mut(&mut value).cast::(), + &raw mut length, + ) + }; + if result != 0 { + return Err(ListenerError::Inspect(std::io::Error::last_os_error())); + } + Ok(value) +} + +/// Inherited listeners are a Linux-only contract in this milestone. +#[cfg(not(target_os = "linux"))] +fn adopt_inherited(_descriptor: i32) -> Result { + Err(ListenerError::InheritanceUnsupported) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn binds_an_ephemeral_loopback_port() { + let listener = acquire(ListenerSource::Bind(SocketAddrV4::new( + Ipv4Addr::LOCALHOST, + 0, + ))) + .unwrap(); + let SocketAddr::V4(address) = listener.local_addr().unwrap() else { + panic!("expected an IPv4 listener"); + }; + assert_eq!(*address.ip(), Ipv4Addr::LOCALHOST); + assert_ne!(address.port(), 0); + } + + #[test] + fn rejects_non_loopback_binds() { + let error = acquire(ListenerSource::Bind(SocketAddrV4::new( + Ipv4Addr::UNSPECIFIED, + 0, + ))) + .unwrap_err(); + assert!(matches!(error, ListenerError::NotLoopback(_))); + } + + #[test] + fn rejects_a_negative_descriptor() { + let error = acquire(ListenerSource::Inherited(-1)).unwrap_err(); + assert!(matches!( + error, + ListenerError::InvalidDescriptor | ListenerError::InheritanceUnsupported + )); + } + + #[cfg(target_os = "linux")] + #[test] + fn adopts_an_inherited_loopback_listener() { + use std::os::fd::IntoRawFd; + + let bound = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); + let expected = bound.local_addr().unwrap(); + let descriptor = bound.into_raw_fd(); + + let adopted = acquire(ListenerSource::Inherited(descriptor)).unwrap(); + assert_eq!(adopted.local_addr().unwrap(), expected); + } + + #[cfg(target_os = "linux")] + #[test] + fn rejects_a_connected_socket() { + use std::os::fd::IntoRawFd; + + let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); + let client = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap(); + let error = acquire(ListenerSource::Inherited(client.into_raw_fd())).unwrap_err(); + assert!(matches!(error, ListenerError::NotAnIpv4Listener)); + } + + #[cfg(target_os = "linux")] + #[test] + fn rejects_a_datagram_socket() { + use std::os::fd::IntoRawFd; + + let socket = std::net::UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); + let error = acquire(ListenerSource::Inherited(socket.into_raw_fd())).unwrap_err(); + assert!(matches!(error, ListenerError::NotAnIpv4Listener)); + } +} diff --git a/litebox_egress_proxy/src/main.rs b/litebox_egress_proxy/src/main.rs new file mode 100644 index 000000000..a79ae5abd --- /dev/null +++ b/litebox_egress_proxy/src/main.rs @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! The `litebox_egress_proxy` executable. +//! +//! Startup failures are reported on standard error and produce a nonzero exit +//! status; the readiness line on standard output is written only once the proxy +//! is fully configured and listening. + +use std::process::ExitCode; + +use clap::Parser; +use litebox_egress_proxy::config::Cli; +use litebox_egress_proxy::run; + +fn main() -> ExitCode { + let config = match Cli::parse().into_config() { + Ok(config) => config, + Err(error) => return fail(&error), + }; + + // A single-threaded runtime keeps the standalone TCB deterministic; all + // request and tunnel work is asynchronous and explicitly bounded. + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + { + Ok(runtime) => runtime, + Err(error) => return fail(&error), + }; + + match runtime.block_on(run(&config)) { + Ok(()) => ExitCode::SUCCESS, + Err(error) => fail(&error), + } +} + +/// Reports a fatal error on standard error, including its causes. +fn fail(error: &dyn core::error::Error) -> ExitCode { + eprint!("litebox_egress_proxy: {error}"); + let mut source = error.source(); + while let Some(cause) = source { + eprint!(": {cause}"); + source = cause.source(); + } + eprintln!(); + ExitCode::FAILURE +} diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs new file mode 100644 index 000000000..0ea258534 --- /dev/null +++ b/litebox_egress_proxy/src/proxy.rs @@ -0,0 +1,341 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Connection acceptance, authorization, and CONNECT tunnelling. +//! +//! Every raw-validated CONNECT request is authorized before DNS or upstream +//! activity. A successful request consumes its client connection by upgrading +//! it to one bounded bidirectional tunnel. + +use core::convert::Infallible; +use core::error::Error as StdError; +use std::io; +use std::sync::Arc; + +use bytes::Bytes; +use http_body_util::combinators::BoxBody; +use http_body_util::{BodyExt, Empty}; +use hyper::body::Incoming; +use hyper::header::{self, HeaderValue}; +use hyper::http::uri::Authority; +use hyper::server::conn::http1 as server_http1; +use hyper::service::service_fn; +use hyper::{Method, Request, Response, StatusCode, Uri}; +use hyper_util::rt::{TokioIo, TokioTimer}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{Mutex, OwnedSemaphorePermit, Semaphore}; +use tokio::time::timeout; + +use crate::authority::{RequestAuthority, parse_authority}; +use crate::headers::validate_connect_framing; +use crate::limits::{ + CLIENT_CLOSE_DRAIN_TIMEOUT, IDLE_TIMEOUT, MAX_CLIENT_CLOSE_DRAIN_BYTES, + MAX_CONCURRENT_CLIENT_CONNECTIONS, MAX_HEADER_FIELDS, MAX_REQUEST_HEADER_BYTES, + REQUEST_HEADER_READ_TIMEOUT, UPSTREAM_CONNECT_TIMEOUT, +}; +use crate::policy::HostPolicy; +use crate::request_head::read_validated_request_prefix; +use crate::stream::{LimitedStream, PrefixedStream, share_tcp_read}; +use crate::upstream::{BoxedUpstreamStream, UpstreamConnector}; + +/// Boxed error type used by response bodies. +type BoxError = Box; + +/// Response body type produced by the proxy. +type ProxyBody = BoxBody; + +/// Immutable state shared by every connection. +pub struct ProxyState { + policy: HostPolicy, + connector: Arc, +} + +impl ProxyState { + /// Builds shared state from a validated policy and upstream connector. + pub fn new(policy: HostPolicy, connector: Arc) -> Self { + Self { policy, connector } + } +} + +/// Serves client connections until `listener` fails. +pub async fn serve(listener: TcpListener, state: Arc) -> io::Result<()> { + let slots = Arc::new(Semaphore::new(MAX_CONCURRENT_CLIENT_CONNECTIONS)); + + loop { + let Ok(permit) = Arc::clone(&slots).acquire_owned().await else { + return Err(io::Error::other("connection slot semaphore closed")); + }; + + let stream = match listener.accept().await { + Ok((stream, _peer)) => stream, + Err(error) + if matches!( + error.kind(), + io::ErrorKind::ConnectionAborted | io::ErrorKind::Interrupted + ) => + { + continue; + } + Err(error) => return Err(error), + }; + + let state = Arc::clone(&state); + tokio::spawn(async move { + serve_connection(state, stream, permit).await; + }); + } +} + +/// Serves one CONNECT tunnel or rejection on a client connection. +async fn serve_connection(state: Arc, stream: TcpStream, permit: OwnedSemaphorePermit) { + if stream.set_nodelay(true).is_err() { + return; + } + + let (stream, mut drain_handle) = share_tcp_read(stream); + let mut stream = LimitedStream::new(stream, IDLE_TIMEOUT); + let prefix = match read_validated_request_prefix(&mut stream).await { + Ok(prefix) => prefix.into_bytes(), + Err(error) => { + if let Some(response) = error.response() { + if let Err(write_error) = stream.write_all(response).await { + diagnostic(format_args!( + "failed to write request rejection after {error}: {write_error}" + )); + } else { + let _ = stream.shutdown().await; + drain_client_input(&mut stream).await; + } + } + return; + } + }; + + let io = TokioIo::new(PrefixedStream::new(prefix, stream)); + let connection_slot = Arc::new(Mutex::new(Some(permit))); + let service_connection_slot = Arc::clone(&connection_slot); + let service = service_fn(move |request: Request| { + let state = Arc::clone(&state); + let connection_slot = Arc::clone(&service_connection_slot); + async move { Ok::<_, Infallible>(handle_request(state, connection_slot, request).await) } + }); + + let mut builder = server_http1::Builder::new(); + builder + .timer(TokioTimer::new()) + .header_read_timeout(Some(REQUEST_HEADER_READ_TIMEOUT)) + .max_buf_size(MAX_REQUEST_HEADER_BYTES) + .max_headers(MAX_HEADER_FIELDS) + .keep_alive(true) + .half_close(true); + + let result = builder.serve_connection(io, service).with_upgrades().await; + let upgraded = connection_slot.lock().await.is_none(); + if !upgraded { + drain_client_input(&mut drain_handle).await; + } + if let Err(error) = result { + diagnostic(format_args!("client connection ended: {error}")); + } +} + +/// Drains bounded client input so unread bytes cannot reset a rejection. +async fn drain_client_input(stream: &mut S) +where + S: tokio::io::AsyncRead + Unpin, +{ + let drain = async { + let mut remaining = MAX_CLIENT_CLOSE_DRAIN_BYTES; + let mut buffer = [0_u8; 1024]; + while remaining != 0 { + let capacity = remaining.min(buffer.len()); + let read = stream.read(&mut buffer[..capacity]).await?; + if read == 0 { + break; + } + remaining -= read; + } + Ok::<(), io::Error>(()) + }; + let _ = timeout(CLIENT_CLOSE_DRAIN_TIMEOUT, drain).await; +} + +/// Dispatches one request. +async fn handle_request( + state: Arc, + connection_slot: Arc>>, + request: Request, +) -> Response { + let is_connect = request.method() == Method::CONNECT; + let mut response = if is_connect { + handle_connect(&state, connection_slot, request).await + } else { + status_response(StatusCode::NOT_IMPLEMENTED) + }; + + if !is_connect || !response.status().is_success() { + response + .headers_mut() + .insert(header::CONNECTION, HeaderValue::from_static("close")); + } + response +} + +/// Handles one CONNECT tunnel request. +async fn handle_connect( + state: &ProxyState, + connection_slot: Arc>>, + mut request: Request, +) -> Response { + if request.headers().contains_key(header::UPGRADE) { + return status_response(StatusCode::NOT_IMPLEMENTED); + } + if validate_connect_framing(request.headers()).is_err() { + return status_response(StatusCode::BAD_REQUEST); + } + + let Some(authority) = connect_authority(request.uri()) else { + return status_response(StatusCode::BAD_REQUEST); + }; + + if !host_header_is_consistent(&request, &authority) { + return status_response(StatusCode::BAD_REQUEST); + } + + if !state.policy.allows(authority.host(), authority.port()) { + return status_response(StatusCode::FORBIDDEN); + } + + let upstream = match connect_upstream(state, &authority).await { + Ok(stream) => LimitedStream::new(stream, IDLE_TIMEOUT), + Err(failure) => return status_response(failure.status()), + }; + + let Some(permit) = connection_slot.lock().await.take() else { + return status_response(StatusCode::SERVICE_UNAVAILABLE); + }; + let upgrade = hyper::upgrade::on(&mut request); + tokio::spawn(async move { + let _permit = permit; + match upgrade.await { + Ok(upgraded) => { + let mut client = TokioIo::new(upgraded); + let mut upstream = upstream; + if let Err(error) = tokio::io::copy_bidirectional(&mut client, &mut upstream).await + { + diagnostic(format_args!("tunnel ended: {error}")); + } + } + Err(error) => diagnostic(format_args!("tunnel upgrade failed: {error}")), + } + }); + + let mut response = Response::new(empty_body()); + *response.status_mut() = StatusCode::OK; + response +} + +/// Canonicalizes the authority-form target of a CONNECT request. +fn connect_authority(uri: &Uri) -> Option { + if uri.scheme_str().is_some() || !uri.path().is_empty() || uri.query().is_some() { + return None; + } + let raw = uri.authority().map(Authority::as_str)?; + parse_authority(raw, None).ok() +} + +/// Returns whether a CONNECT Host header matches the request target. +fn host_header_is_consistent(request: &Request, authority: &RequestAuthority) -> bool { + let Some(value) = request.headers().get(header::HOST) else { + return true; + }; + value + .to_str() + .ok() + .and_then(|raw| parse_authority(raw, None).ok()) + .is_some_and(|host_header| &host_header == authority) +} + +/// Reason no upstream connection could be established. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum UpstreamFailure { + Failed, + TimedOut, +} + +impl UpstreamFailure { + fn status(self) -> StatusCode { + match self { + Self::Failed => StatusCode::BAD_GATEWAY, + Self::TimedOut => StatusCode::GATEWAY_TIMEOUT, + } + } +} + +/// Resolves and connects to an authorized hostname. +async fn connect_upstream( + state: &ProxyState, + authority: &RequestAuthority, +) -> Result { + timeout( + UPSTREAM_CONNECT_TIMEOUT, + state + .connector + .connect(authority.host().clone(), authority.port()), + ) + .await + .map_err(|_elapsed| UpstreamFailure::TimedOut)? + .map_err(|_error| UpstreamFailure::Failed) +} + +fn empty_body() -> ProxyBody { + Empty::::new() + .map_err(|never| match never {}) + .boxed() +} + +fn status_response(status: StatusCode) -> Response { + let mut response = Response::new(empty_body()); + *response.status_mut() = status; + response + .headers_mut() + .insert(header::CONNECTION, HeaderValue::from_static("close")); + response +} + +fn diagnostic(message: core::fmt::Arguments<'_>) { + eprintln!("litebox_egress_proxy: {message}"); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn uri(raw: &str) -> Uri { + raw.parse().unwrap() + } + + #[test] + fn connect_targets_must_be_authority_form() { + let authority = connect_authority(&uri("example.com:443")).unwrap(); + assert_eq!(authority.host().as_str(), "example.com"); + assert_eq!(authority.port(), 443); + + assert!(connect_authority(&uri("http://example.com:443")).is_none()); + assert!(connect_authority(&uri("example.com")).is_none()); + assert!(connect_authority(&uri("example.com:0")).is_none()); + } + + #[test] + fn upstream_failures_map_to_statuses() { + assert_eq!( + UpstreamFailure::Failed.status(), + StatusCode::BAD_GATEWAY + ); + assert_eq!( + UpstreamFailure::TimedOut.status(), + StatusCode::GATEWAY_TIMEOUT + ); + } +} diff --git a/litebox_egress_proxy/src/request_head.rs b/litebox_egress_proxy/src/request_head.rs new file mode 100644 index 000000000..a94462561 --- /dev/null +++ b/litebox_egress_proxy/src/request_head.rs @@ -0,0 +1,108 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Raw request-head validation before HTTP framing normalization. + +use std::io; + +use bytes::Bytes; +use thiserror::Error; +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::time::timeout; + +use crate::headers::validate_raw_request_framing; +use crate::limits::{MAX_HEADER_FIELDS, MAX_REQUEST_HEADER_BYTES, REQUEST_HEADER_READ_TIMEOUT}; + +/// A complete, raw-validated request prefix, including bytes read ahead. +pub(crate) struct ValidatedRequestPrefix(Bytes); + +impl ValidatedRequestPrefix { + pub(crate) fn into_bytes(self) -> Bytes { + self.0 + } +} + +/// Reason the first request head could not be accepted. +#[derive(Debug, Error)] +pub(crate) enum RequestHeadError { + #[error("client closed before sending a complete request head")] + Closed, + #[error("request head exceeded the read timeout")] + TimedOut, + #[error("request head exceeded a configured limit")] + TooLarge, + #[error("request head is malformed or framing-ambiguous")] + Malformed, + #[error("request-head read failed: {0}")] + Io(#[from] io::Error), +} + +impl RequestHeadError { + /// A complete HTTP rejection for errors caused by client input. + pub(crate) fn response(&self) -> Option<&'static [u8]> { + match self { + Self::Closed | Self::Io(_) => None, + Self::TimedOut => Some( + b"HTTP/1.1 408 Request Timeout\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ), + Self::TooLarge => Some( + b"HTTP/1.1 431 Request Header Fields Too Large\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ), + Self::Malformed => Some( + b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ), + } + } +} + +/// Reads and validates exactly the first raw HTTP/1 request head. +pub(crate) async fn read_validated_request_prefix( + stream: &mut S, +) -> Result +where + S: AsyncRead + Unpin, +{ + timeout(REQUEST_HEADER_READ_TIMEOUT, read_request_prefix(stream)) + .await + .map_err(|_| RequestHeadError::TimedOut)? +} + +async fn read_request_prefix(stream: &mut S) -> Result +where + S: AsyncRead + Unpin, +{ + let mut prefix = Vec::with_capacity(1024); + let mut chunk = [0_u8; 1024]; + + loop { + if prefix.len() == MAX_REQUEST_HEADER_BYTES { + return Err(RequestHeadError::TooLarge); + } + let remaining = MAX_REQUEST_HEADER_BYTES - prefix.len(); + let read_capacity = remaining.min(chunk.len()); + let read = stream.read(&mut chunk[..read_capacity]).await?; + if read == 0 { + return Err(RequestHeadError::Closed); + } + prefix.extend_from_slice(&chunk[..read]); + + let mut headers = [httparse::EMPTY_HEADER; MAX_HEADER_FIELDS]; + let mut request = httparse::Request::new(&mut headers); + match request.parse(&prefix) { + Ok(httparse::Status::Partial) => {} + Ok(httparse::Status::Complete(_)) => { + let method = request.method.ok_or(RequestHeadError::Malformed)?; + let target = request.path.ok_or(RequestHeadError::Malformed)?; + let version = request.version.ok_or(RequestHeadError::Malformed)?; + if target.as_bytes().contains(&b'#') { + return Err(RequestHeadError::Malformed); + } + validate_raw_request_framing(method, version, request.headers) + .map_err(|_| RequestHeadError::Malformed)?; + return Ok(ValidatedRequestPrefix(Bytes::from(prefix))); + } + Err(httparse::Error::TooManyHeaders) => return Err(RequestHeadError::TooLarge), + Err(_) => return Err(RequestHeadError::Malformed), + } + } +} diff --git a/litebox_egress_proxy/src/stream.rs b/litebox_egress_proxy/src/stream.rs new file mode 100644 index 000000000..6c69fe168 --- /dev/null +++ b/litebox_egress_proxy/src/stream.rs @@ -0,0 +1,264 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Idle-timeout stream wrapper. + +use core::future::Future; +use core::pin::Pin; +use core::task::{Context, Poll}; +use core::time::Duration; +use std::io; +use std::sync::{Arc, Mutex}; + +use bytes::Bytes; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::net::TcpStream; +use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; +use tokio::time::{Instant, Sleep, sleep_until}; + +/// A clonable handle to the read half of a TCP stream. +/// +/// The proxy retains one handle so it can perform a bounded drain after Hyper +/// flushes a non-upgraded response and releases its stream. +#[derive(Clone)] +pub(crate) struct SharedTcpRead { + inner: Arc>, +} + +impl AsyncRead for SharedTcpRead { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let Ok(mut inner) = self.inner.lock() else { + return Poll::Ready(Err(io::Error::other("TCP read half mutex poisoned"))); + }; + Pin::new(&mut *inner).poll_read(cx, buf) + } +} + +/// A split TCP stream whose read half can be retained for bounded closing. +pub(crate) struct SharedTcpStream { + read: SharedTcpRead, + write: OwnedWriteHalf, +} + +/// Splits a stream while retaining a clonable handle to its read half. +pub(crate) fn share_tcp_read(stream: TcpStream) -> (SharedTcpStream, SharedTcpRead) { + let (read, write) = stream.into_split(); + let read = SharedTcpRead { + inner: Arc::new(Mutex::new(read)), + }; + ( + SharedTcpStream { + read: read.clone(), + write, + }, + read, + ) +} + +impl AsyncRead for SharedTcpStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().read).poll_read(cx, buf) + } +} + +impl AsyncWrite for SharedTcpStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().write).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().write).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().write).poll_shutdown(cx) + } +} + +/// A stream that replays a prefix before reading from its inner stream. +pub struct PrefixedStream { + prefix: Bytes, + inner: S, +} + +impl PrefixedStream { + /// Creates a stream that yields `prefix` before bytes from `inner`. + pub fn new(prefix: Bytes, inner: S) -> Self { + Self { prefix, inner } + } +} + +impl AsyncRead for PrefixedStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if !this.prefix.is_empty() && buf.remaining() != 0 { + let length = this.prefix.len().min(buf.remaining()); + let bytes = this.prefix.split_to(length); + buf.put_slice(&bytes); + return Poll::Ready(Ok(())); + } + Pin::new(&mut this.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for PrefixedStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + +/// A stream that fails once it stalls for too long, or outlives its deadline. +pub struct LimitedStream { + inner: S, + idle: Duration, + idle_timer: Pin>, +} + +impl LimitedStream { + /// Wraps `inner` with an idle timeout. + pub fn new(inner: S, idle: Duration) -> Self { + Self { + inner, + idle, + idle_timer: Box::pin(sleep_until(Instant::now() + idle)), + } + } + + /// Restarts the idle timeout after observable progress. + fn touch(&mut self) { + let deadline = Instant::now() + self.idle; + self.idle_timer.as_mut().reset(deadline); + } + + /// Returns `true` when the stream has been idle for too long. + fn idle_expired(&mut self, cx: &mut Context<'_>) -> bool { + self.idle_timer.as_mut().poll(cx).is_ready() + } +} + +fn timed_out(reason: &'static str) -> io::Error { + io::Error::new(io::ErrorKind::TimedOut, reason) +} + +impl AsyncRead for LimitedStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_read(cx, buf) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.idle_expired(cx) => { + Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncWrite for LimitedStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_write(cx, buf) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.idle_expired(cx) => { + Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) + } + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_flush(cx) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.idle_expired(cx) => { + Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) + } + Poll::Pending => Poll::Pending, + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_shutdown(cx) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.idle_expired(cx) => { + Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) + } + Poll::Pending => Poll::Pending, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use tokio::io::{AsyncReadExt, duplex}; + + fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + } + + #[test] + fn idle_stream_times_out() { + runtime().block_on(async { + tokio::time::pause(); + let (client, _server) = duplex(64); + let mut limited = LimitedStream::new(client, Duration::from_secs(60)); + let mut buffer = [0_u8; 8]; + let error = limited.read(&mut buffer).await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + }); + } + +} diff --git a/litebox_egress_proxy/src/upstream.rs b/litebox_egress_proxy/src/upstream.rs new file mode 100644 index 000000000..2c3602547 --- /dev/null +++ b/litebox_egress_proxy/src/upstream.rs @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Upstream hostname resolution and connection. + +use core::future::Future; +use core::pin::Pin; +use std::io; + +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::net::{TcpStream, lookup_host}; + +use crate::policy::Hostname; + +/// A bidirectional upstream byte stream. +pub trait UpstreamStream: AsyncRead + AsyncWrite + Send + Unpin {} + +impl UpstreamStream for T {} + +/// An owned upstream byte stream. +pub type BoxedUpstreamStream = Box; + +/// A future returned by an [`UpstreamConnector`]. +pub type ConnectFuture<'a> = + Pin> + Send + 'a>>; + +/// Resolves and connects to an authorized hostname and port. +pub trait UpstreamConnector: Send + Sync + 'static { + /// Opens an upstream connection. + fn connect(&self, host: Hostname, port: u16) -> ConnectFuture<'_>; +} + +/// The production connector, using the trusted host resolver. +#[derive(Clone, Copy, Debug, Default)] +pub struct TcpUpstreamConnector; + +impl UpstreamConnector for TcpUpstreamConnector { + fn connect(&self, host: Hostname, port: u16) -> ConnectFuture<'_> { + Box::pin(async move { + let addresses = lookup_host((host.as_str(), port)).await?; + let mut last_error = None; + + for address in addresses { + match TcpStream::connect(address).await { + Ok(stream) => { + stream.set_nodelay(true)?; + return Ok(Box::new(stream) as BoxedUpstreamStream); + } + Err(error) => last_error = Some(error), + } + } + + Err(last_error.unwrap_or_else(|| { + io::Error::new(io::ErrorKind::NotFound, "hostname resolved to no addresses") + })) + }) + } +} diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs new file mode 100644 index 000000000..bdffe0473 --- /dev/null +++ b/litebox_egress_proxy/tests/loopback.rs @@ -0,0 +1,380 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Hermetic loopback tests for the CONNECT proxy. + +use std::collections::HashMap; +use std::io; +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use litebox_egress_proxy::listener::{ListenerSource, acquire}; +use litebox_egress_proxy::policy::{HostPolicy, Hostname}; +use litebox_egress_proxy::proxy::{ProxyState, serve}; +use litebox_egress_proxy::upstream::{BoxedUpstreamStream, ConnectFuture, UpstreamConnector}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::time::timeout; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +struct LoopbackConnector { + routes: HashMap<(Hostname, u16), SocketAddr>, + attempts: Arc, +} + +impl UpstreamConnector for LoopbackConnector { + fn connect(&self, host: Hostname, port: u16) -> ConnectFuture<'_> { + self.attempts.fetch_add(1, Ordering::SeqCst); + let route = self.routes.get(&(host, port)).copied(); + Box::pin(async move { + let Some(route) = route else { + return Err(io::Error::from(io::ErrorKind::ConnectionRefused)); + }; + let stream = TcpStream::connect(route).await?; + Ok(Box::new(stream) as BoxedUpstreamStream) + }) + } +} + +struct TestProxy { + address: SocketAddrV4, + attempts: Arc, +} + +impl TestProxy { + fn start(rules: &[&str], routes: &[(&str, u16, SocketAddr)]) -> Self { + let policy = HostPolicy::from_rules(rules.iter().copied()).expect("valid policy"); + + let mut mapped = HashMap::new(); + for (host, port, address) in routes { + let host = Hostname::parse(host).expect("valid hostname"); + mapped.insert((host, *port), *address); + } + + let attempts = Arc::new(AtomicUsize::new(0)); + let connector = LoopbackConnector { + routes: mapped, + attempts: Arc::clone(&attempts), + }; + let state = Arc::new(ProxyState::new(policy, Arc::new(connector))); + + let listener = acquire(ListenerSource::Bind(SocketAddrV4::new( + Ipv4Addr::LOCALHOST, + 0, + ))) + .expect("loopback listener"); + let listener = TcpListener::from_std(listener).expect("async listener"); + let SocketAddr::V4(address) = listener.local_addr().expect("listener address") else { + panic!("expected an IPv4 listener"); + }; + + tokio::spawn(async move { + let _ = serve(listener, state).await; + }); + + Self { address, attempts } + } + + fn upstream_attempts(&self) -> usize { + self.attempts.load(Ordering::SeqCst) + } + + async fn connect(&self) -> ProxyClient { + let stream = timeout(TEST_TIMEOUT, TcpStream::connect(self.address)) + .await + .expect("connect did not time out") + .expect("client connects to the proxy"); + ProxyClient { + stream, + buffer: Vec::new(), + } + } + + async fn request(&self, raw: &str) -> HttpResponse { + let mut client = self.connect().await; + client.send(raw.as_bytes()).await; + client.read_response().await + } +} + +struct ProxyClient { + stream: TcpStream, + buffer: Vec, +} + +impl ProxyClient { + async fn send(&mut self, bytes: &[u8]) { + timeout(TEST_TIMEOUT, self.stream.write_all(bytes)) + .await + .expect("write did not time out") + .expect("write succeeds"); + } + + async fn fill(&mut self) -> bool { + let mut chunk = [0_u8; 4096]; + let read = timeout(TEST_TIMEOUT, self.stream.read(&mut chunk)) + .await + .expect("read did not time out") + .expect("read succeeds"); + if read == 0 { + return false; + } + self.buffer.extend_from_slice(&chunk[..read]); + true + } + + async fn read_response(&mut self) -> HttpResponse { + let head_end = loop { + if let Some(index) = find_subslice(&self.buffer, b"\r\n\r\n") { + break index + 4; + } + assert!( + self.fill().await, + "connection closed before a response head" + ); + }; + + let head = String::from_utf8(self.buffer[..head_end].to_vec()).expect("ASCII head"); + self.buffer.drain(..head_end); + let status = head + .split_whitespace() + .nth(1) + .and_then(|code| code.parse::().ok()) + .expect("status code"); + HttpResponse { status, head } + } + + async fn read_exact(&mut self, length: usize) -> Vec { + while self.buffer.len() < length { + assert!( + self.fill().await, + "connection closed before the tunnel data" + ); + } + self.buffer.drain(..length).collect() + } +} + +struct HttpResponse { + status: u16, + head: String, +} + +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +async fn echo_upstream() -> SocketAddr { + let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)) + .await + .expect("echo listener"); + let address = listener.local_addr().expect("echo address"); + + tokio::spawn(async move { + while let Ok((mut stream, _peer)) = listener.accept().await { + tokio::spawn(async move { + let (mut reader, mut writer) = stream.split(); + let _ = tokio::io::copy(&mut reader, &mut writer).await; + }); + } + }); + + address +} + +#[tokio::test] +async fn connect_tunnel_relays_bytes() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + let mut client = proxy.connect().await; + client + .send( + concat!( + "CONNECT allowed.example:443 HTTP/1.1\r\n", + "Host: allowed.example:443\r\n", + "\r\n", + "early" + ) + .as_bytes(), + ) + .await; + + let response = client.read_response().await; + assert_eq!(response.status, 200); + assert!( + !response + .head + .to_ascii_lowercase() + .contains("connection: close") + ); + assert_eq!(client.read_exact(5).await, b"early"); + + client.send(b"tunnelled").await; + assert_eq!(client.read_exact(9).await, b"tunnelled"); +} + +#[tokio::test] +async fn allowed_hostname_is_connected_for_each_request() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + for _ in 0..2 { + let response = proxy + .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\n\r\n") + .await; + assert_eq!(response.status, 200); + } + + assert_eq!(proxy.upstream_attempts(), 2); +} + +#[tokio::test] +async fn denied_host_and_port_do_not_trigger_network_activity() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + let denied_host = proxy + .request("CONNECT denied.example:443 HTTP/1.1\r\nHost: denied.example:443\r\n\r\n") + .await; + assert_eq!(denied_host.status, 403); + + let denied_port = proxy + .request("CONNECT allowed.example:8443 HTTP/1.1\r\nHost: allowed.example:8443\r\n\r\n") + .await; + assert_eq!(denied_port.status, 403); + + assert_eq!(proxy.upstream_attempts(), 0); +} + +#[tokio::test] +async fn connect_authority_and_framing_are_validated() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + for request in [ + "CONNECT allowed.example HTTP/1.1\r\nHost: allowed.example\r\n\r\n", + "CONNECT 93.184.216.1:443 HTTP/1.1\r\nHost: 93.184.216.1:443\r\n\r\n", + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:80\r\n\r\n", + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nContent-Length: 0\r\n\r\n", + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nTransfer-Encoding: chunked\r\n\r\n", + ] { + assert_eq!(proxy.request(request).await.status, 400); + } + + assert_eq!(proxy.upstream_attempts(), 0); +} + +#[tokio::test] +async fn explicitly_allowed_dns_port_is_forwarded() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:53"], + &[("allowed.example", 53, upstream)], + ); + + let response = proxy + .request("CONNECT allowed.example:53 HTTP/1.1\r\nHost: allowed.example:53\r\n\r\n") + .await; + assert_eq!(response.status, 200); + assert_eq!(proxy.upstream_attempts(), 1); +} + +#[tokio::test] +async fn unreachable_upstream_yields_bad_gateway() { + let proxy = TestProxy::start(&["allowed.example:443"], &[]); + + let response = proxy + .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\n\r\n") + .await; + assert_eq!(response.status, 502); + assert_eq!(proxy.upstream_attempts(), 1); +} + +#[tokio::test] +async fn denied_connect_early_bytes_are_drained_before_close() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + let early = "x".repeat(32 * 1024); + let request = format!( + "CONNECT denied.example:443 HTTP/1.1\r\n\ + Host: denied.example:443\r\n\ + \r\n\ + {early}" + ); + + let mut client = proxy.connect().await; + client.send(request.as_bytes()).await; + assert_eq!(client.read_response().await.status, 403); + assert!(!client.fill().await); + assert_eq!(proxy.upstream_attempts(), 0); +} + +#[tokio::test] +async fn unsupported_methods_are_rejected_without_network_activity() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + let response = proxy + .request("GET http://allowed.example/ HTTP/1.1\r\nHost: allowed.example\r\n\r\n") + .await; + assert_eq!(response.status, 501); + assert_eq!(proxy.upstream_attempts(), 0); +} + +#[tokio::test] +async fn malformed_and_oversized_heads_are_rejected() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ); + + let spaced = proxy + .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost : allowed.example:443\r\n\r\n") + .await; + assert_eq!(spaced.status, 400); + + let folded = proxy + .request(concat!( + "CONNECT allowed.example:443 HTTP/1.1\r\n", + "Host: allowed.example:443\r\n", + "X-Folded: one\r\n two\r\n", + "\r\n" + )) + .await; + assert_eq!(folded.status, 400); + + let mut oversized = String::from( + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nX-Big: ", + ); + oversized.push_str(&"a".repeat(32 * 1024)); + oversized.push_str("\r\n\r\n"); + assert_eq!(proxy.request(&oversized).await.status, 431); + + assert_eq!(proxy.upstream_attempts(), 0); +} From d3fa5d25f6e6ac07ebc61ecd5ac58fae6a63f6fb Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 11:57:06 -0700 Subject: [PATCH 02/11] Restore alloc import for proxy library Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/lib.rs | 2 ++ litebox_egress_proxy/src/proxy.rs | 5 +---- litebox_egress_proxy/src/stream.rs | 1 - 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index b1740aad2..b47bb64b5 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -7,6 +7,8 @@ //! resolver, then connected using the returned numeric addresses. CONNECT //! tunnels relay bytes without inspecting or terminating TLS. +extern crate alloc; + pub mod authority; pub mod config; pub mod headers; diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 0ea258534..59e0a7686 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -329,10 +329,7 @@ mod tests { #[test] fn upstream_failures_map_to_statuses() { - assert_eq!( - UpstreamFailure::Failed.status(), - StatusCode::BAD_GATEWAY - ); + assert_eq!(UpstreamFailure::Failed.status(), StatusCode::BAD_GATEWAY); assert_eq!( UpstreamFailure::TimedOut.status(), StatusCode::GATEWAY_TIMEOUT diff --git a/litebox_egress_proxy/src/stream.rs b/litebox_egress_proxy/src/stream.rs index 6c69fe168..c48807c1e 100644 --- a/litebox_egress_proxy/src/stream.rs +++ b/litebox_egress_proxy/src/stream.rs @@ -260,5 +260,4 @@ mod tests { assert_eq!(error.kind(), io::ErrorKind::TimedOut); }); } - } From cce7fb3052009294e9b6952c56420f2ce9071c8f Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 12:49:30 -0700 Subject: [PATCH 03/11] Simplify CONNECT proxy request path Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- Cargo.lock | 1 - litebox_egress_proxy/Cargo.toml | 1 - litebox_egress_proxy/src/headers.rs | 99 ------------ litebox_egress_proxy/src/lib.rs | 80 ++-------- litebox_egress_proxy/src/limits.rs | 37 ----- litebox_egress_proxy/src/listener.rs | 141 +++++----------- litebox_egress_proxy/src/proxy.rs | 176 +++++--------------- litebox_egress_proxy/src/request_head.rs | 108 ------------- litebox_egress_proxy/src/stream.rs | 195 +++-------------------- litebox_egress_proxy/tests/loopback.rs | 64 +------- 10 files changed, 121 insertions(+), 781 deletions(-) delete mode 100644 litebox_egress_proxy/src/headers.rs delete mode 100644 litebox_egress_proxy/src/limits.rs delete mode 100644 litebox_egress_proxy/src/request_head.rs diff --git a/Cargo.lock b/Cargo.lock index 1f23a61c5..686475231 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1545,7 +1545,6 @@ dependencies = [ "clap", "hashbrown", "http-body-util", - "httparse", "hyper", "hyper-util", "libc", diff --git a/litebox_egress_proxy/Cargo.toml b/litebox_egress_proxy/Cargo.toml index 2ded1634f..b39bd65e8 100644 --- a/litebox_egress_proxy/Cargo.toml +++ b/litebox_egress_proxy/Cargo.toml @@ -14,7 +14,6 @@ clap = { version = "4.5", default-features = false, features = [ ] } hashbrown = "0.15.2" http-body-util = { version = "0.1.3", default-features = false } -httparse = { version = "1.10.1", default-features = false } hyper = { version = "1.8", default-features = false, features = [ "http1", "server", diff --git a/litebox_egress_proxy/src/headers.rs b/litebox_egress_proxy/src/headers.rs deleted file mode 100644 index f7f4db4f5..000000000 --- a/litebox_egress_proxy/src/headers.rs +++ /dev/null @@ -1,99 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -//! CONNECT request framing validation. - -use hyper::header; -use hyper::header::HeaderMap; -use thiserror::Error; - -/// Reason a CONNECT request was rejected as framing-ambiguous. -#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)] -pub enum FramingError { - /// More than one `Host` header was present. - #[error("message carries more than one Host header")] - DuplicateHost, - /// A CONNECT request carried body framing. - #[error("CONNECT request carries a body")] - BodyOnConnect, -} - -/// Validates request framing before `hyper` normalizes the raw header list. -pub(crate) fn validate_raw_request_framing( - method: &str, - _version: u8, - headers: &[httparse::Header<'_>], -) -> Result<(), FramingError> { - if raw_header_count(headers, b"host") > 1 { - return Err(FramingError::DuplicateHost); - } - if method == "CONNECT" - && (raw_header_count(headers, b"transfer-encoding") != 0 - || raw_header_count(headers, b"content-length") != 0) - { - return Err(FramingError::BodyOnConnect); - } - Ok(()) -} - -/// Validates that a parsed CONNECT request is unambiguously bodyless. -pub fn validate_connect_framing(headers: &HeaderMap) -> Result<(), FramingError> { - if headers.contains_key(header::TRANSFER_ENCODING) - || headers.contains_key(header::CONTENT_LENGTH) - { - return Err(FramingError::BodyOnConnect); - } - if headers.get_all(header::HOST).iter().count() > 1 { - return Err(FramingError::DuplicateHost); - } - Ok(()) -} - -fn raw_header_count(headers: &[httparse::Header<'_>], name: &[u8]) -> usize { - headers - .iter() - .filter(|header| header.name.as_bytes().eq_ignore_ascii_case(name)) - .count() -} - -#[cfg(test)] -mod tests { - use super::*; - - use hyper::header::{HeaderName, HeaderValue}; - - fn headers(pairs: &[(&str, &str)]) -> HeaderMap { - let mut map = HeaderMap::new(); - for (name, value) in pairs { - map.append( - HeaderName::from_bytes(name.as_bytes()).unwrap(), - HeaderValue::from_str(value).unwrap(), - ); - } - map - } - - #[test] - fn connect_requests_must_be_bodyless() { - assert!(validate_connect_framing(&headers(&[("host", "example.com:443")])).is_ok()); - assert_eq!( - validate_connect_framing(&headers(&[("content-length", "0")])), - Err(FramingError::BodyOnConnect) - ); - assert_eq!( - validate_connect_framing(&headers(&[("transfer-encoding", "chunked")])), - Err(FramingError::BodyOnConnect) - ); - } - - #[test] - fn duplicate_host_is_rejected() { - assert_eq!( - validate_connect_framing(&headers(&[ - ("host", "example.com:443"), - ("host", "example.com:443"), - ])), - Err(FramingError::DuplicateHost) - ); - } -} diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index b47bb64b5..f47e1f895 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -11,18 +11,14 @@ extern crate alloc; pub mod authority; pub mod config; -pub mod headers; -pub mod limits; pub mod listener; pub mod policy; pub mod proxy; -pub mod stream; pub mod upstream; -mod request_head; +mod stream; use std::io::{self, Write}; -use std::net::{SocketAddr, SocketAddrV4}; use std::sync::Arc; use thiserror::Error; @@ -44,75 +40,21 @@ pub enum StartupError { Io(#[from] io::Error), } -/// A proxy that has completed startup and is ready to serve. -pub struct StartedProxy { - listener: TcpListener, - state: Arc, - local_address: SocketAddrV4, -} - -impl StartedProxy { - /// Returns the loopback address the proxy listens on. - pub fn local_address(&self) -> SocketAddrV4 { - self.local_address - } - - /// Serves client connections until the listener fails. - pub async fn serve(self) -> io::Result<()> { - proxy::serve(self.listener, self.state).await - } -} - -/// Acquires the listener and prepares the shared state. -pub fn start(config: &ProxyConfig) -> Result { - let listener = listener::acquire(config.listener)?; +/// Runs the proxy: startup, readiness announcement, then serving. +pub async fn run(config: &ProxyConfig) -> Result<(), StartupError> { + let (listener, local_address) = listener::acquire(config.listener)?; let listener = TcpListener::from_std(listener)?; - let SocketAddr::V4(local_address) = listener.local_addr()? else { - return Err(StartupError::Listener(ListenerError::NotLoopback( - listener.local_addr()?, - ))); - }; - let state = Arc::new(ProxyState::new( config.policy.clone(), - Arc::new(TcpUpstreamConnector), + Box::new(TcpUpstreamConnector), )); - Ok(StartedProxy { - listener, - state, - local_address, - }) -} - -/// Writes the single readiness line. -pub fn write_readiness(writer: &mut impl Write, address: SocketAddrV4) -> io::Result<()> { - writeln!(writer, "READY {address}")?; - writer.flush() -} - -/// Runs the proxy: startup, readiness announcement, then serving. -pub async fn run(config: &ProxyConfig) -> Result<(), StartupError> { - let started = start(config)?; - - let mut stdout = io::stdout().lock(); - write_readiness(&mut stdout, started.local_address())?; - drop(stdout); + { + let mut stdout = io::stdout().lock(); + writeln!(stdout, "READY {local_address}")?; + stdout.flush()?; + } - started.serve().await?; + proxy::serve(listener, state).await?; Ok(()) } - -#[cfg(test)] -mod tests { - use super::*; - - use std::net::Ipv4Addr; - - #[test] - fn readiness_line_is_exactly_one_line() { - let mut output = Vec::new(); - write_readiness(&mut output, SocketAddrV4::new(Ipv4Addr::LOCALHOST, 34567)).unwrap(); - assert_eq!(output, b"READY 127.0.0.1:34567\n"); - } -} diff --git a/litebox_egress_proxy/src/limits.rs b/litebox_egress_proxy/src/limits.rs deleted file mode 100644 index 028a1d7d8..000000000 --- a/litebox_egress_proxy/src/limits.rs +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -//! Fixed resource limits for the egress proxy. -//! -//! None of these limits is caller-configurable: the proxy is a trusted -//! component whose behaviour must be identical for every sandbox. - -use core::time::Duration; - -/// Maximum number of client connections served concurrently. -/// -/// Additional connections stay in the listener backlog until a slot frees up. -pub const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; - -/// Maximum number of bytes buffered for a client request head. -pub const MAX_REQUEST_HEADER_BYTES: usize = 16 * 1024; - -/// Maximum number of individual header fields parsed per message. -pub const MAX_HEADER_FIELDS: usize = 100; - -/// Total timeout for hostname resolution and connection attempts. -pub const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); - -/// Idle timeout applied to HTTP bodies and CONNECT tunnels. -/// -/// A stream that makes no read or write progress for this long is torn down. -pub const IDLE_TIMEOUT: Duration = Duration::from_secs(60); - -/// Maximum time a client may take to send a complete request head. -pub const REQUEST_HEADER_READ_TIMEOUT: Duration = Duration::from_secs(30); - -/// Maximum time spent draining client input after a non-upgraded response. -pub const CLIENT_CLOSE_DRAIN_TIMEOUT: Duration = Duration::from_secs(1); - -/// Maximum client input discarded while closing a non-upgraded connection. -pub const MAX_CLIENT_CLOSE_DRAIN_BYTES: usize = 64 * 1024; diff --git a/litebox_egress_proxy/src/listener.rs b/litebox_egress_proxy/src/listener.rs index ab7994592..0a8b33528 100644 --- a/litebox_egress_proxy/src/listener.rs +++ b/litebox_egress_proxy/src/listener.rs @@ -1,13 +1,9 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT license. -//! Listener acquisition for the standalone and broker modes. -//! -//! Both modes end with the same invariant: the proxy only ever serves a bound, -//! listening IPv4 loopback TCP socket. The standalone mode binds it itself; the -//! broker mode adopts a listener that a launcher bound and inherited to this -//! process. +//! Loopback listener acquisition. +use std::io; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener}; use thiserror::Error; @@ -17,114 +13,60 @@ use thiserror::Error; pub enum ListenerSource { /// Bind a fresh loopback listener. Port zero requests an ephemeral port. Bind(SocketAddrV4), - /// Adopt an inherited, already-bound listener by descriptor number. + /// Adopt an inherited listener by descriptor number. Inherited(i32), } /// Reason a listener could not be acquired. #[derive(Debug, Error)] pub enum ListenerError { - /// Binding the requested address failed. - #[error("failed to bind {address}: {source}")] - Bind { - /// The address that could not be bound. - address: SocketAddrV4, - /// The underlying failure. - source: std::io::Error, - }, - /// Reading the listener's local address failed. - #[error("failed to read the listener address: {0}")] - LocalAddress(#[source] std::io::Error), - /// The listener was not bound to canonical IPv4 loopback. - #[error("listener is bound to {0}, which is not 127.0.0.1")] - NotLoopback(SocketAddr), - /// The listener was bound to port zero, which an inherited listener never - /// is once it has been bound. - #[error("inherited listener is not bound to a concrete port")] - UnboundPort, - /// The descriptor number was negative. - #[error("inherited listener descriptor is not a valid descriptor number")] - InvalidDescriptor, - /// The descriptor did not refer to an open file. - #[error("inherited listener descriptor is not open")] - DescriptorNotOpen, - /// Inspecting the socket failed. - #[error("failed to inspect the inherited listener: {0}")] - Inspect(#[source] std::io::Error), - /// The descriptor was not an IPv4 stream socket in the listening state. - #[error("inherited descriptor is not a listening IPv4 TCP socket")] - NotAnIpv4Listener, - /// The platform has no inherited-listener contract. - #[error("--listener-fd is only supported on Linux")] - InheritanceUnsupported, - /// Configuring the listener for asynchronous use failed. - #[error("failed to configure the listener: {0}")] - Configure(#[source] std::io::Error), + /// A listener operation failed. + #[error("listener operation failed: {0}")] + Io(#[from] io::Error), + /// The listener was not a bound IPv4 loopback TCP listener. + #[error("listener must be a bound IPv4 loopback TCP listener")] + Invalid, + /// This platform cannot adopt inherited listeners. + #[error("inherited listeners are unsupported on this platform")] + Unsupported, } -/// Acquires the listener described by `source`. -/// -/// The returned listener is non-blocking and validated to be bound to -/// canonical IPv4 loopback. -pub fn acquire(source: ListenerSource) -> Result { +/// Acquires a nonblocking IPv4 loopback listener and its bound address. +pub fn acquire(source: ListenerSource) -> Result<(TcpListener, SocketAddrV4), ListenerError> { let listener = match source { - ListenerSource::Bind(address) => { - TcpListener::bind(address).map_err(|source| ListenerError::Bind { address, source })? - } + ListenerSource::Bind(address) => TcpListener::bind(address)?, ListenerSource::Inherited(descriptor) => adopt_inherited(descriptor)?, }; - let local = listener.local_addr().map_err(ListenerError::LocalAddress)?; - let SocketAddr::V4(local) = local else { - return Err(ListenerError::NotLoopback(local)); + let SocketAddr::V4(address) = listener.local_addr()? else { + return Err(ListenerError::Invalid); }; - if *local.ip() != Ipv4Addr::LOCALHOST { - return Err(ListenerError::NotLoopback(SocketAddr::V4(local))); - } - if local.port() == 0 { - return Err(ListenerError::UnboundPort); + if *address.ip() != Ipv4Addr::LOCALHOST || address.port() == 0 { + return Err(ListenerError::Invalid); } - listener - .set_nonblocking(true) - .map_err(ListenerError::Configure)?; - Ok(listener) + listener.set_nonblocking(true)?; + Ok((listener, address)) } -/// Adopts an inherited descriptor after validating that it really is a bound, -/// listening IPv4 TCP socket. #[cfg(target_os = "linux")] fn adopt_inherited(descriptor: i32) -> Result { use std::os::fd::FromRawFd; - if descriptor < 0 { - return Err(ListenerError::InvalidDescriptor); - } - - // SAFETY: `fcntl(F_GETFD)` only reads the descriptor flags of `descriptor`. - // It neither takes ownership nor mutates process state, and it reports an - // invalid descriptor as `-1` instead of causing undefined behaviour. - let flags = unsafe { libc::fcntl(descriptor, libc::F_GETFD) }; - if flags < 0 { - return Err(ListenerError::DescriptorNotOpen); - } - - if socket_option(descriptor, libc::SO_DOMAIN)? != libc::AF_INET + if descriptor < 0 + || socket_option(descriptor, libc::SO_DOMAIN)? != libc::AF_INET || socket_option(descriptor, libc::SO_TYPE)? != libc::SOCK_STREAM || socket_option(descriptor, libc::SO_ACCEPTCONN)? != 1 { - return Err(ListenerError::NotAnIpv4Listener); + return Err(ListenerError::Invalid); } - // SAFETY: the checks above established that `descriptor` is an open, - // listening IPv4 stream socket. The launcher contract for `--listener-fd` - // transfers ownership of that descriptor to this process, and nothing else - // in this process holds or closes it, so wrapping it in a `TcpListener` - // gives a single unique owner. + // SAFETY: the checks above establish that `descriptor` is an open, + // listening IPv4 stream socket. The launcher transfers ownership of the + // descriptor to this process, so the returned listener is its sole owner. Ok(unsafe { TcpListener::from_raw_fd(descriptor) }) } -/// Reads a `SOL_SOCKET` integer option. #[cfg(target_os = "linux")] fn socket_option(descriptor: i32, option: libc::c_int) -> Result { let mut value: libc::c_int = 0; @@ -132,9 +74,7 @@ fn socket_option(descriptor: i32, option: libc::c_int) -> Result Result Result { - Err(ListenerError::InheritanceUnsupported) + Err(ListenerError::Unsupported) } #[cfg(test)] @@ -162,14 +102,11 @@ mod tests { #[test] fn binds_an_ephemeral_loopback_port() { - let listener = acquire(ListenerSource::Bind(SocketAddrV4::new( + let (_listener, address) = acquire(ListenerSource::Bind(SocketAddrV4::new( Ipv4Addr::LOCALHOST, 0, ))) .unwrap(); - let SocketAddr::V4(address) = listener.local_addr().unwrap() else { - panic!("expected an IPv4 listener"); - }; assert_eq!(*address.ip(), Ipv4Addr::LOCALHOST); assert_ne!(address.port(), 0); } @@ -181,7 +118,7 @@ mod tests { 0, ))) .unwrap_err(); - assert!(matches!(error, ListenerError::NotLoopback(_))); + assert!(matches!(error, ListenerError::Invalid)); } #[test] @@ -189,7 +126,7 @@ mod tests { let error = acquire(ListenerSource::Inherited(-1)).unwrap_err(); assert!(matches!( error, - ListenerError::InvalidDescriptor | ListenerError::InheritanceUnsupported + ListenerError::Invalid | ListenerError::Unsupported )); } @@ -202,7 +139,7 @@ mod tests { let expected = bound.local_addr().unwrap(); let descriptor = bound.into_raw_fd(); - let adopted = acquire(ListenerSource::Inherited(descriptor)).unwrap(); + let (adopted, _) = acquire(ListenerSource::Inherited(descriptor)).unwrap(); assert_eq!(adopted.local_addr().unwrap(), expected); } @@ -214,7 +151,7 @@ mod tests { let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); let client = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap(); let error = acquire(ListenerSource::Inherited(client.into_raw_fd())).unwrap_err(); - assert!(matches!(error, ListenerError::NotAnIpv4Listener)); + assert!(matches!(error, ListenerError::Invalid)); } #[cfg(target_os = "linux")] @@ -224,6 +161,6 @@ mod tests { let socket = std::net::UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); let error = acquire(ListenerSource::Inherited(socket.into_raw_fd())).unwrap_err(); - assert!(matches!(error, ListenerError::NotAnIpv4Listener)); + assert!(matches!(error, ListenerError::Invalid)); } } diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 59e0a7686..e97444188 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -3,18 +3,17 @@ //! Connection acceptance, authorization, and CONNECT tunnelling. //! -//! Every raw-validated CONNECT request is authorized before DNS or upstream -//! activity. A successful request consumes its client connection by upgrading -//! it to one bounded bidirectional tunnel. +//! Every CONNECT request is authorized before DNS or upstream activity. A +//! successful request consumes its client connection by upgrading it to one +//! bounded bidirectional tunnel. use core::convert::Infallible; -use core::error::Error as StdError; +use core::time::Duration; use std::io; use std::sync::Arc; use bytes::Bytes; -use http_body_util::combinators::BoxBody; -use http_body_util::{BodyExt, Empty}; +use http_body_util::Empty; use hyper::body::Incoming; use hyper::header::{self, HeaderValue}; use hyper::http::uri::Authority; @@ -22,38 +21,34 @@ use hyper::server::conn::http1 as server_http1; use hyper::service::service_fn; use hyper::{Method, Request, Response, StatusCode, Uri}; use hyper_util::rt::{TokioIo, TokioTimer}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{Mutex, OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::time::timeout; use crate::authority::{RequestAuthority, parse_authority}; -use crate::headers::validate_connect_framing; -use crate::limits::{ - CLIENT_CLOSE_DRAIN_TIMEOUT, IDLE_TIMEOUT, MAX_CLIENT_CLOSE_DRAIN_BYTES, - MAX_CONCURRENT_CLIENT_CONNECTIONS, MAX_HEADER_FIELDS, MAX_REQUEST_HEADER_BYTES, - REQUEST_HEADER_READ_TIMEOUT, UPSTREAM_CONNECT_TIMEOUT, -}; use crate::policy::HostPolicy; -use crate::request_head::read_validated_request_prefix; -use crate::stream::{LimitedStream, PrefixedStream, share_tcp_read}; +use crate::stream::LimitedStream; use crate::upstream::{BoxedUpstreamStream, UpstreamConnector}; -/// Boxed error type used by response bodies. -type BoxError = Box; +const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; +const MAX_REQUEST_HEADER_BYTES: usize = 16 * 1024; +const MAX_HEADER_FIELDS: usize = 100; +const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const IDLE_TIMEOUT: Duration = Duration::from_secs(60); +const REQUEST_HEADER_READ_TIMEOUT: Duration = Duration::from_secs(30); /// Response body type produced by the proxy. -type ProxyBody = BoxBody; +type ProxyBody = Empty; /// Immutable state shared by every connection. pub struct ProxyState { policy: HostPolicy, - connector: Arc, + connector: Box, } impl ProxyState { /// Builds shared state from a validated policy and upstream connector. - pub fn new(policy: HostPolicy, connector: Arc) -> Self { + pub fn new(policy: HostPolicy, connector: Box) -> Self { Self { policy, connector } } } @@ -93,32 +88,12 @@ async fn serve_connection(state: Arc, stream: TcpStream, permit: Own return; } - let (stream, mut drain_handle) = share_tcp_read(stream); - let mut stream = LimitedStream::new(stream, IDLE_TIMEOUT); - let prefix = match read_validated_request_prefix(&mut stream).await { - Ok(prefix) => prefix.into_bytes(), - Err(error) => { - if let Some(response) = error.response() { - if let Err(write_error) = stream.write_all(response).await { - diagnostic(format_args!( - "failed to write request rejection after {error}: {write_error}" - )); - } else { - let _ = stream.shutdown().await; - drain_client_input(&mut stream).await; - } - } - return; - } - }; - - let io = TokioIo::new(PrefixedStream::new(prefix, stream)); - let connection_slot = Arc::new(Mutex::new(Some(permit))); - let service_connection_slot = Arc::clone(&connection_slot); + let io = TokioIo::new(LimitedStream::new(stream, IDLE_TIMEOUT)); + let permit = Arc::new(permit); let service = service_fn(move |request: Request| { let state = Arc::clone(&state); - let connection_slot = Arc::clone(&service_connection_slot); - async move { Ok::<_, Infallible>(handle_request(state, connection_slot, request).await) } + let permit = Arc::clone(&permit); + async move { Ok::<_, Infallible>(handle_request(state, permit, request).await) } }); let mut builder = server_http1::Builder::new(); @@ -126,72 +101,38 @@ async fn serve_connection(state: Arc, stream: TcpStream, permit: Own .timer(TokioTimer::new()) .header_read_timeout(Some(REQUEST_HEADER_READ_TIMEOUT)) .max_buf_size(MAX_REQUEST_HEADER_BYTES) - .max_headers(MAX_HEADER_FIELDS) - .keep_alive(true) - .half_close(true); - - let result = builder.serve_connection(io, service).with_upgrades().await; - let upgraded = connection_slot.lock().await.is_none(); - if !upgraded { - drain_client_input(&mut drain_handle).await; - } - if let Err(error) = result { + .max_headers(MAX_HEADER_FIELDS); + + if let Err(error) = builder.serve_connection(io, service).with_upgrades().await { diagnostic(format_args!("client connection ended: {error}")); } } -/// Drains bounded client input so unread bytes cannot reset a rejection. -async fn drain_client_input(stream: &mut S) -where - S: tokio::io::AsyncRead + Unpin, -{ - let drain = async { - let mut remaining = MAX_CLIENT_CLOSE_DRAIN_BYTES; - let mut buffer = [0_u8; 1024]; - while remaining != 0 { - let capacity = remaining.min(buffer.len()); - let read = stream.read(&mut buffer[..capacity]).await?; - if read == 0 { - break; - } - remaining -= read; - } - Ok::<(), io::Error>(()) - }; - let _ = timeout(CLIENT_CLOSE_DRAIN_TIMEOUT, drain).await; -} - /// Dispatches one request. async fn handle_request( state: Arc, - connection_slot: Arc>>, + permit: Arc, request: Request, ) -> Response { - let is_connect = request.method() == Method::CONNECT; - let mut response = if is_connect { - handle_connect(&state, connection_slot, request).await - } else { - status_response(StatusCode::NOT_IMPLEMENTED) - }; - - if !is_connect || !response.status().is_success() { - response - .headers_mut() - .insert(header::CONNECTION, HeaderValue::from_static("close")); + if request.method() != Method::CONNECT { + return status_response(StatusCode::NOT_IMPLEMENTED); } - response + handle_connect(&state, permit, request).await } /// Handles one CONNECT tunnel request. async fn handle_connect( state: &ProxyState, - connection_slot: Arc>>, + permit: Arc, mut request: Request, ) -> Response { if request.headers().contains_key(header::UPGRADE) { return status_response(StatusCode::NOT_IMPLEMENTED); } - if validate_connect_framing(request.headers()).is_err() { + if request.headers().contains_key(header::TRANSFER_ENCODING) + || request.headers().contains_key(header::CONTENT_LENGTH) + || request.headers().get_all(header::HOST).iter().count() > 1 + { return status_response(StatusCode::BAD_REQUEST); } @@ -209,12 +150,9 @@ async fn handle_connect( let upstream = match connect_upstream(state, &authority).await { Ok(stream) => LimitedStream::new(stream, IDLE_TIMEOUT), - Err(failure) => return status_response(failure.status()), + Err(status) => return status_response(status), }; - let Some(permit) = connection_slot.lock().await.take() else { - return status_response(StatusCode::SERVICE_UNAVAILABLE); - }; let upgrade = hyper::upgrade::on(&mut request); tokio::spawn(async move { let _permit = permit; @@ -231,7 +169,7 @@ async fn handle_connect( } }); - let mut response = Response::new(empty_body()); + let mut response = Response::new(Empty::new()); *response.status_mut() = StatusCode::OK; response } @@ -257,46 +195,27 @@ fn host_header_is_consistent(request: &Request, authority: &RequestAut .is_some_and(|host_header| &host_header == authority) } -/// Reason no upstream connection could be established. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum UpstreamFailure { - Failed, - TimedOut, -} - -impl UpstreamFailure { - fn status(self) -> StatusCode { - match self { - Self::Failed => StatusCode::BAD_GATEWAY, - Self::TimedOut => StatusCode::GATEWAY_TIMEOUT, - } - } -} - /// Resolves and connects to an authorized hostname. async fn connect_upstream( state: &ProxyState, authority: &RequestAuthority, -) -> Result { - timeout( +) -> Result { + let result = timeout( UPSTREAM_CONNECT_TIMEOUT, state .connector .connect(authority.host().clone(), authority.port()), ) - .await - .map_err(|_elapsed| UpstreamFailure::TimedOut)? - .map_err(|_error| UpstreamFailure::Failed) -} - -fn empty_body() -> ProxyBody { - Empty::::new() - .map_err(|never| match never {}) - .boxed() + .await; + match result { + Ok(Ok(stream)) => Ok(stream), + Ok(Err(_error)) => Err(StatusCode::BAD_GATEWAY), + Err(_elapsed) => Err(StatusCode::GATEWAY_TIMEOUT), + } } fn status_response(status: StatusCode) -> Response { - let mut response = Response::new(empty_body()); + let mut response = Response::new(Empty::new()); *response.status_mut() = status; response .headers_mut() @@ -326,13 +245,4 @@ mod tests { assert!(connect_authority(&uri("example.com")).is_none()); assert!(connect_authority(&uri("example.com:0")).is_none()); } - - #[test] - fn upstream_failures_map_to_statuses() { - assert_eq!(UpstreamFailure::Failed.status(), StatusCode::BAD_GATEWAY); - assert_eq!( - UpstreamFailure::TimedOut.status(), - StatusCode::GATEWAY_TIMEOUT - ); - } } diff --git a/litebox_egress_proxy/src/request_head.rs b/litebox_egress_proxy/src/request_head.rs deleted file mode 100644 index a94462561..000000000 --- a/litebox_egress_proxy/src/request_head.rs +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -//! Raw request-head validation before HTTP framing normalization. - -use std::io; - -use bytes::Bytes; -use thiserror::Error; -use tokio::io::{AsyncRead, AsyncReadExt}; -use tokio::time::timeout; - -use crate::headers::validate_raw_request_framing; -use crate::limits::{MAX_HEADER_FIELDS, MAX_REQUEST_HEADER_BYTES, REQUEST_HEADER_READ_TIMEOUT}; - -/// A complete, raw-validated request prefix, including bytes read ahead. -pub(crate) struct ValidatedRequestPrefix(Bytes); - -impl ValidatedRequestPrefix { - pub(crate) fn into_bytes(self) -> Bytes { - self.0 - } -} - -/// Reason the first request head could not be accepted. -#[derive(Debug, Error)] -pub(crate) enum RequestHeadError { - #[error("client closed before sending a complete request head")] - Closed, - #[error("request head exceeded the read timeout")] - TimedOut, - #[error("request head exceeded a configured limit")] - TooLarge, - #[error("request head is malformed or framing-ambiguous")] - Malformed, - #[error("request-head read failed: {0}")] - Io(#[from] io::Error), -} - -impl RequestHeadError { - /// A complete HTTP rejection for errors caused by client input. - pub(crate) fn response(&self) -> Option<&'static [u8]> { - match self { - Self::Closed | Self::Io(_) => None, - Self::TimedOut => Some( - b"HTTP/1.1 408 Request Timeout\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", - ), - Self::TooLarge => Some( - b"HTTP/1.1 431 Request Header Fields Too Large\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", - ), - Self::Malformed => Some( - b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", - ), - } - } -} - -/// Reads and validates exactly the first raw HTTP/1 request head. -pub(crate) async fn read_validated_request_prefix( - stream: &mut S, -) -> Result -where - S: AsyncRead + Unpin, -{ - timeout(REQUEST_HEADER_READ_TIMEOUT, read_request_prefix(stream)) - .await - .map_err(|_| RequestHeadError::TimedOut)? -} - -async fn read_request_prefix(stream: &mut S) -> Result -where - S: AsyncRead + Unpin, -{ - let mut prefix = Vec::with_capacity(1024); - let mut chunk = [0_u8; 1024]; - - loop { - if prefix.len() == MAX_REQUEST_HEADER_BYTES { - return Err(RequestHeadError::TooLarge); - } - let remaining = MAX_REQUEST_HEADER_BYTES - prefix.len(); - let read_capacity = remaining.min(chunk.len()); - let read = stream.read(&mut chunk[..read_capacity]).await?; - if read == 0 { - return Err(RequestHeadError::Closed); - } - prefix.extend_from_slice(&chunk[..read]); - - let mut headers = [httparse::EMPTY_HEADER; MAX_HEADER_FIELDS]; - let mut request = httparse::Request::new(&mut headers); - match request.parse(&prefix) { - Ok(httparse::Status::Partial) => {} - Ok(httparse::Status::Complete(_)) => { - let method = request.method.ok_or(RequestHeadError::Malformed)?; - let target = request.path.ok_or(RequestHeadError::Malformed)?; - let version = request.version.ok_or(RequestHeadError::Malformed)?; - if target.as_bytes().contains(&b'#') { - return Err(RequestHeadError::Malformed); - } - validate_raw_request_framing(method, version, request.headers) - .map_err(|_| RequestHeadError::Malformed)?; - return Ok(ValidatedRequestPrefix(Bytes::from(prefix))); - } - Err(httparse::Error::TooManyHeaders) => return Err(RequestHeadError::TooLarge), - Err(_) => return Err(RequestHeadError::Malformed), - } - } -} diff --git a/litebox_egress_proxy/src/stream.rs b/litebox_egress_proxy/src/stream.rs index c48807c1e..2a189731c 100644 --- a/litebox_egress_proxy/src/stream.rs +++ b/litebox_egress_proxy/src/stream.rs @@ -8,181 +8,52 @@ use core::pin::Pin; use core::task::{Context, Poll}; use core::time::Duration; use std::io; -use std::sync::{Arc, Mutex}; -use bytes::Bytes; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -use tokio::net::TcpStream; -use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; use tokio::time::{Instant, Sleep, sleep_until}; -/// A clonable handle to the read half of a TCP stream. -/// -/// The proxy retains one handle so it can perform a bounded drain after Hyper -/// flushes a non-upgraded response and releases its stream. -#[derive(Clone)] -pub(crate) struct SharedTcpRead { - inner: Arc>, -} - -impl AsyncRead for SharedTcpRead { - fn poll_read( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - let Ok(mut inner) = self.inner.lock() else { - return Poll::Ready(Err(io::Error::other("TCP read half mutex poisoned"))); - }; - Pin::new(&mut *inner).poll_read(cx, buf) - } -} - -/// A split TCP stream whose read half can be retained for bounded closing. -pub(crate) struct SharedTcpStream { - read: SharedTcpRead, - write: OwnedWriteHalf, -} - -/// Splits a stream while retaining a clonable handle to its read half. -pub(crate) fn share_tcp_read(stream: TcpStream) -> (SharedTcpStream, SharedTcpRead) { - let (read, write) = stream.into_split(); - let read = SharedTcpRead { - inner: Arc::new(Mutex::new(read)), - }; - ( - SharedTcpStream { - read: read.clone(), - write, - }, - read, - ) -} - -impl AsyncRead for SharedTcpStream { - fn poll_read( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - Pin::new(&mut self.get_mut().read).poll_read(cx, buf) - } -} - -impl AsyncWrite for SharedTcpStream { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - Pin::new(&mut self.get_mut().write).poll_write(cx, buf) - } - - fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().write).poll_flush(cx) - } - - fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().write).poll_shutdown(cx) - } -} - -/// A stream that replays a prefix before reading from its inner stream. -pub struct PrefixedStream { - prefix: Bytes, - inner: S, -} - -impl PrefixedStream { - /// Creates a stream that yields `prefix` before bytes from `inner`. - pub fn new(prefix: Bytes, inner: S) -> Self { - Self { prefix, inner } - } -} - -impl AsyncRead for PrefixedStream { - fn poll_read( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - let this = self.get_mut(); - if !this.prefix.is_empty() && buf.remaining() != 0 { - let length = this.prefix.len().min(buf.remaining()); - let bytes = this.prefix.split_to(length); - buf.put_slice(&bytes); - return Poll::Ready(Ok(())); - } - Pin::new(&mut this.inner).poll_read(cx, buf) - } -} - -impl AsyncWrite for PrefixedStream { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) - } - - fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_flush(cx) - } - - fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) - } -} - -/// A stream that fails once it stalls for too long, or outlives its deadline. -pub struct LimitedStream { +/// A stream that fails after making no progress for its idle timeout. +pub(crate) struct LimitedStream { inner: S, idle: Duration, - idle_timer: Pin>, + timer: Pin>, } impl LimitedStream { - /// Wraps `inner` with an idle timeout. - pub fn new(inner: S, idle: Duration) -> Self { + pub(crate) fn new(inner: S, idle: Duration) -> Self { Self { inner, idle, - idle_timer: Box::pin(sleep_until(Instant::now() + idle)), + timer: Box::pin(sleep_until(Instant::now() + idle)), } } - /// Restarts the idle timeout after observable progress. fn touch(&mut self) { - let deadline = Instant::now() + self.idle; - self.idle_timer.as_mut().reset(deadline); + self.timer.as_mut().reset(Instant::now() + self.idle); } - /// Returns `true` when the stream has been idle for too long. - fn idle_expired(&mut self, cx: &mut Context<'_>) -> bool { - self.idle_timer.as_mut().poll(cx).is_ready() + fn expired(&mut self, cx: &mut Context<'_>) -> bool { + self.timer.as_mut().poll(cx).is_ready() } } -fn timed_out(reason: &'static str) -> io::Error { - io::Error::new(io::ErrorKind::TimedOut, reason) +fn timed_out() -> io::Error { + io::Error::new(io::ErrorKind::TimedOut, "stream idle timeout exceeded") } impl AsyncRead for LimitedStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, + buffer: &mut ReadBuf<'_>, ) -> Poll> { let this = self.get_mut(); - match Pin::new(&mut this.inner).poll_read(cx, buf) { + match Pin::new(&mut this.inner).poll_read(cx, buffer) { Poll::Ready(result) => { this.touch(); Poll::Ready(result) } - Poll::Pending if this.idle_expired(cx) => { - Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) - } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), Poll::Pending => Poll::Pending, } } @@ -192,17 +63,15 @@ impl AsyncWrite for LimitedStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, - buf: &[u8], + buffer: &[u8], ) -> Poll> { let this = self.get_mut(); - match Pin::new(&mut this.inner).poll_write(cx, buf) { + match Pin::new(&mut this.inner).poll_write(cx, buffer) { Poll::Ready(result) => { this.touch(); Poll::Ready(result) } - Poll::Pending if this.idle_expired(cx) => { - Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) - } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), Poll::Pending => Poll::Pending, } } @@ -214,9 +83,7 @@ impl AsyncWrite for LimitedStream { this.touch(); Poll::Ready(result) } - Poll::Pending if this.idle_expired(cx) => { - Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) - } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), Poll::Pending => Poll::Pending, } } @@ -228,9 +95,7 @@ impl AsyncWrite for LimitedStream { this.touch(); Poll::Ready(result) } - Poll::Pending if this.idle_expired(cx) => { - Poll::Ready(Err(timed_out("stream idle timeout exceeded"))) - } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), Poll::Pending => Poll::Pending, } } @@ -242,22 +107,12 @@ mod tests { use tokio::io::{AsyncReadExt, duplex}; - fn runtime() -> tokio::runtime::Runtime { - tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap() - } - - #[test] - fn idle_stream_times_out() { - runtime().block_on(async { - tokio::time::pause(); - let (client, _server) = duplex(64); - let mut limited = LimitedStream::new(client, Duration::from_secs(60)); - let mut buffer = [0_u8; 8]; - let error = limited.read(&mut buffer).await.unwrap_err(); - assert_eq!(error.kind(), io::ErrorKind::TimedOut); - }); + #[tokio::test(start_paused = true)] + async fn idle_stream_times_out() { + let (client, _server) = duplex(64); + let mut limited = LimitedStream::new(client, Duration::from_secs(60)); + let mut buffer = [0_u8; 8]; + let error = limited.read(&mut buffer).await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); } } diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs index bdffe0473..50a84f869 100644 --- a/litebox_egress_proxy/tests/loopback.rs +++ b/litebox_egress_proxy/tests/loopback.rs @@ -59,17 +59,14 @@ impl TestProxy { routes: mapped, attempts: Arc::clone(&attempts), }; - let state = Arc::new(ProxyState::new(policy, Arc::new(connector))); + let state = Arc::new(ProxyState::new(policy, Box::new(connector))); - let listener = acquire(ListenerSource::Bind(SocketAddrV4::new( + let (listener, address) = acquire(ListenerSource::Bind(SocketAddrV4::new( Ipv4Addr::LOCALHOST, 0, ))) .expect("loopback listener"); let listener = TcpListener::from_std(listener).expect("async listener"); - let SocketAddr::V4(address) = listener.local_addr().expect("listener address") else { - panic!("expected an IPv4 listener"); - }; tokio::spawn(async move { let _ = serve(listener, state).await; @@ -273,6 +270,7 @@ async fn connect_authority_and_framing_are_validated() { "CONNECT allowed.example HTTP/1.1\r\nHost: allowed.example\r\n\r\n", "CONNECT 93.184.216.1:443 HTTP/1.1\r\nHost: 93.184.216.1:443\r\n\r\n", "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:80\r\n\r\n", + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nHost: allowed.example:443\r\n\r\n", "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nContent-Length: 0\r\n\r\n", "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nTransfer-Encoding: chunked\r\n\r\n", ] { @@ -308,29 +306,6 @@ async fn unreachable_upstream_yields_bad_gateway() { assert_eq!(proxy.upstream_attempts(), 1); } -#[tokio::test] -async fn denied_connect_early_bytes_are_drained_before_close() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:443"], - &[("allowed.example", 443, upstream)], - ); - - let early = "x".repeat(32 * 1024); - let request = format!( - "CONNECT denied.example:443 HTTP/1.1\r\n\ - Host: denied.example:443\r\n\ - \r\n\ - {early}" - ); - - let mut client = proxy.connect().await; - client.send(request.as_bytes()).await; - assert_eq!(client.read_response().await.status, 403); - assert!(!client.fill().await); - assert_eq!(proxy.upstream_attempts(), 0); -} - #[tokio::test] async fn unsupported_methods_are_rejected_without_network_activity() { let upstream = echo_upstream().await; @@ -345,36 +320,3 @@ async fn unsupported_methods_are_rejected_without_network_activity() { assert_eq!(response.status, 501); assert_eq!(proxy.upstream_attempts(), 0); } - -#[tokio::test] -async fn malformed_and_oversized_heads_are_rejected() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:443"], - &[("allowed.example", 443, upstream)], - ); - - let spaced = proxy - .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost : allowed.example:443\r\n\r\n") - .await; - assert_eq!(spaced.status, 400); - - let folded = proxy - .request(concat!( - "CONNECT allowed.example:443 HTTP/1.1\r\n", - "Host: allowed.example:443\r\n", - "X-Folded: one\r\n two\r\n", - "\r\n" - )) - .await; - assert_eq!(folded.status, 400); - - let mut oversized = String::from( - "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nX-Big: ", - ); - oversized.push_str(&"a".repeat(32 * 1024)); - oversized.push_str("\r\n\r\n"); - assert_eq!(proxy.request(&oversized).await.status, 431); - - assert_eq!(proxy.upstream_attempts(), 0); -} From 7062a887aedf8cf3af0f03a67068d9606ce9df23 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 12:56:34 -0700 Subject: [PATCH 04/11] Use only self-bound proxy listeners Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- Cargo.lock | 1 - litebox_egress_proxy/Cargo.toml | 3 - litebox_egress_proxy/src/config.rs | 49 ++------ litebox_egress_proxy/src/lib.rs | 20 +-- litebox_egress_proxy/src/listener.rs | 166 ------------------------- litebox_egress_proxy/tests/loopback.rs | 35 +++--- 6 files changed, 32 insertions(+), 242 deletions(-) delete mode 100644 litebox_egress_proxy/src/listener.rs diff --git a/Cargo.lock b/Cargo.lock index 686475231..eccc4f5a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1547,7 +1547,6 @@ dependencies = [ "http-body-util", "hyper", "hyper-util", - "libc", "thiserror", "tokio", ] diff --git a/litebox_egress_proxy/Cargo.toml b/litebox_egress_proxy/Cargo.toml index b39bd65e8..542e5b058 100644 --- a/litebox_egress_proxy/Cargo.toml +++ b/litebox_egress_proxy/Cargo.toml @@ -42,8 +42,5 @@ tokio = { version = "1.50", default-features = false, features = [ "time", ] } -[target.'cfg(target_os = "linux")'.dependencies] -libc = { version = "0.2", default-features = false } - [lints] workspace = true diff --git a/litebox_egress_proxy/src/config.rs b/litebox_egress_proxy/src/config.rs index aa70f8fdf..ce9efa44e 100644 --- a/litebox_egress_proxy/src/config.rs +++ b/litebox_egress_proxy/src/config.rs @@ -5,27 +5,21 @@ use std::net::{Ipv4Addr, SocketAddrV4}; -use clap::{ArgGroup, Parser}; +use clap::Parser; use thiserror::Error; -use crate::listener::ListenerSource; use crate::policy::{HostPolicy, PolicyError}; /// The standalone egress proxy for LiteBox sandboxes. #[derive(Debug, Parser)] #[command( name = "litebox_egress_proxy", - about = "Hostname-filtering CONNECT egress proxy", - group(ArgGroup::new("listener").required(true).args(["listen", "listener_fd"])) + about = "Hostname-filtering CONNECT egress proxy" )] pub struct Cli { /// Loopback address to bind, for example `127.0.0.1:0`. #[arg(long, value_name = "IPV4:PORT")] - listen: Option, - - /// Inherited, already-bound loopback listener descriptor. - #[arg(long, value_name = "FD", conflicts_with = "listen")] - listener_fd: Option, + listen: String, /// Allowed hostname and destination ports, repeatable. #[arg(long = "allow-host", value_name = "HOST:PORT[-PORT]")] @@ -49,8 +43,8 @@ pub enum ConfigError { /// The validated configuration of one proxy process. #[derive(Clone, Debug)] pub struct ProxyConfig { - /// Where the listener comes from. - pub listener: ListenerSource, + /// IPv4 loopback address to bind. + pub listen: SocketAddrV4, /// The immutable hostname policy. pub policy: HostPolicy, } @@ -58,14 +52,10 @@ pub struct ProxyConfig { impl Cli { /// Converts parsed arguments into a validated configuration. pub fn into_config(self) -> Result { - let listener = match (self.listen, self.listener_fd) { - (Some(address), _) => ListenerSource::Bind(parse_listen_address(&address)?), - (None, Some(descriptor)) => ListenerSource::Inherited(descriptor), - (None, None) => return Err(ConfigError::ListenAddress), - }; + let listen = parse_listen_address(&self.listen)?; let policy = HostPolicy::from_rules(&self.allow_host)?; - Ok(ProxyConfig { listener, policy }) + Ok(ProxyConfig { listen, policy }) } } @@ -101,10 +91,7 @@ mod tests { ]) .unwrap(); - assert_eq!( - config.listener, - ListenerSource::Bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)) - ); + assert_eq!(config.listen, SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)); let host = Hostname::parse("example.com").unwrap(); assert!(config.policy.allows(&host, 443)); @@ -113,25 +100,7 @@ mod tests { } #[test] - fn parses_an_inherited_listener_configuration() { - let config = parse(&["--listener-fd", "7"]).unwrap(); - let host = Hostname::parse("example.com").unwrap(); - assert_eq!(config.listener, ListenerSource::Inherited(7)); - assert!(!config.policy.allows(&host, 443)); - } - - #[test] - fn listener_modes_are_mutually_exclusive_and_required() { - assert!( - Cli::try_parse_from([ - "litebox_egress_proxy", - "--listen", - "127.0.0.1:0", - "--listener-fd", - "3", - ]) - .is_err() - ); + fn listen_is_required() { assert!(Cli::try_parse_from(["litebox_egress_proxy"]).is_err()); } diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index f47e1f895..463ea691c 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -11,7 +11,6 @@ extern crate alloc; pub mod authority; pub mod config; -pub mod listener; pub mod policy; pub mod proxy; pub mod upstream; @@ -21,29 +20,16 @@ mod stream; use std::io::{self, Write}; use std::sync::Arc; -use thiserror::Error; use tokio::net::TcpListener; use crate::config::ProxyConfig; -use crate::listener::ListenerError; use crate::proxy::ProxyState; use crate::upstream::TcpUpstreamConnector; -/// Reason the proxy could not start. -#[derive(Debug, Error)] -pub enum StartupError { - /// The listener could not be acquired or validated. - #[error(transparent)] - Listener(#[from] ListenerError), - /// An I/O operation failed during startup or while serving. - #[error(transparent)] - Io(#[from] io::Error), -} - /// Runs the proxy: startup, readiness announcement, then serving. -pub async fn run(config: &ProxyConfig) -> Result<(), StartupError> { - let (listener, local_address) = listener::acquire(config.listener)?; - let listener = TcpListener::from_std(listener)?; +pub async fn run(config: &ProxyConfig) -> io::Result<()> { + let listener = TcpListener::bind(config.listen).await?; + let local_address = listener.local_addr()?; let state = Arc::new(ProxyState::new( config.policy.clone(), Box::new(TcpUpstreamConnector), diff --git a/litebox_egress_proxy/src/listener.rs b/litebox_egress_proxy/src/listener.rs deleted file mode 100644 index 0a8b33528..000000000 --- a/litebox_egress_proxy/src/listener.rs +++ /dev/null @@ -1,166 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -//! Loopback listener acquisition. - -use std::io; -use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener}; - -use thiserror::Error; - -/// Where the proxy's listener comes from. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ListenerSource { - /// Bind a fresh loopback listener. Port zero requests an ephemeral port. - Bind(SocketAddrV4), - /// Adopt an inherited listener by descriptor number. - Inherited(i32), -} - -/// Reason a listener could not be acquired. -#[derive(Debug, Error)] -pub enum ListenerError { - /// A listener operation failed. - #[error("listener operation failed: {0}")] - Io(#[from] io::Error), - /// The listener was not a bound IPv4 loopback TCP listener. - #[error("listener must be a bound IPv4 loopback TCP listener")] - Invalid, - /// This platform cannot adopt inherited listeners. - #[error("inherited listeners are unsupported on this platform")] - Unsupported, -} - -/// Acquires a nonblocking IPv4 loopback listener and its bound address. -pub fn acquire(source: ListenerSource) -> Result<(TcpListener, SocketAddrV4), ListenerError> { - let listener = match source { - ListenerSource::Bind(address) => TcpListener::bind(address)?, - ListenerSource::Inherited(descriptor) => adopt_inherited(descriptor)?, - }; - - let SocketAddr::V4(address) = listener.local_addr()? else { - return Err(ListenerError::Invalid); - }; - if *address.ip() != Ipv4Addr::LOCALHOST || address.port() == 0 { - return Err(ListenerError::Invalid); - } - - listener.set_nonblocking(true)?; - Ok((listener, address)) -} - -#[cfg(target_os = "linux")] -fn adopt_inherited(descriptor: i32) -> Result { - use std::os::fd::FromRawFd; - - if descriptor < 0 - || socket_option(descriptor, libc::SO_DOMAIN)? != libc::AF_INET - || socket_option(descriptor, libc::SO_TYPE)? != libc::SOCK_STREAM - || socket_option(descriptor, libc::SO_ACCEPTCONN)? != 1 - { - return Err(ListenerError::Invalid); - } - - // SAFETY: the checks above establish that `descriptor` is an open, - // listening IPv4 stream socket. The launcher transfers ownership of the - // descriptor to this process, so the returned listener is its sole owner. - Ok(unsafe { TcpListener::from_raw_fd(descriptor) }) -} - -#[cfg(target_os = "linux")] -fn socket_option(descriptor: i32, option: libc::c_int) -> Result { - let mut value: libc::c_int = 0; - let mut length = libc::socklen_t::try_from(size_of::()) - .expect("the size of a C int fits in socklen_t"); - - // SAFETY: `value` and `length` are valid, correctly sized and aligned - // locals. `getsockopt` writes at most `length` bytes into `value`. - let result = unsafe { - libc::getsockopt( - descriptor, - libc::SOL_SOCKET, - option, - std::ptr::from_mut(&mut value).cast::(), - &raw mut length, - ) - }; - if result == 0 { - Ok(value) - } else { - Err(io::Error::last_os_error().into()) - } -} - -#[cfg(not(target_os = "linux"))] -fn adopt_inherited(_descriptor: i32) -> Result { - Err(ListenerError::Unsupported) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn binds_an_ephemeral_loopback_port() { - let (_listener, address) = acquire(ListenerSource::Bind(SocketAddrV4::new( - Ipv4Addr::LOCALHOST, - 0, - ))) - .unwrap(); - assert_eq!(*address.ip(), Ipv4Addr::LOCALHOST); - assert_ne!(address.port(), 0); - } - - #[test] - fn rejects_non_loopback_binds() { - let error = acquire(ListenerSource::Bind(SocketAddrV4::new( - Ipv4Addr::UNSPECIFIED, - 0, - ))) - .unwrap_err(); - assert!(matches!(error, ListenerError::Invalid)); - } - - #[test] - fn rejects_a_negative_descriptor() { - let error = acquire(ListenerSource::Inherited(-1)).unwrap_err(); - assert!(matches!( - error, - ListenerError::Invalid | ListenerError::Unsupported - )); - } - - #[cfg(target_os = "linux")] - #[test] - fn adopts_an_inherited_loopback_listener() { - use std::os::fd::IntoRawFd; - - let bound = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); - let expected = bound.local_addr().unwrap(); - let descriptor = bound.into_raw_fd(); - - let (adopted, _) = acquire(ListenerSource::Inherited(descriptor)).unwrap(); - assert_eq!(adopted.local_addr().unwrap(), expected); - } - - #[cfg(target_os = "linux")] - #[test] - fn rejects_a_connected_socket() { - use std::os::fd::IntoRawFd; - - let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); - let client = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap(); - let error = acquire(ListenerSource::Inherited(client.into_raw_fd())).unwrap_err(); - assert!(matches!(error, ListenerError::Invalid)); - } - - #[cfg(target_os = "linux")] - #[test] - fn rejects_a_datagram_socket() { - use std::os::fd::IntoRawFd; - - let socket = std::net::UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); - let error = acquire(ListenerSource::Inherited(socket.into_raw_fd())).unwrap_err(); - assert!(matches!(error, ListenerError::Invalid)); - } -} diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs index 50a84f869..a17134de2 100644 --- a/litebox_egress_proxy/tests/loopback.rs +++ b/litebox_egress_proxy/tests/loopback.rs @@ -10,7 +10,6 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use litebox_egress_proxy::listener::{ListenerSource, acquire}; use litebox_egress_proxy::policy::{HostPolicy, Hostname}; use litebox_egress_proxy::proxy::{ProxyState, serve}; use litebox_egress_proxy::upstream::{BoxedUpstreamStream, ConnectFuture, UpstreamConnector}; @@ -45,7 +44,7 @@ struct TestProxy { } impl TestProxy { - fn start(rules: &[&str], routes: &[(&str, u16, SocketAddr)]) -> Self { + async fn start(rules: &[&str], routes: &[(&str, u16, SocketAddr)]) -> Self { let policy = HostPolicy::from_rules(rules.iter().copied()).expect("valid policy"); let mut mapped = HashMap::new(); @@ -61,12 +60,12 @@ impl TestProxy { }; let state = Arc::new(ProxyState::new(policy, Box::new(connector))); - let (listener, address) = acquire(ListenerSource::Bind(SocketAddrV4::new( - Ipv4Addr::LOCALHOST, - 0, - ))) - .expect("loopback listener"); - let listener = TcpListener::from_std(listener).expect("async listener"); + let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)) + .await + .expect("loopback listener"); + let SocketAddr::V4(address) = listener.local_addr().expect("listener address") else { + panic!("expected IPv4 loopback"); + }; tokio::spawn(async move { let _ = serve(listener, state).await; @@ -190,7 +189,8 @@ async fn connect_tunnel_relays_bytes() { let proxy = TestProxy::start( &["allowed.example:443"], &[("allowed.example", 443, upstream)], - ); + ) + .await; let mut client = proxy.connect().await; client @@ -225,7 +225,8 @@ async fn allowed_hostname_is_connected_for_each_request() { let proxy = TestProxy::start( &["allowed.example:443"], &[("allowed.example", 443, upstream)], - ); + ) + .await; for _ in 0..2 { let response = proxy @@ -243,7 +244,8 @@ async fn denied_host_and_port_do_not_trigger_network_activity() { let proxy = TestProxy::start( &["allowed.example:443"], &[("allowed.example", 443, upstream)], - ); + ) + .await; let denied_host = proxy .request("CONNECT denied.example:443 HTTP/1.1\r\nHost: denied.example:443\r\n\r\n") @@ -264,7 +266,8 @@ async fn connect_authority_and_framing_are_validated() { let proxy = TestProxy::start( &["allowed.example:443"], &[("allowed.example", 443, upstream)], - ); + ) + .await; for request in [ "CONNECT allowed.example HTTP/1.1\r\nHost: allowed.example\r\n\r\n", @@ -286,7 +289,8 @@ async fn explicitly_allowed_dns_port_is_forwarded() { let proxy = TestProxy::start( &["allowed.example:53"], &[("allowed.example", 53, upstream)], - ); + ) + .await; let response = proxy .request("CONNECT allowed.example:53 HTTP/1.1\r\nHost: allowed.example:53\r\n\r\n") @@ -297,7 +301,7 @@ async fn explicitly_allowed_dns_port_is_forwarded() { #[tokio::test] async fn unreachable_upstream_yields_bad_gateway() { - let proxy = TestProxy::start(&["allowed.example:443"], &[]); + let proxy = TestProxy::start(&["allowed.example:443"], &[]).await; let response = proxy .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\n\r\n") @@ -312,7 +316,8 @@ async fn unsupported_methods_are_rejected_without_network_activity() { let proxy = TestProxy::start( &["allowed.example:443"], &[("allowed.example", 443, upstream)], - ); + ) + .await; let response = proxy .request("GET http://allowed.example/ HTTP/1.1\r\nHost: allowed.example\r\n\r\n") From 2c656a9d6a859e17d8ffacc652ee925113700e9d Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 13:11:33 -0700 Subject: [PATCH 05/11] Trim proxy dependency features Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- Cargo.lock | 1 - litebox_egress_proxy/Cargo.toml | 9 +-------- litebox_egress_proxy/src/proxy.rs | 3 +-- 3 files changed, 2 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index eccc4f5a1..0ea43ff1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1541,7 +1541,6 @@ dependencies = [ name = "litebox_egress_proxy" version = "0.1.0" dependencies = [ - "bytes", "clap", "hashbrown", "http-body-util", diff --git a/litebox_egress_proxy/Cargo.toml b/litebox_egress_proxy/Cargo.toml index 542e5b058..2defd9c33 100644 --- a/litebox_egress_proxy/Cargo.toml +++ b/litebox_egress_proxy/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition = "2024" [dependencies] -bytes = { version = "1.10", default-features = false, features = ["std"] } clap = { version = "4.5", default-features = false, features = [ "derive", "error-context", @@ -21,7 +20,7 @@ hyper = { version = "1.8", default-features = false, features = [ hyper-util = { version = "0.1.20", default-features = false, features = [ "tokio", ] } -thiserror = { version = "2.0", default-features = false, features = ["std"] } +thiserror = { version = "2.0", default-features = false } tokio = { version = "1.50", default-features = false, features = [ "io-util", "net", @@ -32,14 +31,8 @@ tokio = { version = "1.50", default-features = false, features = [ [dev-dependencies] tokio = { version = "1.50", default-features = false, features = [ - "io-util", "macros", - "net", - "rt", - "rt-multi-thread", - "sync", "test-util", - "time", ] } [lints] diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index e97444188..8edd8aa8c 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -12,9 +12,8 @@ use core::time::Duration; use std::io; use std::sync::Arc; -use bytes::Bytes; use http_body_util::Empty; -use hyper::body::Incoming; +use hyper::body::{Bytes, Incoming}; use hyper::header::{self, HeaderValue}; use hyper::http::uri::Authority; use hyper::server::conn::http1 as server_http1; From fa9741f8c389109a27d3f5082821248c0441258f Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 13:38:16 -0700 Subject: [PATCH 06/11] Simplify CONNECT proxy tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/config.rs | 5 - litebox_egress_proxy/tests/loopback.rs | 147 ++++++++----------------- 2 files changed, 45 insertions(+), 107 deletions(-) diff --git a/litebox_egress_proxy/src/config.rs b/litebox_egress_proxy/src/config.rs index ce9efa44e..8430ed9e8 100644 --- a/litebox_egress_proxy/src/config.rs +++ b/litebox_egress_proxy/src/config.rs @@ -99,11 +99,6 @@ mod tests { assert!(!config.policy.allows(&host, 80)); } - #[test] - fn listen_is_required() { - assert!(Cli::try_parse_from(["litebox_egress_proxy"]).is_err()); - } - #[test] fn rejects_non_loopback_listen_addresses() { assert!(matches!( diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs index a17134de2..7eadd9774 100644 --- a/litebox_egress_proxy/tests/loopback.rs +++ b/litebox_egress_proxy/tests/loopback.rs @@ -89,7 +89,7 @@ impl TestProxy { } } - async fn request(&self, raw: &str) -> HttpResponse { + async fn request(&self, raw: &str) -> u16 { let mut client = self.connect().await; client.send(raw.as_bytes()).await; client.read_response().await @@ -122,7 +122,7 @@ impl ProxyClient { true } - async fn read_response(&mut self) -> HttpResponse { + async fn read_response(&mut self) -> u16 { let head_end = loop { if let Some(index) = find_subslice(&self.buffer, b"\r\n\r\n") { break index + 4; @@ -135,12 +135,10 @@ impl ProxyClient { let head = String::from_utf8(self.buffer[..head_end].to_vec()).expect("ASCII head"); self.buffer.drain(..head_end); - let status = head - .split_whitespace() + head.split_whitespace() .nth(1) .and_then(|code| code.parse::().ok()) - .expect("status code"); - HttpResponse { status, head } + .expect("status code") } async fn read_exact(&mut self, length: usize) -> Vec { @@ -154,11 +152,6 @@ impl ProxyClient { } } -struct HttpResponse { - status: u16, - head: String, -} - fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { haystack .windows(needle.len()) @@ -205,14 +198,7 @@ async fn connect_tunnel_relays_bytes() { ) .await; - let response = client.read_response().await; - assert_eq!(response.status, 200); - assert!( - !response - .head - .to_ascii_lowercase() - .contains("connection: close") - ); + assert_eq!(client.read_response().await, 200); assert_eq!(client.read_exact(5).await, b"early"); client.send(b"tunnelled").await; @@ -220,26 +206,7 @@ async fn connect_tunnel_relays_bytes() { } #[tokio::test] -async fn allowed_hostname_is_connected_for_each_request() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:443"], - &[("allowed.example", 443, upstream)], - ) - .await; - - for _ in 0..2 { - let response = proxy - .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\n\r\n") - .await; - assert_eq!(response.status, 200); - } - - assert_eq!(proxy.upstream_attempts(), 2); -} - -#[tokio::test] -async fn denied_host_and_port_do_not_trigger_network_activity() { +async fn firewall_rejects_without_network_activity() { let upstream = echo_upstream().await; let proxy = TestProxy::start( &["allowed.example:443"], @@ -247,58 +214,50 @@ async fn denied_host_and_port_do_not_trigger_network_activity() { ) .await; - let denied_host = proxy - .request("CONNECT denied.example:443 HTTP/1.1\r\nHost: denied.example:443\r\n\r\n") - .await; - assert_eq!(denied_host.status, 403); - - let denied_port = proxy - .request("CONNECT allowed.example:8443 HTTP/1.1\r\nHost: allowed.example:8443\r\n\r\n") - .await; - assert_eq!(denied_port.status, 403); - - assert_eq!(proxy.upstream_attempts(), 0); -} - -#[tokio::test] -async fn connect_authority_and_framing_are_validated() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:443"], - &[("allowed.example", 443, upstream)], - ) - .await; - - for request in [ - "CONNECT allowed.example HTTP/1.1\r\nHost: allowed.example\r\n\r\n", - "CONNECT 93.184.216.1:443 HTTP/1.1\r\nHost: 93.184.216.1:443\r\n\r\n", - "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:80\r\n\r\n", - "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nHost: allowed.example:443\r\n\r\n", - "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nContent-Length: 0\r\n\r\n", - "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nTransfer-Encoding: chunked\r\n\r\n", + for (request, expected_status) in [ + ( + "CONNECT denied.example:443 HTTP/1.1\r\nHost: denied.example:443\r\n\r\n", + 403, + ), + ( + "CONNECT allowed.example:8443 HTTP/1.1\r\nHost: allowed.example:8443\r\n\r\n", + 403, + ), + ( + "CONNECT allowed.example HTTP/1.1\r\nHost: allowed.example\r\n\r\n", + 400, + ), + ( + "CONNECT 93.184.216.1:443 HTTP/1.1\r\nHost: 93.184.216.1:443\r\n\r\n", + 400, + ), + ( + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:80\r\n\r\n", + 400, + ), + ( + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nHost: allowed.example:443\r\n\r\n", + 400, + ), + ( + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nContent-Length: 0\r\n\r\n", + 400, + ), + ( + "CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\nTransfer-Encoding: chunked\r\n\r\n", + 400, + ), + ( + "GET http://allowed.example/ HTTP/1.1\r\nHost: allowed.example\r\n\r\n", + 501, + ), ] { - assert_eq!(proxy.request(request).await.status, 400); + assert_eq!(proxy.request(request).await, expected_status, "{request:?}"); } assert_eq!(proxy.upstream_attempts(), 0); } -#[tokio::test] -async fn explicitly_allowed_dns_port_is_forwarded() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:53"], - &[("allowed.example", 53, upstream)], - ) - .await; - - let response = proxy - .request("CONNECT allowed.example:53 HTTP/1.1\r\nHost: allowed.example:53\r\n\r\n") - .await; - assert_eq!(response.status, 200); - assert_eq!(proxy.upstream_attempts(), 1); -} - #[tokio::test] async fn unreachable_upstream_yields_bad_gateway() { let proxy = TestProxy::start(&["allowed.example:443"], &[]).await; @@ -306,22 +265,6 @@ async fn unreachable_upstream_yields_bad_gateway() { let response = proxy .request("CONNECT allowed.example:443 HTTP/1.1\r\nHost: allowed.example:443\r\n\r\n") .await; - assert_eq!(response.status, 502); + assert_eq!(response, 502); assert_eq!(proxy.upstream_attempts(), 1); } - -#[tokio::test] -async fn unsupported_methods_are_rejected_without_network_activity() { - let upstream = echo_upstream().await; - let proxy = TestProxy::start( - &["allowed.example:443"], - &[("allowed.example", 443, upstream)], - ) - .await; - - let response = proxy - .request("GET http://allowed.example/ HTTP/1.1\r\nHost: allowed.example\r\n\r\n") - .await; - assert_eq!(response.status, 501); - assert_eq!(proxy.upstream_attempts(), 0); -} From 96e3efa3b760c16401d9a985d8b4023226cf7f71 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 13:45:09 -0700 Subject: [PATCH 07/11] Rename idle timeout stream Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/proxy.rs | 6 +++--- litebox_egress_proxy/src/stream.rs | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 8edd8aa8c..185ffe18e 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -26,7 +26,7 @@ use tokio::time::timeout; use crate::authority::{RequestAuthority, parse_authority}; use crate::policy::HostPolicy; -use crate::stream::LimitedStream; +use crate::stream::IdleTimeoutStream; use crate::upstream::{BoxedUpstreamStream, UpstreamConnector}; const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; @@ -87,7 +87,7 @@ async fn serve_connection(state: Arc, stream: TcpStream, permit: Own return; } - let io = TokioIo::new(LimitedStream::new(stream, IDLE_TIMEOUT)); + let io = TokioIo::new(IdleTimeoutStream::new(stream, IDLE_TIMEOUT)); let permit = Arc::new(permit); let service = service_fn(move |request: Request| { let state = Arc::clone(&state); @@ -148,7 +148,7 @@ async fn handle_connect( } let upstream = match connect_upstream(state, &authority).await { - Ok(stream) => LimitedStream::new(stream, IDLE_TIMEOUT), + Ok(stream) => IdleTimeoutStream::new(stream, IDLE_TIMEOUT), Err(status) => return status_response(status), }; diff --git a/litebox_egress_proxy/src/stream.rs b/litebox_egress_proxy/src/stream.rs index 2a189731c..38508a569 100644 --- a/litebox_egress_proxy/src/stream.rs +++ b/litebox_egress_proxy/src/stream.rs @@ -13,13 +13,13 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::time::{Instant, Sleep, sleep_until}; /// A stream that fails after making no progress for its idle timeout. -pub(crate) struct LimitedStream { +pub(crate) struct IdleTimeoutStream { inner: S, idle: Duration, timer: Pin>, } -impl LimitedStream { +impl IdleTimeoutStream { pub(crate) fn new(inner: S, idle: Duration) -> Self { Self { inner, @@ -41,7 +41,7 @@ fn timed_out() -> io::Error { io::Error::new(io::ErrorKind::TimedOut, "stream idle timeout exceeded") } -impl AsyncRead for LimitedStream { +impl AsyncRead for IdleTimeoutStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, @@ -59,7 +59,7 @@ impl AsyncRead for LimitedStream { } } -impl AsyncWrite for LimitedStream { +impl AsyncWrite for IdleTimeoutStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, @@ -110,7 +110,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn idle_stream_times_out() { let (client, _server) = duplex(64); - let mut limited = LimitedStream::new(client, Duration::from_secs(60)); + let mut limited = IdleTimeoutStream::new(client, Duration::from_secs(60)); let mut buffer = [0_u8; 8]; let error = limited.read(&mut buffer).await.unwrap_err(); assert_eq!(error.kind(), io::ErrorKind::TimedOut); From e5399c12a68c6f568e03338a139d7893aaa14be3 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 13:49:24 -0700 Subject: [PATCH 08/11] Remove leftover proxy cloning Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/config.rs | 2 +- litebox_egress_proxy/src/lib.rs | 4 ++-- litebox_egress_proxy/src/main.rs | 2 +- litebox_egress_proxy/src/upstream.rs | 1 - 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/litebox_egress_proxy/src/config.rs b/litebox_egress_proxy/src/config.rs index 8430ed9e8..5863c3d99 100644 --- a/litebox_egress_proxy/src/config.rs +++ b/litebox_egress_proxy/src/config.rs @@ -41,7 +41,7 @@ pub enum ConfigError { } /// The validated configuration of one proxy process. -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct ProxyConfig { /// IPv4 loopback address to bind. pub listen: SocketAddrV4, diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index 463ea691c..022f8ee7a 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -27,11 +27,11 @@ use crate::proxy::ProxyState; use crate::upstream::TcpUpstreamConnector; /// Runs the proxy: startup, readiness announcement, then serving. -pub async fn run(config: &ProxyConfig) -> io::Result<()> { +pub async fn run(config: ProxyConfig) -> io::Result<()> { let listener = TcpListener::bind(config.listen).await?; let local_address = listener.local_addr()?; let state = Arc::new(ProxyState::new( - config.policy.clone(), + config.policy, Box::new(TcpUpstreamConnector), )); diff --git a/litebox_egress_proxy/src/main.rs b/litebox_egress_proxy/src/main.rs index a79ae5abd..185fe6586 100644 --- a/litebox_egress_proxy/src/main.rs +++ b/litebox_egress_proxy/src/main.rs @@ -30,7 +30,7 @@ fn main() -> ExitCode { Err(error) => return fail(&error), }; - match runtime.block_on(run(&config)) { + match runtime.block_on(run(config)) { Ok(()) => ExitCode::SUCCESS, Err(error) => fail(&error), } diff --git a/litebox_egress_proxy/src/upstream.rs b/litebox_egress_proxy/src/upstream.rs index 2c3602547..353c690ae 100644 --- a/litebox_egress_proxy/src/upstream.rs +++ b/litebox_egress_proxy/src/upstream.rs @@ -31,7 +31,6 @@ pub trait UpstreamConnector: Send + Sync + 'static { } /// The production connector, using the trusted host resolver. -#[derive(Clone, Copy, Debug, Default)] pub struct TcpUpstreamConnector; impl UpstreamConnector for TcpUpstreamConnector { From 50257280e94f86bb1caf8a382495ad3ad13fb5a3 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 13:58:30 -0700 Subject: [PATCH 09/11] Clarify proxy module names Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/{upstream.rs => connector.rs} | 0 litebox_egress_proxy/src/{stream.rs => idle_timeout.rs} | 0 litebox_egress_proxy/src/lib.rs | 6 +++--- litebox_egress_proxy/src/proxy.rs | 4 ++-- litebox_egress_proxy/tests/loopback.rs | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) rename litebox_egress_proxy/src/{upstream.rs => connector.rs} (100%) rename litebox_egress_proxy/src/{stream.rs => idle_timeout.rs} (100%) diff --git a/litebox_egress_proxy/src/upstream.rs b/litebox_egress_proxy/src/connector.rs similarity index 100% rename from litebox_egress_proxy/src/upstream.rs rename to litebox_egress_proxy/src/connector.rs diff --git a/litebox_egress_proxy/src/stream.rs b/litebox_egress_proxy/src/idle_timeout.rs similarity index 100% rename from litebox_egress_proxy/src/stream.rs rename to litebox_egress_proxy/src/idle_timeout.rs diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index 022f8ee7a..657730c14 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -11,11 +11,11 @@ extern crate alloc; pub mod authority; pub mod config; +pub mod connector; pub mod policy; pub mod proxy; -pub mod upstream; -mod stream; +mod idle_timeout; use std::io::{self, Write}; use std::sync::Arc; @@ -23,8 +23,8 @@ use std::sync::Arc; use tokio::net::TcpListener; use crate::config::ProxyConfig; +use crate::connector::TcpUpstreamConnector; use crate::proxy::ProxyState; -use crate::upstream::TcpUpstreamConnector; /// Runs the proxy: startup, readiness announcement, then serving. pub async fn run(config: ProxyConfig) -> io::Result<()> { diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 185ffe18e..2f8588f65 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -25,9 +25,9 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::time::timeout; use crate::authority::{RequestAuthority, parse_authority}; +use crate::connector::{BoxedUpstreamStream, UpstreamConnector}; +use crate::idle_timeout::IdleTimeoutStream; use crate::policy::HostPolicy; -use crate::stream::IdleTimeoutStream; -use crate::upstream::{BoxedUpstreamStream, UpstreamConnector}; const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; const MAX_REQUEST_HEADER_BYTES: usize = 16 * 1024; diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs index 7eadd9774..277c14aeb 100644 --- a/litebox_egress_proxy/tests/loopback.rs +++ b/litebox_egress_proxy/tests/loopback.rs @@ -10,9 +10,9 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; +use litebox_egress_proxy::connector::{BoxedUpstreamStream, ConnectFuture, UpstreamConnector}; use litebox_egress_proxy::policy::{HostPolicy, Hostname}; use litebox_egress_proxy::proxy::{ProxyState, serve}; -use litebox_egress_proxy::upstream::{BoxedUpstreamStream, ConnectFuture, UpstreamConnector}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::time::timeout; From 8e9c1e7b4cfbca6659921bdb74af0672431685a3 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 14:20:24 -0700 Subject: [PATCH 10/11] Bound upstream address attempts Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/connector.rs | 21 ++++++++++++++++++--- litebox_egress_proxy/src/proxy.rs | 3 +++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/litebox_egress_proxy/src/connector.rs b/litebox_egress_proxy/src/connector.rs index 353c690ae..49aebac1e 100644 --- a/litebox_egress_proxy/src/connector.rs +++ b/litebox_egress_proxy/src/connector.rs @@ -5,13 +5,17 @@ use core::future::Future; use core::pin::Pin; +use core::time::Duration; use std::io; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::net::{TcpStream, lookup_host}; +use tokio::time::timeout; use crate::policy::Hostname; +const UPSTREAM_ADDRESS_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); + /// A bidirectional upstream byte stream. pub trait UpstreamStream: AsyncRead + AsyncWrite + Send + Unpin {} @@ -40,12 +44,23 @@ impl UpstreamConnector for TcpUpstreamConnector { let mut last_error = None; for address in addresses { - match TcpStream::connect(address).await { - Ok(stream) => { + match timeout( + UPSTREAM_ADDRESS_CONNECT_TIMEOUT, + TcpStream::connect(address), + ) + .await + { + Ok(Ok(stream)) => { stream.set_nodelay(true)?; return Ok(Box::new(stream) as BoxedUpstreamStream); } - Err(error) => last_error = Some(error), + Ok(Err(error)) => last_error = Some(error), + Err(_elapsed) => { + last_error = Some(io::Error::new( + io::ErrorKind::TimedOut, + "upstream address connection timed out", + )); + } } } diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 2f8588f65..563716c4c 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -208,6 +208,9 @@ async fn connect_upstream( .await; match result { Ok(Ok(stream)) => Ok(stream), + Ok(Err(error)) if error.kind() == io::ErrorKind::TimedOut => { + Err(StatusCode::GATEWAY_TIMEOUT) + } Ok(Err(_error)) => Err(StatusCode::BAD_GATEWAY), Err(_elapsed) => Err(StatusCode::GATEWAY_TIMEOUT), } From bb2414382dcfcba150a9e5c0b627da96ed781070 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Tue, 1 Sep 2026 14:30:07 -0700 Subject: [PATCH 11/11] Share connect budget across addresses Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 239ec5f7-870a-4259-bcae-4ca85fb913a0 --- litebox_egress_proxy/src/connector.rs | 21 +++++++++++---------- litebox_egress_proxy/src/proxy.rs | 3 +-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litebox_egress_proxy/src/connector.rs b/litebox_egress_proxy/src/connector.rs index 49aebac1e..5c6d4d05a 100644 --- a/litebox_egress_proxy/src/connector.rs +++ b/litebox_egress_proxy/src/connector.rs @@ -10,11 +10,11 @@ use std::io; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::net::{TcpStream, lookup_host}; -use tokio::time::timeout; +use tokio::time::{Instant, timeout}; use crate::policy::Hostname; -const UPSTREAM_ADDRESS_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +pub(crate) const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// A bidirectional upstream byte stream. pub trait UpstreamStream: AsyncRead + AsyncWrite + Send + Unpin {} @@ -40,16 +40,17 @@ pub struct TcpUpstreamConnector; impl UpstreamConnector for TcpUpstreamConnector { fn connect(&self, host: Hostname, port: u16) -> ConnectFuture<'_> { Box::pin(async move { - let addresses = lookup_host((host.as_str(), port)).await?; + let deadline = Instant::now() + UPSTREAM_CONNECT_TIMEOUT; + let addresses: Vec<_> = lookup_host((host.as_str(), port)).await?.collect(); + let address_count = addresses.len(); let mut last_error = None; - for address in addresses { - match timeout( - UPSTREAM_ADDRESS_CONNECT_TIMEOUT, - TcpStream::connect(address), - ) - .await - { + for (index, address) in addresses.into_iter().enumerate() { + let attempts_left = u32::try_from(address_count - index).unwrap_or(u32::MAX); + let attempt_timeout = + deadline.saturating_duration_since(Instant::now()) / attempts_left; + + match timeout(attempt_timeout, TcpStream::connect(address)).await { Ok(Ok(stream)) => { stream.set_nodelay(true)?; return Ok(Box::new(stream) as BoxedUpstreamStream); diff --git a/litebox_egress_proxy/src/proxy.rs b/litebox_egress_proxy/src/proxy.rs index 563716c4c..17f961a19 100644 --- a/litebox_egress_proxy/src/proxy.rs +++ b/litebox_egress_proxy/src/proxy.rs @@ -25,14 +25,13 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::time::timeout; use crate::authority::{RequestAuthority, parse_authority}; -use crate::connector::{BoxedUpstreamStream, UpstreamConnector}; +use crate::connector::{BoxedUpstreamStream, UPSTREAM_CONNECT_TIMEOUT, UpstreamConnector}; use crate::idle_timeout::IdleTimeoutStream; use crate::policy::HostPolicy; const MAX_CONCURRENT_CLIENT_CONNECTIONS: usize = 256; const MAX_REQUEST_HEADER_BYTES: usize = 16 * 1024; const MAX_HEADER_FIELDS: usize = 100; -const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const IDLE_TIMEOUT: Duration = Duration::from_secs(60); const REQUEST_HEADER_READ_TIMEOUT: Duration = Duration::from_secs(30);