From 694c0272cd201581c813cbf3a82edf9a4679dd82 Mon Sep 17 00:00:00 2001 From: Gregory Schier Date: Wed, 14 Oct 2026 13:10:51 -0700 Subject: [PATCH] Proxy engine: persistent CA, capped capture, upgrade tunnel (#750) --- Cargo.lock | 6 + crates-proxy/yaak-proxy-lib/src/lib.rs | 48 ++- crates/yaak-proxy/Cargo.toml | 11 +- crates/yaak-proxy/src/body.rs | 143 +++---- crates/yaak-proxy/src/cert.rs | 180 ++++++-- crates/yaak-proxy/src/connection.rs | 38 +- crates/yaak-proxy/src/lib.rs | 238 ++++++++--- crates/yaak-proxy/src/request.rs | 530 ++++++++++++------------ crates/yaak-proxy/src/tunnel.rs | 85 ++++ crates/yaak-proxy/tests/proxy.rs | 548 +++++++++++++++++++++++++ 10 files changed, 1379 insertions(+), 448 deletions(-) create mode 100644 crates/yaak-proxy/src/tunnel.rs create mode 100644 crates/yaak-proxy/tests/proxy.rs diff --git a/Cargo.lock b/Cargo.lock index fec1893e1..304da4fb3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11494,16 +11494,22 @@ name = "yaak-proxy" version = "0.1.0" dependencies = [ "bytes", + "futures-util", "http", "http-body-util", "hyper", "hyper-util", + "log 0.4.34", "pem", "rcgen", + "reqwest", "rustls", "rustls-native-certs", + "tempfile", + "time", "tokio", "tokio-rustls", + "tokio-tungstenite", ] [[package]] diff --git a/crates-proxy/yaak-proxy-lib/src/lib.rs b/crates-proxy/yaak-proxy-lib/src/lib.rs index f8cace44d..de186e97d 100644 --- a/crates-proxy/yaak-proxy-lib/src/lib.rs +++ b/crates-proxy/yaak-proxy-lib/src/lib.rs @@ -8,11 +8,11 @@ use crate::models::{HttpExchange, ModelPayload, ProxyHeader}; use log::warn; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use ts_rs::TS; use yaak_database::{ModelChangeEvent, UpdateSource}; -use yaak_proxy::{CapturedRequest, ProxyEvent, ProxyHandle, RequestState}; +use yaak_proxy::{CaStore, CapturedRequest, ProxyConfig, ProxyEvent, ProxyHandle, RequestState}; use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc}; // -- Context -- @@ -21,13 +21,20 @@ use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc}; #[derive(Clone)] pub struct ProxyCtx { handle: Arc>>, + ca_dir: PathBuf, pub db: ProxyQueryManager, pub events: RpcEventEmitter, } impl ProxyCtx { pub fn new(db_path: &Path, events: RpcEventEmitter) -> Self { - Self { handle: Arc::new(Mutex::new(None)), db: ProxyQueryManager::new(db_path), events } + let ca_dir = db_path.parent().unwrap_or(Path::new(".")).join("ca"); + Self { + handle: Arc::new(Mutex::new(None)), + ca_dir, + db: ProxyQueryManager::new(db_path), + events, + } } } @@ -93,8 +100,11 @@ fn execute_action(ctx: &ProxyCtx, invocation: ActionInvocation) -> Result { + ProxyEvent::RequestStart { id, method, url, http_version, .. } => { in_flight.insert( id, CapturedRequest { @@ -182,12 +192,15 @@ fn run_event_loop( r.request_headers.push((name, value)); } } - ProxyEvent::RequestBody { id, body } => { - if let Some(r) = in_flight.get_mut(&id) { - r.request_body = Some(body); + ProxyEvent::RequestBodyChunk { id, data, .. } => { + if let Some(r) = in_flight.get_mut(&id) + && !data.is_empty() + { + r.request_body.get_or_insert_default().extend_from_slice(&data); } } - ProxyEvent::ResponseStart { id, status, http_version, elapsed_ms } => { + ProxyEvent::RequestBodyComplete { .. } => {} + ProxyEvent::ResponseStart { id, status, http_version, elapsed_ms, .. } => { if let Some(r) = in_flight.get_mut(&id) { r.status = Some(status); r.remote_http_version = Some(http_version); @@ -200,19 +213,24 @@ fn run_event_loop( r.response_headers.push((name, value)); } } - ProxyEvent::ResponseBodyChunk { .. } => { - // Progress only — no action needed + ProxyEvent::ResponseBodyChunk { id, data, .. } => { + if let Some(r) = in_flight.get_mut(&id) + && !data.is_empty() + { + r.response_body.get_or_insert_default().extend_from_slice(&data); + } } - ProxyEvent::ResponseBodyComplete { id, body, size, elapsed_ms } => { + ProxyEvent::ResponseBodyComplete { id, size, elapsed_ms, error, .. } => { if let Some(mut r) = in_flight.remove(&id) { - r.response_body = body; r.response_body_size = size; r.elapsed_ms = r.elapsed_ms.or(Some(elapsed_ms)); - r.state = RequestState::Complete; + r.state = + if error.is_some() { RequestState::Error } else { RequestState::Complete }; + r.error = error; write_entry(&db, &events, &r); } } - ProxyEvent::Error { id, error } => { + ProxyEvent::Error { id, error, .. } => { if let Some(mut r) = in_flight.remove(&id) { r.error = Some(error); r.state = RequestState::Error; diff --git a/crates/yaak-proxy/Cargo.toml b/crates/yaak-proxy/Cargo.toml index 3cc4e725d..58a52bc50 100644 --- a/crates/yaak-proxy/Cargo.toml +++ b/crates/yaak-proxy/Cargo.toml @@ -6,10 +6,11 @@ publish = false [dependencies] hyper = { version = "1", features = ["http1", "http2", "server", "client"] } -hyper-util = { version = "0.1", features = ["tokio", "server-auto", "client-legacy"] } +hyper-util = { version = "0.1", features = ["tokio", "http1", "client-legacy"] } http-body-util = "0.1" http = "1" bytes = "1" +log = { workspace = true } tokio = { workspace = true, features = [ "rt-multi-thread", "net", @@ -23,3 +24,11 @@ rustls = { workspace = true, features = ["ring"] } rustls-native-certs = "0.8" tokio-rustls = "0.26" pem = "3" +time = "0.3" + +[dev-dependencies] +futures-util = "0.3" +hyper-util = { version = "0.1", features = ["server-auto"] } +reqwest = { workspace = true } +tempfile = "3" +tokio-tungstenite = { version = "0.30.0", default-features = false, features = ["handshake"] } diff --git a/crates/yaak-proxy/src/body.rs b/crates/yaak-proxy/src/body.rs index 5eef207aa..0eae11e7a 100644 --- a/crates/yaak-proxy/src/body.rs +++ b/crates/yaak-proxy/src/body.rs @@ -1,113 +1,116 @@ use std::pin::Pin; -use std::sync::mpsc as std_mpsc; +use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Instant; use bytes::Bytes; -use hyper::body::{Body, Frame}; +use hyper::body::{Body, Frame, Incoming, SizeHint}; -use crate::ProxyEvent; +use crate::{Engine, ProxyEvent}; -/// A body wrapper that emits `ResponseBodyChunk` per frame and -/// `ResponseBodyComplete` when the stream finishes. -pub struct MeasuredBody { - inner: B, - request_id: u64, - bytes_count: u64, - chunks: Vec, - event_tx: std_mpsc::Sender, +#[derive(Clone, Copy, PartialEq)] +pub(crate) enum Direction { + Request, + Response, +} + +/// Streams a body through untouched, reporting each frame as a chunk event (with data up to +/// `capture_body_limit`) and the end as a completion event. +pub(crate) struct CaptureBody { + inner: Incoming, + direction: Direction, + id: u64, start: Instant, + engine: Arc, + captured: usize, + size: u64, finished: bool, } -impl MeasuredBody { - pub fn new( - inner: B, - request_id: u64, +impl CaptureBody { + pub(crate) fn new( + inner: Incoming, + direction: Direction, + id: u64, start: Instant, - event_tx: std_mpsc::Sender, + engine: Arc, ) -> Self { - Self { - inner, - request_id, - bytes_count: 0, - chunks: Vec::new(), - event_tx, - start, - finished: false, - } + Self { inner, direction, id, start, engine, captured: 0, size: 0, finished: false } } - fn send_complete(&mut self) { - if !self.finished { - self.finished = true; - let body = if self.chunks.is_empty() { - None - } else { - let mut buf = Vec::with_capacity(self.bytes_count as usize); - for chunk in self.chunks.drain(..) { - buf.extend_from_slice(&chunk); - } - Some(buf) - }; - let _ = self.event_tx.send(ProxyEvent::ResponseBodyComplete { - id: self.request_id, - body, - size: self.bytes_count, - elapsed_ms: self.start.elapsed().as_millis() as u64, - }); + fn finish(&mut self, error: Option) { + if self.finished { + return; } + self.finished = true; + let truncated = self.size > self.captured as u64; + let event = match self.direction { + Direction::Request if self.size == 0 => return, + Direction::Request => { + ProxyEvent::RequestBodyComplete { id: self.id, size: self.size, truncated } + } + Direction::Response => ProxyEvent::ResponseBodyComplete { + id: self.id, + size: self.size, + truncated, + elapsed_ms: self.start.elapsed().as_millis() as u64, + error, + }, + }; + self.engine.emit(event); } } -impl Body for MeasuredBody -where - B: Body + Unpin, - B::Error: std::error::Error + Send + Sync + 'static, -{ +impl Body for CaptureBody { type Data = Bytes; - type Error = B::Error; + type Error = hyper::Error; fn poll_frame( mut self: Pin<&mut Self>, cx: &mut Context<'_>, - ) -> Poll, Self::Error>>> { - let inner = Pin::new(&mut self.inner); - match inner.poll_frame(cx) { + ) -> Poll, hyper::Error>>> { + let result = Pin::new(&mut self.inner).poll_frame(cx); + match &result { Poll::Ready(Some(Ok(frame))) => { - if let Some(data) = frame.data_ref() { - let len = data.len(); - self.bytes_count += len as u64; - self.chunks.push(data.clone()); - let _ = self - .event_tx - .send(ProxyEvent::ResponseBodyChunk { id: self.request_id, bytes: len }); + if let Some(frame_data) = frame.data_ref() { + let len = frame_data.len(); + let room = self.engine.capture_body_limit.saturating_sub(self.captured); + let data = frame_data.slice(..len.min(room)); + self.captured += data.len(); + self.size += len as u64; + let id = self.id; + self.engine.emit(match self.direction { + Direction::Request => ProxyEvent::RequestBodyChunk { id, len, data }, + Direction::Response => ProxyEvent::ResponseBodyChunk { id, len, data }, + }); + } + if self.inner.is_end_stream() { + self.finish(None); } - Poll::Ready(Some(Ok(frame))) } Poll::Ready(Some(Err(e))) => { - self.send_complete(); - Poll::Ready(Some(Err(e))) + let error = e.to_string(); + self.finish(Some(error)); } - Poll::Ready(None) => { - self.send_complete(); - Poll::Ready(None) - } - Poll::Pending => Poll::Pending, + Poll::Ready(None) => self.finish(None), + Poll::Pending => {} } + result } fn is_end_stream(&self) -> bool { self.inner.is_end_stream() } - fn size_hint(&self) -> hyper::body::SizeHint { + fn size_hint(&self) -> SizeHint { self.inner.size_hint() } } -impl Drop for MeasuredBody { +impl Drop for CaptureBody { fn drop(&mut self) { - self.send_complete(); + let error = (self.direction == Direction::Response && !self.inner.is_end_stream()) + .then(|| "Client closed the connection before the response finished".to_string()); + self.finish(error); } } diff --git a/crates/yaak-proxy/src/cert.rs b/crates/yaak-proxy/src/cert.rs index e9105e02d..cfdd73a44 100644 --- a/crates/yaak-proxy/src/cert.rs +++ b/crates/yaak-proxy/src/cert.rs @@ -1,69 +1,175 @@ use std::collections::HashMap; +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; -use rcgen::{BasicConstraints, Certificate, CertificateParams, IsCa, KeyPair, KeyUsagePurpose}; +use rcgen::{ + BasicConstraints, Certificate, CertificateParams, DnType, ExtendedKeyUsagePurpose, IsCa, + KeyPair, KeyUsagePurpose, +}; use rustls::ServerConfig; use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}; +use time::{Duration, OffsetDateTime}; -pub struct CertificateAuthority { - ca_cert: Certificate, - ca_cert_der: CertificateDer<'static>, - ca_key: KeyPair, - cache: Mutex>>, +const CA_COMMON_NAME: &str = "Yaak Proxy CA"; +const CA_VALIDITY_DAYS: i64 = 3650; +const LEAF_VALIDITY_DAYS: i64 = 365; + +/// The proxy's root CA, kept on disk as `ca.pem` and `ca.key` so trust survives restarts. +pub struct CaStore { + cert_path: PathBuf, + key_path: PathBuf, + pem: String, + cert_der: CertificateDer<'static>, + issuer: Certificate, + key: KeyPair, + leaf_configs: Mutex>>, } -impl CertificateAuthority { - pub fn new() -> Result> { - let mut params = CertificateParams::default(); - params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); - params.key_usages.push(KeyUsagePurpose::KeyCertSign); - params.key_usages.push(KeyUsagePurpose::CrlSign); - params.distinguished_name.push(rcgen::DnType::CommonName, "Debug Proxy CA"); - params.distinguished_name.push(rcgen::DnType::OrganizationName, "Debug Proxy"); +impl CaStore { + /// Loads the CA from `dir`, or creates one there (key mode 0600) if either file is missing. + /// Safe to call from several processes at once: `ca.lock` serializes them. + pub fn load_or_create(dir: &Path) -> io::Result { + let cert_path = dir.join("ca.pem"); + let key_path = dir.join("ca.key"); - let key = KeyPair::generate()?; - let ca_cert = params.self_signed(&key)?; - let ca_cert_der = ca_cert.der().clone(); + fs::create_dir_all(dir)?; + let lock = fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(dir.join("ca.lock"))?; + lock.lock()?; - Ok(Self { ca_cert, ca_cert_der, ca_key: key, cache: Mutex::new(HashMap::new()) }) + let (pem, key) = if cert_path.exists() && key_path.exists() { + let key = + KeyPair::from_pem(&fs::read_to_string(&key_path)?).map_err(io::Error::other)?; + (fs::read_to_string(&cert_path)?, key) + } else { + log::info!("Creating proxy CA in {}", dir.display()); + let key = KeyPair::generate().map_err(io::Error::other)?; + let now = OffsetDateTime::now_utc(); + let mut params = ca_params(); + params.not_before = now - Duration::days(1); + params.not_after = now + Duration::days(CA_VALIDITY_DAYS); + let pem = params.self_signed(&key).map_err(io::Error::other)?.pem(); + remove_if_exists(&cert_path)?; + write_atomic(&key_path, key.serialize_pem().as_bytes(), true)?; + write_atomic(&cert_path, pem.as_bytes(), false)?; + (pem, key) + }; + drop(lock); + + let cert_der = pem::parse(&pem) + .map_err(io::Error::other) + .map(|p| CertificateDer::from(p.into_contents()))?; + let issuer = ca_params().self_signed(&key).map_err(io::Error::other)?; + + Ok(Self { + cert_path, + key_path, + pem, + cert_der, + issuer, + key, + leaf_configs: Mutex::new(HashMap::new()), + }) } - pub fn ca_pem(&self) -> String { - pem::encode(&pem::Pem::new("CERTIFICATE", self.ca_cert_der.to_vec())) + /// The CA certificate as PEM, for installing into a trust store. + pub fn pem(&self) -> &str { + &self.pem } - pub fn server_config( + /// Path of `ca.pem`. + pub fn cert_path(&self) -> &Path { + &self.cert_path + } + + /// Path of `ca.key`. + pub fn key_path(&self) -> &Path { + &self.key_path + } + + pub(crate) fn server_config( &self, - domain: &str, + host: &str, ) -> Result, Box> { - { - let cache = self.cache.lock().unwrap(); - if let Some(config) = cache.get(domain) { - return Ok(config.clone()); - } + if let Some(config) = self.leaf_configs.lock().unwrap().get(host) { + return Ok(config.clone()); } - let mut params = CertificateParams::new(vec![domain.to_string()])?; - params.distinguished_name.push(rcgen::DnType::CommonName, domain); + let now = OffsetDateTime::now_utc(); + let mut params = CertificateParams::new(vec![host.to_string()])?; + params.distinguished_name.push(DnType::CommonName, host); + params.key_usages = vec![ + KeyUsagePurpose::DigitalSignature, + KeyUsagePurpose::KeyEncipherment, + ]; + params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth]; + params.use_authority_key_identifier_extension = true; + params.not_before = now - Duration::days(1); + params.not_after = now + Duration::days(LEAF_VALIDITY_DAYS); let leaf_key = KeyPair::generate()?; - let leaf_cert = params.signed_by(&leaf_key, &self.ca_cert, &self.ca_key)?; - - let cert_der = leaf_cert.der().clone(); - let key_der = leaf_key.serialize_der(); + let leaf_cert = params.signed_by(&leaf_key, &self.issuer, &self.key)?; let mut config = ServerConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) .with_safe_default_protocol_versions()? .with_no_client_auth() .with_single_cert( - vec![cert_der, self.ca_cert_der.clone()], - PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key_der)), + vec![leaf_cert.der().clone(), self.cert_der.clone()], + PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(leaf_key.serialize_der())), )?; - config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + config.alpn_protocols = vec![b"http/1.1".to_vec()]; let config = Arc::new(config); - self.cache.lock().unwrap().insert(domain.to_string(), config.clone()); + self.leaf_configs.lock().unwrap().insert(host.to_string(), config.clone()); Ok(config) } } + +/// Issuer fields for the CA. Leaf certificates are signed against these, so they must stay +/// identical to what was written to `ca.pem`. +fn ca_params() -> CertificateParams { + let mut params = CertificateParams::default(); + params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + params.key_usages = vec![KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign]; + params.distinguished_name.push(DnType::CommonName, CA_COMMON_NAME); + params.distinguished_name.push(DnType::OrganizationName, "Yaak"); + params +} + +fn remove_if_exists(path: &Path) -> io::Result<()> { + match fs::remove_file(path) { + Err(e) if e.kind() != io::ErrorKind::NotFound => Err(e), + _ => Ok(()), + } +} + +/// Writes through a temporary file and a rename, so readers never see a partial file. +fn write_atomic(path: &Path, contents: &[u8], private: bool) -> io::Result<()> { + use std::io::Write; + let mut tmp_name = path.file_name().unwrap_or_default().to_os_string(); + tmp_name.push(".tmp"); + let tmp = path.with_file_name(tmp_name); + + let mut options = fs::OpenOptions::new(); + options.write(true).create(true).truncate(true); + #[cfg(unix)] + if private { + std::os::unix::fs::OpenOptionsExt::mode(&mut options, 0o600); + } + #[cfg(not(unix))] + let _ = private; + let mut file = options.open(&tmp)?; + #[cfg(unix)] + if private { + file.set_permissions(std::os::unix::fs::PermissionsExt::from_mode(0o600))?; + } + file.write_all(contents)?; + file.sync_all()?; + fs::rename(&tmp, path) +} diff --git a/crates/yaak-proxy/src/connection.rs b/crates/yaak-proxy/src/connection.rs index 9719876eb..27776dae0 100644 --- a/crates/yaak-proxy/src/connection.rs +++ b/crates/yaak-proxy/src/connection.rs @@ -1,32 +1,32 @@ use std::sync::Arc; -use std::sync::mpsc as std_mpsc; use hyper::server::conn::http1; use hyper::service::service_fn; -use tokio::net::TcpStream; +use hyper_util::rt::TokioIo; +use tokio::io::{AsyncRead, AsyncWrite}; -use crate::ProxyEvent; -use crate::cert::CertificateAuthority; +use crate::Engine; use crate::request::handle_request; -pub(crate) async fn handle_connection( - stream: TcpStream, - event_tx: std_mpsc::Sender, - ca: Arc, -) -> Result<(), Box> { - let tx = event_tx.clone(); +pub(crate) async fn serve_http1(io: S, handler: F) -> Result<(), hyper::Error> +where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + F: Fn(hyper::Request) -> Fut + Send + Sync + 'static, + Fut: Future, std::convert::Infallible>> + + Send + + 'static, +{ http1::Builder::new() .preserve_header_case(true) .title_case_headers(true) - .serve_connection( - hyper_util::rt::TokioIo::new(stream), - service_fn(move |req| { - let tx = tx.clone(); - let ca = ca.clone(); - async move { handle_request(req, tx, ca).await } - }), - ) + .serve_connection(TokioIo::new(io), service_fn(handler)) .with_upgrades() .await - .map_err(|e| Box::new(e) as Box) +} + +pub(crate) async fn handle_connection( + stream: tokio::net::TcpStream, + engine: Arc, +) -> Result<(), hyper::Error> { + serve_http1(stream, move |req| handle_request(req, engine.clone())).await } diff --git a/crates/yaak-proxy/src/lib.rs b/crates/yaak-proxy/src/lib.rs index a83a2d7c9..43a8a5733 100644 --- a/crates/yaak-proxy/src/lib.rs +++ b/crates/yaak-proxy/src/lib.rs @@ -1,63 +1,131 @@ -pub mod body; -pub mod cert; +mod body; +mod cert; mod connection; mod request; +mod tunnel; +use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; -use std::sync::Arc; -use std::sync::atomic::AtomicU64; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::mpsc as std_mpsc; +use std::sync::{Arc, Mutex}; +use std::time::SystemTime; -use cert::CertificateAuthority; +use bytes::Bytes; +use hyper_util::client::legacy::Client; +use hyper_util::client::legacy::connect::HttpConnector; +use hyper_util::rt::TokioExecutor; +use rustls::ClientConfig; +use rustls::pki_types::CertificateDer; use tokio::net::TcpListener; +pub use cert::CaStore; use connection::handle_connection; +use request::BoxBody; -static REQUEST_ID: AtomicU64 = AtomicU64::new(1); +const PASSTHROUGH_AFTER_FAILURES: u32 = 2; -/// Granular events emitted during request/response lifecycle. -/// Each event carries a request `id` so consumers can correlate events. +/// Default for [`ProxyConfig::capture_body_limit`]: 10 MB. +pub const DEFAULT_CAPTURE_BODY_LIMIT: usize = 10 * 1024 * 1024; + +/// Settings for one run of the engine. +pub struct ProxyConfig { + /// Port to listen on at 127.0.0.1; 0 picks a free one (see [`ProxyHandle::port`]). + pub port: u16, + /// Signs the per-host certificates presented to clients. + pub ca: Arc, + /// Bytes of each request and response body carried in chunk events. The rest still + /// streams through, counted but not captured. + pub capture_body_limit: usize, + /// Roots trusted for upstream TLS in addition to the platform's. + pub upstream_root_certs: Vec>, +} + +impl ProxyConfig { + pub fn new(port: u16, ca: Arc) -> Self { + Self { + port, + ca, + capture_body_limit: DEFAULT_CAPTURE_BODY_LIMIT, + upstream_root_certs: Vec::new(), + } + } +} + +/// Lifecycle of each proxied exchange (a flow), keyed by `id`. A flow starts with +/// `RequestStart` and ends with exactly one `ResponseBodyComplete` or `Error`. Traffic the +/// engine tunnels without decrypting (pinned clients, non-TLS CONNECT payloads, the frames +/// after a `101 Switching Protocols`) produces no events. #[derive(Debug, Clone)] pub enum ProxyEvent { - /// A new request has been received from the client. + /// Request headers arrived from the client. RequestStart { id: u64, + started_at: SystemTime, method: String, url: String, + /// Version spoken by the client. http_version: String, }, - /// A request header sent to the upstream server. RequestHeader { id: u64, name: String, value: String, }, - /// The full request body (buffered before forwarding). - RequestBody { id: u64, body: Vec }, - /// Response headers received from upstream. + /// A frame of the request body was forwarded upstream. + RequestBodyChunk { + id: u64, + /// Bytes in the frame. + len: usize, + /// The captured part of the frame; empty once `capture_body_limit` is reached. + data: Bytes, + }, + /// The request body finished streaming upstream. Not sent when the request has no body. + RequestBodyComplete { + id: u64, + size: u64, + /// Whether `size` exceeded `capture_body_limit`. + truncated: bool, + }, + /// Response headers arrived from upstream. ResponseStart { id: u64, status: u16, + /// Version spoken with the upstream server. http_version: String, + remote_addr: Option, + /// Milliseconds from `started_at` to the response headers. elapsed_ms: u64, }, - /// A response header received from the upstream server. ResponseHeader { id: u64, name: String, value: String, }, - /// A chunk of the response body was received (emitted per-frame). - ResponseBodyChunk { id: u64, bytes: usize }, - /// The response body stream has completed. + /// A frame of the response body was forwarded to the client. + ResponseBodyChunk { + id: u64, + /// Bytes in the frame. + len: usize, + /// The captured part of the frame; empty once `capture_body_limit` is reached. + data: Bytes, + }, + /// The response body ended, normally or with `error` (upstream failed or client left). ResponseBodyComplete { id: u64, - body: Option>, size: u64, + /// Whether `size` exceeded `capture_body_limit`. + truncated: bool, + /// Milliseconds from `started_at` to the end of the body. + elapsed_ms: u64, + error: Option, + }, + /// No response: connecting to or requesting from upstream failed. The client got a 502. + Error { + id: u64, + error: String, elapsed_ms: u64, }, - /// The upstream request failed. - Error { id: u64, error: String }, } /// Accumulated view of a proxied request, built from `ProxyEvent`s. @@ -87,12 +155,13 @@ pub enum RequestState { Error, } +/// A running engine. Dropping it stops the listener and joins its thread. pub struct ProxyHandle { shutdown_tx: Option>, thread_handle: Option>, event_rx: Option>, + /// The bound port. pub port: u16, - pub ca_pem: String, } impl ProxyHandle { @@ -113,12 +182,88 @@ impl Drop for ProxyHandle { } } -pub fn start_proxy(port: u16) -> Result { - let ca = CertificateAuthority::new().map_err(|e| format!("Failed to create CA: {e}"))?; - let ca_pem = ca.ca_pem(); - let ca = Arc::new(ca); +pub(crate) struct Engine { + ca: Arc, + capture_body_limit: usize, + events: std_mpsc::Sender, + next_id: AtomicU64, + http_client: Client, + upstream_tls: Arc, + upstream_tls_h1: Arc, + passthrough: Mutex>, + handshake_failures: Mutex>, +} + +impl Engine { + fn emit(&self, event: ProxyEvent) { + let _ = self.events.send(event); + } + + fn next_flow_id(&self) -> u64 { + self.next_id.fetch_add(1, Ordering::Relaxed) + } + + fn is_passthrough(&self, addr: &str) -> bool { + self.passthrough.lock().unwrap().contains(addr) + } + + /// Returns whether `addr` is now passed through. macOS clients reject the CA with a bare + /// reset, indistinguishable from an abandoned handshake, so those need two in a row. + fn record_handshake_failure(&self, addr: &str, certificate_rejected: bool) -> bool { + let mut failures = self.handshake_failures.lock().unwrap(); + let count = failures.entry(addr.to_string()).or_default(); + *count += 1; + if !certificate_rejected && *count < PASSTHROUGH_AFTER_FAILURES { + return false; + } + failures.remove(addr); + self.passthrough.lock().unwrap().insert(addr.to_string()); + true + } + + fn record_handshake_success(&self, addr: &str) { + self.handshake_failures.lock().unwrap().remove(addr); + } +} + +fn upstream_tls_config( + roots: &rustls::RootCertStore, + alpn: &[&[u8]], +) -> Result, String> { + let mut config = + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map_err(|e| e.to_string())? + .with_root_certificates(roots.clone()) + .with_no_client_auth(); + config.alpn_protocols = alpn.iter().map(|p| p.to_vec()).collect(); + Ok(Arc::new(config)) +} + +/// Starts the engine on its own thread and Tokio runtime, returning once the port is bound. +pub fn start_proxy(config: ProxyConfig) -> Result { + let mut roots = rustls::RootCertStore::empty(); + for cert in rustls_native_certs::load_native_certs().certs { + let _ = roots.add(cert); + } + for cert in config.upstream_root_certs { + roots.add(cert).map_err(|e| format!("Invalid upstream root certificate: {e}"))?; + } let (event_tx, event_rx) = std_mpsc::channel(); + let engine = Arc::new(Engine { + ca: config.ca, + capture_body_limit: config.capture_body_limit, + events: event_tx, + next_id: AtomicU64::new(1), + http_client: Client::builder(TokioExecutor::new()).build_http(), + upstream_tls: upstream_tls_config(&roots, &[b"h2", b"http/1.1"])?, + upstream_tls_h1: upstream_tls_config(&roots, &[b"http/1.1"])?, + passthrough: Mutex::new(HashSet::new()), + handshake_failures: Mutex::new(HashMap::new()), + }); + + let port = config.port; let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); let (ready_tx, ready_rx) = std_mpsc::channel(); @@ -132,50 +277,43 @@ pub fn start_proxy(port: u16) -> Result { }; rt.block_on(async move { - let addr = SocketAddr::from(([127, 0, 0, 1], port)); - let listener = match TcpListener::bind(addr).await { + let listener = match TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], port))).await { Ok(l) => l, Err(e) => { - let _ = ready_tx.send(Err(format!("Failed to bind: {e}"))); + let _ = ready_tx.send(Err(format!("Failed to bind port {port}: {e}"))); return; } }; - - let bound_port = listener.local_addr().unwrap().port(); - let _ = ready_tx.send(Ok(bound_port)); + let bound = listener.local_addr().map(|a| a.port()).map_err(|e| e.to_string()); + let _ = ready_tx.send(bound); + log::info!("Proxy listening on {:?}", listener.local_addr()); let mut shutdown_rx = shutdown_rx; loop { tokio::select! { - result = listener.accept() => { - match result { - Ok((stream, _addr)) => { - let tx = event_tx.clone(); - let ca = ca.clone(); - tokio::spawn(async move { - if let Err(e) = handle_connection(stream, tx, ca).await { - eprintln!("Connection error: {e}"); - } - }); - } - Err(e) => eprintln!("Accept error: {e}"), + result = listener.accept() => match result { + Ok((stream, _)) => { + let engine = engine.clone(); + tokio::spawn(async move { + if let Err(e) = handle_connection(stream, engine).await { + log::debug!("Client connection error: {e}"); + } + }); } - } - _ = &mut shutdown_rx => { - break; - } + Err(e) => log::warn!("Accept error: {e}"), + }, + _ = &mut shutdown_rx => break, } } }); }); match ready_rx.recv() { - Ok(Ok(bound_port)) => Ok(ProxyHandle { + Ok(Ok(port)) => Ok(ProxyHandle { shutdown_tx: Some(shutdown_tx), thread_handle: Some(thread_handle), event_rx: Some(event_rx), - port: bound_port, - ca_pem, + port, }), Ok(Err(e)) => Err(e), Err(_) => Err("Proxy thread died before binding".into()), diff --git a/crates/yaak-proxy/src/request.rs b/crates/yaak-proxy/src/request.rs index 1de6bb632..1c14029a1 100644 --- a/crates/yaak-proxy/src/request.rs +++ b/crates/yaak-proxy/src/request.rs @@ -1,40 +1,38 @@ use std::convert::Infallible; +use std::error::Error; +use std::net::SocketAddr; use std::sync::Arc; -use std::sync::atomic::Ordering; -use std::sync::mpsc as std_mpsc; -use std::time::Instant; +use std::time::{Duration, Instant, SystemTime}; use bytes::Bytes; use http_body_util::{BodyExt, Full}; use hyper::body::Incoming; -use hyper::header::HeaderMap; -use hyper::service::service_fn; +use hyper::header::{HOST, UPGRADE}; +use hyper::upgrade::OnUpgrade; use hyper::{Method, Request, Response, StatusCode, Uri}; -use hyper_util::client::legacy::Client; -use hyper_util::rt::TokioExecutor; -use hyper_util::server::conn::auto; -use rustls::ClientConfig; +use hyper_util::client::legacy::connect::HttpInfo; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use log::{debug, info, warn}; use rustls::pki_types::ServerName; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite}; use tokio::net::TcpStream; -use tokio_rustls::TlsAcceptor; +use tokio_rustls::{LazyConfigAcceptor, TlsConnector}; -use crate::body::MeasuredBody; -use crate::cert::CertificateAuthority; -use crate::{ProxyEvent, REQUEST_ID}; +use crate::body::{CaptureBody, Direction}; +use crate::connection::serve_http1; +use crate::tunnel::{Rewind, splice, tunnel}; +use crate::{Engine, ProxyEvent}; -type BoxBody = http_body_util::combinators::BoxBody; +pub(crate) type BoxBody = http_body_util::combinators::BoxBody; +type BoxError = Box; -fn full_body(bytes: Bytes) -> BoxBody { - Full::new(bytes).map_err(|never| match never {}).boxed() -} +const TLS_HANDSHAKE_RECORD: u8 = 0x16; +/// TLS clients send their ClientHello as soon as CONNECT succeeds; protocols where the server +/// speaks first (SMTP, for one) send nothing, so after this they are tunneled blind. +const CLIENT_FIRST_BYTES_TIMEOUT: Duration = Duration::from_secs(1); -fn measured_incoming( - incoming: Incoming, - id: u64, - start: Instant, - tx: std_mpsc::Sender, -) -> BoxBody { - MeasuredBody::new(incoming, id, start, tx).boxed() +fn text_body(text: impl Into) -> BoxBody { + Full::new(text.into()).map_err(|never| match never {}).boxed() } fn version_str(v: hyper::Version) -> String { @@ -49,239 +47,241 @@ fn version_str(v: hyper::Version) -> String { .to_string() } -fn emit_request_events( - tx: &std_mpsc::Sender, - id: u64, - headers: &HeaderMap, - body: &Option>, -) { - for (name, value) in headers.iter() { - let _ = tx.send(ProxyEvent::RequestHeader { - id, - name: name.to_string(), - value: value.to_str().unwrap_or("").to_string(), - }); +fn error_chain(e: &(dyn Error + 'static)) -> String { + let mut message = e.to_string(); + let mut source = e.source(); + while let Some(s) = source { + message.push_str(": "); + message.push_str(&s.to_string()); + source = s.source(); } - if let Some(body) = body { - let _ = tx.send(ProxyEvent::RequestBody { id, body: body.clone() }); + message +} + +/// Per-flow bookkeeping shared by the plain HTTP and the intercepted HTTPS paths. +struct Flow { + id: u64, + start: Instant, + engine: Arc, +} + +impl Flow { + fn begin(engine: Arc, req: &Request, url: String) -> Self { + let flow = Self { id: engine.next_flow_id(), start: Instant::now(), engine }; + flow.engine.emit(ProxyEvent::RequestStart { + id: flow.id, + started_at: SystemTime::now(), + method: req.method().to_string(), + url, + http_version: version_str(req.version()), + }); + for (name, value) in req.headers() { + flow.engine.emit(ProxyEvent::RequestHeader { + id: flow.id, + name: name.to_string(), + value: String::from_utf8_lossy(value.as_bytes()).into_owned(), + }); + } + flow + } + + fn capture_request(&self, body: Incoming) -> BoxBody { + CaptureBody::new(body, Direction::Request, self.id, self.start, self.engine.clone()).boxed() + } + + fn respond( + self, + resp: Response, + remote_addr: Option, + ) -> Response { + self.engine.emit(ProxyEvent::ResponseStart { + id: self.id, + status: resp.status().as_u16(), + http_version: version_str(resp.version()), + remote_addr, + elapsed_ms: self.start.elapsed().as_millis() as u64, + }); + for (name, value) in resp.headers() { + self.engine.emit(ProxyEvent::ResponseHeader { + id: self.id, + name: name.to_string(), + value: String::from_utf8_lossy(value.as_bytes()).into_owned(), + }); + } + let (parts, body) = resp.into_parts(); + let body = CaptureBody::new(body, Direction::Response, self.id, self.start, self.engine); + Response::from_parts(parts, body.boxed()) + } + + fn fail(self, error: String) -> Response { + debug!("Flow {} failed: {error}", self.id); + let body = format!("Yaak Proxy could not reach the server: {error}"); + self.engine.emit(ProxyEvent::Error { + id: self.id, + error, + elapsed_ms: self.start.elapsed().as_millis() as u64, + }); + bad_gateway(body) } } -fn emit_response_events( - tx: &std_mpsc::Sender, - id: u64, - resp: &Response, - start: &Instant, -) { - let _ = tx.send(ProxyEvent::ResponseStart { - id, - status: resp.status().as_u16(), - http_version: version_str(resp.version()), - elapsed_ms: start.elapsed().as_millis() as u64, - }); - for (name, value) in resp.headers().iter() { - let _ = tx.send(ProxyEvent::ResponseHeader { - id, - name: name.to_string(), - value: value.to_str().unwrap_or("").to_string(), - }); - } +fn bad_gateway(body: String) -> Response { + Response::builder().status(StatusCode::BAD_GATEWAY).body(text_body(body)).unwrap() } pub(crate) async fn handle_request( req: Request, - event_tx: std_mpsc::Sender, - ca: Arc, + engine: Arc, ) -> Result, Infallible> { - let result = if req.method() == Method::CONNECT { - handle_connect(req, event_tx, ca).await - } else { - handle_http(req, event_tx).await - }; - match result { - Ok(resp) => Ok(resp), - Err(e) => { - eprintln!("Proxy error: {e}"); - Ok(Response::builder() - .status(StatusCode::BAD_GATEWAY) - .body(full_body(Bytes::from(format!("Proxy error: {e}")))) - .unwrap()) - } + if req.method() == Method::CONNECT { + return Ok(handle_connect(req, engine)); } -} - -async fn handle_http( - req: Request, - event_tx: std_mpsc::Sender, -) -> Result, Box> { - let id = REQUEST_ID.fetch_add(1, Ordering::Relaxed); - let method = req.method().to_string(); - let uri = req.uri().to_string(); - let http_version = version_str(req.version()); - let start = Instant::now(); - - let _ = event_tx.send(ProxyEvent::RequestStart { id, method, url: uri.clone(), http_version }); - - let client: Client<_, Full> = Client::builder(TokioExecutor::new()).build_http(); + let flow = Flow::begin(engine.clone(), &req, req.uri().to_string()); let (parts, body) = req.into_parts(); - let body_bytes = body.collect().await?.to_bytes(); - let request_body = if body_bytes.is_empty() { None } else { Some(body_bytes.to_vec()) }; - emit_request_events(&event_tx, id, &parts.headers, &request_body); - - let outgoing_req = Request::from_parts(parts, Full::new(body_bytes)); - - match client.request(outgoing_req).await { + let outgoing = Request::from_parts(parts, flow.capture_request(body)); + Ok(match engine.http_client.request(outgoing).await { Ok(resp) => { - emit_response_events(&event_tx, id, &resp, &start); - - let (parts, body) = resp.into_parts(); - Ok(Response::from_parts(parts, measured_incoming(body, id, start, event_tx))) + let remote_addr = resp.extensions().get::().map(|i| i.remote_addr()); + flow.respond(resp, remote_addr) } - Err(e) => { - let _ = event_tx.send(ProxyEvent::Error { id, error: e.to_string() }); - Err(Box::new(e) as Box) - } - } + Err(e) => flow.fail(error_chain(&e)), + }) } -async fn handle_connect( - req: Request, - event_tx: std_mpsc::Sender, - ca: Arc, -) -> Result, Box> { - let authority = req.uri().authority().map(|a| a.to_string()).unwrap_or_default(); - let (host, port) = parse_host_port(&authority); +/// Where an intercepted CONNECT tunnel leads. +struct Target { + /// Host for the leaf certificate and upstream SNI. + server_name: String, + /// `host[:port]` as it appears in captured URLs. + authority: String, + /// `host:port` to dial. + addr: String, +} - let server_config = ca.server_config(&host)?; - let acceptor = TlsAcceptor::from(server_config); - - let target_addr = format!("{host}:{port}"); +fn handle_connect(req: Request, engine: Arc) -> Response { + let Some(authority) = req.uri().authority().cloned() else { + return Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(text_body("CONNECT requires host:port")) + .unwrap(); + }; + let host = authority.host().to_string(); + let port = authority.port_u16().unwrap_or(443); tokio::spawn(async move { let upgraded = match hyper::upgrade::on(req).await { - Ok(u) => u, - Err(e) => { - eprintln!("CONNECT upgrade failed: {e}"); - return; - } + Ok(u) => TokioIo::new(u), + Err(e) => return warn!("CONNECT upgrade failed for {host}: {e}"), }; - - let tls_stream = match acceptor.accept(hyper_util::rt::TokioIo::new(upgraded)).await { - Ok(s) => s, - Err(e) => { - eprintln!("TLS accept failed for {host}: {e}"); - return; - } + let target = Target { + server_name: host.trim_start_matches('[').trim_end_matches(']').to_string(), + authority: if port == 443 { host.clone() } else { format!("{host}:{port}") }, + addr: format!("{host}:{port}"), }; - - let tx = event_tx.clone(); - let host_for_requests = host.clone(); - let mut builder = auto::Builder::new(TokioExecutor::new()); - builder.http1().preserve_header_case(true).title_case_headers(true); - if let Err(e) = builder - .serve_connection_with_upgrades( - hyper_util::rt::TokioIo::new(tls_stream), - service_fn(move |req| { - let tx = tx.clone(); - let host = host_for_requests.clone(); - let target_addr = target_addr.clone(); - async move { handle_tunneled_request(req, tx, &host, &target_addr).await } - }), - ) - .await - { - eprintln!("MITM connection error for {host}: {e}"); + if let Err(e) = intercept(upgraded, target, engine).await { + debug!("Tunnel to {host}:{port} ended: {}", error_chain(e.as_ref())); } }); - Ok(Response::new(full_body(Bytes::new()))) + Response::new(text_body(Bytes::new())) } -async fn handle_tunneled_request( - req: Request, - event_tx: std_mpsc::Sender, - host: &str, - target_addr: &str, -) -> Result, Infallible> { - let result = forward_https(req, event_tx, host, target_addr).await; - match result { - Ok(resp) => Ok(resp), +async fn intercept(mut client: S, target: Target, engine: Arc) -> Result<(), BoxError> +where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + if engine.is_passthrough(&target.addr) { + return Ok(tunnel(client, &target.addr, &[]).await?); + } + + let mut buf = vec![0u8; 4096]; + let Ok(n) = tokio::time::timeout(CLIENT_FIRST_BYTES_TIMEOUT, client.read(&mut buf)).await + else { + debug!("Client sent nothing, tunneling server-first traffic to {}", target.addr); + return Ok(tunnel(client, &target.addr, &[]).await?); + }; + let n = n?; + if n == 0 { + return Ok(()); + } + buf.truncate(n); + if buf[0] != TLS_HANDSHAKE_RECORD { + debug!("Tunneling non-TLS traffic to {}", target.addr); + return Ok(tunnel(client, &target.addr, &buf).await?); + } + + let acceptor = rustls::server::Acceptor::default(); + let start = LazyConfigAcceptor::new(acceptor, Rewind::new(Bytes::from(buf), client)).await?; + let server_name = start + .client_hello() + .server_name() + .map(str::to_string) + .unwrap_or_else(|| target.server_name.clone()); + let config = engine.ca.server_config(&server_name)?; + let tls = match start.into_stream(config).await { + Ok(tls) => tls, Err(e) => { - eprintln!("HTTPS forward error: {e:?}"); - Ok(Response::builder() - .status(StatusCode::BAD_GATEWAY) - .body(full_body(Bytes::from(format!("Proxy error: {e}")))) - .unwrap()) + if engine.record_handshake_failure(&target.addr, is_certificate_rejection(&e)) { + info!("Client rejected the proxy for {}, passing it through: {e}", target.addr); + } else { + debug!("TLS handshake with client failed for {}: {e}", target.addr); + } + return Ok(()); } - } + }; + engine.record_handshake_success(&target.addr); + + let target = Arc::new(Target { server_name, ..target }); + serve_http1(tls, move |req| forward_https(req, target.clone(), engine.clone())).await?; + Ok(()) } -enum HttpSender { - H1(hyper::client::conn::http1::SendRequest>), - H2(hyper::client::conn::http2::SendRequest>), +fn is_certificate_rejection(e: &std::io::Error) -> bool { + use rustls::AlertDescription::{ + BadCertificate, CertificateExpired, CertificateRevoked, CertificateUnknown, UnknownCA, + UnsupportedCertificate, + }; + matches!( + e.get_ref().and_then(|inner| inner.downcast_ref::()), + Some(rustls::Error::AlertReceived( + BadCertificate + | UnsupportedCertificate + | CertificateRevoked + | CertificateExpired + | CertificateUnknown + | UnknownCA + )) + ) } -impl HttpSender { - async fn send_request( - &mut self, - req: Request>, - ) -> Result, hyper::Error> { - match self { - HttpSender::H1(s) => s.send_request(req).await, - HttpSender::H2(s) => s.send_request(req).await, - } - } +enum Upstream { + H1(hyper::client::conn::http1::SendRequest), + H2(hyper::client::conn::http2::SendRequest), } -async fn forward_https( - req: Request, - event_tx: std_mpsc::Sender, - host: &str, - target_addr: &str, -) -> Result, Box> { - let id = REQUEST_ID.fetch_add(1, Ordering::Relaxed); - let method = req.method().to_string(); - let http_version = version_str(req.version()); - let path = req.uri().path_and_query().map(|pq| pq.to_string()).unwrap_or_else(|| "/".into()); - let uri_str = format!("https://{host}{path}"); - let start = Instant::now(); +async fn connect_upstream( + engine: &Engine, + target: &Target, + h1_only: bool, +) -> Result<(Upstream, Option), BoxError> { + let tcp = TcpStream::connect(&target.addr).await?; + let remote_addr = tcp.peer_addr().ok(); + let config = if h1_only { &engine.upstream_tls_h1 } else { &engine.upstream_tls }; + let server_name = ServerName::try_from(target.server_name.clone())?; + let tls = TlsConnector::from(config.clone()).connect(server_name, tcp).await?; + let is_h2 = tls.get_ref().1.alpn_protocol() == Some(b"h2"); + let io = TokioIo::new(tls); - let _ = - event_tx.send(ProxyEvent::RequestStart { id, method, url: uri_str.clone(), http_version }); - - // Connect to upstream with TLS - let tcp_stream = TcpStream::connect(target_addr).await?; - - let mut root_store = rustls::RootCertStore::empty(); - for cert in rustls_native_certs::load_native_certs().certs { - let _ = root_store.add(cert); - } - - let mut tls_config = - ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) - .with_safe_default_protocol_versions()? - .with_root_certificates(root_store) - .with_no_client_auth(); - tls_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; - - let connector = tokio_rustls::TlsConnector::from(Arc::new(tls_config)); - let server_name = ServerName::try_from(host.to_string())?; - let tls_stream = connector.connect(server_name, tcp_stream).await?; - - let negotiated_h2 = tls_stream.get_ref().1.alpn_protocol().map_or(false, |p| p == b"h2"); - - let io = hyper_util::rt::TokioIo::new(tls_stream); - - let mut sender = if negotiated_h2 { + let upstream = if is_h2 { let (sender, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()).handshake(io).await?; tokio::spawn(async move { if let Err(e) = conn.await { - eprintln!("Upstream h2 connection error: {e}"); + debug!("Upstream h2 connection error: {e}"); } }); - HttpSender::H2(sender) + Upstream::H2(sender) } else { let (sender, conn) = hyper::client::conn::http1::Builder::new() .preserve_header_case(true) @@ -289,51 +289,69 @@ async fn forward_https( .handshake(io) .await?; tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("Upstream h1 connection error: {e}"); + if let Err(e) = conn.with_upgrades().await { + debug!("Upstream h1 connection error: {e}"); } }); - HttpSender::H1(sender) + Upstream::H1(sender) + }; + Ok((upstream, remote_addr)) +} + +async fn forward_https( + mut req: Request, + target: Arc, + engine: Arc, +) -> Result, Infallible> { + let path = req.uri().path_and_query().map(|pq| pq.to_string()).unwrap_or_else(|| "/".into()); + let url = format!("https://{}{path}", target.authority); + let flow = Flow::begin(engine.clone(), &req, url.clone()); + + let client_upgrade: Option = + req.headers().contains_key(UPGRADE).then(|| hyper::upgrade::on(&mut req)); + + let (mut upstream, remote_addr) = + match connect_upstream(&engine, &target, client_upgrade.is_some()).await { + Ok(u) => u, + Err(e) => return Ok(flow.fail(error_chain(e.as_ref()))), + }; + + let (mut parts, body) = req.into_parts(); + let uri = match &upstream { + Upstream::H2(_) => url.parse::(), + Upstream::H1(_) => path.parse::(), + }; + parts.uri = match uri { + Ok(uri) => uri, + Err(e) => return Ok(flow.fail(e.to_string())), + }; + if !parts.headers.contains_key(HOST) + && let Ok(host) = target.authority.parse() + { + parts.headers.insert(HOST, host); + } + let outgoing = Request::from_parts(parts, flow.capture_request(body)); + + let sent = match &mut upstream { + Upstream::H1(s) => s.send_request(outgoing).await, + Upstream::H2(s) => s.send_request(outgoing).await, + }; + let mut resp = match sent { + Ok(resp) => resp, + Err(e) => return Ok(flow.fail(error_chain(&e))), }; - // Capture request metadata - let (mut parts, body) = req.into_parts(); - let body_bytes = body.collect().await?.to_bytes(); - let request_body = if body_bytes.is_empty() { None } else { Some(body_bytes.to_vec()) }; - emit_request_events(&event_tx, id, &parts.headers, &request_body); - - if negotiated_h2 { - // HTTP/2 requires absolute-form URI with scheme + authority - parts.uri = uri_str.parse::()?; - } else { - parts.uri = path.parse::()?; + if resp.status() == StatusCode::SWITCHING_PROTOCOLS + && let Some(client_upgrade) = client_upgrade + { + let upstream_upgrade = hyper::upgrade::on(&mut resp); + let addr = target.addr.clone(); + tokio::spawn(async move { + if let Err(e) = splice(client_upgrade, upstream_upgrade).await { + debug!("Upgraded tunnel to {addr} ended: {e}"); + } + }); } - if !parts.headers.contains_key(hyper::header::HOST) { - parts.headers.insert(hyper::header::HOST, host.parse()?); - } - - let outgoing = Request::from_parts(parts, Full::new(body_bytes)); - - match sender.send_request(outgoing).await { - Ok(resp) => { - emit_response_events(&event_tx, id, &resp, &start); - - let (parts, body) = resp.into_parts(); - Ok(Response::from_parts(parts, measured_incoming(body, id, start, event_tx))) - } - Err(e) => { - let _ = event_tx.send(ProxyEvent::Error { id, error: e.to_string() }); - Err(Box::new(e) as Box) - } - } -} - -fn parse_host_port(authority: &str) -> (String, u16) { - if let Some((host, port_str)) = authority.rsplit_once(':') { - if let Ok(port) = port_str.parse::() { - return (host.to_string(), port); - } - } - (authority.to_string(), 443) + Ok(flow.respond(resp, remote_addr)) } diff --git a/crates/yaak-proxy/src/tunnel.rs b/crates/yaak-proxy/src/tunnel.rs new file mode 100644 index 000000000..5fe5a93d9 --- /dev/null +++ b/crates/yaak-proxy/src/tunnel.rs @@ -0,0 +1,85 @@ +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use bytes::Bytes; +use hyper::upgrade::OnUpgrade; +use hyper_util::rt::TokioIo; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_bidirectional}; +use tokio::net::TcpStream; + +/// Copies bytes both ways between `client` and `target`, after sending `prefix` (bytes already +/// read from the client) upstream. +pub(crate) async fn tunnel(mut client: S, target: &str, prefix: &[u8]) -> io::Result<()> +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let mut upstream = TcpStream::connect(target).await?; + upstream.write_all(prefix).await?; + copy_bidirectional(&mut client, &mut upstream).await?; + Ok(()) +} + +/// Joins the two halves of an HTTP upgrade once both sides have switched protocols. +pub(crate) async fn splice(client: OnUpgrade, upstream: OnUpgrade) -> io::Result<()> { + let (client, upstream) = tokio::try_join!(client, upstream).map_err(io::Error::other)?; + copy_bidirectional(&mut TokioIo::new(client), &mut TokioIo::new(upstream)).await?; + Ok(()) +} + +/// A stream that replays bytes already read from it before reading further. +pub(crate) struct Rewind { + prefix: Bytes, + inner: S, +} + +impl Rewind { + pub(crate) fn new(prefix: Bytes, inner: S) -> Self { + Self { prefix, inner } + } +} + +impl AsyncRead for Rewind { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if !self.prefix.is_empty() { + let n = self.prefix.len().min(buf.remaining()); + buf.put_slice(&self.prefix.split_to(n)); + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for Rewind { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} diff --git a/crates/yaak-proxy/tests/proxy.rs b/crates/yaak-proxy/tests/proxy.rs new file mode 100644 index 000000000..a3d0c3b72 --- /dev/null +++ b/crates/yaak-proxy/tests/proxy.rs @@ -0,0 +1,548 @@ +use std::convert::Infallible; +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::mpsc::Receiver; +use std::time::{Duration, SystemTime}; + +use bytes::Bytes; +use futures_util::{SinkExt, StreamExt}; +use http_body_util::{BodyExt, Full}; +use hyper::body::Incoming; +use hyper::header::{CONNECTION, UPGRADE}; +use hyper::service::service_fn; +use hyper::{Request, Response, StatusCode}; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto; +use rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; +use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName}; +use rustls::{ClientConfig, RootCertStore, ServerConfig}; +use tempfile::TempDir; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio_rustls::{TlsAcceptor, TlsConnector}; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::handshake::derive_accept_key; +use tokio_tungstenite::tungstenite::protocol::Role; +use yaak_proxy::{CaStore, ProxyConfig, ProxyEvent, ProxyHandle, start_proxy}; + +const BIG_BODY: usize = 64 * 1024; + +macro_rules! find { + ($events:expr, $pat:pat => $out:expr) => { + $events + .iter() + .find_map(|e| match e { + $pat => Some($out), + _ => None, + }) + .unwrap_or_else(|| panic!("no {} in {:#?}", stringify!($pat), $events)) + }; +} + +fn ring() -> Arc { + Arc::new(rustls::crypto::ring::default_provider()) +} + +/// A throwaway CA and a `localhost` certificate it signed, standing in for a real server. +struct TestPki { + ca_pem: String, + ca_der: CertificateDer<'static>, + server: Arc, +} + +fn test_pki() -> TestPki { + let ca_key = KeyPair::generate().unwrap(); + let mut ca_params = CertificateParams::default(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + ca_params.distinguished_name.push(rcgen::DnType::CommonName, "Upstream Test CA"); + let ca = ca_params.self_signed(&ca_key).unwrap(); + + let leaf_key = KeyPair::generate().unwrap(); + let leaf_params = CertificateParams::new(vec!["localhost".into()]).unwrap(); + let leaf = leaf_params.signed_by(&leaf_key, &ca, &ca_key).unwrap(); + + let mut server = ServerConfig::builder_with_provider(ring()) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert( + vec![leaf.der().clone()], + PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(leaf_key.serialize_der())), + ) + .unwrap(); + server.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + + TestPki { ca_pem: ca.pem(), ca_der: ca.der().clone(), server: Arc::new(server) } +} + +async fn route(req: Request) -> Result>, Infallible> { + Ok(match req.uri().path() { + "/hello" => Response::new(Full::from("hello")), + "/echo" => Response::new(Full::new(req.into_body().collect().await.unwrap().to_bytes())), + "/big" => Response::new(Full::from(vec![b'x'; BIG_BODY])), + "/ws" => websocket_echo(req), + _ => { + Response::builder().status(StatusCode::NOT_FOUND).body(Full::from("not found")).unwrap() + } + }) +} + +fn websocket_echo(mut req: Request) -> Response> { + let accept = derive_accept_key(req.headers()["sec-websocket-key"].as_bytes()); + let on_upgrade = hyper::upgrade::on(&mut req); + tokio::spawn(async move { + let io = TokioIo::new(on_upgrade.await.unwrap()); + let mut ws = WebSocketStream::from_raw_socket(io, Role::Server, None).await; + while let Some(Ok(msg)) = ws.next().await { + if msg.is_text() { + ws.send(msg).await.unwrap(); + } + } + }); + Response::builder() + .status(StatusCode::SWITCHING_PROTOCOLS) + .header(UPGRADE, "websocket") + .header(CONNECTION, "Upgrade") + .header("sec-websocket-accept", accept) + .body(Full::default()) + .unwrap() +} + +async fn spawn_upstream(tls: Option>) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + loop { + let (tcp, _) = listener.accept().await.unwrap(); + let tls = tls.clone(); + tokio::spawn(async move { + let builder = auto::Builder::new(TokioExecutor::new()); + let svc = service_fn(route); + let _ = match tls { + Some(config) => { + let Ok(stream) = TlsAcceptor::from(config).accept(tcp).await else { + return; + }; + builder.serve_connection_with_upgrades(TokioIo::new(stream), svc).await + } + None => builder.serve_connection_with_upgrades(TokioIo::new(tcp), svc).await, + }; + }); + } + }); + addr +} + +struct Harness { + proxy: ProxyHandle, + events: Receiver, + ca: Arc, + pki: TestPki, + _dir: TempDir, +} + +impl Harness { + fn start(capture_body_limit: usize) -> Self { + let dir = TempDir::new().unwrap(); + let ca_dir = dir.path().join("ca"); + let created = CaStore::load_or_create(&ca_dir).unwrap(); + // Run against a reloaded CA so leaf certificates must chain to the saved ca.pem. + let ca = Arc::new(CaStore::load_or_create(&ca_dir).unwrap()); + assert_eq!(created.pem(), ca.pem()); + + let pki = test_pki(); + let mut config = ProxyConfig::new(0, ca.clone()); + config.capture_body_limit = capture_body_limit; + config.upstream_root_certs = vec![pki.ca_der.clone()]; + let mut proxy = start_proxy(config).unwrap(); + let events = proxy.take_event_rx().unwrap(); + Self { proxy, events, ca, pki, _dir: dir } + } + + fn proxy_url(&self) -> String { + format!("http://127.0.0.1:{}", self.proxy.port) + } + + fn client_trusting(&self, root_pem: &str) -> reqwest::Client { + reqwest::Client::builder() + .proxy(reqwest::Proxy::all(self.proxy_url()).unwrap()) + .add_root_certificate(reqwest::Certificate::from_pem(root_pem.as_bytes()).unwrap()) + .build() + .unwrap() + } + + fn client(&self) -> reqwest::Client { + self.client_trusting(self.ca.pem()) + } + + /// A raw connection tunneled through the proxy with CONNECT to `localhost:port`. + async fn connect_tunnel(&self, port: u16) -> TcpStream { + let mut tcp = TcpStream::connect(("127.0.0.1", self.proxy.port)).await.unwrap(); + let connect = + format!("CONNECT localhost:{port} HTTP/1.1\r\nHost: localhost:{port}\r\n\r\n"); + tcp.write_all(connect.as_bytes()).await.unwrap(); + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + head.push(tcp.read_u8().await.unwrap()); + } + assert!(head.starts_with(b"HTTP/1.1 200"), "{}", String::from_utf8_lossy(&head)); + tcp + } + + fn tls_client_config(&self) -> Arc { + let mut roots = RootCertStore::empty(); + roots + .add(CertificateDer::from(pem::parse(self.ca.pem()).unwrap().into_contents())) + .unwrap(); + let config = ClientConfig::builder_with_provider(ring()) + .with_safe_default_protocol_versions() + .unwrap() + .with_root_certificates(roots) + .with_no_client_auth(); + Arc::new(config) + } + + /// Events of the next flow, through its terminal event plus any stragglers. + fn flow_events(&self) -> Vec { + let mut events = Vec::new(); + loop { + let event = + self.events.recv_timeout(Duration::from_secs(10)).expect("flow did not end"); + let done = + matches!(event, ProxyEvent::ResponseBodyComplete { .. } | ProxyEvent::Error { .. }); + events.push(event); + if done { + break; + } + } + while let Ok(event) = self.events.recv_timeout(Duration::from_millis(100)) { + events.push(event); + } + events + } +} + +/// The captured bytes of one side of a flow, and how many bytes passed through. +fn body(events: &[ProxyEvent], request: bool) -> (Vec, usize) { + let mut captured = Vec::new(); + let mut passed = 0; + for event in events { + match event { + ProxyEvent::RequestBodyChunk { len, data, .. } if request => { + captured.extend_from_slice(data); + passed += len; + } + ProxyEvent::ResponseBodyChunk { len, data, .. } if !request => { + captured.extend_from_slice(data); + passed += len; + } + _ => {} + } + } + (captured, passed) +} + +#[test] +fn ca_is_persisted_with_private_key() { + let dir = TempDir::new().unwrap(); + let first = CaStore::load_or_create(dir.path()).unwrap(); + let second = CaStore::load_or_create(dir.path()).unwrap(); + + assert_eq!(first.pem(), second.pem()); + assert_eq!(std::fs::read_to_string(first.cert_path()).unwrap(), first.pem()); + let der = pem::parse(first.pem()).unwrap().into_contents(); + assert!(der.windows(13).any(|w| w == b"Yaak Proxy CA")); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = std::fs::metadata(first.key_path()).unwrap().permissions().mode(); + assert_eq!(mode & 0o777, 0o600); + } +} + +#[test] +fn concurrent_creators_agree_on_one_ca() { + for _ in 0..20 { + let dir = TempDir::new().unwrap(); + let barrier = Arc::new(std::sync::Barrier::new(8)); + let pems: Vec = (0..8) + .map(|_| { + let (dir, barrier) = (dir.path().to_path_buf(), barrier.clone()); + std::thread::spawn(move || { + barrier.wait(); + CaStore::load_or_create(&dir).unwrap().pem().to_string() + }) + }) + .collect::>() + .into_iter() + .map(|t| t.join().unwrap()) + .collect(); + + let on_disk = std::fs::read_to_string(dir.path().join("ca.pem")).unwrap(); + assert!(pems.iter().all(|p| *p == on_disk), "creators returned different CAs"); + let key = KeyPair::from_pem(&std::fs::read_to_string(dir.path().join("ca.key")).unwrap()) + .unwrap(); + let cert = pem::parse(&on_disk).unwrap().into_contents(); + let spki = key.public_key_der(); + assert!(cert.windows(spki.len()).any(|w| w == spki), "ca.key does not match ca.pem"); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn https_get() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let url = format!("https://localhost:{}/hello", upstream.port()); + + let resp = h.client().get(&url).send().await.unwrap(); + assert_eq!(resp.status(), 200); + assert_eq!(resp.text().await.unwrap(), "hello"); + + let events = h.flow_events(); + let (method, req_url, req_version, started_at) = find!(events, + ProxyEvent::RequestStart { method, url, http_version, started_at, .. } => + (method.clone(), url.clone(), http_version.clone(), *started_at)); + assert_eq!(method, "GET"); + assert_eq!(req_url, url); + assert_eq!(req_version, "HTTP/1.1"); + assert!(SystemTime::now().duration_since(started_at).unwrap() < Duration::from_secs(30)); + + let (status, res_version, remote_addr, headers_ms) = find!(events, + ProxyEvent::ResponseStart { status, http_version, remote_addr, elapsed_ms, .. } => + (*status, http_version.clone(), *remote_addr, *elapsed_ms)); + assert_eq!(status, 200); + assert_eq!(res_version, "HTTP/2"); + assert_eq!(remote_addr, Some(upstream)); + + let (size, truncated, end_ms, error) = find!(events, + ProxyEvent::ResponseBodyComplete { size, truncated, elapsed_ms, error, .. } => + (*size, *truncated, *elapsed_ms, error.clone())); + assert_eq!(body(&events, false), (b"hello".to_vec(), 5)); + assert_eq!((size, truncated, error), (5, false, None)); + assert!(end_ms >= headers_ms); + assert_eq!(body(&events, true), (Vec::new(), 0)); + assert!(!events.iter().any(|e| matches!(e, ProxyEvent::RequestBodyComplete { .. }))); +} + +#[tokio::test(flavor = "multi_thread")] +async fn https_post_with_body() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let url = format!("https://localhost:{}/echo", upstream.port()); + + let resp = h.client().post(&url).body("name=yaak").send().await.unwrap(); + assert_eq!(resp.text().await.unwrap(), "name=yaak"); + + let events = h.flow_events(); + let method = find!(events, ProxyEvent::RequestStart { method, .. } => method.clone()); + assert_eq!(method, "POST"); + let (size, truncated) = find!(events, + ProxyEvent::RequestBodyComplete { size, truncated, .. } => (*size, *truncated)); + assert_eq!((size, truncated), (9, false)); + assert_eq!(body(&events, true), (b"name=yaak".to_vec(), 9)); + assert_eq!(body(&events, false), (b"name=yaak".to_vec(), 9)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn https_not_found() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let url = format!("https://localhost:{}/missing", upstream.port()); + + let resp = h.client().get(&url).send().await.unwrap(); + assert_eq!(resp.status(), 404); + assert_eq!(resp.text().await.unwrap(), "not found"); + + let events = h.flow_events(); + assert_eq!(find!(events, ProxyEvent::ResponseStart { status, .. } => *status), 404); + let error = find!(events, ProxyEvent::ResponseBodyComplete { error, .. } => error.clone()); + assert_eq!(error, None); + assert_eq!(body(&events, false), (b"not found".to_vec(), 9)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn bodies_over_the_cap_are_truncated_but_delivered() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let client = h.client(); + + let url = format!("https://localhost:{}/big", upstream.port()); + let received = client.get(&url).send().await.unwrap().bytes().await.unwrap(); + assert_eq!(received.len(), BIG_BODY); + + let events = h.flow_events(); + let (size, truncated) = find!(events, + ProxyEvent::ResponseBodyComplete { size, truncated, .. } => (*size, *truncated)); + assert_eq!((size, truncated), (BIG_BODY as u64, true)); + assert_eq!(body(&events, false), (vec![b'x'; 1024], BIG_BODY)); + + let url = format!("https://localhost:{}/echo", upstream.port()); + let sent = vec![b'y'; 4096]; + let echoed = client.post(&url).body(sent.clone()).send().await.unwrap().bytes().await.unwrap(); + assert_eq!(echoed.as_ref(), sent.as_slice()); + + let events = h.flow_events(); + let (size, truncated) = find!(events, + ProxyEvent::RequestBodyComplete { size, truncated, .. } => (*size, *truncated)); + assert_eq!((size, truncated), (4096, true)); + assert_eq!(body(&events, true), (vec![b'y'; 1024], 4096)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn websocket_upgrade_passes_through() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let port = upstream.port(); + + let tcp = h.connect_tunnel(port).await; + let tls = TlsConnector::from(h.tls_client_config()) + .connect(ServerName::try_from("localhost").unwrap(), tcp) + .await + .unwrap(); + + let url = format!("wss://localhost:{port}/ws"); + let (mut ws, _) = tokio_tungstenite::client_async(url.as_str(), tls).await.unwrap(); + for text in ["ping", "pong"] { + ws.send(Message::text(text)).await.unwrap(); + let reply = ws.next().await.unwrap().unwrap(); + assert_eq!(reply.into_text().unwrap().as_str(), text); + } + ws.close(None).await.unwrap(); + + let events = h.flow_events(); + let req_url = find!(events, ProxyEvent::RequestStart { url, .. } => url.clone()); + assert_eq!(req_url, format!("https://localhost:{port}/ws")); + let (status, version) = find!(events, + ProxyEvent::ResponseStart { status, http_version, .. } => (*status, http_version.clone())); + assert_eq!((status, version.as_str()), (101, "HTTP/1.1")); + let (size, error) = find!(events, ProxyEvent::ResponseBodyComplete { size, error, .. } => (*size, error.clone())); + assert_eq!((size, error), (0, None)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn pinned_client_falls_back_to_passthrough() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let url = format!("https://localhost:{}/hello", upstream.port()); + let pinned = h.client_trusting(&h.pki.ca_pem); + + assert!(pinned.get(&url).send().await.is_err()); + + let mut body = None; + for _ in 0..40 { + if let Ok(resp) = pinned.get(&url).send().await { + body = Some(resp.text().await.unwrap()); + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!(body.as_deref(), Some("hello")); + assert!(h.events.try_recv().is_err(), "passthrough traffic must not produce flows"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn certificate_alert_passes_through_on_the_next_connection() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let port = upstream.port(); + let mut roots = RootCertStore::empty(); + roots.add(h.pki.ca_der.clone()).unwrap(); + let config = ClientConfig::builder_with_provider(ring()) + .with_safe_default_protocol_versions() + .unwrap() + .with_root_certificates(roots) + .with_no_client_auth(); + let pinned = TlsConnector::from(Arc::new(config)); + let name = ServerName::try_from("localhost").unwrap(); + + assert!(pinned.connect(name.clone(), h.connect_tunnel(port).await).await.is_err()); + tokio::time::sleep(Duration::from_millis(100)).await; + pinned + .connect(name, h.connect_tunnel(port).await) + .await + .expect("one certificate alert should be enough to pass the host through"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn abandoned_handshake_does_not_disable_capture() { + let h = Harness::start(1024); + let upstream = spawn_upstream(Some(h.pki.server.clone())).await; + let port = upstream.port(); + + let mut tcp = h.connect_tunnel(port).await; + let mut hello = Vec::new(); + rustls::ClientConnection::new( + h.tls_client_config(), + ServerName::try_from("localhost").unwrap(), + ) + .unwrap() + .write_tls(&mut hello) + .unwrap(); + tcp.write_all(&hello).await.unwrap(); + let mut server_hello = [0u8; 1]; + tcp.read_exact(&mut server_hello).await.unwrap(); + drop(tcp); + tokio::time::sleep(Duration::from_millis(100)).await; + + let url = format!("https://localhost:{port}/hello"); + assert_eq!(h.client().get(&url).send().await.unwrap().text().await.unwrap(), "hello"); + let events = h.flow_events(); + assert_eq!(find!(events, ProxyEvent::ResponseStart { status, .. } => *status), 200); +} + +#[tokio::test(flavor = "multi_thread")] +async fn server_first_tunnel_does_not_wait_for_the_client() { + let h = Harness::start(1024); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + socket.write_all(b"220 ready\r\n").await.unwrap(); + let mut buf = [0u8; 64]; + let n = socket.read(&mut buf).await.unwrap(); + socket.write_all(&buf[..n]).await.unwrap(); + }); + + let mut tcp = h.connect_tunnel(port).await; + let mut greeting = [0u8; 11]; + tokio::time::timeout(Duration::from_secs(5), tcp.read_exact(&mut greeting)) + .await + .expect("the server's greeting never arrived") + .unwrap(); + assert_eq!(&greeting, b"220 ready\r\n"); + tcp.write_all(b"EHLO yaak\r\n").await.unwrap(); + let mut echo = [0u8; 11]; + tcp.read_exact(&mut echo).await.unwrap(); + assert_eq!(&echo, b"EHLO yaak\r\n"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn plain_http_get() { + let h = Harness::start(1024); + let upstream = spawn_upstream(None).await; + let url = format!("http://127.0.0.1:{}/hello", upstream.port()); + + assert_eq!(h.client().get(&url).send().await.unwrap().text().await.unwrap(), "hello"); + + let events = h.flow_events(); + assert_eq!(find!(events, ProxyEvent::RequestStart { url, .. } => url.clone()), url); + let (status, remote_addr) = find!(events, + ProxyEvent::ResponseStart { status, remote_addr, .. } => (*status, *remote_addr)); + assert_eq!((status, remote_addr), (200, Some(upstream))); +} + +#[tokio::test(flavor = "multi_thread")] +async fn unreachable_upstream_is_an_error() { + let h = Harness::start(1024); + let closed = TcpListener::bind("127.0.0.1:0").await.unwrap().local_addr().unwrap(); + let url = format!("http://{closed}/"); + + let resp = h.client().get(&url).send().await.unwrap(); + assert_eq!(resp.status(), 502); + + let events = h.flow_events(); + let error = find!(events, ProxyEvent::Error { error, .. } => error.clone()); + assert!(error.to_lowercase().contains("connect"), "{error}"); + assert!(!events.iter().any(|e| matches!(e, ProxyEvent::ResponseStart { .. }))); +}