Add bounded commit-pinned Hugging Face source streaming

crates/remora-hf-source/tests/hf_source.rs

@@ -0,0 +1,650 @@
+use std::collections::HashMap;
+use std::io;
+use std::net::IpAddr;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+
+use http::{HeaderMap, HeaderValue};
+use ipnet::IpNet;
+use remora_hf_source::{
+    HfMember, HfMemberStorage, HfRepoType, HfSelection, HfSource, HfSourceError, HfSourceRequest,
+};
+use remora_outbound::{OutboundPolicyBuilder, Resolver};
+use remora_stream::{RetryTiming, SourceBody, StreamError};
+use sha1::{Digest, Sha1};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
+
+const COMMIT: &str = "0123456789abcdef0123456789abcdef01234567";
+
+#[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(Clone, Copy)]
+enum Fixture {
+    Normal,
+    Collision,
+    Empty,
+    Symlink,
+    Submodule,
+    Changing,
+    BadDigests,
+    OutOfOrderNormalized,
+    MismatchedCommit,
+    XetRedirect(u16),
+    Content,
+}
+
+struct ObservedRequest {
+    path: String,
+    headers: HashMap<String, String>,
+}
+
+struct TestServer {
+    port: u16,
+    observed: Arc<Mutex<Vec<ObservedRequest>>>,
+    task: JoinHandle<()>,
+}
+
+impl Drop for TestServer {
+    fn drop(&mut self) {
+        self.task.abort();
+    }
+}
+
+async fn serve(fixture: Fixture) -> TestServer {
+    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
+    let port = listener.local_addr().expect("address").port();
+    let observed = Arc::new(Mutex::new(Vec::new()));
+    let task_observed = Arc::clone(&observed);
+    let scans = Arc::new(AtomicUsize::new(0));
+    let task = tokio::spawn(async move {
+        loop {
+            let (mut stream, _) = listener.accept().await.expect("accept");
+            let observed = Arc::clone(&task_observed);
+            let scans = Arc::clone(&scans);
+            tokio::spawn(async move {
+                let request = read_request(&mut stream).await;
+                let response = fixture_response(fixture, port, &request.path, &scans);
+                observed.lock().expect("observed").push(request);
+                stream
+                    .write_all(response.as_bytes())
+                    .await
+                    .expect("response");
+            });
+        }
+    });
+    TestServer {
+        port,
+        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("request");
+        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 fixture_response(fixture: Fixture, port: u16, path: &str, scans: &AtomicUsize) -> String {
+    if path.contains("/revision/missing")
+        || path.contains("/revision/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+            && !matches!(fixture, Fixture::MismatchedCommit)
+    {
+        return response("404 Not Found", &[], b"missing");
+    }
+    if path.contains("/revision/") {
+        let repo_id = if path.contains("/datasets/") {
+            "acme/dataset"
+        } else if path.contains("/spaces/") {
+            "acme/space"
+        } else {
+            "acme/model"
+        };
+        let body = format!(r#"{{"id":"{repo_id}","sha":"{COMMIT}"}}"#);
+        return response("200 OK", &[], body.as_bytes());
+    }
+    if path.contains("/tree/") {
+        return tree_response(fixture, port, path, scans);
+    }
+    if path.ends_with("/z.bin")
+        && let Fixture::XetRedirect(content_port) = fixture
+    {
+        return response(
+            "302 Found",
+            &[(
+                "Location",
+                &format!("http://content.test:{content_port}/z.bin"),
+            )],
+            b"",
+        );
+    }
+    let body = if path.ends_with("/a.txt") {
+        b"alpha".as_slice()
+    } else if path.ends_with("/dir/lfs.bin") {
+        b"lfs-body".as_slice()
+    } else if path.ends_with("/z.bin") {
+        b"xet-body".as_slice()
+    } else {
+        return response("404 Not Found", &[], b"missing");
+    };
+    response(
+        "200 OK",
+        &[("Content-Type", "application/octet-stream")],
+        body,
+    )
+}
+
+fn tree_response(fixture: Fixture, port: u16, path: &str, scans: &AtomicUsize) -> String {
+    let page_two = path.contains("page=2");
+    let scan = if page_two {
+        scans.load(Ordering::SeqCst)
+    } else {
+        scans.fetch_add(1, Ordering::SeqCst) + 1
+    };
+    let body = match fixture {
+        Fixture::Normal
+        | Fixture::Changing
+        | Fixture::BadDigests
+        | Fixture::MismatchedCommit
+        | Fixture::XetRedirect(_)
+            if page_two =>
+        {
+            let size = if matches!(fixture, Fixture::Changing) && scan > 1 {
+                9
+            } else {
+                8
+            };
+            format!(
+                r#"[{{"type":"file","oid":"3333333333333333333333333333333333333333","size":{size},"path":"z.bin","lfs":null,"xetHash":"xet-333"}}]"#
+            )
+        }
+        Fixture::Normal
+        | Fixture::Changing
+        | Fixture::MismatchedCommit
+        | Fixture::XetRedirect(_) => format!(
+            r#"[{{"type":"file","oid":"{}","size":5,"path":"a.txt","lfs":null}},{{"type":"directory","oid":"1111111111111111111111111111111111111111","path":"dir"}},{{"type":"file","oid":"2222222222222222222222222222222222222222","size":8,"path":"dir/lfs.bin","lfs":{{"size":8,"sha256":"{}","pointer_size":64}}}}]"#,
+            git_blob_sha1(b"alpha"),
+            sha256_hex(b"lfs-body")
+        ),
+        Fixture::BadDigests => format!(
+            r#"[{{"type":"file","oid":"ffffffffffffffffffffffffffffffffffffffff","size":5,"path":"a.txt","lfs":null}},{{"type":"directory","oid":"1111111111111111111111111111111111111111","path":"dir"}},{{"type":"file","oid":"2222222222222222222222222222222222222222","size":8,"path":"dir/lfs.bin","lfs":{{"size":8,"sha256":"{}","pointer_size":64}}}}]"#,
+            "f".repeat(64)
+        ),
+        Fixture::OutOfOrderNormalized => format!(
+            r#"[{{"type":"file","oid":"{}","size":5,"path":"Å.txt","lfs":null}},{{"type":"file","oid":"{}","size":5,"path":"Ä.txt","lfs":null}}]"#,
+            git_blob_sha1(b"alpha"),
+            git_blob_sha1(b"alpha")
+        ),
+        Fixture::Collision => format!(
+            r#"[{{"type":"file","oid":"{}","size":5,"path":"café.txt","lfs":null}},{{"type":"file","oid":"{}","size":5,"path":"café.txt","lfs":null}}]"#,
+            git_blob_sha1(b"alpha"),
+            git_blob_sha1(b"alpha")
+        ),
+        Fixture::Empty => format!(
+            r#"[{{"type":"file","oid":"{}","size":5,"path":"a.txt","lfs":null}}]"#,
+            git_blob_sha1(b"alpha")
+        ),
+        Fixture::Symlink => {
+            r#"[{"type":"symlink","oid":"1111111111111111111111111111111111111111","size":4,"path":"link"}]"#.to_owned()
+        }
+        Fixture::Submodule => {
+            r#"[{"type":"submodule","oid":"1111111111111111111111111111111111111111","path":"child"}]"#.to_owned()
+        }
+        Fixture::Content => unreachable!("content fixture has no tree"),
+    };
+    let headers = if !page_two
+        && matches!(
+            fixture,
+            Fixture::Normal
+                | Fixture::Changing
+                | Fixture::BadDigests
+                | Fixture::MismatchedCommit
+                | Fixture::XetRedirect(_)
+        ) {
+        vec![(
+            "Link",
+            format!("<http://source.test:{port}{}&page=2>; rel=\"next\"", path),
+        )]
+    } else {
+        Vec::new()
+    };
+    let borrowed = headers
+        .iter()
+        .map(|(name, value)| (*name, value.as_str()))
+        .collect::<Vec<_>>();
+    response("200 OK", &borrowed, body.as_bytes())
+}
+
+fn response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> String {
+    let mut response = format!(
+        "HTTP/1.1 {status}\r\nConnection: close\r\nContent-Length: {}\r\n",
+        body.len()
+    );
+    for (name, value) in headers {
+        response.push_str(&format!("{name}: {value}\r\n"));
+    }
+    response.push_str("\r\n");
+    response.push_str(&String::from_utf8_lossy(body));
+    response
+}
+
+fn source(port: u16) -> HfSource<LocalResolver> {
+    let policy = OutboundPolicyBuilder::new()
+        .allow_scheme_port("http", port)
+        .allow_non_global("127.0.0.0/8".parse::<IpNet>().expect("network"))
+        .build();
+    HfSource::new(LocalResolver, policy, RetryTiming::default_bounds())
+}
+
+fn source_with_ports(ports: &[u16]) -> HfSource<LocalResolver> {
+    let mut policy = OutboundPolicyBuilder::new()
+        .allow_non_global("127.0.0.0/8".parse::<IpNet>().expect("network"));
+    for &port in ports {
+        policy = policy.allow_scheme_port("http", port);
+    }
+    HfSource::new(LocalResolver, policy.build(), RetryTiming::default_bounds())
+}
+
+fn request(
+    port: u16,
+    repo_type: HfRepoType,
+    revision: &str,
+    selection: HfSelection,
+) -> HfSourceRequest {
+    let repo_id = match repo_type {
+        HfRepoType::Model => "acme/model",
+        HfRepoType::Dataset => "acme/dataset",
+        HfRepoType::Space => "acme/space",
+    };
+    HfSourceRequest::new(
+        &format!("http://source.test:{port}"),
+        repo_type,
+        repo_id,
+        revision,
+        selection,
+    )
+    .expect("request")
+}
+
+async fn members(
+    source: &HfSource<LocalResolver>,
+    identity: &remora_hf_source::HfSnapshotIdentity,
+) -> Result<Vec<HfMember>, HfSourceError> {
+    let mut inventory = source.inventory(identity, HeaderMap::new())?;
+    let mut members = Vec::new();
+    while let Some(member) = inventory.next().await? {
+        members.push(member);
+    }
+    Ok(members)
+}
+
+async fn read_body(
+    mut body: remora_hf_source::HfSourceBody,
+) -> Result<(Vec<u8>, String), StreamError> {
+    let mut bytes = Vec::new();
+    let mut chunk = [0_u8; 3];
+    loop {
+        let count = body.read(&mut chunk).await?;
+        if count == 0 {
+            break;
+        }
+        bytes.extend_from_slice(&chunk[..count]);
+    }
+    Ok((
+        bytes,
+        body.verified_sha256().expect("verified digest").to_owned(),
+    ))
+}
+
+#[tokio::test]
+async fn repeated_model_dataset_space_inventories_are_pinned_paginated_and_streamed() {
+    let server = serve(Fixture::Normal).await;
+    let source = source(server.port);
+    let mut identities = Vec::new();
+    for repo_type in [HfRepoType::Model, HfRepoType::Dataset, HfRepoType::Space] {
+        let identity = source
+            .resolve(request(
+                server.port,
+                repo_type,
+                "main",
+                HfSelection::default(),
+            ))
+            .await
+            .expect("resolve");
+        assert_eq!(identity.commit(), COMMIT);
+        assert_eq!(identity.inventory().member_count, 3);
+        identities.push(identity);
+    }
+    let repeated = source
+        .resolve(request(
+            server.port,
+            HfRepoType::Model,
+            "main",
+            HfSelection::default(),
+        ))
+        .await
+        .expect("repeat");
+    assert_eq!(identities[0].inventory(), repeated.inventory());
+    assert_eq!(
+        members(&source, &identities[0])
+            .await
+            .expect("inventory")
+            .iter()
+            .map(HfMember::path)
+            .collect::<Vec<_>>(),
+        ["a.txt", "dir/lfs.bin", "z.bin"]
+    );
+    assert!(
+        server
+            .observed
+            .lock()
+            .expect("observed")
+            .iter()
+            .any(|request| request.path.contains("page=2"))
+    );
+}
+
+#[tokio::test]
+async fn exact_and_segment_prefix_selection_collapses_duplicates() {
+    let server = serve(Fixture::Normal).await;
+    let source = source(server.port);
+    let selection = HfSelection::new(["dir/lfs.bin", "dir/lfs.bin"], ["a.txt", "dir", "dir"])
+        .expect("selection");
+    let identity = source
+        .resolve(request(server.port, HfRepoType::Model, "main", selection))
+        .await
+        .expect("resolve");
+    assert_eq!(identity.inventory().member_count, 2);
+    assert_eq!(
+        members(&source, &identity)
+            .await
+            .expect("inventory")
+            .iter()
+            .map(HfMember::path)
+            .collect::<Vec<_>>(),
+        ["a.txt", "dir/lfs.bin"]
+    );
+}
+
+#[tokio::test]
+async fn git_lfs_and_xet_members_stream_with_required_identity_checks() {
+    let server = serve(Fixture::Normal).await;
+    let source = source(server.port);
+    let identity = source
+        .resolve(request(
+            server.port,
+            HfRepoType::Model,
+            "main",
+            HfSelection::default(),
+        ))
+        .await
+        .expect("resolve");
+    let members = members(&source, &identity).await.expect("inventory");
+    assert!(matches!(members[0].storage(), HfMemberStorage::Git { .. }));
+    assert!(matches!(members[1].storage(), HfMemberStorage::Lfs { .. }));
+    assert!(matches!(members[2].storage(), HfMemberStorage::Xet { .. }));
+    for member in &members {
+        let body = source
+            .open_member(&identity, member, HeaderMap::new())
+            .await
+            .expect("open");
+        let (bytes, digest) = read_body(body).await.expect("stream");
+        assert_eq!(digest, sha256_hex(&bytes));
+    }
+}
+
+#[tokio::test]
+async fn xet_download_follows_policy_checked_redirect_without_forwarding_credentials() {
+    let content = serve(Fixture::Content).await;
+    let hub = serve(Fixture::XetRedirect(content.port)).await;
+    let source = source_with_ports(&[hub.port, content.port]);
+    let mut credentials = HeaderMap::new();
+    credentials.insert(
+        http::header::AUTHORIZATION,
+        HeaderValue::from_static("Bearer fixture-secret"),
+    );
+    let identity = source
+        .resolve(
+            request(hub.port, HfRepoType::Model, "main", HfSelection::default())
+                .with_credential_headers(credentials.clone())
+                .expect("credentials"),
+        )
+        .await
+        .expect("resolve");
+    let mut inventory = source
+        .inventory(&identity, credentials.clone())
+        .expect("inventory");
+    let mut xet = None;
+    while let Some(member) = inventory.next().await.expect("member") {
+        if matches!(member.storage(), HfMemberStorage::Xet { .. }) {
+            xet = Some(member);
+        }
+    }
+    let body = source
+        .open_member(&identity, &xet.expect("xet"), credentials)
+        .await
+        .expect("open");
+    assert_eq!(read_body(body).await.expect("stream").0, b"xet-body");
+    assert!(
+        content
+            .observed
+            .lock()
+            .expect("content observations")
+            .iter()
+            .all(|request| !request.headers.contains_key("authorization"))
+    );
+}
+
+#[tokio::test]
+async fn credentials_are_not_exposed_by_debug_or_sent_without_explicit_grant() {
+    let server = serve(Fixture::Normal).await;
+    let source = source(server.port);
+    let mut headers = HeaderMap::new();
+    headers.insert(
+        http::header::AUTHORIZATION,
+        HeaderValue::from_static("Bearer fixture-secret"),
+    );
+    let request = request(
+        server.port,
+        HfRepoType::Model,
+        "main",
+        HfSelection::default(),
+    )
+    .with_credential_headers(headers)
+    .expect("credentials");
+    assert!(!format!("{request:?}").contains("fixture-secret"));
+    source.resolve(request).await.expect("resolve");
+    assert!(
+        server
+            .observed
+            .lock()
+            .expect("observed")
+            .iter()
+            .all(|request| request.headers.contains_key("authorization"))
+    );
+}
+
+#[tokio::test]
+async fn collision_empty_selection_symlink_and_submodule_fail_closed() {
+    for (fixture, expected) in [
+        (Fixture::Collision, HfSourceError::InvalidInventory),
+        (Fixture::Empty, HfSourceError::InvalidInventory),
+        (Fixture::Symlink, HfSourceError::UnsupportedEntry),
+        (Fixture::Submodule, HfSourceError::UnsupportedEntry),
+    ] {
+        let server = serve(fixture).await;
+        let source = source(server.port);
+        let selection = if matches!(fixture, Fixture::Empty) {
+            HfSelection::new(["missing.bin"], std::iter::empty::<&str>()).expect("selection")
+        } else {
+            HfSelection::default()
+        };
+        assert_eq!(
+            source
+                .resolve(request(server.port, HfRepoType::Model, "main", selection,))
+                .await
+                .expect_err("must fail"),
+            expected
+        );
+    }
+}
+
+#[tokio::test]
+async fn missing_pinned_commits_and_changed_pinned_inventories_are_source_changed() {
+    let missing = serve(Fixture::Normal).await;
+    let missing_source = source(missing.port);
+    assert_eq!(
+        missing_source
+            .resolve(request(
+                missing.port,
+                HfRepoType::Model,
+                "missing",
+                HfSelection::default(),
+            ))
+            .await
+            .expect_err("missing symbolic ref"),
+        HfSourceError::NotFound
+    );
+    assert_eq!(
+        missing_source
+            .resolve(request(
+                missing.port,
+                HfRepoType::Model,
+                "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                HfSelection::default(),
+            ))
+            .await
+            .expect_err("missing commit"),
+        HfSourceError::SourceChanged
+    );
+
+    let mismatched = serve(Fixture::MismatchedCommit).await;
+    let mismatched_source = source(mismatched.port);
+    assert_eq!(
+        mismatched_source
+            .resolve(request(
+                mismatched.port,
+                HfRepoType::Model,
+                "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                HfSelection::default(),
+            ))
+            .await
+            .expect_err("mismatched commit"),
+        HfSourceError::SourceChanged
+    );
+
+    let changing = serve(Fixture::Changing).await;
+    let source = source(changing.port);
+    let identity = source
+        .resolve(request(
+            changing.port,
+            HfRepoType::Model,
+            "main",
+            HfSelection::default(),
+        ))
+        .await
+        .expect("initial inventory");
+    assert_eq!(
+        members(&source, &identity)
+            .await
+            .expect_err("changed inventory"),
+        HfSourceError::SourceChanged
+    );
+}
+
+#[tokio::test]
+async fn changed_git_or_lfs_payload_fails_during_streaming() {
+    let server = serve(Fixture::BadDigests).await;
+    let source = source(server.port);
+    let identity = source
+        .resolve(request(
+            server.port,
+            HfRepoType::Model,
+            "main",
+            HfSelection::default(),
+        ))
+        .await
+        .expect("resolve");
+    let members = members(&source, &identity).await.expect("inventory");
+    for member in &members[..2] {
+        let body = source
+            .open_member(&identity, member, HeaderMap::new())
+            .await
+            .expect("open");
+        assert!(matches!(
+            read_body(body).await,
+            Err(StreamError::SourceChanged)
+        ));
+    }
+}
+
+#[tokio::test]
+async fn canonical_inventory_sorts_normalized_paths_independently_of_hub_order() {
+    let server = serve(Fixture::OutOfOrderNormalized).await;
+    let source = source(server.port);
+    let identity = source
+        .resolve(request(
+            server.port,
+            HfRepoType::Model,
+            "main",
+            HfSelection::default(),
+        ))
+        .await
+        .expect("resolve");
+    assert_eq!(
+        members(&source, &identity)
+            .await
+            .expect("inventory")
+            .iter()
+            .map(HfMember::path)
+            .collect::<Vec<_>>(),
+        ["Ä.txt", "Å.txt"]
+    );
+}
+
+fn git_blob_sha1(body: &[u8]) -> String {
+    let mut hasher = Sha1::new();
+    hasher.update(format!("blob {}\0", body.len()).as_bytes());
+    hasher.update(body);
+    hex(hasher.finalize())
+}
+
+fn sha256_hex(body: &[u8]) -> String {
+    use sha2::Digest as _;
+    hex(sha2::Sha256::digest(body))
+}
+
+fn hex(bytes: impl AsRef<[u8]>) -> String {
+    const DIGITS: &[u8; 16] = b"0123456789abcdef";
+    let mut output = String::with_capacity(bytes.as_ref().len() * 2);
+    for &byte in bytes.as_ref() {
+        output.push(DIGITS[(byte >> 4) as usize] as char);
+        output.push(DIGITS[(byte & 0x0f) as usize] as char);
+    }
+    output
+}