Skip to main content

dig_download/
plan.rs

1//! Range planning: turn a resource's chunk layout into the chunk-aligned byte ranges a download fans
2//! across providers, and track each range's scheduling state.
3//!
4//! A resource's ciphertext is a sequence of chunks whose per-chunk lengths are `chunk_lens` (L7 §9,
5//! from the first `dig.fetchRange` frame / the availability answer). [`ChunkLayout`] turns that into
6//! byte offsets; [`plan_ranges`] partitions the resource into contiguous, **chunk-aligned** ranges of
7//! at most the node window (so a single range always maps to whole chunk(s) and is independently
8//! verifiable — L7 §9 "a requested range maps to whole chunk(s)"). Each range is scheduled
9//! independently: fetched from some provider, verified, and marked done — and a done range is never
10//! re-fetched (the basis of resume).
11
12use crate::error::VerifyError;
13
14/// The chunk boundaries of a resource: the per-chunk ciphertext lengths and their cumulative byte
15/// offsets. Built from the `chunk_lens` a peer reports in the first range frame (or the availability
16/// answer), this is the map from a byte range to the chunk(s) it covers.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct ChunkLayout {
19    /// Per-chunk ciphertext lengths, in order (`chunk_lens` on the wire).
20    chunk_lens: Vec<u64>,
21    /// Cumulative start offset of each chunk (len = `chunk_lens.len() + 1`; last = total length).
22    offsets: Vec<u64>,
23}
24
25/// The hard upper bound on the number of chunks a resource layout may declare.
26///
27/// `chunk_lens` arrives from an untrusted peer's first frame and its COUNT sizes the layout's own
28/// vectors, so an absurd count is a one-message allocation attack — bounded here, before any
29/// allocation. 1 Mi chunks covers any real resource at any sane chunk size (an independently-chosen
30/// equal ceiling for the module puller's [`MAX_MODULE_CHUNK_COUNT`](crate::module::MAX_MODULE_CHUNK_COUNT)
31/// — NOT wired to the same source; they bound different wire messages, so they may legally diverge).
32///
33/// Re-exported from `dig_nat` (the wire's own bound, #2231) rather than redefined, so the two
34/// crates can never drift onto different ceilings for the same wire limit.
35pub use dig_nat::MAX_RESOURCE_CHUNK_COUNT;
36
37impl ChunkLayout {
38    /// Build a layout from TRUSTED per-chunk lengths. Zero-length chunks are permitted (an empty
39    /// resource has no chunks; a resource with content has ≥1).
40    ///
41    /// Overflow SATURATES here, which is safe only because the caller already knows the lengths are
42    /// real. For lengths off the wire use [`try_new`](Self::try_new): a saturating sum lets a hostile
43    /// `[1, u64::MAX]` land on exactly `u64::MAX` and pass a consistency check against a declared
44    /// total (#1608).
45    pub fn new(chunk_lens: Vec<u64>) -> Self {
46        let mut offsets = Vec::with_capacity(chunk_lens.len() + 1);
47        let mut acc = 0u64;
48        offsets.push(0);
49        for &len in &chunk_lens {
50            acc = acc.saturating_add(len);
51            offsets.push(acc);
52        }
53        ChunkLayout {
54            chunk_lens,
55            offsets,
56        }
57    }
58
59    /// Build a layout from UNTRUSTED per-chunk lengths — the peer-facing constructor.
60    ///
61    /// Every step is total over a hostile input: the chunk COUNT is bounded before any allocation, the
62    /// offset reservation is FALLIBLE (an infallible `Vec::with_capacity` sized by wire input aborts the
63    /// process via the uncatchable `handle_alloc_error`), and the cumulative offsets are CHECKED so an
64    /// overflow is a typed rejection instead of a silent wrap or saturation. A library must never
65    /// delegate this to `[profile.release] overflow-checks`: only the ROOT package's profile applies, so
66    /// in a consumer build a wrapping add is a silent ACCEPT, not a panic (#1608).
67    ///
68    /// # Errors
69    /// [`VerifyError::Metadata`] when the declared chunk count exceeds
70    /// [`MAX_RESOURCE_CHUNK_COUNT`], the offsets cannot be allocated, or the cumulative length
71    /// overflows `u64`.
72    pub fn try_new(chunk_lens: Vec<u64>) -> Result<Self, VerifyError> {
73        if chunk_lens.len() > MAX_RESOURCE_CHUNK_COUNT {
74            return Err(VerifyError::Metadata(format!(
75                "declared chunk_lens count {} exceeds the maximum {MAX_RESOURCE_CHUNK_COUNT}",
76                chunk_lens.len()
77            )));
78        }
79        let mut offsets: Vec<u64> = Vec::new();
80        offsets
81            .try_reserve_exact(chunk_lens.len() + 1)
82            .map_err(|e| {
83                VerifyError::Metadata(format!(
84                    "cannot allocate a {}-entry chunk layout: {e}",
85                    chunk_lens.len() + 1
86                ))
87            })?;
88        let mut acc = 0u64;
89        offsets.push(0);
90        for &len in &chunk_lens {
91            acc = acc.checked_add(len).ok_or_else(|| {
92                VerifyError::Metadata(
93                    "chunk_lens cumulative length overflows u64 (hostile metadata)".into(),
94                )
95            })?;
96            offsets.push(acc);
97        }
98        Ok(ChunkLayout {
99            chunk_lens,
100            offsets,
101        })
102    }
103
104    /// The number of chunks.
105    pub fn chunk_count(&self) -> usize {
106        self.chunk_lens.len()
107    }
108
109    /// The per-chunk lengths.
110    pub fn chunk_lens(&self) -> &[u64] {
111        &self.chunk_lens
112    }
113
114    /// The total ciphertext length (sum of all chunk lengths).
115    pub fn total_length(&self) -> u64 {
116        *self.offsets.last().unwrap_or(&0)
117    }
118
119    /// The byte start offset of chunk `index`, or `None` if out of range.
120    pub fn chunk_offset(&self, index: usize) -> Option<u64> {
121        self.offsets.get(index).copied()
122    }
123
124    /// The length of chunk `index`, or `None` if out of range.
125    pub fn chunk_len(&self, index: usize) -> Option<u64> {
126        self.chunk_lens.get(index).copied()
127    }
128
129    /// The chunk index range `[start, end)` that a byte range `[offset, offset+length)` covers,
130    /// requiring the byte range to be **chunk-aligned** (start on a chunk boundary, end on a chunk
131    /// boundary). Returns [`VerifyError::Alignment`] otherwise — an unaligned range is not a
132    /// verifiable unit.
133    pub fn chunks_for_range(
134        &self,
135        offset: u64,
136        length: u64,
137    ) -> Result<(usize, usize), VerifyError> {
138        let end = offset.saturating_add(length);
139        let start_idx = self
140            .offsets
141            .iter()
142            .position(|&o| o == offset)
143            .ok_or_else(|| {
144                VerifyError::Alignment(format!("offset {offset} is not a chunk boundary"))
145            })?;
146        let end_idx =
147            self.offsets.iter().position(|&o| o == end).ok_or_else(|| {
148                VerifyError::Alignment(format!("end {end} is not a chunk boundary"))
149            })?;
150        if end_idx < start_idx {
151            return Err(VerifyError::Alignment(format!(
152                "range end {end} precedes start {offset}"
153            )));
154        }
155        Ok((start_idx, end_idx))
156    }
157}
158
159/// One planned byte range of the resource — a contiguous set of whole chunks fetched as a unit.
160///
161/// A range is the scheduling atom: it is fetched from a single provider at a time, verified against
162/// the resource commitment, and marked done. Ranges are independent, so different ranges of the same
163/// resource are fetched from different providers concurrently and a failed range is re-fetched
164/// elsewhere without disturbing the others.
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166pub struct Range {
167    /// Stable index of this range in the plan (0-based, ascending by offset).
168    pub index: usize,
169    /// Byte start offset within the resource ciphertext.
170    pub offset: u64,
171    /// Byte length (sum of the lengths of the chunks it covers).
172    pub length: u64,
173    /// First chunk index this range covers (into `chunk_lens`).
174    pub chunk_start: usize,
175    /// One-past-the-last chunk index this range covers.
176    pub chunk_end: usize,
177}
178
179impl Range {
180    /// The `[chunk_start, chunk_end)` chunk index range this byte range covers.
181    pub fn chunk_range(&self) -> std::ops::Range<usize> {
182        self.chunk_start..self.chunk_end
183    }
184}
185
186/// Partition a resource into chunk-aligned ranges of at most `window` bytes each.
187///
188/// Chunks are packed greedily into ranges: a chunk is added to the current range while the range
189/// stays within `window`; a chunk larger than `window` becomes its own range (a range is always ≥ one
190/// whole chunk, since a chunk is the smallest verifiable unit). The result tiles the whole resource
191/// exactly, in ascending offset order.
192pub fn plan_ranges(layout: &ChunkLayout, window: u64) -> Vec<Range> {
193    let window = window.max(1);
194    let mut ranges = Vec::new();
195    let mut i = 0usize;
196    let n = layout.chunk_count();
197    while i < n {
198        let chunk_start = i;
199        let offset = layout.chunk_offset(i).unwrap_or(0);
200        let mut length = 0u64;
201        // Always take at least one chunk; keep adding whole chunks while within the window.
202        while i < n {
203            let clen = layout.chunk_len(i).unwrap_or(0);
204            if length > 0 && length.saturating_add(clen) > window {
205                break;
206            }
207            length = length.saturating_add(clen);
208            i += 1;
209        }
210        ranges.push(Range {
211            index: ranges.len(),
212            offset,
213            length,
214            chunk_start,
215            chunk_end: i,
216        });
217    }
218    ranges
219}
220
221/// The scheduling state of one range in a running download.
222#[derive(Debug, Clone, PartialEq, Eq)]
223pub enum RangeState {
224    /// Not yet started (or re-queued after a failure) — awaiting assignment to a provider.
225    Pending,
226    /// Currently being fetched from the provider with this `peer_id` (64-hex).
227    InFlight(String),
228    /// Fetched and verified — will never be fetched again (the resume invariant).
229    Done,
230}
231
232impl RangeState {
233    /// Whether this range still needs work (pending or in-flight).
234    pub fn is_incomplete(&self) -> bool {
235        !matches!(self, RangeState::Done)
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn layout_offsets_and_total() {
245        let l = ChunkLayout::new(vec![10, 20, 5]);
246        assert_eq!(l.chunk_count(), 3);
247        assert_eq!(l.total_length(), 35);
248        assert_eq!(l.chunk_offset(0), Some(0));
249        assert_eq!(l.chunk_offset(1), Some(10));
250        assert_eq!(l.chunk_offset(2), Some(30));
251        assert_eq!(l.chunk_offset(3), Some(35));
252        assert_eq!(l.chunk_offset(4), None);
253        assert_eq!(l.chunk_len(1), Some(20));
254        assert_eq!(l.chunk_len(9), None);
255    }
256
257    #[test]
258    fn chunks_for_aligned_range() {
259        let l = ChunkLayout::new(vec![10, 20, 5]);
260        assert_eq!(l.chunks_for_range(0, 30).unwrap(), (0, 2));
261        assert_eq!(l.chunks_for_range(10, 25).unwrap(), (1, 3));
262        assert_eq!(l.chunks_for_range(30, 5).unwrap(), (2, 3));
263        // Whole resource.
264        assert_eq!(l.chunks_for_range(0, 35).unwrap(), (0, 3));
265    }
266
267    #[test]
268    fn unaligned_range_rejected() {
269        let l = ChunkLayout::new(vec![10, 20, 5]);
270        assert!(matches!(
271            l.chunks_for_range(5, 10),
272            Err(VerifyError::Alignment(_))
273        ));
274        assert!(matches!(
275            l.chunks_for_range(0, 15),
276            Err(VerifyError::Alignment(_))
277        ));
278    }
279
280    #[test]
281    fn plan_packs_chunks_into_windows() {
282        let l = ChunkLayout::new(vec![10, 10, 10, 10]);
283        // window 25 → [10,10] (20), [10,10] (20)
284        let ranges = plan_ranges(&l, 25);
285        assert_eq!(ranges.len(), 2);
286        assert_eq!(ranges[0].offset, 0);
287        assert_eq!(ranges[0].length, 20);
288        assert_eq!(ranges[0].chunk_range(), 0..2);
289        assert_eq!(ranges[1].offset, 20);
290        assert_eq!(ranges[1].length, 20);
291        assert_eq!(ranges[1].chunk_range(), 2..4);
292        // The plan tiles the whole resource exactly.
293        assert_eq!(
294            ranges.iter().map(|r| r.length).sum::<u64>(),
295            l.total_length()
296        );
297    }
298
299    #[test]
300    fn plan_oversized_chunk_is_its_own_range() {
301        let l = ChunkLayout::new(vec![100, 5]);
302        let ranges = plan_ranges(&l, 25);
303        assert_eq!(ranges.len(), 2);
304        assert_eq!(ranges[0].length, 100); // one big chunk, over the window, alone
305        assert_eq!(ranges[0].chunk_range(), 0..1);
306        assert_eq!(ranges[1].length, 5);
307        assert_eq!(ranges[1].chunk_range(), 1..2);
308    }
309
310    #[test]
311    fn plan_single_range_when_window_large() {
312        let l = ChunkLayout::new(vec![10, 20, 5]);
313        let ranges = plan_ranges(&l, 1_000_000);
314        assert_eq!(ranges.len(), 1);
315        assert_eq!(ranges[0].offset, 0);
316        assert_eq!(ranges[0].length, 35);
317        assert_eq!(ranges[0].index, 0);
318    }
319
320    #[test]
321    fn plan_empty_resource_has_no_ranges() {
322        let l = ChunkLayout::new(vec![]);
323        assert!(plan_ranges(&l, 100).is_empty());
324        assert_eq!(l.total_length(), 0);
325    }
326
327    #[test]
328    fn range_state_incompleteness() {
329        assert!(RangeState::Pending.is_incomplete());
330        assert!(RangeState::InFlight("p".into()).is_incomplete());
331        assert!(!RangeState::Done.is_incomplete());
332    }
333}