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

async_http_range_reader/
lib.rs
