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..0ea43ff1b 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,13 @@ dependencies = [ name = "litebox_egress_proxy" version = "0.1.0" dependencies = [ + "clap", "hashbrown", + "http-body-util", + "hyper", + "hyper-util", "thiserror", + "tokio", ] [[package]] diff --git a/litebox_egress_proxy/Cargo.toml b/litebox_egress_proxy/Cargo.toml index b1630eb63..2defd9c33 100644 --- a/litebox_egress_proxy/Cargo.toml +++ b/litebox_egress_proxy/Cargo.toml @@ -4,8 +4,36 @@ version = "0.1.0" edition = "2024" [dependencies] +clap = { version = "4.5", default-features = false, features = [ + "derive", + "error-context", + "help", + "std", + "usage", +] } hashbrown = "0.15.2" +http-body-util = { version = "0.1.3", 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 } +tokio = { version = "1.50", default-features = false, features = [ + "io-util", + "net", + "rt", + "sync", + "time", +] } + +[dev-dependencies] +tokio = { version = "1.50", default-features = false, features = [ + "macros", + "test-util", +] } [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..5863c3d99 --- /dev/null +++ b/litebox_egress_proxy/src/config.rs @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Executable configuration. + +use std::net::{Ipv4Addr, SocketAddrV4}; + +use clap::Parser; +use thiserror::Error; + +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" +)] +pub struct Cli { + /// Loopback address to bind, for example `127.0.0.1:0`. + #[arg(long, value_name = "IPV4:PORT")] + listen: String, + + /// 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(Debug)] +pub struct ProxyConfig { + /// IPv4 loopback address to bind. + pub listen: SocketAddrV4, + /// The immutable hostname policy. + pub policy: HostPolicy, +} + +impl Cli { + /// Converts parsed arguments into a validated configuration. + pub fn into_config(self) -> Result { + let listen = parse_listen_address(&self.listen)?; + let policy = HostPolicy::from_rules(&self.allow_host)?; + + Ok(ProxyConfig { listen, 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.listen, 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 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/connector.rs b/litebox_egress_proxy/src/connector.rs new file mode 100644 index 000000000..5c6d4d05a --- /dev/null +++ b/litebox_egress_proxy/src/connector.rs @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Upstream hostname resolution and connection. + +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::{Instant, timeout}; + +use crate::policy::Hostname; + +pub(crate) const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +/// 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. +pub struct TcpUpstreamConnector; + +impl UpstreamConnector for TcpUpstreamConnector { + fn connect(&self, host: Hostname, port: u16) -> ConnectFuture<'_> { + Box::pin(async move { + 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 (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); + } + Ok(Err(error)) => last_error = Some(error), + Err(_elapsed) => { + last_error = Some(io::Error::new( + io::ErrorKind::TimedOut, + "upstream address connection timed out", + )); + } + } + } + + Err(last_error.unwrap_or_else(|| { + io::Error::new(io::ErrorKind::NotFound, "hostname resolved to no addresses") + })) + }) + } +} diff --git a/litebox_egress_proxy/src/idle_timeout.rs b/litebox_egress_proxy/src/idle_timeout.rs new file mode 100644 index 000000000..38508a569 --- /dev/null +++ b/litebox_egress_proxy/src/idle_timeout.rs @@ -0,0 +1,118 @@ +// 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 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 IdleTimeoutStream { + inner: S, + idle: Duration, + timer: Pin>, +} + +impl IdleTimeoutStream { + pub(crate) fn new(inner: S, idle: Duration) -> Self { + Self { + inner, + idle, + timer: Box::pin(sleep_until(Instant::now() + idle)), + } + } + + fn touch(&mut self) { + self.timer.as_mut().reset(Instant::now() + self.idle); + } + + fn expired(&mut self, cx: &mut Context<'_>) -> bool { + self.timer.as_mut().poll(cx).is_ready() + } +} + +fn timed_out() -> io::Error { + io::Error::new(io::ErrorKind::TimedOut, "stream idle timeout exceeded") +} + +impl AsyncRead for IdleTimeoutStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_read(cx, buffer) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncWrite for IdleTimeoutStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.inner).poll_write(cx, buffer) { + Poll::Ready(result) => { + this.touch(); + Poll::Ready(result) + } + Poll::Pending if this.expired(cx) => Poll::Ready(Err(timed_out())), + 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.expired(cx) => Poll::Ready(Err(timed_out())), + 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.expired(cx) => Poll::Ready(Err(timed_out())), + Poll::Pending => Poll::Pending, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use tokio::io::{AsyncReadExt, duplex}; + + #[tokio::test(start_paused = true)] + async fn idle_stream_times_out() { + let (client, _server) = duplex(64); + 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); + } +} diff --git a/litebox_egress_proxy/src/lib.rs b/litebox_egress_proxy/src/lib.rs index 77af2d3ed..657730c14 100644 --- a/litebox_egress_proxy/src/lib.rs +++ b/litebox_egress_proxy/src/lib.rs @@ -1,15 +1,46 @@ // 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] +//! 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. extern crate alloc; pub mod authority; +pub mod config; +pub mod connector; pub mod policy; +pub mod proxy; + +mod idle_timeout; + +use std::io::{self, Write}; +use std::sync::Arc; + +use tokio::net::TcpListener; + +use crate::config::ProxyConfig; +use crate::connector::TcpUpstreamConnector; +use crate::proxy::ProxyState; + +/// Runs the proxy: startup, readiness announcement, then serving. +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, + Box::new(TcpUpstreamConnector), + )); + + { + let mut stdout = io::stdout().lock(); + writeln!(stdout, "READY {local_address}")?; + stdout.flush()?; + } + + proxy::serve(listener, state).await?; + Ok(()) +} diff --git a/litebox_egress_proxy/src/main.rs b/litebox_egress_proxy/src/main.rs new file mode 100644 index 000000000..185fe6586 --- /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..17f961a19 --- /dev/null +++ b/litebox_egress_proxy/src/proxy.rs @@ -0,0 +1,249 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Connection acceptance, authorization, and CONNECT tunnelling. +//! +//! 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::time::Duration; +use std::io; +use std::sync::Arc; + +use http_body_util::Empty; +use hyper::body::{Bytes, 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::net::{TcpListener, TcpStream}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::time::timeout; + +use crate::authority::{RequestAuthority, parse_authority}; +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 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 = Empty; + +/// Immutable state shared by every connection. +pub struct ProxyState { + policy: HostPolicy, + connector: Box, +} + +impl ProxyState { + /// Builds shared state from a validated policy and upstream connector. + pub fn new(policy: HostPolicy, connector: Box) -> 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 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); + let permit = Arc::clone(&permit); + async move { Ok::<_, Infallible>(handle_request(state, permit, 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); + + if let Err(error) = builder.serve_connection(io, service).with_upgrades().await { + diagnostic(format_args!("client connection ended: {error}")); + } +} + +/// Dispatches one request. +async fn handle_request( + state: Arc, + permit: Arc, + request: Request, +) -> Response { + if request.method() != Method::CONNECT { + return status_response(StatusCode::NOT_IMPLEMENTED); + } + handle_connect(&state, permit, request).await +} + +/// Handles one CONNECT tunnel request. +async fn handle_connect( + state: &ProxyState, + permit: Arc, + mut request: Request, +) -> Response { + if request.headers().contains_key(header::UPGRADE) { + return status_response(StatusCode::NOT_IMPLEMENTED); + } + 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); + } + + 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) => IdleTimeoutStream::new(stream, IDLE_TIMEOUT), + Err(status) => return status_response(status), + }; + + 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::new()); + *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) +} + +/// Resolves and connects to an authorized hostname. +async fn connect_upstream( + state: &ProxyState, + authority: &RequestAuthority, +) -> Result { + let result = timeout( + UPSTREAM_CONNECT_TIMEOUT, + state + .connector + .connect(authority.host().clone(), authority.port()), + ) + .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), + } +} + +fn status_response(status: StatusCode) -> Response { + let mut response = Response::new(Empty::new()); + *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()); + } +} diff --git a/litebox_egress_proxy/tests/loopback.rs b/litebox_egress_proxy/tests/loopback.rs new file mode 100644 index 000000000..277c14aeb --- /dev/null +++ b/litebox_egress_proxy/tests/loopback.rs @@ -0,0 +1,270 @@ +// 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::connector::{BoxedUpstreamStream, ConnectFuture, UpstreamConnector}; +use litebox_egress_proxy::policy::{HostPolicy, Hostname}; +use litebox_egress_proxy::proxy::{ProxyState, serve}; +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 { + 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(); + 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, Box::new(connector))); + + 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; + }); + + 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) -> u16 { + 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) -> u16 { + 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); + head.split_whitespace() + .nth(1) + .and_then(|code| code.parse::().ok()) + .expect("status code") + } + + 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() + } +} + +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)], + ) + .await; + + 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; + + assert_eq!(client.read_response().await, 200); + 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 firewall_rejects_without_network_activity() { + let upstream = echo_upstream().await; + let proxy = TestProxy::start( + &["allowed.example:443"], + &[("allowed.example", 443, upstream)], + ) + .await; + + 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, expected_status, "{request:?}"); + } + + assert_eq!(proxy.upstream_attempts(), 0); +} + +#[tokio::test] +async fn unreachable_upstream_yields_bad_gateway() { + 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") + .await; + assert_eq!(response, 502); + assert_eq!(proxy.upstream_attempts(), 1); +}