1use std::io::{self, Read, Seek, SeekFrom};
47use std::time::Duration;
48
49use serde::Serialize;
50
51use crate::chunked;
52use crate::error::FetchError;
53
54pub const TAIL_PREFETCH_BYTES: u64 = 64 * 1024;
64
65pub const READAHEAD_BYTES: u64 = 4 * 1024;
71
72pub const MAX_TRANSFER_BUDGET: u64 = 32 * 1024 * 1024;
79
80pub const MAX_RANGE_REQUESTS: u32 = 256;
86
87const RANGE_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
92
93pub trait RangeFetcher {
105 fn fetch(&mut self, start: u64, end_inclusive: u64) -> Result<Vec<u8>, FetchError>;
112
113 fn total_size(&self) -> u64;
115
116 fn extra_requests(&self) -> u32 {
121 0
122 }
123}
124
125#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
131#[non_exhaustive]
132pub struct RangeStats {
133 pub requests: u32,
136 pub bytes_fetched: u64,
140}
141
142struct Extent {
148 start: u64,
150 data: Vec<u8>,
152}
153
154impl Extent {
155 fn end(&self) -> u64 {
157 #[allow(clippy::as_conversions)]
159 self.start.saturating_add(self.data.len() as u64)
160 }
161
162 fn contains(&self, pos: u64) -> bool {
164 pos >= self.start && pos < self.end()
165 }
166
167 fn copy_at(&self, pos: u64, buf: &mut [u8]) -> Option<usize> {
170 if !self.contains(pos) {
171 return None;
172 }
173 let offset = usize::try_from(pos.checked_sub(self.start)?).ok()?;
174 let available = self.data.len().checked_sub(offset)?;
175 let n = buf.len().min(available);
176 let src = self.data.get(offset..offset.checked_add(n)?)?;
177 let dst = buf.get_mut(..n)?;
178 dst.copy_from_slice(src);
179 Some(n)
180 }
181}
182
183pub struct RangeReader<F: RangeFetcher> {
188 fetcher: F,
190 pos: u64,
193 window: Option<Extent>,
195 tail: Option<Extent>,
197 requests: u32,
199 bytes_fetched: u64,
201 max_requests: u32,
203 max_transfer_bytes: u64,
205 last_error: Option<FetchError>,
207}
208
209impl<F: RangeFetcher> RangeReader<F> {
210 #[must_use]
213 pub const fn new(fetcher: F) -> Self {
214 Self::with_limits(fetcher, MAX_RANGE_REQUESTS, MAX_TRANSFER_BUDGET)
215 }
216
217 #[must_use]
223 pub const fn with_limits(fetcher: F, max_requests: u32, max_transfer_bytes: u64) -> Self {
224 Self {
225 fetcher,
226 pos: 0,
227 window: None,
228 tail: None,
229 requests: 0,
230 bytes_fetched: 0,
231 max_requests,
232 max_transfer_bytes,
233 last_error: None,
234 }
235 }
236
237 #[must_use]
239 pub fn stats(&self) -> RangeStats {
240 RangeStats {
241 requests: self.requests.saturating_add(self.fetcher.extra_requests()),
242 bytes_fetched: self.bytes_fetched,
243 }
244 }
245
246 #[must_use]
248 pub fn total_size(&self) -> u64 {
249 self.fetcher.total_size()
250 }
251
252 #[must_use]
258 pub const fn take_last_error(&mut self) -> Option<FetchError> {
259 self.last_error.take()
260 }
261
262 fn checked_fetch(&mut self, start: u64, end_inclusive: u64) -> io::Result<Vec<u8>> {
265 let len = end_inclusive
266 .checked_sub(start)
267 .and_then(|d| d.checked_add(1))
268 .ok_or_else(|| {
269 io::Error::new(
270 io::ErrorKind::InvalidInput,
271 format!("invalid range {start}..={end_inclusive}"),
272 )
273 })?;
274
275 if self.requests >= self.max_requests {
276 return Err(io::Error::other(format!(
277 "range request cap exceeded ({} requests): pathological archive \
278 layout or seek storm; refusing further fetches",
279 self.max_requests
280 )));
281 }
282 if self.bytes_fetched.saturating_add(len) > self.max_transfer_bytes {
283 return Err(io::Error::other(format!(
284 "range transfer budget exceeded ({} bytes fetched, {len} more \
285 requested, cap {}): metadata inspection should never read \
286 this much; refusing further fetches",
287 self.bytes_fetched, self.max_transfer_bytes
288 )));
289 }
290
291 self.requests = self.requests.saturating_add(1);
292 match self.fetcher.fetch(start, end_inclusive) {
293 Ok(data) => {
294 #[allow(clippy::as_conversions)]
296 let got = data.len() as u64;
297 if got != len {
298 return Err(io::Error::other(format!(
299 "range fetcher returned {got} bytes for a {len}-byte \
300 range ({start}..={end_inclusive})"
301 )));
302 }
303 self.bytes_fetched = self.bytes_fetched.saturating_add(len);
304 Ok(data)
305 }
306 Err(fetch_err) => {
307 let io_err = io::Error::other(fetch_err.to_string());
308 self.last_error = Some(fetch_err);
309 Err(io_err)
310 }
311 }
312 }
313
314 fn copy_cached(&self, buf: &mut [u8]) -> Option<usize> {
316 if let Some(w) = &self.window {
317 if let Some(n) = w.copy_at(self.pos, buf) {
318 return Some(n);
319 }
320 }
321 if let Some(t) = &self.tail {
322 if let Some(n) = t.copy_at(self.pos, buf) {
323 return Some(n);
324 }
325 }
326 None
327 }
328
329 fn advance_pos(&mut self, n: usize) {
331 #[allow(clippy::as_conversions)]
333 {
334 self.pos = self.pos.saturating_add(n as u64);
335 }
336 }
337
338 fn serve_from_cache_after_fetch(&mut self, buf: &mut [u8]) -> io::Result<usize> {
343 let n = self.copy_cached(buf).ok_or_else(|| {
344 io::Error::other(format!(
345 "internal range cache inconsistency at offset {}",
346 self.pos
347 ))
348 })?;
349 self.advance_pos(n);
350 Ok(n)
351 }
352}
353
354impl<F: RangeFetcher> Read for RangeReader<F> {
355 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
356 if buf.is_empty() {
357 return Ok(0);
358 }
359 let total = self.fetcher.total_size();
360 if self.pos >= total {
361 return Ok(0); }
363
364 if let Some(n) = self.copy_cached(buf) {
366 self.advance_pos(n);
367 return Ok(n);
368 }
369
370 let tail_start = total.saturating_sub(TAIL_PREFETCH_BYTES);
372 if self.tail.is_none() && self.pos >= tail_start {
373 let data = self.checked_fetch(tail_start, total.saturating_sub(1))?;
374 self.tail = Some(Extent {
375 start: tail_start,
376 data,
377 });
378 return self.serve_from_cache_after_fetch(buf);
379 }
380
381 #[allow(clippy::as_conversions)]
384 let want = (buf.len() as u64).max(READAHEAD_BYTES);
385 let mut end = self
386 .pos
387 .saturating_add(want)
388 .saturating_sub(1)
389 .min(total.saturating_sub(1));
390 if let Some(t) = &self.tail {
392 if self.pos < t.start {
393 end = end.min(t.start.saturating_sub(1));
394 }
395 }
396 let data = self.checked_fetch(self.pos, end)?;
397 self.window = Some(Extent {
398 start: self.pos,
399 data,
400 });
401 self.serve_from_cache_after_fetch(buf)
402 }
403}
404
405impl<F: RangeFetcher> Seek for RangeReader<F> {
406 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
407 let total = self.fetcher.total_size();
408 let target: i128 = match pos {
409 SeekFrom::Start(p) => i128::from(p),
410 SeekFrom::End(delta) => i128::from(total).saturating_add(i128::from(delta)),
411 SeekFrom::Current(delta) => i128::from(self.pos).saturating_add(i128::from(delta)),
412 };
413 if target < 0 {
414 return Err(io::Error::new(
415 io::ErrorKind::InvalidInput,
416 format!("seek to negative offset {target}"),
417 ));
418 }
419 let new_pos = u64::try_from(target).map_err(|_| {
420 io::Error::new(
421 io::ErrorKind::InvalidInput,
422 format!("seek offset {target} exceeds u64::MAX"),
423 )
424 })?;
425 self.pos = new_pos;
426 Ok(new_pos)
427 }
428}
429
430pub type HttpRangeReader = RangeReader<HttpRangeFetcher>;
437
438pub struct HttpRangeFetcher {
468 handle: tokio::runtime::Handle,
471 client: reqwest::Client,
475 hf_url: String,
477 filename: String,
479 response_etag: Option<String>,
482 total_size: u64,
484 extra: u32,
486}
487
488impl HttpRangeReader {
489 pub async fn open(
508 repo_id: &str,
509 revision: Option<&str>,
510 filename: &str,
511 token: Option<&str>,
512 ) -> Result<Self, FetchError> {
513 let fetcher = HttpRangeFetcher::open(repo_id, revision, filename, token).await?;
514 Ok(RangeReader::new(fetcher))
515 }
516}
517
518impl HttpRangeFetcher {
519 async fn open(
527 repo_id: &str,
528 revision: Option<&str>,
529 filename: &str,
530 token: Option<&str>,
531 ) -> Result<Self, FetchError> {
532 let rev = revision.unwrap_or("main");
533 let hf_url = chunked::build_download_url(repo_id, rev, filename);
534 let client = chunked::build_client(token)?;
535
536 let info = chunked::probe_range_support(
537 client.clone(),
538 hf_url.clone(),
539 token.map(String::from),
541 )
542 .await?;
543 let Some(info) = info else {
544 return Err(Self::classify_no_range_support(&client, &hf_url, filename).await);
545 };
546
547 Ok(Self {
548 handle: tokio::runtime::Handle::current(),
549 client,
550 hf_url,
551 filename: filename.to_owned(),
553 response_etag: None,
554 total_size: info.content_length,
555 extra: 2, })
557 }
558
559 async fn classify_no_range_support(
567 client: &reqwest::Client,
568 url: &str,
569 filename: &str,
570 ) -> FetchError {
571 let result = client
572 .get(url)
573 .header(reqwest::header::RANGE, "bytes=0-0")
574 .timeout(RANGE_REQUEST_TIMEOUT)
575 .send()
576 .await;
577 match result {
578 Ok(resp) => {
579 let status = resp.status();
580 if status.is_client_error() || status.is_server_error() {
581 FetchError::Http(format!(
582 "Range request for {filename} returned status {status}"
583 ))
584 } else {
585 FetchError::Http(format!(
586 "server does not support Range requests for {filename}"
587 ))
588 }
589 }
590 Err(e) => FetchError::Http(format!("failed to probe {filename}: {e}")),
591 }
592 }
593
594 fn fetch_once(
600 &self,
601 start: u64,
602 end_inclusive: u64,
603 ) -> Result<(Vec<u8>, Option<String>), FetchError> {
604 let expected_len = end_inclusive
605 .checked_sub(start)
606 .and_then(|d| d.checked_add(1))
607 .ok_or_else(|| {
608 FetchError::Http(format!(
609 "invalid range {start}..={end_inclusive} for {}",
610 self.filename
611 ))
612 })?;
613 let expected_usize = usize::try_from(expected_len).map_err(|_| {
614 FetchError::Http(format!(
615 "range length {expected_len} exceeds addressable memory for {}",
616 self.filename
617 ))
618 })?;
619
620 let range_value = format!("bytes={start}-{end_inclusive}");
621 let filename = self.filename.as_str();
623 let total_size = self.total_size;
624
625 self.handle.block_on(async {
626 let resp = self
627 .client
628 .get(self.hf_url.as_str())
629 .header(reqwest::header::RANGE, range_value.as_str())
631 .timeout(RANGE_REQUEST_TIMEOUT)
632 .send()
633 .await
634 .map_err(|e| {
635 FetchError::Http(format!("failed to send Range request for {filename}: {e}"))
636 })?;
637
638 let status = resp.status();
639 if status == reqwest::StatusCode::OK {
640 return Err(FetchError::Http(format!(
643 "server ignored the Range header for {filename} (status 200 \
644 for bytes={start}-{end_inclusive}); refusing to read the full file"
645 )));
646 }
647 if status != reqwest::StatusCode::PARTIAL_CONTENT {
648 return Err(FetchError::Http(format!(
649 "Range request for {filename} returned status {status}"
650 )));
651 }
652
653 let content_range = resp
654 .headers()
655 .get(reqwest::header::CONTENT_RANGE)
656 .and_then(|v| v.to_str().ok())
657 .map(str::to_owned)
659 .ok_or_else(|| {
660 FetchError::Http(format!("missing Content-Range header for {filename}"))
661 })?;
662 let (cr_start, cr_end, cr_total) = parse_content_range(content_range.as_str())
663 .ok_or_else(|| {
664 FetchError::Http(format!(
665 "invalid Content-Range header for {filename}: {content_range}"
666 ))
667 })?;
668 if cr_start != start || cr_end != end_inclusive || cr_total != total_size {
669 return Err(FetchError::Http(format!(
670 "Content-Range mismatch for {filename}: requested \
671 bytes={start}-{end_inclusive} of {total_size}, server answered {content_range}"
672 )));
673 }
674
675 let etag = resp
676 .headers()
677 .get(reqwest::header::ETAG)
678 .and_then(|v| v.to_str().ok())
679 .map(clean_etag);
680
681 let mut data: Vec<u8> = Vec::with_capacity(expected_usize);
683 let mut resp = resp;
684 while let Some(chunk) = resp.chunk().await.map_err(|e| {
685 FetchError::Http(format!("failed to read Range response for {filename}: {e}"))
686 })? {
687 if data.len().saturating_add(chunk.len()) > expected_usize {
688 return Err(FetchError::Http(format!(
689 "server sent more than the requested {expected_len} bytes \
690 for {filename} (bytes={start}-{end_inclusive}); aborting"
691 )));
692 }
693 data.extend_from_slice(&chunk);
694 }
695 if data.len() != expected_usize {
696 return Err(FetchError::Http(format!(
697 "server returned {} bytes for a {expected_len}-byte range \
698 of {filename} (bytes={start}-{end_inclusive})",
699 data.len()
700 )));
701 }
702
703 Ok((data, etag))
704 })
705 }
706
707 fn check_response_etag(&mut self, etag: Option<String>) -> Result<(), FetchError> {
709 if let Some(current) = etag {
710 match &self.response_etag {
711 Some(previous) if *previous != current => {
712 return Err(FetchError::Http(format!(
713 "{} changed upstream during inspect (etag {previous} \
714 became {current})",
715 self.filename
716 )));
717 }
718 Some(_) => {} None => self.response_etag = Some(current),
720 }
721 }
722 Ok(())
723 }
724}
725
726impl RangeFetcher for HttpRangeFetcher {
727 fn fetch(&mut self, start: u64, end_inclusive: u64) -> Result<Vec<u8>, FetchError> {
728 let (data, etag) = self.fetch_once(start, end_inclusive)?;
732 self.check_response_etag(etag)?;
733 Ok(data)
734 }
735
736 fn total_size(&self) -> u64 {
737 self.total_size
738 }
739
740 fn extra_requests(&self) -> u32 {
741 self.extra
742 }
743}
744
745fn parse_content_range(value: &str) -> Option<(u64, u64, u64)> {
750 let rest = value.strip_prefix("bytes ")?;
751 let (range, total) = rest.split_once('/')?;
752 let (start, end) = range.split_once('-')?;
753 Some((
754 start.trim().parse().ok()?,
755 end.trim().parse().ok()?,
756 total.trim().parse().ok()?,
757 ))
758}
759
760fn clean_etag(raw: &str) -> String {
765 raw.strip_prefix("W/").unwrap_or(raw).replace('"', "")
766}
767
768#[cfg(test)]
773mod tests {
774 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::indexing_slicing)]
775
776 use super::*;
777
778 struct InMemoryFetcher {
780 data: Vec<u8>,
781 calls: Vec<(u64, u64)>,
782 }
783
784 impl InMemoryFetcher {
785 fn new(data: Vec<u8>) -> Self {
786 Self {
787 data,
788 calls: Vec::new(),
789 }
790 }
791 }
792
793 impl RangeFetcher for InMemoryFetcher {
794 fn fetch(&mut self, start: u64, end_inclusive: u64) -> Result<Vec<u8>, FetchError> {
795 self.calls.push((start, end_inclusive));
796 let s = usize::try_from(start).unwrap();
797 let e = usize::try_from(end_inclusive).unwrap();
798 self.data
799 .get(s..=e)
800 .map(<[u8]>::to_vec)
801 .ok_or_else(|| FetchError::Http(format!("bad range {start}..={end_inclusive}")))
802 }
803
804 fn total_size(&self) -> u64 {
805 u64::try_from(self.data.len()).unwrap()
806 }
807 }
808
809 struct FailingFetcher {
811 size: u64,
812 }
813
814 impl RangeFetcher for FailingFetcher {
815 fn fetch(&mut self, _start: u64, _end_inclusive: u64) -> Result<Vec<u8>, FetchError> {
816 Err(FetchError::Http(
817 "Range request for x.npz returned status 403 Forbidden".to_owned(),
818 ))
819 }
820
821 fn total_size(&self) -> u64 {
822 self.size
823 }
824 }
825
826 struct ShortFetcher {
828 size: u64,
829 }
830
831 impl RangeFetcher for ShortFetcher {
832 fn fetch(&mut self, _start: u64, _end_inclusive: u64) -> Result<Vec<u8>, FetchError> {
833 Ok(vec![0u8; 1])
834 }
835
836 fn total_size(&self) -> u64 {
837 self.size
838 }
839 }
840
841 fn sample_data(len: usize) -> Vec<u8> {
842 #[allow(clippy::as_conversions, clippy::cast_possible_truncation)]
844 (0..len).map(|i| (i % 251) as u8).collect()
845 }
846
847 #[test]
850 fn seek_start_end_current_semantics() {
851 let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(1000)));
852 assert_eq!(r.seek(SeekFrom::Start(10)).unwrap(), 10);
853 assert_eq!(r.seek(SeekFrom::Current(5)).unwrap(), 15);
854 assert_eq!(r.seek(SeekFrom::Current(-15)).unwrap(), 0);
855 assert_eq!(r.seek(SeekFrom::End(0)).unwrap(), 1000);
856 assert_eq!(r.seek(SeekFrom::End(-1000)).unwrap(), 0);
857 assert_eq!(r.seek(SeekFrom::End(50)).unwrap(), 1050);
859 let mut buf = [0u8; 4];
860 assert_eq!(r.read(&mut buf).unwrap(), 0);
861 }
862
863 #[test]
864 fn seek_negative_is_invalid_input() {
865 let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(100)));
866 let err = r.seek(SeekFrom::Current(-1)).unwrap_err();
867 assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
868 let err = r.seek(SeekFrom::End(-101)).unwrap_err();
869 assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
870 }
871
872 #[test]
873 fn empty_file_reads_zero_with_no_requests() {
874 let mut r = RangeReader::new(InMemoryFetcher::new(Vec::new()));
875 let mut buf = [0u8; 8];
876 assert_eq!(r.read(&mut buf).unwrap(), 0);
877 assert_eq!(r.stats().requests, 0);
878 }
879
880 #[test]
883 fn sequential_small_reads_coalesce_into_one_window_fetch() {
884 let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(100 * 1024)));
886 let mut buf = [0u8; 16];
887 for i in 0..10 {
888 r.read_exact(&mut buf).unwrap();
889 assert_eq!(buf[0], sample_data(100 * 1024)[i * 16]);
890 }
891 assert_eq!(r.stats().requests, 1);
893 assert_eq!(r.fetcher.calls[0], (0, READAHEAD_BYTES - 1));
894 }
895
896 #[test]
897 fn tail_region_read_prefetches_tail_once() {
898 let size: u64 = 1024 * 1024; let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(
900 usize::try_from(size).unwrap(),
901 )));
902 r.seek(SeekFrom::End(-22)).unwrap();
904 let mut buf = [0u8; 22];
905 r.read_exact(&mut buf).unwrap();
906 assert_eq!(r.stats().requests, 1);
907 assert_eq!(r.fetcher.calls[0], (size - TAIL_PREFETCH_BYTES, size - 1));
908 r.seek(SeekFrom::End(-4096)).unwrap();
910 let mut big = [0u8; 4096];
911 r.read_exact(&mut big).unwrap();
912 assert_eq!(r.stats().requests, 1);
913 assert_eq!(r.stats().bytes_fetched, TAIL_PREFETCH_BYTES);
914 }
915
916 #[test]
917 fn window_reuse_after_seek_back() {
918 let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(64 * 1024)));
919 let mut buf = [0u8; 128];
920 r.read_exact(&mut buf).unwrap();
921 assert_eq!(r.stats().requests, 1);
922 r.seek(SeekFrom::Start(32)).unwrap();
924 r.read_exact(&mut buf).unwrap();
925 assert_eq!(r.stats().requests, 1);
926 let expected = sample_data(64 * 1024);
927 assert_eq!(&buf[..], &expected[32..160]);
928 }
929
930 #[test]
931 fn window_fetch_never_overlaps_cached_tail() {
932 let size: u64 = 200 * 1024;
933 let mut r = RangeReader::new(InMemoryFetcher::new(sample_data(
934 usize::try_from(size).unwrap(),
935 )));
936 r.seek(SeekFrom::End(-10)).unwrap();
938 let mut small = [0u8; 10];
939 r.read_exact(&mut small).unwrap();
940 let tail_start = size - TAIL_PREFETCH_BYTES;
943 r.seek(SeekFrom::Start(tail_start - 100)).unwrap();
944 let mut big = vec![0u8; 4096];
945 r.read_exact(&mut big).unwrap();
946 assert_eq!(r.fetcher.calls[1], (tail_start - 100, tail_start - 1));
947 let expected = sample_data(usize::try_from(size).unwrap());
948 let from = usize::try_from(tail_start - 100).unwrap();
949 assert_eq!(&big[..], &expected[from..from + 4096]);
950 }
951
952 #[test]
955 fn request_cap_is_enforced() {
956 let data = sample_data(10 * 1024 * 1024);
957 let mut r = RangeReader::with_limits(InMemoryFetcher::new(data), 3, u64::MAX);
958 let mut buf = [0u8; 8];
959 for i in 0u64..3 {
962 r.seek(SeekFrom::Start(i * 100 * 1024)).unwrap();
963 r.read_exact(&mut buf).unwrap();
964 }
965 r.seek(SeekFrom::Start(1024 * 1024)).unwrap();
966 let err = r.read(&mut buf).unwrap_err();
967 assert!(
968 err.to_string()
969 .contains("range request cap exceeded (3 requests)"),
970 "unexpected error: {err}"
971 );
972 }
973
974 #[test]
975 fn transfer_budget_is_enforced() {
976 let data = sample_data(10 * 1024 * 1024);
977 let mut r = RangeReader::with_limits(InMemoryFetcher::new(data), u32::MAX, 6 * 1024);
979 let mut buf = [0u8; 8];
980 r.read_exact(&mut buf).unwrap();
981 r.seek(SeekFrom::Start(1024 * 1024)).unwrap();
982 let err = r.read(&mut buf).unwrap_err();
983 assert!(
984 err.to_string().contains("range transfer budget exceeded"),
985 "unexpected error: {err}"
986 );
987 }
988
989 #[test]
992 fn fetch_error_surfaces_as_io_and_is_recoverable_typed() {
993 let mut r = RangeReader::new(FailingFetcher { size: 1024 });
994 let mut buf = [0u8; 8];
995 let err = r.read(&mut buf).unwrap_err();
996 assert!(err.to_string().contains("returned status 403"));
997 let typed = r.take_last_error().expect("typed error must be stored");
998 assert!(matches!(typed, FetchError::Http(msg)
999 if msg.contains("returned status 403 Forbidden")));
1000 assert!(r.take_last_error().is_none());
1002 }
1003
1004 #[test]
1005 fn short_fetch_is_a_contract_error() {
1006 let mut r = RangeReader::new(ShortFetcher { size: 1024 * 1024 });
1007 let mut buf = [0u8; 8];
1008 let err = r.read(&mut buf).unwrap_err();
1009 assert!(
1010 err.to_string().contains("bytes for a"),
1011 "unexpected error: {err}"
1012 );
1013 }
1014
1015 #[test]
1018 fn parse_content_range_accepts_the_exact_206_form() {
1019 assert_eq!(parse_content_range("bytes 0-7/1234"), Some((0, 7, 1234)));
1020 assert_eq!(
1021 parse_content_range("bytes 100-199/200"),
1022 Some((100, 199, 200))
1023 );
1024 }
1025
1026 #[test]
1027 fn parse_content_range_rejects_deviant_forms() {
1028 assert_eq!(parse_content_range("bytes */1234"), None);
1029 assert_eq!(parse_content_range("bytes 0-7/*"), None);
1030 assert_eq!(parse_content_range("0-7/1234"), None);
1031 assert_eq!(parse_content_range("bytes 7/1234"), None);
1032 assert_eq!(parse_content_range(""), None);
1033 }
1034
1035 #[test]
1036 fn clean_etag_strips_quotes_and_weak_prefix() {
1037 assert_eq!(clean_etag("\"abc123\""), "abc123");
1038 assert_eq!(clean_etag("W/\"abc123\""), "abc123");
1039 assert_eq!(clean_etag("abc123"), "abc123");
1040 }
1041
1042 fn npy_bytes(descr: &str, shape_literal: &str, data_len: usize) -> Vec<u8> {
1047 let dict =
1048 format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape_literal}, }}");
1049 let unpadded = 10 + dict.len() + 1; let padding = (64 - unpadded % 64) % 64;
1052 let header_len = dict.len() + padding + 1;
1053 let mut out = Vec::with_capacity(10 + header_len + data_len);
1054 out.extend_from_slice(b"\x93NUMPY\x01\x00");
1055 out.extend_from_slice(&u16::try_from(header_len).unwrap().to_le_bytes());
1056 out.extend_from_slice(dict.as_bytes());
1057 out.extend(std::iter::repeat_n(b' ', padding));
1058 out.push(b'\n');
1059 out.extend(std::iter::repeat_n(0u8, data_len));
1060 out
1061 }
1062
1063 fn stored_zip(entries: &[(&str, Vec<u8>)]) -> Vec<u8> {
1067 let mut out = Vec::new();
1068 let mut offsets = Vec::new();
1069 for (name, payload) in entries {
1070 offsets.push(u32::try_from(out.len()).unwrap());
1071 let size = u32::try_from(payload.len()).unwrap();
1072 out.extend_from_slice(&0x0403_4b50u32.to_le_bytes()); out.extend_from_slice(&20u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&size.to_le_bytes()); out.extend_from_slice(&size.to_le_bytes()); out.extend_from_slice(&u16::try_from(name.len()).unwrap().to_le_bytes());
1081 out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(name.as_bytes());
1083 out.extend_from_slice(payload);
1084 }
1085 let cd_offset = u32::try_from(out.len()).unwrap();
1086 for ((name, payload), lfh_offset) in entries.iter().zip(&offsets) {
1087 let size = u32::try_from(payload.len()).unwrap();
1088 out.extend_from_slice(&0x0201_4b50u32.to_le_bytes()); out.extend_from_slice(&20u16.to_le_bytes()); out.extend_from_slice(&20u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&size.to_le_bytes()); out.extend_from_slice(&size.to_le_bytes()); out.extend_from_slice(&u16::try_from(name.len()).unwrap().to_le_bytes());
1098 out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&lfh_offset.to_le_bytes());
1104 out.extend_from_slice(name.as_bytes());
1105 }
1106 let cd_size = u32::try_from(out.len()).unwrap() - cd_offset;
1107 let n = u16::try_from(entries.len()).unwrap();
1108 out.extend_from_slice(&0x0605_4b50u32.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&0u16.to_le_bytes()); out.extend_from_slice(&n.to_le_bytes()); out.extend_from_slice(&n.to_le_bytes()); out.extend_from_slice(&cd_size.to_le_bytes());
1114 out.extend_from_slice(&cd_offset.to_le_bytes());
1115 out.extend_from_slice(&0u16.to_le_bytes()); out
1117 }
1118
1119 #[test]
1120 fn npz_inspect_over_range_reader_reads_metadata_not_data() {
1121 let npz = stored_zip(&[
1124 ("w_enc.npy", npy_bytes("<f4", "(2, 3)", 2 * 3 * 4)),
1125 ("b_dec.npy", npy_bytes("<f4", "(150, 1024)", 600 * 1024)),
1126 ]);
1127 let total = u64::try_from(npz.len()).unwrap();
1128 let mut reader = RangeReader::new(InMemoryFetcher::new(npz));
1129
1130 let info = anamnesis::inspect_npz_from_reader(&mut reader)
1131 .expect("synthetic NPZ must inspect cleanly");
1132
1133 assert_eq!(info.tensors.len(), 2);
1134 let names: Vec<&str> = info.tensors.iter().map(|t| t.name.as_str()).collect();
1135 assert!(names.contains(&"w_enc"), "names: {names:?}");
1136 assert!(names.contains(&"b_dec"), "names: {names:?}");
1137 let b_dec = info.tensors.iter().find(|t| t.name == "b_dec").unwrap();
1138 assert_eq!(b_dec.shape, vec![150, 1024]);
1139
1140 let stats = reader.stats();
1143 assert!(
1144 stats.requests <= 8,
1145 "expected a handful of range requests, got {}",
1146 stats.requests
1147 );
1148 assert!(
1149 stats.bytes_fetched < total / 4,
1150 "fetched {} of {total} bytes — inspect must not read tensor data",
1151 stats.bytes_fetched
1152 );
1153 }
1154
1155 fn synthetic_safetensors(count: usize, data_section_len: usize) -> (Vec<u8>, usize) {
1164 use std::fmt::Write as _;
1165
1166 let mut entries = String::new();
1167 let mut offset: u64 = 0;
1168 for i in 0..count {
1169 if i > 0 {
1170 entries.push(',');
1171 }
1172 let start = offset;
1173 let end = offset + 16;
1174 offset = end;
1175 let _ = write!(
1176 entries,
1177 "\"tensor_{i}\":{{\"dtype\":\"F32\",\"shape\":[4],\"data_offsets\":[{start},{end}]}}"
1178 );
1179 }
1180 let header_bytes = format!("{{{entries}}}").into_bytes();
1181 let header_len = header_bytes.len();
1182
1183 let mut buf = Vec::with_capacity(8 + header_len + data_section_len);
1184 buf.extend_from_slice(&u64::try_from(header_len).unwrap().to_le_bytes());
1185 buf.extend_from_slice(&header_bytes);
1186 buf.extend(std::iter::repeat_n(0u8, data_section_len));
1187 (buf, header_len)
1188 }
1189
1190 #[test]
1191 fn safetensors_inspect_over_range_reader_reads_metadata_not_data() {
1192 let (buf, header_len) = synthetic_safetensors(100, 256 * 1024);
1197 assert!(
1198 header_len > 4096,
1199 "fixture must exceed the 4 KiB read-ahead window to exercise \
1200 the second-fetch path, got {header_len} bytes"
1201 );
1202 let total = u64::try_from(buf.len()).unwrap();
1203 let mut reader = RangeReader::new(InMemoryFetcher::new(buf));
1204
1205 let header = anamnesis::parse_safetensors_header_from_reader(&mut reader)
1206 .expect("synthetic safetensors header must parse cleanly");
1207 assert_eq!(header.tensors.len(), 100);
1208
1209 let stats = reader.stats();
1213 assert_eq!(
1214 stats.requests, 2,
1215 "a {header_len}-byte header should need exactly two range \
1216 requests (first window fetch covers 0..4096, second covers \
1217 the rest); got {}",
1218 stats.requests
1219 );
1220 assert!(
1221 stats.bytes_fetched < total / 4,
1222 "fetched {} of {total} bytes — inspect must not read tensor data",
1223 stats.bytes_fetched
1224 );
1225 }
1226
1227 #[test]
1228 fn safetensors_small_header_fits_in_one_read_ahead_window() {
1229 let (buf, header_len) = synthetic_safetensors(3, 1024);
1233 assert!(
1234 header_len < 4096,
1235 "fixture must stay under the 4 KiB read-ahead window, got {header_len} bytes"
1236 );
1237 let mut reader = RangeReader::new(InMemoryFetcher::new(buf));
1238
1239 let header = anamnesis::parse_safetensors_header_from_reader(&mut reader)
1240 .expect("synthetic safetensors header must parse cleanly");
1241 assert_eq!(header.tensors.len(), 3);
1242 assert_eq!(reader.stats().requests, 1);
1243 }
1244
1245 struct GgufBuilder {
1252 buf: Vec<u8>,
1253 }
1254
1255 impl GgufBuilder {
1256 fn new(tensor_count: u64) -> Self {
1257 let mut buf = Vec::new();
1258 buf.extend_from_slice(b"GGUF");
1259 buf.extend_from_slice(&3u32.to_le_bytes()); buf.extend_from_slice(&tensor_count.to_le_bytes());
1261 buf.extend_from_slice(&0u64.to_le_bytes()); Self { buf }
1263 }
1264
1265 fn push_f32_tensor_info(&mut self, name: &str, shape: &[u64], relative_offset: u64) {
1266 self.buf
1267 .extend_from_slice(&u64::try_from(name.len()).unwrap().to_le_bytes());
1268 self.buf.extend_from_slice(name.as_bytes());
1269 self.buf
1270 .extend_from_slice(&u32::try_from(shape.len()).unwrap().to_le_bytes());
1271 for d in shape {
1272 self.buf.extend_from_slice(&d.to_le_bytes());
1273 }
1274 self.buf.extend_from_slice(&0u32.to_le_bytes()); self.buf.extend_from_slice(&relative_offset.to_le_bytes());
1276 }
1277
1278 fn pad_to_alignment(&mut self, alignment: usize) {
1279 let rem = self.buf.len() % alignment;
1280 if rem != 0 {
1281 self.buf.extend(std::iter::repeat_n(0u8, alignment - rem));
1282 }
1283 }
1284
1285 fn push_zeros(&mut self, n: usize) {
1286 self.buf.extend(std::iter::repeat_n(0u8, n));
1287 }
1288
1289 fn finish(self) -> Vec<u8> {
1290 self.buf
1291 }
1292 }
1293
1294 fn synthetic_gguf() -> Vec<u8> {
1299 let mut b = GgufBuilder::new(2);
1300 b.push_f32_tensor_info("w_enc", &[2, 3], 0);
1301 b.push_f32_tensor_info("b_dec", &[153_600], 32);
1302 b.pad_to_alignment(32); b.push_zeros(24); b.push_zeros(8); b.push_zeros(153_600 * 4); b.finish()
1307 }
1308
1309 #[test]
1310 fn gguf_front_matter_over_range_reader_reads_metadata_not_data() {
1311 let gguf = synthetic_gguf();
1312 let total = u64::try_from(gguf.len()).unwrap();
1313 let mut reader = RangeReader::new(InMemoryFetcher::new(gguf));
1314
1315 let front = anamnesis::parse_gguf_front_matter_from_reader(&mut reader)
1316 .expect("synthetic GGUF must parse cleanly");
1317
1318 assert_eq!(front.tensor_infos.len(), 2);
1319 let names: Vec<&str> = front.tensor_infos.iter().map(|t| t.name.as_str()).collect();
1320 assert!(names.contains(&"w_enc"), "names: {names:?}");
1321 assert!(names.contains(&"b_dec"), "names: {names:?}");
1322 let b_dec = front
1323 .tensor_infos
1324 .iter()
1325 .find(|t| t.name == "b_dec")
1326 .unwrap();
1327 assert_eq!(b_dec.shape, vec![153_600]);
1328
1329 let stats = reader.stats();
1334 assert!(
1335 stats.requests <= 8,
1336 "expected a handful of range requests, got {}",
1337 stats.requests
1338 );
1339 assert!(
1340 stats.bytes_fetched < total / 4,
1341 "fetched {} of {total} bytes — front-matter parse must not read tensor data",
1342 stats.bytes_fetched
1343 );
1344 }
1345}