use std::{ collections::HashMap, io, net::IpAddr, sync::{Arc, Mutex}, time::Duration, }; use http::{HeaderMap, HeaderValue}; use ipnet::IpNet; use remora_domain::MemberPath; use remora_http_source::{ HttpContentEncoding, HttpResumeCheckpoint, HttpSource, HttpSourceError, HttpSourceIdentity, HttpSourceRequest, StrongEtag, }; use remora_outbound::{CanonicalUrl, OutboundPolicyBuilder, Resolver}; use remora_stream::{RetryTiming, Sha256Accountant, SourceBody, StreamError}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, task::JoinHandle, }; #[derive(Clone)] struct LocalResolver; impl Resolver for LocalResolver { async fn resolve(&self, _host: &str, _port: u16) -> io::Result> { Ok(vec!["127.0.0.1".parse().expect("address")]) } } #[derive(Debug)] struct ObservedRequest { path: String, headers: HashMap, } struct ResponseSpec { status: &'static str, headers: Vec<(&'static str, String)>, body: Vec, } impl ResponseSpec { fn new(status: &'static str, body: impl Into>) -> Self { Self { status, headers: Vec::new(), body: body.into(), } } fn header(mut self, name: &'static str, value: impl Into) -> Self { self.headers.push((name, value.into())); self } } async fn serve( connections: usize, handler: impl Fn(usize, &ObservedRequest) -> ResponseSpec + Send + Sync + 'static, ) -> (u16, Arc>>, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); let port = listener.local_addr().expect("address").port(); let (observed, task) = start_server(listener, connections, handler); (port, observed, task) } async fn serve_on( port: u16, connections: usize, handler: impl Fn(usize, &ObservedRequest) -> ResponseSpec + Send + Sync + 'static, ) -> (Arc>>, JoinHandle<()>) { let listener = TcpListener::bind(("127.0.0.1", port)).await.expect("bind"); start_server(listener, connections, handler) } fn start_server( listener: TcpListener, connections: usize, handler: impl Fn(usize, &ObservedRequest) -> ResponseSpec + Send + Sync + 'static, ) -> (Arc>>, JoinHandle<()>) { let observed = Arc::new(Mutex::new(Vec::new())); let task_observed = Arc::clone(&observed); let handler = Arc::new(handler); let task = tokio::spawn(async move { for index in 0..connections { let (mut stream, _) = listener.accept().await.expect("accept"); let request = read_request(&mut stream).await; let response = handler(index, &request); task_observed.lock().expect("observed").push(request); let mut head = format!("HTTP/1.1 {}\r\nConnection: close\r\n", response.status); for (name, value) in response.headers { head.push_str(&format!("{name}: {value}\r\n")); } head.push_str("\r\n"); if let Err(error) = stream.write_all(head.as_bytes()).await { assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); continue; } if let Err(error) = stream.write_all(&response.body).await { assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); } } }); (observed, task) } async fn read_request(stream: &mut tokio::net::TcpStream) -> ObservedRequest { let mut bytes = Vec::new(); let mut chunk = [0_u8; 1024]; while !bytes.windows(4).any(|window| window == b"\r\n\r\n") { let count = stream.read(&mut chunk).await.expect("read"); assert_ne!(count, 0); bytes.extend_from_slice(&chunk[..count]); } let request = String::from_utf8(bytes).expect("request text"); let mut lines = request.split("\r\n"); let path = lines .next() .and_then(|line| line.split_whitespace().nth(1)) .expect("path") .to_owned(); let headers = lines .filter_map(|line| line.split_once(':')) .map(|(name, value)| (name.to_ascii_lowercase(), value.trim().to_owned())) .collect(); ObservedRequest { path, headers } } fn source(port: u16) -> HttpSource { let policy = OutboundPolicyBuilder::new() .allow_scheme_port("http", port) .allow_non_global("127.0.0.0/8".parse::().expect("network")) .build(); HttpSource::new(LocalResolver, policy, RetryTiming::default_bounds()) } fn request(port: u16) -> HttpSourceRequest { HttpSourceRequest::new( &format!("http://source.test:{port}/entry?secret=hidden"), "dir/cafe\u{301}.bin", ) .expect("request") } fn strong(body: &[u8], etag: &str) -> ResponseSpec { ResponseSpec::new("200 OK", body) .header("Content-Length", body.len().to_string()) .header("ETag", etag) } async fn read_all(mut body: remora_http_source::HttpSourceBody) -> Result, StreamError> { let mut output = Vec::new(); let mut buffer = [0_u8; 3]; loop { let count = body.read(&mut buffer).await?; if count == 0 { return Ok(output); } output.extend_from_slice(&buffer[..count]); } } fn digest(bytes: &[u8]) -> remora_stream::IntegrityDigest { let mut accountant = Sha256Accountant::new(); accountant.update(bytes); accountant.finalize().digest().expect("digest").clone() } #[tokio::test] async fn full_stream_exposes_strong_identity_and_explicit_canonical_result_path() { let body = b"representation".to_vec(); let expected = body.clone(); let (port, observed, server) = serve(1, move |_, _| strong(&body, "\"v1\"")).await; let opened = source(port).open(request(port)).await.expect("open"); assert_eq!(opened.result_path().canonical(), "dir/café.bin"); let identity = opened.identity().expect("strong identity"); assert_eq!(identity.etag().as_str(), "\"v1\""); assert_eq!(identity.representation_length(), expected.len() as u64); assert_eq!(identity.content_encoding(), HttpContentEncoding::Identity); assert_eq!(read_all(opened.into_body()).await.expect("body"), expected); server.await.expect("server"); let observed = observed.lock().expect("observed"); assert_eq!(observed[0].path, "/entry?secret=hidden"); assert_eq!( observed[0] .headers .get("accept-encoding") .map(String::as_str), Some("identity") ); } #[tokio::test] async fn matching_conditional_range_resumes_exact_suffix() { let bytes = b"abcdefghij".to_vec(); let served = bytes.clone(); let (port, observed, server) = serve(2, move |index, request| { if index == 0 { strong(&served, "\"v1\"") } else { assert_eq!( request.headers.get("range").map(String::as_str), Some("bytes=4-") ); assert_eq!( request.headers.get("if-match").map(String::as_str), Some("\"v1\"") ); ResponseSpec::new("206 Partial Content", served[4..].to_vec()) .header("Content-Length", (served.len() - 4).to_string()) .header( "Content-Range", format!("bytes 4-{}/{}", served.len() - 1, served.len()), ) .header("ETag", "\"v1\"") } }) .await; let initial = source(port).open(request(port)).await.expect("initial"); let checkpoint = HttpResumeCheckpoint::new( initial.identity().expect("identity").clone(), initial.result_path().clone(), 4, digest(&bytes[..4]), ) .expect("checkpoint"); drop(initial); let resumed = source(port) .resume(&checkpoint, &digest(&bytes[..4]), HeaderMap::new()) .await .expect("resume"); assert_eq!( read_all(resumed.into_body()).await.expect("body"), bytes[4..] ); server.await.expect("server"); assert_eq!(observed.lock().expect("observed").len(), 2); } #[tokio::test] async fn durable_identity_and_path_resume_after_worker_and_server_restart() { let (port, _, first_server) = serve(2, |index, _| match index { 0 => ResponseSpec::new("302 Found", Vec::new()) .header("Content-Length", "0") .header("Location", "/stable"), 1 => strong(b"abcdefghij", "\"v1\""), _ => unreachable!(), }) .await; let initial = source(port).open(request(port)).await.expect("initial"); let identity = initial.identity().expect("identity"); let persisted = ( identity.requested_url().as_url().as_str().to_owned(), identity.final_url().as_url().as_str().to_owned(), identity.etag().as_str().to_owned(), identity.representation_length(), initial.result_path().canonical(), 4, digest(b"abcd"), ); drop(initial); first_server.await.expect("first server"); let (observed, second_server) = serve_on(port, 2, |index, _| match index { 0 => ResponseSpec::new("302 Found", Vec::new()) .header("Content-Length", "0") .header("Location", "/stable"), 1 => ResponseSpec::new("206 Partial Content", b"efghij".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 4-9/10") .header("ETag", "\"v1\""), _ => unreachable!(), }) .await; let rehydrated = HttpSourceIdentity::from_parts( CanonicalUrl::parse(&persisted.0).expect("requested URL"), CanonicalUrl::parse(&persisted.1).expect("final URL"), StrongEtag::parse(&persisted.2).expect("etag"), persisted.3, ); let checkpoint = HttpResumeCheckpoint::new( rehydrated, MemberPath::parse(&persisted.4).expect("result path"), persisted.5, persisted.6, ) .expect("checkpoint"); let resumed = source(port) .resume(&checkpoint, &digest(b"abcd"), HeaderMap::new()) .await .expect("resume"); assert_eq!(resumed.result_path().canonical(), "dir/café.bin"); assert_eq!( read_all(resumed.into_body()).await.expect("body"), b"efghij" ); second_server.await.expect("second server"); let observed = observed.lock().expect("observed"); assert_eq!(observed[0].path, "/entry?secret=hidden"); assert_eq!(observed[1].path, "/stable"); } #[tokio::test] async fn weak_absent_and_modification_time_only_responses_are_restart_only() { for (etag, last_modified) in [ (Some("W/\"weak\""), None), (None, None), (None, Some("Wed, 21 Oct 2015 07:28:00 GMT")), ] { let (port, _, server) = serve(1, move |_, _| { let mut response = ResponseSpec::new("200 OK", b"body".to_vec()).header("Content-Length", "4"); if let Some(etag) = etag { response = response.header("ETag", etag); } if let Some(last_modified) = last_modified { response = response.header("Last-Modified", last_modified); } response }) .await; let opened = source(port).open(request(port)).await.expect("open"); assert!(opened.identity().is_none()); assert_eq!(read_all(opened.into_body()).await.expect("body"), b"body"); server.await.expect("server"); } } #[tokio::test] async fn changed_or_unproven_resume_responses_fail_as_source_changed() { for response in [ ResponseSpec::new("412 Precondition Failed", Vec::new()).header("Content-Length", "0"), ResponseSpec::new("200 OK", b"changed".to_vec()).header("Content-Length", "7"), ResponseSpec::new("206 Partial Content", b"efghij".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 5-9/10") .header("ETag", "\"v1\""), ResponseSpec::new("206 Partial Content", b"efghij".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 4-9/10") .header("ETag", "\"v2\""), ResponseSpec::new("206 Partial Content", b"efghij".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 4-9/10") .header("ETag", "\"v1\"") .header("Content-Encoding", "gzip"), ResponseSpec::new("206 Partial Content", b"efghijk".to_vec()) .header("Content-Length", "7") .header("Content-Range", "bytes 4-10/11") .header("ETag", "\"v1\""), ] { let response = Arc::new(Mutex::new(Some(response))); let task_response = Arc::clone(&response); let (port, _, server) = serve(2, move |index, _| { if index == 0 { strong(b"abcdefghij", "\"v1\"") } else { task_response .lock() .expect("response") .take() .expect("once") } }) .await; let source = source(port); let opened = source.open(request(port)).await.expect("initial"); let identity = opened.identity().expect("identity").clone(); let checkpoint = HttpResumeCheckpoint::new(identity, opened.result_path().clone(), 4, digest(b"abcd")) .expect("checkpoint"); drop(opened); assert_eq!( source .resume(&checkpoint, &digest(b"abcd"), HeaderMap::new()) .await .expect_err("changed"), HttpSourceError::SourceChanged ); server.await.expect("server"); } } #[tokio::test] async fn changed_redirect_target_invalidates_resume_identity() { let (port, _, server) = serve(4, |index, _| match index { 0 => ResponseSpec::new("302 Found", Vec::new()) .header("Content-Length", "0") .header("Location", "/a"), 1 => strong(b"abcdefghij", "\"v1\""), 2 => ResponseSpec::new("302 Found", Vec::new()) .header("Content-Length", "0") .header("Location", "/b"), 3 => ResponseSpec::new("206 Partial Content", b"efghij".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 4-9/10") .header("ETag", "\"v1\""), _ => unreachable!(), }) .await; let source = source(port); let opened = source.open(request(port)).await.expect("initial"); let checkpoint = HttpResumeCheckpoint::new( opened.identity().expect("identity").clone(), opened.result_path().clone(), 4, digest(b"abcd"), ) .expect("checkpoint"); drop(opened); assert_eq!( source .resume(&checkpoint, &digest(b"abcd"), HeaderMap::new()) .await .expect_err("redirect changed"), HttpSourceError::SourceChanged ); server.await.expect("server"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn response_start_timeout_is_retryable() { let (port, _, server) = serve(1, |_, _| { std::thread::sleep(Duration::from_millis(50)); strong(b"body", "\"v1\"") }) .await; let policy = OutboundPolicyBuilder::new() .allow_scheme_port("http", port) .allow_non_global("127.0.0.0/8".parse::().expect("network")) .build(); let timing = RetryTiming { response_start_timeout: Duration::from_millis(5), ..RetryTiming::default_bounds() }; assert_eq!( HttpSource::new(LocalResolver, policy, timing) .open(request(port)) .await .expect_err("timeout"), HttpSourceError::Transient { retry_after: None } ); server.await.expect("server"); } #[tokio::test] async fn durable_prefix_mismatch_fails_before_network_io() { let (port, _, server) = serve(1, |_, _| strong(b"abcdefghij", "\"v1\"")).await; let source = source(port); let opened = source.open(request(port)).await.expect("initial"); let identity = opened.identity().expect("identity").clone(); let checkpoint = HttpResumeCheckpoint::new(identity, opened.result_path().clone(), 4, digest(b"abcd")) .expect("checkpoint"); drop(opened); assert_eq!( source .resume(&checkpoint, &digest(b"abce"), HeaderMap::new()) .await .expect_err("prefix"), HttpSourceError::SourceChanged ); server.await.expect("server"); } #[tokio::test] async fn truncated_stream_is_retryable_and_non_identity_encoding_fails_closed() { let (port, _, server) = serve(1, |_, _| { ResponseSpec::new("200 OK", b"short".to_vec()) .header("Content-Length", "10") .header("ETag", "\"v1\"") .header("Accept-Ranges", "bytes") }) .await; let opened = source(port).open(request(port)).await.expect("open"); assert!(matches!( read_all(opened.into_body()).await.expect_err("truncated"), StreamError::Transient { .. } )); server.await.expect("server"); let (port, _, server) = serve(1, |_, _| { strong(b"compressed", "\"gzip\"").header("Content-Encoding", "gzip") }) .await; assert_eq!( source(port) .open(request(port)) .await .expect_err("encoding"), HttpSourceError::InvalidResponse ); server.await.expect("server"); } #[tokio::test] async fn truncated_resumed_stream_is_retryable() { let (port, _, server) = serve(2, |index, _| { if index == 0 { strong(b"abcdefghij", "\"v1\"") } else { ResponseSpec::new("206 Partial Content", b"efg".to_vec()) .header("Content-Length", "6") .header("Content-Range", "bytes 4-9/10") .header("ETag", "\"v1\"") } }) .await; let initial = source(port).open(request(port)).await.expect("initial"); let checkpoint = HttpResumeCheckpoint::new( initial.identity().expect("identity").clone(), initial.result_path().clone(), 4, digest(b"abcd"), ) .expect("checkpoint"); drop(initial); let resumed = source(port) .resume(&checkpoint, &digest(b"abcd"), HeaderMap::new()) .await .expect("resume"); assert!(matches!( read_all(resumed.into_body()).await.expect_err("truncated"), StreamError::Transient { .. } )); server.await.expect("server"); } #[tokio::test] async fn status_mapping_and_retry_hints_are_bounded() { for (status, retry_after, expected) in [ ("401 Unauthorized", None, HttpSourceError::Authorization), ("404 Not Found", None, HttpSourceError::NotFound), ( "503 Service Unavailable", Some("3"), HttpSourceError::Transient { retry_after: Some(Duration::from_secs(3)), }, ), ( "429 Too Many Requests", Some("600"), HttpSourceError::Transient { retry_after: None }, ), ] { let (port, _, server) = serve(1, move |_, _| { let mut response = ResponseSpec::new(status, Vec::new()).header("Content-Length", "0"); if let Some(value) = retry_after { response = response.header("Retry-After", value); } response }) .await; assert_eq!( source(port).open(request(port)).await.expect_err("status"), expected ); server.await.expect("server"); } } #[test] fn request_debug_and_validation_do_not_expose_sensitive_url_or_headers() { assert!(HttpSourceRequest::new("http://source.test/object", "../bad").is_err()); let mut headers = HeaderMap::new(); headers.insert("x-api-key", HeaderValue::from_static("top-secret")); let request = HttpSourceRequest::new("https://source.test/object?token=secret", "artifact.bin") .expect("request") .with_credential_headers(headers) .expect("credentials"); let debug = format!("{request:?}"); assert!(!debug.contains("token")); assert!(!debug.contains("secret")); assert!(!debug.contains("top-secret")); }