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

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"));
}