async_http_range_reader/
lib.rs

1//! This library provides the [`AsyncHttpRangeReader`] type.
2//!
3//! It allows streaming a file over HTTP while also allow random access. The type implements both
4//! [`AsyncRead`] as well as [`AsyncSeek`]. This is supported through the use of range requests.
5//! Each individual read will request a portion of the file using an HTTP range request.
6//!
7//! Requesting numerous small reads might turn out to be relatively slow because each reads needs to
8//! perform an HTTP request. To alleviate this issue [`AsyncHttpRangeReader::prefetch`] is provided.
9//! Using this method you can *prefect* a number of bytes which will be streamed in on the
10//! background. If a read operation is reading from already (pre)fetched ranges it will stream from
11//! the internal cache instead.
12//!
13//! The primary use-case for this library is to be able to sparsely stream a zip archive over HTTP
14//! but its designed in a generic fashion.
15
16mod error;
17mod sparse_range;
18
19use futures::{FutureExt, Stream, StreamExt};
20use http_content_range::{ContentRange, ContentRangeBytes};
21use reqwest::header::HeaderMap;
22use reqwest::{Response, Url};
23use sparse_range::SparseRange;
24use std::fmt::Debug;
25use std::ops::{Deref, DerefMut};
26use std::{
27    collections::BTreeMap,
28    fmt,
29    io::{self, SeekFrom},
30    ops::Range,
31    pin::Pin,
32    sync::{self, Arc},
33    task::{ready, Context, Poll},
34};
35use tokio::{
36    io::{AsyncRead, AsyncSeek, ReadBuf},
37    sync::watch::Sender,
38    sync::{watch, Mutex},
39};
40use tokio_stream::wrappers::WatchStream;
41use tokio_util::sync::PollSender;
42use tracing::{info_span, Instrument};
43
44pub use error::AsyncHttpRangeReaderError;
45
46/// An `AsyncRangeReader` enables reading from a file over HTTP using range requests.
47///
48/// See the [`crate`] level documentation for more information.
49///
50/// The general entrypoint is [`AsyncHttpRangeReader::new`]. Depending on the
51/// [`CheckSupportMethod`], this will either call [`AsyncHttpRangeReader::initial_tail_request`] or
52/// [`AsyncHttpRangeReader::initial_head_request`] to send the initial request and then
53/// [`AsyncHttpRangeReader::from_range_response`] or [`AsyncHttpRangeReader::from_head_response`] to
54/// initialize the async reader. If you want to apply a caching layer, you can send the initial head
55/// (or tail) request yourself with your cache headers (e.g. through the
56/// [http-cache-semantics](https://docs.rs/http-cache-semantics) crate):
57///
58/// ```rust
59/// # use url::Url;
60/// # use async_http_range_reader::{AsyncHttpRangeReader, AsyncHttpRangeReaderError};
61/// # use reqwest::header::HeaderMap;
62/// async fn get_reader_cached(
63///     url: Url,
64/// ) -> Result<Option<AsyncHttpRangeReader>, AsyncHttpRangeReaderError> {
65///     let etag = "63c550e8-5ae";
66///     let client = reqwest::Client::new();
67///     let response = client
68///         .head(url.clone())
69///         .header(reqwest::header::IF_NONE_MATCH, etag)
70///         .send()
71///         .await?;
72///     if response.status() == reqwest::StatusCode::NOT_MODIFIED {
73///         Ok(None)
74///     } else {
75///         let reader = AsyncHttpRangeReader::from_head_response(client, response, url, HeaderMap::default()).await?;
76///         Ok(Some(reader))
77///     }
78/// }
79/// ```
80#[derive(Debug)]
81pub struct AsyncHttpRangeReader {
82    inner: Mutex<Inner>,
83    len: u64,
84}
85
86#[derive(Default, Clone, Debug)]
87struct StreamerState {
88    resident_range: SparseRange,
89    requested_ranges: Vec<Range<u64>>,
90    error: Option<AsyncHttpRangeReaderError>,
91}
92
93/// Loaded segments of a larger `Vec<u8>` (the remote file), indexed by start position.
94#[derive(Default)]
95struct Chunks(BTreeMap<usize, Vec<u8>>);
96
97impl fmt::Debug for Chunks {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        // Don't leak the (potentially sensitive) bytes on debug, show only the ranges for
100        // debugging.
101        let ranges = self
102            .0
103            .iter()
104            .map(|(start, bytes)| *start..*start + bytes.len());
105        f.debug_list().entries(ranges).finish()
106    }
107}
108
109impl Deref for Chunks {
110    type Target = BTreeMap<usize, Vec<u8>>;
111
112    fn deref(&self) -> &Self::Target {
113        &self.0
114    }
115}
116
117impl DerefMut for Chunks {
118    fn deref_mut(&mut self) -> &mut Self::Target {
119        &mut self.0
120    }
121}
122
123/// Downloaded chunks owned by both the reader and its background download task.
124///
125/// Each requested range grows one buffer as bytes arrive. This releases the HTTP buffers and keeps
126/// fragmented responses compact. Ranges do not overlap. The lock is never held across an await.
127#[derive(Debug)]
128struct SharedCache {
129    chunks: sync::Mutex<Chunks>,
130    /// The length of the remote file.
131    len: usize,
132}
133
134impl SharedCache {
135    fn new(len: u64) -> Result<Self, AsyncHttpRangeReaderError> {
136        // On 32-bit platforms, we only support files up to 4GB (the max of `Vec<u8>`).
137        let len = usize::try_from(len).map_err(|_| AsyncHttpRangeReaderError::FileTooLarge(len))?;
138        Ok(Self {
139            len,
140            chunks: sync::Mutex::new(Chunks::default()),
141        })
142    }
143
144    /// Create or add to a chunk starting at the given offset.
145    fn create_or_extend_chunk(
146        &self,
147        offset: usize,
148        bytes: &[u8],
149    ) -> Result<(), AsyncHttpRangeReaderError> {
150        let Ok(mut chunks) = self.chunks.lock() else {
151            return Err(AsyncHttpRangeReaderError::LockPoisoned);
152        };
153        let _new_len = {
154            let entry = chunks.entry(offset).or_default();
155            entry.extend_from_slice(bytes);
156            entry.len()
157        };
158        // Check that the chunks are non-overlapping.
159        #[cfg(debug_assertions)]
160        {
161            if let Some((next_start, _)) = chunks
162                .range((
163                    std::ops::Bound::Excluded(offset),
164                    std::ops::Bound::Unbounded,
165                ))
166                .next()
167            {
168                assert!(offset + _new_len <= *next_start);
169            }
170        }
171        Ok(())
172    }
173
174    /// Read a range of bytes known to be part of the storage.
175    ///
176    /// This includes reads across chunk boundaries.
177    fn read(&self, range: Range<usize>, buf: &mut ReadBuf<'_>) -> io::Result<()> {
178        let chunks = self
179            .chunks
180            .lock()
181            .map_err(|_| io::Error::other(AsyncHttpRangeReaderError::LockPoisoned))?;
182        let (&first_stored_start, _) = chunks
183            .range(..=range.start)
184            .next_back()
185            .expect("resident range must have a cached chunk");
186        let mut offset = range.start;
187        // Iterate over all chunks in the request range.
188        for (&start, bytes) in chunks.range(first_stored_start..range.end) {
189            let end = range.end.min(start + bytes.len());
190            buf.put_slice(&bytes[offset - start..end - start]);
191            offset = end;
192        }
193        assert_eq!(offset, range.end, "resident range must be fully cached");
194        Ok(())
195    }
196}
197
198#[derive(Debug)]
199struct Inner {
200    /// Downloaded chunks shared with the background task. `resident_range` tracks their coverage.
201    data: Arc<SharedCache>,
202
203    /// The current read position in the stream
204    pos: u64,
205
206    /// The range of bytes that have been requested for download
207    requested_range: SparseRange,
208
209    /// The range of bytes that have actually been downloaded to `data`.
210    streamer_state: StreamerState,
211
212    /// A channel receiver that holds the last downloaded range (or an error) from the background
213    /// task.
214    streamer_state_rx: WatchStream<StreamerState>,
215
216    /// A channel sender to send range requests to the background task
217    ///
218    /// Contract: All ranges sent must be inside the file.
219    request_tx: tokio::sync::mpsc::Sender<Range<u64>>,
220
221    /// An optional object to reserve a slot in the `request_tx` sender. When in the process of
222    /// sending a requests this contains an actual value.
223    poll_request_tx: Option<PollSender<Range<u64>>>,
224}
225
226/// For the initial request, we support either directly requesting N bytes from the end for file
227/// or, if you the server doesn't support negative byte offsets, starting with a HEAD request
228/// instead
229pub enum CheckSupportMethod {
230    /// Perform a range request with a negative byte range. This will return the N bytes from the
231    /// *end* of the file as well as the file-size. This is especially useful to also immediately
232    /// get some bytes from the end of the file.
233    NegativeRangeRequest(u64),
234
235    /// Perform a head request to get the length of the file and check if the server supports range
236    /// requests.
237    Head,
238}
239
240fn error_for_status(response: reqwest::Response) -> reqwest_middleware::Result<Response> {
241    response
242        .error_for_status()
243        .map_err(reqwest_middleware::Error::Reqwest)
244}
245
246impl AsyncHttpRangeReader {
247    /// Construct a new `AsyncHttpRangeReader`.
248    pub async fn new(
249        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
250        url: Url,
251        check_method: CheckSupportMethod,
252        extra_headers: HeaderMap,
253    ) -> Result<(Self, HeaderMap), AsyncHttpRangeReaderError> {
254        let client = client.into();
255        match check_method {
256            CheckSupportMethod::NegativeRangeRequest(initial_chunk_size) => {
257                let response = Self::initial_tail_request(
258                    client.clone(),
259                    url.clone(),
260                    initial_chunk_size,
261                    HeaderMap::default(),
262                )
263                .await?;
264                let response_headers = response.headers().clone();
265                let self_ = Self::from_range_response(client, response, url, extra_headers).await?;
266                Ok((self_, response_headers))
267            }
268            CheckSupportMethod::Head => {
269                let response =
270                    Self::initial_head_request(client.clone(), url.clone(), HeaderMap::default())
271                        .await?;
272                let response_headers = response.headers().clone();
273                let self_ = Self::from_head_response(client, response, url, extra_headers).await?;
274                Ok((self_, response_headers))
275            }
276        }
277    }
278
279    /// Send an initial range request to determine if the remote accepts range
280    /// requests. This will return a number of bytes from the end of the stream. Use the
281    /// `initial_chunk_size` parameter to define how many bytes should be requested from the end.
282    pub async fn initial_tail_request(
283        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
284        url: reqwest::Url,
285        initial_chunk_size: u64,
286        extra_headers: HeaderMap,
287    ) -> Result<Response, AsyncHttpRangeReaderError> {
288        let client = client.into();
289        let tail_response = client
290            .get(url)
291            .header(
292                reqwest::header::RANGE,
293                format!("bytes=-{initial_chunk_size}"),
294            )
295            .headers(extra_headers)
296            .send()
297            .await
298            .and_then(error_for_status)
299            .map_err(Arc::new)
300            .map_err(AsyncHttpRangeReaderError::HttpError)?;
301        Ok(tail_response)
302    }
303
304    #[deprecated(note = "use `from_range_response` instead")]
305    pub async fn from_tail_response(
306        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
307        tail_request_response: Response,
308        url: Url,
309        extra_headers: HeaderMap,
310    ) -> Result<Self, AsyncHttpRangeReaderError> {
311        Self::from_range_response(client, tail_request_response, url, extra_headers).await
312    }
313
314    /// Initialize the reader from [`AsyncHttpRangeReader::initial_tail_request`] (or a user
315    /// provided range response)
316    pub async fn from_range_response(
317        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
318        response: Response,
319        url: Url,
320        extra_headers: HeaderMap,
321    ) -> Result<Self, AsyncHttpRangeReaderError> {
322        let client = client.into();
323
324        // Get the size of the file from this initial request
325        let content_range_header = response
326            .headers()
327            .get(reqwest::header::CONTENT_RANGE)
328            .ok_or(AsyncHttpRangeReaderError::ContentRangeMissing)?
329            .to_str()
330            .map_err(|_err| AsyncHttpRangeReaderError::ContentRangeMissing)?;
331        // The parser ensures finish < complete_length
332        let content_range = ContentRange::parse(content_range_header).ok_or_else(|| {
333            AsyncHttpRangeReaderError::ContentRangeParser(content_range_header.to_string())
334        })?;
335        let (start, end_inclusive, complete_length) = match content_range {
336            ContentRange::Bytes(ContentRangeBytes {
337                first_byte,
338                last_byte,
339                complete_length,
340            }) => (first_byte, last_byte, complete_length),
341            _ => return Err(AsyncHttpRangeReaderError::HttpRangeRequestUnsupported),
342        };
343
344        let data = Arc::new(SharedCache::new(complete_length)?);
345
346        let requested_range = SparseRange::from_range(start..end_inclusive + 1);
347
348        // adding more than 2 entries to the channel would block the sender. I assumed two would
349        // suffice because I would want to 1) prefetch a certain range and 2) read stuff via the
350        // AsyncRead implementation. Any extra would simply have to wait for one of these to
351        // succeed. I eventually used 10 because who cares.
352        let (request_tx, request_rx) = tokio::sync::mpsc::channel(10);
353        let (state_tx, state_rx) = watch::channel(StreamerState::default());
354        tokio::spawn(run_streamer(
355            client,
356            url,
357            extra_headers,
358            Some((response, start, end_inclusive + 1)),
359            Arc::clone(&data),
360            state_tx,
361            request_rx,
362        ));
363
364        // Configure the initial state of the streamer.
365        let mut streamer_state = StreamerState::default();
366        streamer_state
367            .requested_ranges
368            .push(start..end_inclusive + 1);
369
370        let reader = Self {
371            len: data.len as u64,
372            inner: Mutex::new(Inner {
373                data,
374                pos: 0,
375                requested_range,
376                streamer_state,
377                streamer_state_rx: WatchStream::new(state_rx),
378                request_tx,
379                poll_request_tx: None,
380            }),
381        };
382        Ok(reader)
383    }
384
385    /// Send an initial range request to determine if the remote accepts range
386    /// requests and get the content length
387    pub async fn initial_head_request(
388        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
389        url: reqwest::Url,
390        extra_headers: HeaderMap,
391    ) -> Result<Response, AsyncHttpRangeReaderError> {
392        let client = client.into();
393
394        // Perform a HEAD request to get the content-length.
395        let head_response = client
396            .head(url.clone())
397            .headers(extra_headers)
398            .send()
399            .await
400            .and_then(error_for_status)
401            .map_err(Arc::new)
402            .map_err(AsyncHttpRangeReaderError::HttpError)?;
403        Ok(head_response)
404    }
405
406    /// Initialize the reader from [`AsyncHttpRangeReader::initial_head_request`] (or a user
407    /// provided response)
408    pub async fn from_head_response(
409        client: impl Into<reqwest_middleware::ClientWithMiddleware>,
410        head_response: Response,
411        url: Url,
412        extra_headers: HeaderMap,
413    ) -> Result<Self, AsyncHttpRangeReaderError> {
414        let client = client.into();
415
416        // Are range requests supported?
417        if head_response
418            .headers()
419            .get(reqwest::header::ACCEPT_RANGES)
420            .and_then(|h| h.to_str().ok())
421            != Some("bytes")
422        {
423            return Err(AsyncHttpRangeReaderError::HttpRangeRequestUnsupported);
424        }
425
426        let content_length: u64 = head_response
427            .headers()
428            .get(reqwest::header::CONTENT_LENGTH)
429            .ok_or(AsyncHttpRangeReaderError::ContentLengthMissing)?
430            .to_str()
431            .map_err(|_err| AsyncHttpRangeReaderError::ContentLengthMissing)?
432            .parse()
433            .map_err(|_err| AsyncHttpRangeReaderError::ContentLengthMissing)?;
434
435        let data = Arc::new(SharedCache::new(content_length)?);
436
437        let requested_range = SparseRange::default();
438
439        // adding more than 2 entries to the channel would block the sender. I assumed two would
440        // suffice because I would want to 1) prefetch a certain range and 2) read stuff via the
441        // AsyncRead implementation. Any extra would simply have to wait for one of these to
442        // succeed. I eventually used 10 because who cares.
443        let (request_tx, request_rx) = tokio::sync::mpsc::channel(10);
444        let (state_tx, state_rx) = watch::channel(StreamerState::default());
445        tokio::spawn(run_streamer(
446            client,
447            url,
448            extra_headers,
449            None,
450            Arc::clone(&data),
451            state_tx,
452            request_rx,
453        ));
454
455        // Configure the initial state of the streamer.
456        let streamer_state = StreamerState::default();
457
458        let reader = Self {
459            len: data.len as u64,
460            inner: Mutex::new(Inner {
461                data,
462                pos: 0,
463                requested_range,
464                streamer_state,
465                streamer_state_rx: WatchStream::new(state_rx),
466                request_tx,
467                poll_request_tx: None,
468            }),
469        };
470        Ok(reader)
471    }
472
473    /// Returns the ranges that this instance actually performed HTTP requests for.
474    pub async fn requested_ranges(&self) -> Vec<Range<u64>> {
475        let mut inner = self.inner.lock().await;
476        if let Some(Some(new_state)) = inner.streamer_state_rx.next().now_or_never() {
477            inner.streamer_state = new_state;
478        }
479        inner.streamer_state.requested_ranges.clone()
480    }
481
482    /// Prefetches a range of bytes from the remote. When specifying a large range this can
483    /// drastically reduce the number of requests required to the server.
484    pub async fn prefetch(&mut self, bytes: Range<u64>) {
485        let inner = self.inner.get_mut();
486
487        // Ensure the range is withing the file size and non-zero of length.
488        let range = bytes.start..(bytes.end.min(inner.data.len as u64));
489        if range.start >= range.end {
490            return;
491        }
492
493        // Check if the range has been requested or not.
494        let inner = self.inner.get_mut();
495        if let Some((new_range, _)) = inner.requested_range.cover(range.clone()) {
496            let _ = inner.request_tx.send(range).await;
497            inner.requested_range = new_range;
498        }
499    }
500
501    /// Returns the length of the stream in bytes
502    #[allow(clippy::len_without_is_empty)]
503    pub fn len(&self) -> u64 {
504        self.len
505    }
506}
507
508/// A task that will download parts from the remote archive and "send" them to the frontend as they
509/// become available.
510#[tracing::instrument(name = "fetch_ranges", skip_all, fields(url))]
511async fn run_streamer(
512    client: reqwest_middleware::ClientWithMiddleware,
513    url: Url,
514    extra_headers: HeaderMap,
515    response: Option<(Response, u64, u64)>,
516    data: Arc<SharedCache>,
517    mut state_tx: Sender<StreamerState>,
518    mut request_rx: tokio::sync::mpsc::Receiver<Range<u64>>,
519) {
520    let mut state = StreamerState::default();
521
522    if let Some((response, start, end_exclusive)) = response {
523        // Add the initial range to the state
524        state.requested_ranges.push(start..end_exclusive);
525
526        // Stream the initial data in memory
527        if !stream_response(
528            response,
529            start,
530            end_exclusive,
531            &data,
532            &mut state_tx,
533            &mut state,
534        )
535        .await
536        {
537            return;
538        }
539    }
540
541    // Listen for any new incoming requests
542    'outer: loop {
543        let range = match request_rx.recv().await {
544            Some(range) => range,
545            None => {
546                break 'outer;
547            }
548        };
549
550        // Determine the range that we need to cover
551        let uncovered_ranges = match state.resident_range.cover(range) {
552            None => continue,
553            Some((_, uncovered_ranges)) => uncovered_ranges,
554        };
555
556        // Download and stream each range.
557        for range in uncovered_ranges {
558            // Update the requested ranges
559            state
560                .requested_ranges
561                .push(*range.start()..*range.end() + 1);
562
563            // Execute the request
564            let range_string = format!("bytes={}-{}", range.start(), range.end());
565            let span = info_span!("fetch_range", range = range_string.as_str());
566            let response = match client
567                .get(url.clone())
568                .header(reqwest::header::RANGE, range_string)
569                .headers(extra_headers.clone())
570                .send()
571                .instrument(span)
572                .await
573                .and_then(error_for_status)
574                .map_err(std::io::Error::other)
575            {
576                Err(e) => {
577                    state.error = Some(e.into());
578                    let _ = state_tx.send(state);
579                    break 'outer;
580                }
581                Ok(response) => response,
582            };
583
584            if let Err(err) =
585                validate_content_range(&response, *range.start(), *range.end(), data.len)
586            {
587                state.error = Some(err);
588                let _ = state_tx.send(state);
589                break 'outer;
590            }
591
592            // If the server returns a successful, but non-206 response (e.g., 200), then it
593            // doesn't support range requests (even if the `Accept-Ranges` header is set).
594            if response.status() != reqwest::StatusCode::PARTIAL_CONTENT {
595                state.error = Some(AsyncHttpRangeReaderError::HttpRangeRequestUnsupported);
596                let _ = state_tx.send(state);
597                break 'outer;
598            }
599
600            if !stream_response(
601                response,
602                *range.start(),
603                *range.end() + 1,
604                &data,
605                &mut state_tx,
606                &mut state,
607            )
608            .await
609            {
610                break 'outer;
611            }
612        }
613    }
614}
615
616/// Ensure that the response range headers match the request range headers
617fn validate_content_range(
618    response: &Response,
619    expected_start: u64,
620    expected_end_inclusive: u64,
621    expected_complete_length: usize,
622) -> Result<(), AsyncHttpRangeReaderError> {
623    let content_range_header = response
624        .headers()
625        .get(reqwest::header::CONTENT_RANGE)
626        .ok_or(AsyncHttpRangeReaderError::ContentRangeMissing)?
627        .to_str()
628        .map_err(|_err| AsyncHttpRangeReaderError::ContentRangeMissing)?;
629    let content_range = ContentRange::parse(content_range_header).ok_or_else(|| {
630        AsyncHttpRangeReaderError::ContentRangeParser(content_range_header.to_string())
631    })?;
632    let (actual_start, actual_end_inclusive, actual_complete_length) = match content_range {
633        ContentRange::Bytes(ContentRangeBytes {
634            first_byte,
635            last_byte,
636            complete_length,
637        }) => (first_byte, last_byte, complete_length),
638        _ => return Err(AsyncHttpRangeReaderError::HttpRangeRequestUnsupported),
639    };
640    if expected_start != actual_start
641        || expected_end_inclusive != actual_end_inclusive
642        || expected_complete_length as u64 != actual_complete_length
643    {
644        return Err(AsyncHttpRangeReaderError::RangeMismatch {
645            expected_start,
646            expected_end_inclusive,
647            expected_complete_length,
648            actual_start,
649            actual_end_inclusive,
650            actual_complete_length,
651        });
652    }
653
654    Ok(())
655}
656
657/// Appends response chunks to the range's buffer, updating progress in between.
658/// Returns `true` if everything went fine, `false` if anything went wrong. The error state, if any,
659/// is stored in `state_tx` so the "frontend" will consume it.
660///
661/// The response must return bytes for the range of precisely `start..end_exclusive`.
662async fn stream_response(
663    tail_request_response: Response,
664    start: u64,
665    end_exclusive: u64,
666    data: &SharedCache,
667    state_tx: &mut Sender<StreamerState>,
668    state: &mut StreamerState,
669) -> bool {
670    // Enforce request channel contract
671    assert!(
672        end_exclusive <= data.len as u64,
673        "end is outside of file {} > {}",
674        end_exclusive,
675        data.len
676    );
677
678    let mut offset = start;
679    let mut byte_stream = tail_request_response.bytes_stream();
680    while let Some(bytes) = byte_stream.next().await {
681        let bytes = match bytes {
682            Err(e) => {
683                state.error = Some(e.into());
684                let _ = state_tx.send(state.clone());
685                return false;
686            }
687            Ok(bytes) => bytes,
688        };
689
690        if bytes.is_empty() {
691            continue;
692        }
693
694        // Prevent the server from sending more bytes than advertised in a response
695        if bytes.len() as u64 > end_exclusive - offset {
696            state.error = Some(AsyncHttpRangeReaderError::ResponseTooLong {
697                expected: end_exclusive - start,
698            });
699            let _ = state_tx.send(state.clone());
700            return false;
701        }
702
703        offset += bytes.len() as u64;
704
705        // Copy into one growing buffer as we go, reusing the same buffer through identical `start`.
706        match data.create_or_extend_chunk(start as usize, &bytes) {
707            Ok(()) => {}
708            Err(err) => {
709                state.error = Some(err);
710                let _ = state_tx.send(state.clone());
711                return false;
712            }
713        }
714
715        // Update the range of bytes that have been downloaded
716        state.resident_range.update(start..offset);
717
718        // Notify anyone that's listening that we have downloaded some extra data
719        if state_tx.send(state.clone()).is_err() {
720            // If we failed to set the state it means there is no receiver. In that case we should
721            // just exit.
722            return false;
723        }
724    }
725
726    // Prevent the server from sending less bytes than advertised in a response
727    if offset != end_exclusive {
728        state.error = Some(AsyncHttpRangeReaderError::ResponseTooShort {
729            expected: end_exclusive - start,
730            actual: offset - start,
731        });
732        let _ = state_tx.send(state.clone());
733        return false;
734    }
735
736    true
737}
738
739impl AsyncSeek for AsyncHttpRangeReader {
740    fn start_seek(self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
741        let me = self.get_mut();
742        let inner = me.inner.get_mut();
743
744        inner.pos = match position {
745            SeekFrom::Start(pos) => Some(pos),
746            SeekFrom::End(relative) => (inner.data.len as u64).checked_add_signed(relative),
747            SeekFrom::Current(relative) => inner.pos.checked_add_signed(relative),
748        }
749        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "invalid seek position"))?;
750
751        Ok(())
752    }
753
754    fn poll_complete(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
755        let inner = self.inner.get_mut();
756        Poll::Ready(Ok(inner.pos))
757    }
758}
759
760impl AsyncRead for AsyncHttpRangeReader {
761    fn poll_read(
762        self: Pin<&mut Self>,
763        cx: &mut Context<'_>,
764        buf: &mut ReadBuf<'_>,
765    ) -> Poll<io::Result<()>> {
766        let me = self.get_mut();
767        let inner = me.inner.get_mut();
768
769        // If a previous error occurred we return that.
770        if let Some(e) = inner.streamer_state.error.as_ref() {
771            return Poll::Ready(Err(io::Error::other(e.clone())));
772        }
773
774        // Determine the range to be fetched
775        let range = inner.pos
776            ..inner
777                .pos
778                .saturating_add(buf.remaining() as u64)
779                .min(inner.data.len as u64);
780        if range.start >= range.end {
781            return Poll::Ready(Ok(()));
782        }
783
784        // Ensure we requested the required bytes
785        while !inner.requested_range.is_covered(range.clone()) {
786            // If there is an active range request wait for it to complete
787            if let Some(mut poll) = inner.poll_request_tx.take() {
788                match poll.poll_reserve(cx) {
789                    Poll::Ready(_) => {
790                        let _ = poll.send_item(range.clone());
791                        inner.requested_range.update(range.clone());
792                        break;
793                    }
794                    Poll::Pending => {
795                        inner.poll_request_tx = Some(poll);
796                        return Poll::Pending;
797                    }
798                }
799            }
800
801            // Request the range
802            inner.poll_request_tx = Some(PollSender::new(inner.request_tx.clone()));
803        }
804
805        // If there is still a request poll open but there is no need for a request, abort it.
806        if let Some(mut poll) = inner.poll_request_tx.take() {
807            poll.abort_send();
808        }
809
810        loop {
811            // Is the range already available?
812            if inner
813                .streamer_state
814                .resident_range
815                .is_covered(range.clone())
816            {
817                inner
818                    .data
819                    .read(range.start as usize..range.end as usize, buf)?;
820                inner.pos = range.end;
821                return Poll::Ready(Ok(()));
822            }
823
824            // Otherwise wait for new data to come in
825            match ready!(Pin::new(&mut inner.streamer_state_rx).poll_next(cx)) {
826                None => {
827                    let error = AsyncHttpRangeReaderError::DownloadTaskStopped;
828                    inner.streamer_state.error = Some(error.clone());
829                    return Poll::Ready(Err(io::Error::other(error)));
830                }
831                Some(state) => {
832                    inner.streamer_state = state;
833                    if let Some(e) = inner.streamer_state.error.as_ref() {
834                        return Poll::Ready(Err(io::Error::other(e.clone())));
835                    }
836                }
837            }
838        }
839    }
840}
841
842#[cfg(test)]
843mod static_directory_server;
844
845#[cfg(test)]
846mod lifetime_tests;
847
848#[cfg(test)]
849mod test {
850    use super::*;
851    use crate::static_directory_server::StaticDirectoryServer;
852    use assert_matches::assert_matches;
853    use async_zip::tokio::read::seek::ZipFileReader;
854    use axum::body::Body;
855    use axum::extract::Request;
856    use axum::response::IntoResponse;
857    use futures::AsyncReadExt;
858    use reqwest::header;
859    use reqwest::Method;
860    use reqwest::{Client, StatusCode};
861    use rstest::*;
862    use std::path::Path;
863    use std::time::Duration;
864    use tokio::io::AsyncReadExt as _;
865
866    #[test]
867    fn sparse_cache_reads_across_chunks() {
868        // A few bytes at the end of a huge file must not allocate storage for the holes.
869        let cache = SharedCache::new(isize::MAX as u64).unwrap();
870        let start = cache.len - 16;
871        cache.create_or_extend_chunk(start + 8, b"89ab").unwrap();
872        cache.create_or_extend_chunk(0, b"other range").unwrap();
873        cache.create_or_extend_chunk(start, b"0123").unwrap();
874        cache.create_or_extend_chunk(start + 4, b"4567").unwrap();
875
876        let mut output = [b'?'; 10];
877        let mut buf = ReadBuf::new(&mut output);
878        buf.put_slice(b">");
879        cache.read(start + 2..start + 10, &mut buf).unwrap();
880        assert_eq!(buf.filled(), b">23456789");
881        assert_eq!(&output, b">23456789?");
882    }
883
884    #[tokio::test]
885    async fn closed_download_stream_returns_terminal_error() {
886        let response = axum::http::Response::builder()
887            .header(header::ACCEPT_RANGES, "bytes")
888            .header(header::CONTENT_LENGTH, "4")
889            .body("")
890            .unwrap();
891        let mut reader = AsyncHttpRangeReader::from_head_response(
892            Client::builder().no_proxy().build().unwrap(),
893            response.into(),
894            Url::parse("http://localhost/file").unwrap(),
895            HeaderMap::new(),
896        )
897        .await
898        .unwrap();
899
900        let (state_tx, state_rx) = watch::channel(StreamerState::default());
901        drop(state_tx);
902        let inner = reader.inner.get_mut();
903        inner.streamer_state_rx = WatchStream::new(state_rx);
904        // Treat the range as requested so the read does not issue an HTTP request.
905        inner.requested_range = SparseRange::from_range(0..4);
906
907        let mut contents = [0; 4];
908        let error = tokio::time::timeout(Duration::from_secs(5), reader.read_exact(&mut contents))
909            .await
910            .unwrap()
911            .expect_err("a closed download stream must not leave a read pending");
912        assert_matches!(
913            error
914                .get_ref()
915                .and_then(|error| error.downcast_ref::<AsyncHttpRangeReaderError>()),
916            Some(AsyncHttpRangeReaderError::DownloadTaskStopped)
917        );
918
919        // A cached read must still return the recorded error.
920        reader
921            .inner
922            .get_mut()
923            .streamer_state
924            .resident_range
925            .update(0..4);
926        let repeated_error = reader
927            .read_exact(&mut contents)
928            .await
929            .expect_err("the terminal error must be retained");
930        assert_eq!(repeated_error.to_string(), error.to_string());
931    }
932
933    #[rstest]
934    #[case(CheckSupportMethod::Head)]
935    #[case(CheckSupportMethod::NegativeRangeRequest(8192))]
936    #[tokio::test]
937    async fn async_range_reader_zip(#[case] check_method: CheckSupportMethod) {
938        // Spawn a static file server
939        let path = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap()).join("test-data");
940        let server = StaticDirectoryServer::new(&path)
941            .await
942            .expect("could not initialize server");
943
944        // check that file is there and has the right size
945        let filepath = path.join("andes-1.8.3-pyhd8ed1ab_0.conda");
946        assert!(
947            filepath.exists(),
948            "The conda package is not there yet. Did you run `git lfs pull`?"
949        );
950        let file_size = std::fs::metadata(&filepath).unwrap().len();
951        assert_eq!(
952            file_size, 2_463_995,
953            "The conda package is not there yet. Did you run `git lfs pull`?"
954        );
955
956        // Construct an AsyncRangeReader
957        let (mut range, _) = AsyncHttpRangeReader::new(
958            Client::new(),
959            server.url().join("andes-1.8.3-pyhd8ed1ab_0.conda").unwrap(),
960            check_method,
961            HeaderMap::default(),
962        )
963        .await
964        .expect("Could not download range - did you run `git lfs pull`?");
965
966        // Make sure we have read the last couple of bytes
967        range.prefetch(range.len() - 8192..range.len()).await;
968
969        assert_eq!(range.len(), file_size);
970
971        let mut reader = ZipFileReader::with_tokio(tokio::io::BufReader::with_capacity(0, range))
972            .await
973            .unwrap();
974
975        assert_eq!(
976            reader
977                .file()
978                .entries()
979                .iter()
980                .map(|e| e.filename().as_str().unwrap_or(""))
981                .collect::<Vec<_>>(),
982            vec![
983                "metadata.json",
984                "info-andes-1.8.3-pyhd8ed1ab_0.tar.zst",
985                "pkg-andes-1.8.3-pyhd8ed1ab_0.tar.zst",
986            ]
987        );
988
989        // Get the number of performed requests so far
990        let request_ranges = reader
991            .inner_mut()
992            .get_mut()
993            .get_mut()
994            .requested_ranges()
995            .await;
996        assert_eq!(request_ranges.len(), 1);
997        assert_eq!(
998            request_ranges[0].end - request_ranges[0].start,
999            8192,
1000            "first request should be the size of the initial chunk size"
1001        );
1002        assert_eq!(
1003            request_ranges[0].end, file_size,
1004            "first request should be at the end"
1005        );
1006
1007        // Prefetch the data for the metadata.json file
1008        let entry = reader.file().entries().first().unwrap();
1009        let offset = entry.header_offset();
1010        // Get the size of the entry plus the header + size of the filename. We should also actually
1011        // include bytes for the extra fields but we don't have that information.
1012        let size = entry.compressed_size() + 30 + entry.filename().as_bytes().len() as u64;
1013
1014        // The zip archive uses as BufReader which reads in chunks of 8192. To ensure we prefetch
1015        // enough data we round the size up to the nearest multiple of the buffer size.
1016        let buffer_size = 8192;
1017        let size = size.div_ceil(buffer_size) * buffer_size;
1018
1019        // Fetch the bytes from the zip archive that contain the requested file.
1020        reader
1021            .inner_mut()
1022            .get_mut()
1023            .get_mut()
1024            .prefetch(offset..offset + size as u64)
1025            .await;
1026
1027        // Read the contents of the metadata.json file
1028        let mut contents = String::new();
1029        reader
1030            .reader_with_entry(0)
1031            .await
1032            .unwrap()
1033            .read_to_string(&mut contents)
1034            .await
1035            .unwrap();
1036
1037        // Get the number of performed requests
1038        let request_ranges = reader
1039            .inner_mut()
1040            .get_mut()
1041            .get_mut()
1042            .requested_ranges()
1043            .await;
1044
1045        assert_eq!(contents, r#"{"conda_pkg_format_version": 2}"#);
1046        assert_eq!(request_ranges.len(), 2);
1047        assert_eq!(
1048            request_ranges[1],
1049            0..size,
1050            "expected only two range requests"
1051        );
1052    }
1053
1054    #[rstest]
1055    #[case(CheckSupportMethod::Head)]
1056    #[case(CheckSupportMethod::NegativeRangeRequest(8192))]
1057    #[tokio::test]
1058    async fn async_range_reader(#[case] check_method: CheckSupportMethod) {
1059        // Spawn a static file server
1060        let path = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap()).join("test-data");
1061        let server = StaticDirectoryServer::new(&path)
1062            .await
1063            .expect("could not initialize server");
1064
1065        // Construct an AsyncRangeReader
1066        let (mut range, _) = AsyncHttpRangeReader::new(
1067            Client::new(),
1068            server.url().join("andes-1.8.3-pyhd8ed1ab_0.conda").unwrap(),
1069            check_method,
1070            HeaderMap::default(),
1071        )
1072        .await
1073        .expect("bla");
1074
1075        // Also open a simple file reader
1076        let mut file = tokio::fs::File::open(path.join("andes-1.8.3-pyhd8ed1ab_0.conda"))
1077            .await
1078            .unwrap();
1079
1080        // Read until the end and make sure that the contents matches
1081        let mut range_read = vec![0; 64 * 1024];
1082        let mut file_read = vec![0; 64 * 1024];
1083        loop {
1084            // Read with the async reader
1085            let range_read_bytes = range.read(&mut range_read).await.unwrap();
1086
1087            // Read directly from the file
1088            let file_read_bytes = file
1089                .read_exact(&mut file_read[0..range_read_bytes])
1090                .await
1091                .unwrap();
1092
1093            assert_eq!(range_read_bytes, file_read_bytes);
1094            assert_eq!(
1095                range_read[0..range_read_bytes],
1096                file_read[0..file_read_bytes]
1097            );
1098
1099            if file_read_bytes == 0 && range_read_bytes == 0 {
1100                break;
1101            }
1102        }
1103    }
1104
1105    #[tokio::test]
1106    async fn test_not_found() {
1107        let server = StaticDirectoryServer::new(Path::new(env!("CARGO_MANIFEST_DIR")))
1108            .await
1109            .expect("could not initialize server");
1110        let err = AsyncHttpRangeReader::new(
1111            Client::new(),
1112            server.url().join("not-found").unwrap(),
1113            CheckSupportMethod::Head,
1114            HeaderMap::default(),
1115        )
1116        .await
1117        .expect_err("expected an error");
1118
1119        assert_matches!(
1120            err, AsyncHttpRangeReaderError::HttpError(err) if err.status() == Some(StatusCode::NOT_FOUND)
1121        );
1122    }
1123
1124    /// Spawn a server where the HEAD response reports `head_size` bytes, and range requests always
1125    /// claim to be `pretend_size` bytes, while actually serving `actual_size`.
1126    async fn spawn_mismatch_server(
1127        head_content_length: usize,
1128        pretend_size: usize,
1129        actual_size: usize,
1130    ) -> Url {
1131        let app =
1132            axum::Router::new().fallback(async move |request: Request| match *request.method() {
1133                Method::HEAD => {
1134                    let headers = [
1135                        (header::CONTENT_LENGTH, head_content_length.to_string()),
1136                        (header::ACCEPT_RANGES, "bytes".to_string()),
1137                    ];
1138                    (StatusCode::OK, headers).into_response()
1139                }
1140                Method::GET => {
1141                    let range_header = request
1142                        .headers()
1143                        .get(header::RANGE)
1144                        .unwrap()
1145                        .to_str()
1146                        .unwrap()
1147                        .to_string();
1148
1149                    let range_spec = range_header.strip_prefix("bytes=").unwrap();
1150                    let (start_str, _end_str) = range_spec.split_once('-').unwrap();
1151                    let start = start_str.parse::<usize>().unwrap();
1152                    // The end is inclusive
1153                    let end = start + pretend_size - 1;
1154
1155                    axum::response::Response::builder()
1156                        .status(StatusCode::PARTIAL_CONTENT)
1157                        // Note that the client ignores this value currently, it only checks the
1158                        // actual size
1159                        .header(
1160                            header::CONTENT_RANGE,
1161                            format!("bytes {start}-{end}/{head_content_length}"),
1162                        )
1163                        .body(Body::from(vec![1u8; actual_size]))
1164                        .unwrap()
1165                        .into_response()
1166                }
1167                _ => StatusCode::METHOD_NOT_ALLOWED.into_response(),
1168            });
1169
1170        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1171        let local_addr = listener.local_addr().unwrap();
1172        tokio::spawn(async move {
1173            axum::serve(listener, app.into_make_service())
1174                .await
1175                .unwrap();
1176        });
1177
1178        Url::parse(&format!("http://localhost:{}/file", local_addr.port())).unwrap()
1179    }
1180
1181    /// HEAD says 512 bytes, but range responses return 1024 bytes, beyond the file's end.
1182    #[tokio::test]
1183    async fn test_content_length_response_beyond_content_length() {
1184        /// Extract the [`AsyncHttpRangeReaderError`] from an `io::Error` returned by `read`.
1185        fn into_range_error(err: std::io::Error) -> AsyncHttpRangeReaderError {
1186            err.into_inner()
1187                .unwrap()
1188                .downcast::<AsyncHttpRangeReaderError>()
1189                .map(|e| *e)
1190                .unwrap()
1191        }
1192
1193        let cases: Vec<(usize, usize, usize, Option<AsyncHttpRangeReaderError>)> = vec![
1194            // Baseline
1195            (512, 512, 512, None),
1196            // The requested and declared length is 512, while the actual content is 1024
1197            (
1198                512,
1199                512,
1200                1024,
1201                Some(AsyncHttpRangeReaderError::ResponseTooLong { expected: 512 }),
1202            ),
1203            // The declared total length is 512, but it says and sends a range of 1024
1204            (
1205                512,
1206                1024,
1207                1024,
1208                Some(AsyncHttpRangeReaderError::ContentRangeParser(
1209                    "bytes 0-1023/512".to_string(),
1210                )),
1211            ),
1212            // The declared total length is 512, but it says a range of 1024
1213            (
1214                512,
1215                1024,
1216                512,
1217                Some(AsyncHttpRangeReaderError::ContentRangeParser(
1218                    "bytes 0-1023/512".to_string(),
1219                )),
1220            ),
1221            // Baseline
1222            (1024, 512, 512, None),
1223            // We requested 512, but we're getting 1024
1224            (
1225                1024,
1226                512,
1227                1024,
1228                Some(AsyncHttpRangeReaderError::ResponseTooLong { expected: 512 }),
1229            ),
1230            // We requested 512, but we're getting 1024
1231            (
1232                1024,
1233                1024,
1234                1024,
1235                Some(AsyncHttpRangeReaderError::RangeMismatch {
1236                    expected_start: 0,
1237                    expected_end_inclusive: 511,
1238                    expected_complete_length: 1024,
1239                    actual_start: 0,
1240                    actual_end_inclusive: 1023,
1241                    actual_complete_length: 1024,
1242                }),
1243            ),
1244            // We requested 512, but the header says 1024
1245            (
1246                1024,
1247                1024,
1248                512,
1249                Some(AsyncHttpRangeReaderError::RangeMismatch {
1250                    expected_start: 0,
1251                    expected_end_inclusive: 511,
1252                    expected_complete_length: 1024,
1253                    actual_start: 0,
1254                    actual_end_inclusive: 1023,
1255                    actual_complete_length: 1024,
1256                }),
1257            ),
1258        ];
1259        for (head_content_length, range_header_length, range_actual_length, expected_error) in cases
1260        {
1261            let url = spawn_mismatch_server(
1262                head_content_length,
1263                range_header_length,
1264                range_actual_length,
1265            )
1266            .await;
1267
1268            let (mut reader, _) = AsyncHttpRangeReader::new(
1269                Client::new(),
1270                url,
1271                CheckSupportMethod::Head,
1272                HeaderMap::default(),
1273            )
1274            .await
1275            .unwrap();
1276
1277            assert_eq!(reader.len(), head_content_length as u64);
1278            reader.prefetch(0..512).await;
1279
1280            let mut buf = vec![0u8; 512];
1281            let result = reader.read(&mut buf).await;
1282            let label =
1283                format!("{head_content_length} {range_header_length} {range_actual_length}");
1284            match expected_error {
1285                None => {
1286                    assert_matches!(result, Ok(_), "{label}");
1287                }
1288                Some(expected) => {
1289                    // The nested error don't support `PartialEq`
1290                    assert_eq!(
1291                        into_range_error(result.unwrap_err()).to_string(),
1292                        expected.to_string(),
1293                        "{label}"
1294                    );
1295                }
1296            }
1297        }
1298    }
1299}
async_http_range_reader/lib.rs

async_http_range_reader/
lib.rs
