Add exact resumable HTTP source streaming

crates/remora-http-source/tests/http_source.rs

@@ -0,0 +1,579 @@
+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<Vec<IpAddr>> {
+        Ok(vec!["127.0.0.1".parse().expect("address")])
+    }
+}
+
+#[derive(Debug)]
+struct ObservedRequest {
+    path: String,
+    headers: HashMap<String, String>,
+}
+
+struct ResponseSpec {
+    status: &'static str,
+    headers: Vec<(&'static str, String)>,
+    body: Vec<u8>,
+}
+
+impl ResponseSpec {
+    fn new(status: &'static str, body: impl Into<Vec<u8>>) -> Self {
+        Self {
+            status,
+            headers: Vec::new(),
+            body: body.into(),
+        }
+    }
+
+    fn header(mut self, name: &'static str, value: impl Into<String>) -> Self {
+        self.headers.push((name, value.into()));
+        self
+    }
+}
+
+async fn serve(
+    connections: usize,
+    handler: impl Fn(usize, &ObservedRequest) -> ResponseSpec + Send + Sync + 'static,
+) -> (u16, Arc<Mutex<Vec<ObservedRequest>>>, 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<Mutex<Vec<ObservedRequest>>>, 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<Mutex<Vec<ObservedRequest>>>, 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<LocalResolver> {
+    let policy = OutboundPolicyBuilder::new()
+        .allow_scheme_port("http", port)
+        .allow_non_global("127.0.0.0/8".parse::<IpNet>().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<Vec<u8>, 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::<IpNet>().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"));
+}