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

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
}