Add bounded commit-pinned Hugging Face source streaming

This commit is contained in:
ms2k 2026-07-17 06:32:48 +09:00
commit e13c01441c
14 changed files with 3595 additions and 26 deletions

1305
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -9,6 +9,7 @@ members = [
"crates/remora-durable", "crates/remora-durable",
"crates/remora-identity", "crates/remora-identity",
"crates/remora-http-source", "crates/remora-http-source",
"crates/remora-hf-source",
"crates/remora-object-source", "crates/remora-object-source",
"crates/remora-outbound", "crates/remora-outbound",
"crates/remora-result", "crates/remora-result",
@ -26,11 +27,13 @@ version = "0.0.0"
[workspace.dependencies] [workspace.dependencies]
base64 = "0.22.1" base64 = "0.22.1"
hmac = "0.13.0" hmac = "0.13.0"
hf-hub = "=1.0.0"
remora-domain = { path = "crates/remora-domain" } remora-domain = { path = "crates/remora-domain" }
remora-durable = { path = "crates/remora-durable" } remora-durable = { path = "crates/remora-durable" }
serde = { version = "1.0.228", features = ["derive"] } serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150" serde_json = "1.0.150"
serde_json_canonicalizer = "0.3.2" serde_json_canonicalizer = "0.3.2"
sha1 = "0.11.0"
sha2 = "0.11.0" sha2 = "0.11.0"
sqlx = { version = "0.9.0", default-features = false, features = [ sqlx = { version = "0.9.0", default-features = false, features = [
"json", "json",

View file

@ -272,6 +272,35 @@ COPY crates crates
COPY --from=kind-object-source-verified /tmp/remora-object-source-tests.tar.zst /tests/remora-object-source-tests.tar.zst COPY --from=kind-object-source-verified /tmp/remora-object-source-tests.tar.zst /tests/remora-object-source-tests.tar.zst
CMD ["cargo-nextest", "nextest", "run", "--archive-file", "/tests/remora-object-source-tests.tar.zst", "--workspace-remap", "/workspace", "--test-threads=1"] CMD ["cargo-nextest", "nextest", "run", "--archive-file", "/tests/remora-object-source-tests.tar.zst", "--workspace-remap", "/workspace", "--test-threads=1"]
FROM rust-tools AS kind-hf-source-build
WORKDIR /workspace
COPY Cargo.toml Cargo.lock rust-toolchain.toml ./
COPY .config .config
COPY contracts contracts
COPY crates crates
COPY docs/api docs/api
RUN cargo nextest archive \
--package remora-hf-source \
--build-jobs 4 \
--locked \
--archive-file /tmp/remora-hf-source-tests.tar.zst
FROM kind-hf-source-build AS kind-hf-source-verified
RUN cargo clippy \
--package remora-hf-source \
--all-targets \
--all-features \
--locked \
-- -D warnings \
&& cargo test --package remora-hf-source --doc --locked
FROM rust-tools AS kind-hf-source-runner
WORKDIR /workspace
COPY Cargo.toml ./
COPY crates crates
COPY --from=kind-hf-source-verified /tmp/remora-hf-source-tests.tar.zst /tests/remora-hf-source-tests.tar.zst
CMD ["cargo-nextest", "nextest", "run", "--archive-file", "/tests/remora-hf-source-tests.tar.zst", "--workspace-remap", "/workspace"]
FROM ${NODE_IMAGE} AS node-runtime FROM ${NODE_IMAGE} AS node-runtime
FROM ${PLAYWRIGHT_IMAGE} AS web-dependencies FROM ${PLAYWRIGHT_IMAGE} AS web-dependencies

View file

@ -1,4 +1,4 @@
.PHONY: check-dependencies check-fixture-lib check-kind-sample check-durable check-coordination check-identity check-credentials check-admission check-stream check-target check-result check-outbound check-http-source check-object-source check-fast check-full check-fast-host check-full-host rust-check web-install web-check web-test-fast web-test-full .PHONY: check-dependencies check-fixture-lib check-kind-sample check-durable check-coordination check-identity check-credentials check-admission check-stream check-target check-result check-outbound check-http-source check-object-source check-hf-source check-fast check-full check-fast-host check-full-host rust-check web-install web-check web-test-fast web-test-full
CONTAINER_ENGINE ?= nerdctl CONTAINER_ENGINE ?= nerdctl
CONTAINER_RUN_FLAGS ?= --memory=16g --memory-swap=24g --cpus=4 CONTAINER_RUN_FLAGS ?= --memory=16g --memory-swap=24g --cpus=4
@ -91,3 +91,6 @@ check-http-source:
check-object-source: check-object-source:
./tests/kind/run-object-source-job.sh ./tests/kind/run-object-source-job.sh
check-hf-source:
./tests/kind/run-hf-source-job.sh

View file

@ -0,0 +1,29 @@
[package]
name = "remora-hf-source"
edition.workspace = true
publish = false
rust-version.workspace = true
version.workspace = true
[dependencies]
bytes = "1.12.1"
hf-hub.workspace = true
http = "1.4.0"
httpdate = "1.0.3"
remora-domain.workspace = true
remora-outbound = { path = "../remora-outbound" }
remora-stream = { path = "../remora-stream" }
reqwest = { version = "0.12.28", default-features = false }
serde.workspace = true
serde_json.workspace = true
serde_json_canonicalizer.workspace = true
sha1.workspace = true
sha2.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["fs", "io-util"] }
tempfile = "3.27.0"
url = "2.5.8"
[dev-dependencies]
ipnet = "2.12.0"
tokio = { workspace = true, features = ["io-util", "net"] }

View file

@ -0,0 +1,123 @@
use bytes::Bytes;
use remora_stream::{SourceBody, StreamError};
use sha1::{Digest as _, Sha1};
use sha2::Sha256;
use crate::HfMemberStorage;
use crate::digest_hex;
enum ExpectedDigest {
Git(String),
Sha256(String),
SizeOnly,
}
pub struct HfSourceBody {
response: reqwest::Response,
pending: Bytes,
pending_offset: usize,
remaining: u64,
expected: ExpectedDigest,
sha1: Sha1,
sha256: Sha256,
verified_sha256: Option<String>,
finished: bool,
}
impl HfSourceBody {
pub(crate) fn new(response: reqwest::Response, size: u64, storage: &HfMemberStorage) -> Self {
let expected = match storage {
HfMemberStorage::Git { blob_sha1 } => ExpectedDigest::Git(blob_sha1.clone()),
HfMemberStorage::Lfs { content_sha256, .. } => {
ExpectedDigest::Sha256(content_sha256.clone())
}
HfMemberStorage::Xet { .. } => ExpectedDigest::SizeOnly,
};
let mut sha1 = Sha1::new();
sha1.update(format!("blob {size}\0").as_bytes());
Self {
response,
pending: Bytes::new(),
pending_offset: 0,
remaining: size,
expected,
sha1,
sha256: Sha256::new(),
verified_sha256: None,
finished: false,
}
}
#[must_use]
pub fn verified_sha256(&self) -> Option<&str> {
self.verified_sha256.as_deref()
}
fn copy_pending(&mut self, output: &mut [u8]) -> usize {
let available = &self.pending[self.pending_offset..];
let count = available.len().min(output.len());
let copied = &available[..count];
output[..count].copy_from_slice(copied);
self.sha1.update(copied);
self.sha256.update(copied);
self.pending_offset += count;
self.remaining -= count as u64;
if self.pending_offset == self.pending.len() {
self.pending = Bytes::new();
self.pending_offset = 0;
}
count
}
fn verify_eof(&mut self) -> Result<(), StreamError> {
if self.remaining != 0 {
return Err(StreamError::transient_source(
"Hugging Face response ended before its declared length",
));
}
let sha1 = digest_hex(self.sha1.clone().finalize());
let sha256 = digest_hex(self.sha256.clone().finalize());
match &self.expected {
ExpectedDigest::Git(expected) if expected != &sha1 => {
return Err(StreamError::SourceChanged);
}
ExpectedDigest::Sha256(expected) if expected != &sha256 => {
return Err(StreamError::SourceChanged);
}
ExpectedDigest::Git(_) | ExpectedDigest::Sha256(_) | ExpectedDigest::SizeOnly => {}
}
self.verified_sha256 = Some(sha256);
self.finished = true;
Ok(())
}
}
impl SourceBody for HfSourceBody {
async fn read(&mut self, output: &mut [u8]) -> Result<usize, StreamError> {
if output.is_empty() {
return Ok(0);
}
if self.pending_offset < self.pending.len() {
return Ok(self.copy_pending(output));
}
if self.finished {
return Ok(0);
}
match self.response.chunk().await {
Ok(Some(chunk)) => {
if chunk.len() as u64 > self.remaining {
return Err(StreamError::SourceChanged);
}
self.pending = chunk;
Ok(self.copy_pending(output))
}
Ok(None) => {
self.verify_eof()?;
Ok(0)
}
Err(_) => Err(StreamError::transient_source(
"Hugging Face response body failed",
)),
}
}
}

View file

@ -0,0 +1,25 @@
use std::time::Duration;
use thiserror::Error;
#[derive(Clone, Debug, Eq, Error, PartialEq)]
pub enum HfSourceError {
#[error("Hugging Face source request is invalid")]
InvalidRequest,
#[error("Hugging Face source authorization failed")]
Authorization,
#[error("Hugging Face source was not found")]
NotFound,
#[error("Hugging Face source changed")]
SourceChanged,
#[error("Hugging Face source returned an unsupported entry")]
UnsupportedEntry,
#[error("Hugging Face source response is invalid")]
InvalidResponse,
#[error("Hugging Face source inventory is invalid")]
InvalidInventory,
#[error("Hugging Face inventory capacity is exhausted")]
Capacity,
#[error("Hugging Face source failed transiently")]
Transient { retry_after: Option<Duration> },
}

View file

@ -0,0 +1,24 @@
mod body;
mod error;
mod sort;
mod source;
mod types;
pub use body::HfSourceBody;
pub use error::HfSourceError;
pub use source::{HfInventory, HfSource};
pub use types::{
HfInventoryDescriptor, HfMember, HfMemberStorage, HfRepoType, HfSelection, HfSnapshotIdentity,
HfSourceRequest,
};
fn digest_hex(bytes: impl AsRef<[u8]>) -> String {
const DIGITS: &[u8; 16] = b"0123456789abcdef";
let bytes = bytes.as_ref();
let mut encoded = String::with_capacity(bytes.len() * 2);
for &byte in bytes {
encoded.push(DIGITS[(byte >> 4) as usize] as char);
encoded.push(DIGITS[(byte & 0x0f) as usize] as char);
}
encoded
}

View file

@ -0,0 +1,302 @@
use std::path::{Path, PathBuf};
use sha2::{Digest, Sha256};
use tempfile::TempDir;
use tokio::fs::File;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use crate::types::INVENTORY_DESCRIPTOR_VERSION;
use crate::{HfInventoryDescriptor, HfMember, HfSourceError, digest_hex};
const CHUNK_BYTES: usize = 4 * 1024 * 1024;
const MAX_MEMBER_BYTES: usize = 16 * 1024 * 1024;
const RUN_LEVELS: usize = 64;
pub(crate) struct ExternalSorter {
directory: TempDir,
levels: Vec<Option<PathBuf>>,
chunk: Vec<HfMember>,
chunk_bytes: usize,
next_file: u64,
}
pub(crate) struct SortedMembers {
_directory: TempDir,
reader: BufReader<File>,
}
impl ExternalSorter {
pub(crate) fn new() -> Result<Self, HfSourceError> {
Ok(Self {
directory: tempfile::tempdir().map_err(|_| HfSourceError::Capacity)?,
levels: (0..RUN_LEVELS).map(|_| None).collect(),
chunk: Vec::new(),
chunk_bytes: 0,
next_file: 0,
})
}
pub(crate) async fn push(&mut self, member: HfMember) -> Result<(), HfSourceError> {
let encoded_bytes = member.encode()?.len();
if encoded_bytes > MAX_MEMBER_BYTES {
return Err(HfSourceError::InvalidInventory);
}
if !self.chunk.is_empty() && self.chunk_bytes.saturating_add(encoded_bytes) > CHUNK_BYTES {
self.flush_chunk().await?;
}
self.chunk_bytes = self.chunk_bytes.saturating_add(encoded_bytes);
self.chunk.push(member);
Ok(())
}
pub(crate) async fn finish(mut self) -> Result<SortedMembers, HfSourceError> {
self.flush_chunk().await?;
let mut combined: Option<PathBuf> = None;
for level in 0..RUN_LEVELS {
let Some(run) = self.levels[level].take() else {
continue;
};
combined = Some(match combined {
None => run,
Some(previous) => self.merge(&previous, &run).await?,
});
}
let path = match combined {
Some(path) => path,
None => {
let path = self.next_path();
File::create(&path)
.await
.map_err(|_| HfSourceError::Capacity)?;
path
}
};
let reader = BufReader::new(
File::open(&path)
.await
.map_err(|_| HfSourceError::Capacity)?,
);
Ok(SortedMembers {
_directory: self.directory,
reader,
})
}
async fn flush_chunk(&mut self) -> Result<(), HfSourceError> {
if self.chunk.is_empty() {
return Ok(());
}
self.chunk
.sort_unstable_by(|left, right| left.path().cmp(right.path()));
let path = self.next_path();
let mut writer = BufWriter::new(
File::create(&path)
.await
.map_err(|_| HfSourceError::Capacity)?,
);
for member in self.chunk.drain(..) {
write_member(&mut writer, &member).await?;
}
writer.flush().await.map_err(|_| HfSourceError::Capacity)?;
self.chunk_bytes = 0;
self.insert_run(path).await
}
async fn insert_run(&mut self, mut run: PathBuf) -> Result<(), HfSourceError> {
for level in 0..RUN_LEVELS {
match self.levels[level].take() {
None => {
self.levels[level] = Some(run);
return Ok(());
}
Some(existing) => run = self.merge(&existing, &run).await?,
}
}
Err(HfSourceError::Capacity)
}
async fn merge(&mut self, left: &Path, right: &Path) -> Result<PathBuf, HfSourceError> {
let path = self.next_path();
let left_path = left.to_owned();
let right_path = right.to_owned();
let mut left = BufReader::new(
File::open(left)
.await
.map_err(|_| HfSourceError::Capacity)?,
);
let mut right = BufReader::new(
File::open(right)
.await
.map_err(|_| HfSourceError::Capacity)?,
);
let mut writer = BufWriter::new(
File::create(&path)
.await
.map_err(|_| HfSourceError::Capacity)?,
);
let mut left_member = read_member(&mut left).await?;
let mut right_member = read_member(&mut right).await?;
while left_member.is_some() || right_member.is_some() {
let take_left = match (&left_member, &right_member) {
(Some(left), Some(right)) => left.path() <= right.path(),
(Some(_), None) => true,
(None, Some(_)) => false,
(None, None) => break,
};
if take_left {
let member = left_member
.as_ref()
.ok_or(HfSourceError::InvalidInventory)?;
write_member(&mut writer, member).await?;
left_member = read_member(&mut left).await?;
} else {
let member = right_member
.as_ref()
.ok_or(HfSourceError::InvalidInventory)?;
write_member(&mut writer, member).await?;
right_member = read_member(&mut right).await?;
}
}
writer.flush().await.map_err(|_| HfSourceError::Capacity)?;
drop(left);
drop(right);
drop(writer);
tokio::fs::remove_file(left_path)
.await
.map_err(|_| HfSourceError::Capacity)?;
tokio::fs::remove_file(right_path)
.await
.map_err(|_| HfSourceError::Capacity)?;
Ok(path)
}
fn next_path(&mut self) -> PathBuf {
let path = self.directory.path().join(self.next_file.to_string());
self.next_file = self.next_file.saturating_add(1);
path
}
}
impl SortedMembers {
pub(crate) async fn descriptor(&mut self) -> Result<HfInventoryDescriptor, HfSourceError> {
let mut hasher = Sha256::new();
let mut member_count = 0_u64;
let mut previous_path: Option<String> = None;
while let Some(member) = read_member(&mut self.reader).await? {
if previous_path
.as_ref()
.is_some_and(|previous| previous.as_str() >= member.path())
{
return Err(HfSourceError::InvalidInventory);
}
let encoded = member.encode()?;
hasher.update((encoded.len() as u64).to_be_bytes());
hasher.update(encoded);
member_count = member_count.checked_add(1).ok_or(HfSourceError::Capacity)?;
previous_path = Some(member.path().to_owned());
}
if member_count == 0 {
return Err(HfSourceError::InvalidInventory);
}
Ok(HfInventoryDescriptor {
version: INVENTORY_DESCRIPTOR_VERSION,
member_count,
sha256: digest_hex(hasher.finalize()),
})
}
pub(crate) async fn rewind(&mut self) -> Result<(), HfSourceError> {
use tokio::io::{AsyncSeekExt, SeekFrom};
self.reader
.seek(SeekFrom::Start(0))
.await
.map_err(|_| HfSourceError::Capacity)?;
Ok(())
}
pub(crate) async fn next(&mut self) -> Result<Option<HfMember>, HfSourceError> {
read_member(&mut self.reader).await
}
}
async fn write_member(
writer: &mut BufWriter<File>,
member: &HfMember,
) -> Result<(), HfSourceError> {
let encoded = member.encode()?;
writer
.write_all(&(encoded.len() as u64).to_be_bytes())
.await
.map_err(|_| HfSourceError::Capacity)?;
writer
.write_all(&encoded)
.await
.map_err(|_| HfSourceError::Capacity)
}
async fn read_member(reader: &mut BufReader<File>) -> Result<Option<HfMember>, HfSourceError> {
let mut length = [0_u8; 8];
match reader.read(&mut length[..1]).await {
Ok(0) => return Ok(None),
Ok(1) => {}
Ok(_) => return Err(HfSourceError::InvalidInventory),
Err(_) => return Err(HfSourceError::Capacity),
}
reader
.read_exact(&mut length[1..])
.await
.map_err(|_| HfSourceError::InvalidInventory)?;
let length =
usize::try_from(u64::from_be_bytes(length)).map_err(|_| HfSourceError::InvalidInventory)?;
if length > MAX_MEMBER_BYTES {
return Err(HfSourceError::InvalidInventory);
}
let mut encoded = vec![0_u8; length];
reader
.read_exact(&mut encoded)
.await
.map_err(|_| HfSourceError::InvalidInventory)?;
HfMember::decode(&encoded).map(Some)
}
#[cfg(test)]
mod tests {
use crate::{HfMemberStorage, HfRepoType};
use super::*;
#[tokio::test]
async fn merges_multiple_bounded_runs_into_canonical_order() {
let mut sorter = ExternalSorter::new().expect("sorter");
let suffix = "x".repeat(450);
for index in (0..10_000).rev() {
sorter
.push(HfMember::from_parts(
HfRepoType::Model,
"acme/model".to_owned(),
"0123456789abcdef0123456789abcdef01234567".to_owned(),
format!("{index:05}-{suffix}"),
format!("{index:05}-{suffix}"),
1,
HfMemberStorage::Git {
blob_sha1: "1111111111111111111111111111111111111111".to_owned(),
},
))
.await
.expect("push");
}
let mut sorted = sorter.finish().await.expect("finish");
assert_eq!(
sorted.descriptor().await.expect("descriptor").member_count,
10_000
);
sorted.rewind().await.expect("rewind");
for index in 0..10_000 {
assert_eq!(
sorted.next().await.expect("next").expect("member").path(),
format!("{index:05}-{suffix}")
);
}
assert!(sorted.next().await.expect("end").is_none());
}
}

View file

@ -0,0 +1,669 @@
use std::collections::VecDeque;
use std::time::{Duration, SystemTime};
use hf_hub::repository::RepoTreeEntry;
use http::{HeaderMap, HeaderValue, StatusCode, header};
use remora_domain::MemberPath;
use remora_outbound::{
CanonicalUrl, OutboundClient, OutboundError, OutboundPolicy, OutboundRequest, Resolver,
};
use remora_stream::{RetryTiming, capped_server_hint};
use serde::Deserialize;
use serde_json::Value;
use tokio::time::timeout;
use url::Url;
use crate::sort::{ExternalSorter, SortedMembers};
use crate::types::{INVENTORY_DESCRIPTOR_VERSION, is_full_commit, validate_repo_id};
use crate::{
HfInventoryDescriptor, HfMember, HfMemberStorage, HfSnapshotIdentity, HfSourceBody,
HfSourceError, HfSourceRequest,
};
const MAX_METADATA_BYTES: usize = 16 * 1024 * 1024;
const MAX_INVENTORY_PAGES: u32 = 1_000_000;
pub struct HfSource<R> {
client: OutboundClient<R>,
timing: RetryTiming,
}
impl<R: Resolver> HfSource<R> {
#[must_use]
pub fn new(resolver: R, policy: OutboundPolicy, timing: RetryTiming) -> Self {
let timing = timing.sanitized();
Self {
client: OutboundClient::builder(resolver, policy)
.connect_timeout(timing.connect_timeout)
.build(),
timing,
}
}
pub async fn resolve(
&self,
request: HfSourceRequest,
) -> Result<HfSnapshotIdentity, HfSourceError> {
let info_url = info_url(&request)?;
let response = self
.execute(build_request(
&info_url,
request.credential_headers.clone(),
true,
)?)
.await?;
let status = response.response().status();
if status != StatusCode::OK {
return Err(status_error(
status,
response.response().headers(),
is_full_commit(&request.revision),
&self.timing,
));
}
let info: RepoInfo = decode_response(
response.into_response(),
1024 * 1024,
self.timing.no_progress_timeout,
)
.await?;
validate_repo_id(&info.id).map_err(|_| HfSourceError::InvalidResponse)?;
if !same_repo(&request.repo_id, &info.id) {
return Err(HfSourceError::InvalidResponse);
}
let commit = info.sha.to_ascii_lowercase();
if !is_full_commit(&commit) {
return Err(HfSourceError::InvalidResponse);
}
if is_full_commit(&request.revision) && !request.revision.eq_ignore_ascii_case(&commit) {
return Err(HfSourceError::SourceChanged);
}
let mut scanner = InventoryScanner::new(
self,
&request.endpoint,
request.repo_type,
&info.id,
&commit,
&request.selection,
request.credential_headers.clone(),
None,
)?;
scanner.prepare().await?;
let inventory = scanner
.descriptor()
.cloned()
.ok_or(HfSourceError::InvalidInventory)?;
Ok(HfSnapshotIdentity::from_parts(
request.endpoint,
request.repo_type,
info.id,
commit,
request.selection,
inventory,
))
}
pub fn inventory<'a>(
&'a self,
snapshot: &'a HfSnapshotIdentity,
credential_headers: HeaderMap,
) -> Result<HfInventory<'a, R>, HfSourceError> {
validate_snapshot(snapshot)?;
let scanner = InventoryScanner::new(
self,
snapshot.endpoint(),
snapshot.repo_type(),
snapshot.repo_id(),
snapshot.commit(),
snapshot.selection(),
credential_headers,
Some(snapshot.inventory()),
)?;
Ok(HfInventory { scanner })
}
pub async fn open_member(
&self,
snapshot: &HfSnapshotIdentity,
member: &HfMember,
credential_headers: HeaderMap,
) -> Result<HfSourceBody, HfSourceError> {
validate_snapshot(snapshot)?;
if !member.belongs_to(snapshot) {
return Err(HfSourceError::InvalidRequest);
}
validate_member(member)?;
let url = download_url(snapshot, member.source_path())?;
let response = self
.execute(build_request(&url, credential_headers, false)?)
.await?;
if response.response().status() != StatusCode::OK {
return Err(status_error(
response.response().status(),
response.response().headers(),
true,
&self.timing,
));
}
ensure_identity_encoding(response.response().headers())?;
if let Some(length) = content_length(response.response().headers())?
&& length != member.size()
{
return Err(HfSourceError::SourceChanged);
}
Ok(HfSourceBody::new(
response.into_response(),
member.size(),
member.storage(),
))
}
async fn execute(
&self,
request: OutboundRequest,
) -> Result<remora_outbound::OutboundResponse, HfSourceError> {
match timeout(
self.timing.response_start_timeout,
self.client.execute(request),
)
.await
{
Ok(Ok(response)) => Ok(response),
Ok(Err(error)) => Err(map_outbound_error(error)),
Err(_) => Err(HfSourceError::Transient { retry_after: None }),
}
}
}
pub struct HfInventory<'a, R> {
scanner: InventoryScanner<'a, R>,
}
impl<R: Resolver> HfInventory<'_, R> {
pub async fn next(&mut self) -> Result<Option<HfMember>, HfSourceError> {
self.scanner.next().await
}
}
struct InventoryScanner<'a, R> {
source: &'a HfSource<R>,
endpoint_origin: remora_outbound::Origin,
next_url: Option<CanonicalUrl>,
page_count: u32,
pending: VecDeque<Value>,
selection: &'a crate::HfSelection,
repo_type: crate::HfRepoType,
repo_id: String,
commit: String,
credentials: HeaderMap,
expected: Option<&'a HfInventoryDescriptor>,
sorter: Option<ExternalSorter>,
sorted: Option<SortedMembers>,
descriptor: Option<HfInventoryDescriptor>,
}
impl<'a, R: Resolver> InventoryScanner<'a, R> {
#[allow(clippy::too_many_arguments)]
fn new(
source: &'a HfSource<R>,
endpoint: &str,
repo_type: crate::HfRepoType,
repo_id: &str,
commit: &str,
selection: &'a crate::HfSelection,
credentials: HeaderMap,
expected: Option<&'a HfInventoryDescriptor>,
) -> Result<Self, HfSourceError> {
let endpoint = CanonicalUrl::parse(endpoint).map_err(|_| HfSourceError::InvalidRequest)?;
let next_url = tree_url(endpoint.as_url(), repo_type.plural(), repo_id, commit)?;
Ok(Self {
source,
endpoint_origin: endpoint
.origin()
.map_err(|_| HfSourceError::InvalidRequest)?,
next_url: Some(next_url),
page_count: 0,
pending: VecDeque::new(),
selection,
repo_type,
repo_id: repo_id.to_owned(),
commit: commit.to_owned(),
credentials,
expected,
sorter: Some(ExternalSorter::new()?),
sorted: None,
descriptor: None,
})
}
async fn next(&mut self) -> Result<Option<HfMember>, HfSourceError> {
self.prepare().await?;
self.sorted
.as_mut()
.ok_or(HfSourceError::InvalidInventory)?
.next()
.await
}
async fn prepare(&mut self) -> Result<(), HfSourceError> {
if self.sorted.is_some() {
return Ok(());
}
loop {
if let Some(value) = self.pending.pop_front() {
if let Some(member) = self.decode_entry(value)? {
self.sorter
.as_mut()
.ok_or(HfSourceError::InvalidInventory)?
.push(member)
.await?;
}
continue;
}
let Some(url) = self.next_url.take() else {
break;
};
self.page_count = self
.page_count
.checked_add(1)
.ok_or(HfSourceError::Capacity)?;
if self.page_count > MAX_INVENTORY_PAGES {
return Err(HfSourceError::InvalidResponse);
}
let response = self
.source
.execute(build_request(&url, self.credentials.clone(), true)?)
.await?;
if response.response().status() != StatusCode::OK {
return Err(status_error(
response.response().status(),
response.response().headers(),
true,
&self.source.timing,
));
}
self.next_url = next_page(response.response().headers(), &url, &self.endpoint_origin)?;
let values: Vec<Value> = decode_response(
response.into_response(),
MAX_METADATA_BYTES,
self.source.timing.no_progress_timeout,
)
.await?;
self.pending = values.into();
}
let mut sorted = self
.sorter
.take()
.ok_or(HfSourceError::InvalidInventory)?
.finish()
.await?;
let descriptor = sorted.descriptor().await?;
if self
.expected
.is_some_and(|expected| expected != &descriptor)
{
return Err(HfSourceError::SourceChanged);
}
sorted.rewind().await?;
self.descriptor = Some(descriptor);
self.sorted = Some(sorted);
Ok(())
}
fn decode_entry(&self, value: Value) -> Result<Option<HfMember>, HfSourceError> {
match value.get("type").and_then(Value::as_str) {
Some("directory") => {
serde_json::from_value::<RepoTreeEntry>(value)
.map_err(|_| HfSourceError::InvalidResponse)?;
Ok(None)
}
Some("file") => {
let entry = serde_json::from_value::<RepoTreeEntry>(value)
.map_err(|_| HfSourceError::InvalidResponse)?;
let RepoTreeEntry::File {
oid,
size,
path: source_path,
lfs,
xet_hash,
..
} = entry
else {
return Err(HfSourceError::InvalidResponse);
};
let canonical_path = MemberPath::parse(&source_path)
.map_err(|_| HfSourceError::InvalidInventory)?
.canonical()
.to_owned();
if !self.selection.matches(&canonical_path) {
return Ok(None);
}
let storage = member_storage(&oid, size, lfs.as_ref(), xet_hash.as_deref())?;
let member = HfMember::from_parts(
self.repo_type,
self.repo_id.clone(),
self.commit.clone(),
canonical_path.clone(),
source_path,
size,
storage,
);
Ok(Some(member))
}
Some(_) => Err(HfSourceError::UnsupportedEntry),
None => Err(HfSourceError::InvalidResponse),
}
}
fn descriptor(&self) -> Option<&HfInventoryDescriptor> {
self.descriptor.as_ref()
}
}
fn member_storage(
oid: &str,
size: u64,
lfs: Option<&hf_hub::repository::BlobLfsInfo>,
xet_hash: Option<&str>,
) -> Result<HfMemberStorage, HfSourceError> {
let oid = oid.to_ascii_lowercase();
if !is_hex(&oid, 40) || lfs.is_some() && xet_hash.is_some() {
return Err(HfSourceError::InvalidInventory);
}
if let Some(lfs) = lfs {
if lfs.size.is_some_and(|reported| reported != size) {
return Err(HfSourceError::InvalidInventory);
}
let content_sha256 = lfs
.sha256
.as_deref()
.map(str::to_ascii_lowercase)
.filter(|digest| is_hex(digest, 64))
.ok_or(HfSourceError::InvalidInventory)?;
return Ok(HfMemberStorage::Lfs {
pointer_blob_sha1: oid,
content_sha256,
});
}
if let Some(hash) = xet_hash {
if hash.is_empty() || hash.chars().any(char::is_control) {
return Err(HfSourceError::InvalidInventory);
}
return Ok(HfMemberStorage::Xet {
hash: hash.to_owned(),
});
}
Ok(HfMemberStorage::Git { blob_sha1: oid })
}
fn info_url(request: &HfSourceRequest) -> Result<CanonicalUrl, HfSourceError> {
let mut url = Url::parse(&request.endpoint).map_err(|_| HfSourceError::InvalidRequest)?;
append_segments(
&mut url,
["api", request.repo_type.plural()]
.into_iter()
.chain(request.repo_id.split('/'))
.chain(["revision", request.revision.as_str()]),
)?;
CanonicalUrl::parse(url.as_str()).map_err(|_| HfSourceError::InvalidRequest)
}
fn tree_url(
endpoint: &Url,
plural: &str,
repo_id: &str,
commit: &str,
) -> Result<CanonicalUrl, HfSourceError> {
let mut url = endpoint.clone();
append_segments(
&mut url,
["api", plural]
.into_iter()
.chain(repo_id.split('/'))
.chain(["tree", commit]),
)?;
url.query_pairs_mut()
.append_pair("recursive", "true")
.append_pair("expand", "true");
CanonicalUrl::parse(url.as_str()).map_err(|_| HfSourceError::InvalidRequest)
}
fn download_url(snapshot: &HfSnapshotIdentity, path: &str) -> Result<CanonicalUrl, HfSourceError> {
let mut url = Url::parse(snapshot.endpoint()).map_err(|_| HfSourceError::InvalidRequest)?;
let prefix = snapshot.repo_type().url_prefix().trim_end_matches('/');
let leading = (!prefix.is_empty()).then_some(prefix).into_iter();
append_segments(
&mut url,
leading
.chain(snapshot.repo_id().split('/'))
.chain(["resolve", snapshot.commit()])
.chain(path.split('/')),
)?;
CanonicalUrl::parse(url.as_str()).map_err(|_| HfSourceError::InvalidRequest)
}
fn append_segments<'a>(
url: &mut Url,
segments: impl IntoIterator<Item = &'a str>,
) -> Result<(), HfSourceError> {
url.path_segments_mut()
.map_err(|_| HfSourceError::InvalidRequest)?
.pop_if_empty()
.extend(segments);
Ok(())
}
fn build_request(
url: &CanonicalUrl,
credentials: HeaderMap,
json: bool,
) -> Result<OutboundRequest, HfSourceError> {
let mut public = HeaderMap::new();
public.insert(
header::ACCEPT_ENCODING,
HeaderValue::from_static("identity"),
);
if json {
public.insert(header::ACCEPT, HeaderValue::from_static("application/json"));
}
OutboundRequest::get(url.as_url().as_str())
.and_then(|request| request.with_public_headers(public))
.and_then(|request| request.with_credential_headers(credentials))
.map_err(|_| HfSourceError::InvalidRequest)
}
async fn decode_response<T: for<'de> Deserialize<'de>>(
mut response: reqwest::Response,
limit: usize,
no_progress_timeout: Duration,
) -> Result<T, HfSourceError> {
let mut bytes = Vec::new();
loop {
let chunk = timeout(no_progress_timeout, response.chunk())
.await
.map_err(|_| HfSourceError::Transient { retry_after: None })?
.map_err(|_| HfSourceError::Transient { retry_after: None })?;
let Some(chunk) = chunk else {
break;
};
if bytes.len().saturating_add(chunk.len()) > limit {
return Err(HfSourceError::InvalidResponse);
}
bytes.extend_from_slice(&chunk);
}
serde_json::from_slice(&bytes).map_err(|_| HfSourceError::InvalidResponse)
}
fn next_page(
headers: &HeaderMap,
current: &CanonicalUrl,
endpoint_origin: &remora_outbound::Origin,
) -> Result<Option<CanonicalUrl>, HfSourceError> {
let Some(link) = headers.get(header::LINK) else {
return Ok(None);
};
let link = link.to_str().map_err(|_| HfSourceError::InvalidResponse)?;
for part in link.split(',').map(str::trim) {
if !part
.split(';')
.skip(1)
.any(|parameter| parameter.trim().eq_ignore_ascii_case("rel=\"next\""))
{
continue;
}
let target = part
.strip_prefix('<')
.and_then(|value| value.split_once('>'))
.map(|(target, _)| target)
.ok_or(HfSourceError::InvalidResponse)?;
let next = current
.join(target)
.map_err(|_| HfSourceError::InvalidResponse)?;
if next.origin().map_err(|_| HfSourceError::InvalidResponse)? != *endpoint_origin {
return Err(HfSourceError::InvalidResponse);
}
return Ok(Some(next));
}
Ok(None)
}
fn ensure_identity_encoding(headers: &HeaderMap) -> Result<(), HfSourceError> {
match headers.get(header::CONTENT_ENCODING) {
None => Ok(()),
Some(value) if value.as_bytes().eq_ignore_ascii_case(b"identity") => Ok(()),
Some(_) => Err(HfSourceError::InvalidResponse),
}
}
fn content_length(headers: &HeaderMap) -> Result<Option<u64>, HfSourceError> {
headers
.get(header::CONTENT_LENGTH)
.map(|value| {
value
.to_str()
.map_err(|_| HfSourceError::InvalidResponse)?
.parse::<u64>()
.map_err(|_| HfSourceError::InvalidResponse)
})
.transpose()
}
fn status_error(
status: StatusCode,
headers: &HeaderMap,
pinned: bool,
timing: &RetryTiming,
) -> HfSourceError {
match status {
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => HfSourceError::Authorization,
StatusCode::NOT_FOUND | StatusCode::GONE if pinned => HfSourceError::SourceChanged,
StatusCode::NOT_FOUND | StatusCode::GONE => HfSourceError::NotFound,
status if status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error() => {
HfSourceError::Transient {
retry_after: parse_retry_after(headers, timing, SystemTime::now()),
}
}
_ => HfSourceError::InvalidResponse,
}
}
fn parse_retry_after(
headers: &HeaderMap,
timing: &RetryTiming,
now: SystemTime,
) -> Option<Duration> {
let value = headers.get(header::RETRY_AFTER)?.to_str().ok()?;
let hint = value
.parse::<u64>()
.ok()
.map(Duration::from_secs)
.or_else(|| {
httpdate::parse_http_date(value)
.ok()?
.duration_since(now)
.ok()
});
capped_server_hint(hint, timing)
}
fn map_outbound_error(error: OutboundError) -> HfSourceError {
match error {
OutboundError::Transport | OutboundError::ResolutionFailed => {
HfSourceError::Transient { retry_after: None }
}
OutboundError::InvalidUrl | OutboundError::InvalidHeaders => HfSourceError::InvalidRequest,
OutboundError::EmptyResolution
| OutboundError::AddressDenied
| OutboundError::SchemePortDenied
| OutboundError::HttpsDowngrade => HfSourceError::Authorization,
OutboundError::InvalidRedirect | OutboundError::RedirectLimit => {
HfSourceError::InvalidResponse
}
}
}
fn is_hex(value: &str, length: usize) -> bool {
value.len() == length && value.bytes().all(|byte| byte.is_ascii_hexdigit())
}
fn is_lower_hex(value: &str, length: usize) -> bool {
value.len() == length
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
fn same_repo(requested: &str, canonical: &str) -> bool {
if requested.contains('/') {
requested.eq_ignore_ascii_case(canonical)
} else {
canonical
.rsplit_once('/')
.map_or(canonical, |(_, name)| name)
.eq_ignore_ascii_case(requested)
}
}
fn validate_snapshot(snapshot: &HfSnapshotIdentity) -> Result<(), HfSourceError> {
let endpoint =
CanonicalUrl::parse(snapshot.endpoint()).map_err(|_| HfSourceError::InvalidRequest)?;
if endpoint.as_url().query().is_some()
|| endpoint.as_url().fragment().is_some()
|| !is_lower_hex(snapshot.commit(), 40)
|| validate_repo_id(snapshot.repo_id()).is_err()
|| snapshot.selection().validate().is_err()
|| snapshot.inventory().version != INVENTORY_DESCRIPTOR_VERSION
|| snapshot.inventory().member_count == 0
|| !is_lower_hex(&snapshot.inventory().sha256, 64)
{
return Err(HfSourceError::InvalidRequest);
}
Ok(())
}
fn validate_member(member: &HfMember) -> Result<(), HfSourceError> {
let source_path =
MemberPath::parse(member.source_path()).map_err(|_| HfSourceError::InvalidRequest)?;
if source_path.canonical() != member.path() || !is_lower_hex(member.commit(), 40) {
return Err(HfSourceError::InvalidRequest);
}
match member.storage() {
HfMemberStorage::Git { blob_sha1 } if is_lower_hex(blob_sha1, 40) => Ok(()),
HfMemberStorage::Lfs {
pointer_blob_sha1,
content_sha256,
} if is_lower_hex(pointer_blob_sha1, 40) && is_lower_hex(content_sha256, 64) => Ok(()),
HfMemberStorage::Xet { hash }
if !hash.is_empty() && !hash.chars().any(char::is_control) =>
{
Ok(())
}
_ => Err(HfSourceError::InvalidRequest),
}
}
#[derive(Deserialize)]
struct RepoInfo {
id: String,
sha: String,
}

View file

@ -0,0 +1,382 @@
use std::collections::BTreeSet;
use hf_hub::repository::{RepoType, RepoTypeAny};
use http::HeaderMap;
use remora_domain::MemberPath;
use serde::{Deserialize, Serialize};
use crate::HfSourceError;
pub const INVENTORY_DESCRIPTOR_VERSION: u16 = 1;
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum HfRepoType {
Model,
Dataset,
Space,
}
impl HfRepoType {
pub(crate) fn plural(self) -> &'static str {
self.as_hf_type().plural()
}
pub(crate) fn url_prefix(self) -> &'static str {
self.as_hf_type().url_prefix()
}
fn as_hf_type(self) -> RepoTypeAny {
match self {
Self::Model => RepoTypeAny::Model,
Self::Dataset => RepoTypeAny::Dataset,
Self::Space => RepoTypeAny::Space,
}
}
}
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub struct HfSelection {
exact_paths: BTreeSet<String>,
path_prefixes: BTreeSet<String>,
}
impl HfSelection {
pub fn new(
exact_paths: impl IntoIterator<Item = impl AsRef<str>>,
path_prefixes: impl IntoIterator<Item = impl AsRef<str>>,
) -> Result<Self, HfSourceError> {
Ok(Self {
exact_paths: canonical_paths(exact_paths)?,
path_prefixes: canonical_paths(path_prefixes)?,
})
}
#[must_use]
pub fn is_complete(&self) -> bool {
self.exact_paths.is_empty() && self.path_prefixes.is_empty()
}
pub(crate) fn matches(&self, path: &str) -> bool {
self.is_complete()
|| self.exact_paths.contains(path)
|| self.path_prefixes.iter().any(|prefix| {
path == prefix
|| path
.strip_prefix(prefix)
.is_some_and(|suffix| suffix.starts_with('/'))
})
}
pub(crate) fn validate(&self) -> Result<(), HfSourceError> {
for path in self.exact_paths.iter().chain(&self.path_prefixes) {
let canonical = MemberPath::parse(path).map_err(|_| HfSourceError::InvalidRequest)?;
if canonical.canonical() != path.as_str() {
return Err(HfSourceError::InvalidRequest);
}
}
Ok(())
}
}
fn canonical_paths(
paths: impl IntoIterator<Item = impl AsRef<str>>,
) -> Result<BTreeSet<String>, HfSourceError> {
paths
.into_iter()
.map(|path| {
MemberPath::parse(path.as_ref())
.map(|path| path.canonical().to_owned())
.map_err(|_| HfSourceError::InvalidRequest)
})
.collect()
}
#[derive(Clone)]
pub struct HfSourceRequest {
pub(crate) endpoint: String,
pub(crate) repo_type: HfRepoType,
pub(crate) repo_id: String,
pub(crate) revision: String,
pub(crate) selection: HfSelection,
pub(crate) credential_headers: HeaderMap,
}
impl HfSourceRequest {
pub fn new(
endpoint: &str,
repo_type: HfRepoType,
repo_id: &str,
revision: &str,
selection: HfSelection,
) -> Result<Self, HfSourceError> {
validate_repo_id(repo_id)?;
validate_revision(revision)?;
let endpoint = canonical_endpoint(endpoint)?;
Ok(Self {
endpoint,
repo_type,
repo_id: repo_id.to_owned(),
revision: revision.to_owned(),
selection,
credential_headers: HeaderMap::new(),
})
}
pub fn with_credential_headers(mut self, headers: HeaderMap) -> Result<Self, HfSourceError> {
remora_outbound::OutboundRequest::get(&self.endpoint)
.map_err(|_| HfSourceError::InvalidRequest)?
.with_credential_headers(headers.clone())
.map_err(|_| HfSourceError::InvalidRequest)?;
self.credential_headers = headers;
Ok(self)
}
}
impl std::fmt::Debug for HfSourceRequest {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("HfSourceRequest")
.field("repo_type", &self.repo_type)
.field("repo_id", &self.repo_id)
.field("revision", &self.revision)
.field("selection", &self.selection)
.field("credential_header_count", &self.credential_headers.len())
.finish_non_exhaustive()
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct HfInventoryDescriptor {
pub version: u16,
pub member_count: u64,
pub sha256: String,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct HfSnapshotIdentity {
endpoint: String,
repo_type: HfRepoType,
repo_id: String,
commit: String,
selection: HfSelection,
inventory: HfInventoryDescriptor,
}
impl HfSnapshotIdentity {
#[must_use]
pub fn endpoint(&self) -> &str {
&self.endpoint
}
#[must_use]
pub fn repo_type(&self) -> HfRepoType {
self.repo_type
}
#[must_use]
pub fn repo_id(&self) -> &str {
&self.repo_id
}
#[must_use]
pub fn commit(&self) -> &str {
&self.commit
}
#[must_use]
pub fn selection(&self) -> &HfSelection {
&self.selection
}
#[must_use]
pub fn inventory(&self) -> &HfInventoryDescriptor {
&self.inventory
}
pub(crate) fn from_parts(
endpoint: String,
repo_type: HfRepoType,
repo_id: String,
commit: String,
selection: HfSelection,
inventory: HfInventoryDescriptor,
) -> Self {
Self {
endpoint,
repo_type,
repo_id,
commit,
selection,
inventory,
}
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum HfMemberStorage {
Git {
blob_sha1: String,
},
Lfs {
pointer_blob_sha1: String,
content_sha256: String,
},
Xet {
hash: String,
},
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct HfMember {
repo_type: HfRepoType,
repo_id: String,
commit: String,
path: String,
source_path: String,
size: u64,
storage: HfMemberStorage,
}
impl HfMember {
#[must_use]
pub fn repo_type(&self) -> HfRepoType {
self.repo_type
}
#[must_use]
pub fn repo_id(&self) -> &str {
&self.repo_id
}
#[must_use]
pub fn commit(&self) -> &str {
&self.commit
}
#[must_use]
pub fn path(&self) -> &str {
&self.path
}
#[must_use]
pub fn source_path(&self) -> &str {
&self.source_path
}
#[must_use]
pub fn size(&self) -> u64 {
self.size
}
#[must_use]
pub fn storage(&self) -> &HfMemberStorage {
&self.storage
}
pub(crate) fn from_parts(
repo_type: HfRepoType,
repo_id: String,
commit: String,
path: String,
source_path: String,
size: u64,
storage: HfMemberStorage,
) -> Self {
Self {
repo_type,
repo_id,
commit,
path,
source_path,
size,
storage,
}
}
pub(crate) fn belongs_to(&self, snapshot: &HfSnapshotIdentity) -> bool {
self.repo_type == snapshot.repo_type
&& self.repo_id == snapshot.repo_id
&& self.commit == snapshot.commit
&& snapshot.selection.matches(&self.path)
}
pub(crate) fn encode(&self) -> Result<Vec<u8>, HfSourceError> {
serde_json_canonicalizer::to_vec(self).map_err(|_| HfSourceError::InvalidInventory)
}
pub(crate) fn decode(encoded: &[u8]) -> Result<Self, HfSourceError> {
let wire: HfMemberWire =
serde_json::from_slice(encoded).map_err(|_| HfSourceError::InvalidInventory)?;
Ok(wire.into())
}
}
#[derive(Deserialize)]
struct HfMemberWire {
repo_type: HfRepoType,
repo_id: String,
commit: String,
path: String,
source_path: String,
size: u64,
storage: HfMemberStorage,
}
impl From<HfMemberWire> for HfMember {
fn from(wire: HfMemberWire) -> Self {
Self {
repo_type: wire.repo_type,
repo_id: wire.repo_id,
commit: wire.commit,
path: wire.path,
source_path: wire.source_path,
size: wire.size,
storage: wire.storage,
}
}
}
pub(crate) fn canonical_endpoint(endpoint: &str) -> Result<String, HfSourceError> {
let canonical = remora_outbound::CanonicalUrl::parse(endpoint)
.map_err(|_| HfSourceError::InvalidRequest)?;
if canonical.as_url().query().is_some() {
return Err(HfSourceError::InvalidRequest);
}
let mut url = canonical.as_url().clone();
if url.path() != "/" && !url.path().ends_with('/') {
url.set_path(&format!("{}/", url.path()));
}
Ok(url.to_string().trim_end_matches('/').to_owned())
}
pub(crate) fn validate_repo_id(repo_id: &str) -> Result<(), HfSourceError> {
let mut segments = repo_id.split('/');
let first = segments.next().unwrap_or_default();
let second = segments.next();
if first.is_empty()
|| second.is_some_and(str::is_empty)
|| segments.next().is_some()
|| [Some(first), second].into_iter().flatten().any(|segment| {
matches!(segment, "." | "..")
|| !segment
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
})
{
return Err(HfSourceError::InvalidRequest);
}
Ok(())
}
pub(crate) fn validate_revision(revision: &str) -> Result<(), HfSourceError> {
if revision.is_empty() || revision.chars().any(char::is_control) {
return Err(HfSourceError::InvalidRequest);
}
Ok(())
}
pub(crate) fn is_full_commit(value: &str) -> bool {
value.len() == 40 && value.bytes().all(|byte| byte.is_ascii_hexdigit())
}

View file

@ -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
}

View file

@ -0,0 +1,21 @@
apiVersion: batch/v1
kind: Job
metadata:
name: @@JOB_NAME@@
namespace: remora-testbed
spec:
activeDeadlineSeconds: 300
backoffLimit: 0
ttlSecondsAfterFinished: 300
template:
spec:
automountServiceAccountToken: false
enableServiceLinks: false
restartPolicy: Never
containers:
- name: tests
image: @@TEST_IMAGE@@
imagePullPolicy: Never
resources:
requests: {cpu: 100m, memory: 256Mi}
limits: {cpu: "2", memory: 2Gi}

50
tests/kind/run-hf-source-job.sh Executable file
View file

@ -0,0 +1,50 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
source "$ROOT/tests/kind/lib.sh"
USER_KUBECONFIG_BEFORE="$(user_kubeconfig_digest)"
SUCCEEDED=false
OWNS_FIXTURE=false
cleanup() {
local status=$?
trap - EXIT INT TERM
if [[ "${REMORA_TESTBED_RETAIN:-0}" = 1 ]] && [[ "$SUCCEEDED" != true ]]; then
printf 'retained kubeconfig: %s\n' "$KUBECONFIG_PATH"
printf 'retained context: %s\n' "$CONTEXT"
exit "$status"
fi
if [[ "$OWNS_FIXTURE" = true ]]; then
destroy_fixture || status=1
fi
if [[ "$(user_kubeconfig_digest)" != "$USER_KUBECONFIG_BEFORE" ]]; then
printf 'user kubeconfig changed unexpectedly\n' >&2
status=1
fi
exit "$status"
}
trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
acquire_fixture_lock
REMORA_REAL_NERDCTL="$(command -v nerdctl)"
export REMORA_REAL_NERDCTL
"$REMORA_REAL_NERDCTL" --namespace default build \
--file "$ROOT/Dockerfile.test" \
--target kind-hf-source-runner \
--tag remora-kind-hf-source:local \
"$ROOT"
OWNS_FIXTURE=true
setup_fixture base
load_local_image remora-kind-hf-source:local
run_job_matrix \
hf-source \
"$TESTBED/hf-source-job.yaml.in" \
remora-kind-hf-source:local \
base
SUCCEEDED=true
printf 'Hugging Face source adapter package test Job passed\n'