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::combinators::BoxBody; use http_body_util::{BodyExt, Full, StreamBody}; use hyper::body::{Frame, Incoming}; use hyper::header::{CONNECTION, CONTENT_LENGTH, 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) } } type TestBody = BoxBody; fn full(data: impl Into) -> TestBody { Full::new(data.into()).boxed() } async fn route(req: Request) -> Result, Infallible> { Ok(match req.uri().path() { "/hello" => Response::new(full("hello")), "/echo" => Response::new(full(req.into_body().collect().await.unwrap().to_bytes())), "/big" => Response::new(full(vec![b'x'; BIG_BODY])), "/late-end" => late_end(5), // The stream ends with half the promised bytes, so the upstream resets it mid-body. "/cut" => late_end(10), "/ws" => websocket_echo(req), _ => Response::builder().status(StatusCode::NOT_FOUND).body(full("not found")).unwrap(), }) } /// "hello" under a Content-Length of `content_length`, its end of stream sent a while after the /// last byte, as an HTTP/2 server does when its final DATA frame has no END_STREAM. fn late_end(content_length: usize) -> Response { let frames = futures_util::stream::unfold(false, |sent| async move { if sent { tokio::time::sleep(Duration::from_millis(300)).await; return None; } Some((Ok(Frame::data(Bytes::from("hello"))), true)) }); Response::builder() .header(CONTENT_LENGTH, content_length) .body(BodyExt::boxed(StreamBody::new(frames))) .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(Bytes::new())) .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) } /// Waits until the engine reports `state` (failures counted, passed through) for `addr`. async fn wait_for_handshake_state(&self, addr: &str, state: (u32, bool)) { let deadline = tokio::time::Instant::now() + Duration::from_secs(10); while self.proxy.handshake_state(addr) != state { assert!( tokio::time::Instant::now() < deadline, "{addr} stayed at {:?}, expected {state:?}", self.proxy.handshake_state(addr) ); tokio::time::sleep(Duration::from_millis(10)).await; } } /// 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)); } /// The client is sent every byte the Content-Length promised, and hyper drops the body there, /// before an HTTP/2 upstream's end of stream arrives. That is a complete response, not an error. #[tokio::test(flavor = "multi_thread")] async fn https_body_complete_before_end_of_stream() { let h = Harness::start(1024); let upstream = spawn_upstream(Some(h.pki.server.clone())).await; let url = format!("https://localhost:{}/late-end", upstream.port()); let resp = h.client().get(&url).send().await.unwrap(); assert_eq!(resp.text().await.unwrap(), "hello"); let events = h.flow_events(); let version = find!(events, ProxyEvent::ResponseStart { http_version, .. } => http_version.clone()); assert_eq!(version, "HTTP/2"); let (size, error) = find!(events, ProxyEvent::ResponseBodyComplete { size, error, .. } => (*size, error.clone())); assert_eq!((size, error), (5, None)); assert_eq!(body(&events, false), (b"hello".to_vec(), 5)); } #[tokio::test(flavor = "multi_thread")] async fn https_body_cut_short_is_an_error() { let h = Harness::start(1024); let upstream = spawn_upstream(Some(h.pki.server.clone())).await; let url = format!("https://localhost:{}/cut", upstream.port()); let resp = h.client().get(&url).send().await.unwrap(); assert!(resp.bytes().await.is_err()); let events = h.flow_events(); let error = find!(events, ProxyEvent::ResponseBodyComplete { error, .. } => error.clone()); assert!(error.is_some(), "{events:#?}"); } #[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"); } /// SChannel (native-tls on Windows) checks the certificate only after sending its Finished, so /// the proxy sees a completed handshake and then a hang-up instead of an alert. #[tokio::test(flavor = "multi_thread")] async fn hang_up_after_the_handshake_passes_through() { let h = Harness::start(1024); let upstream = spawn_upstream(Some(h.pki.server.clone())).await; let port = upstream.port(); let name = ServerName::try_from("localhost").unwrap(); let addr = format!("localhost:{port}"); let hangs_up = TlsConnector::from(h.tls_client_config()); for state in [(1, false), (0, true)] { let tls = hangs_up.connect(name.clone(), h.connect_tunnel(port).await).await.unwrap(); drop(tls); h.wait_for_handshake_state(&addr, state).await; } 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(); TlsConnector::from(Arc::new(config)) .connect(name, h.connect_tunnel(port).await) .await .expect("two hang-ups after the handshake should 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 { .. }))); }