Skip to main content

heddle_object_model/object/
tree_stream.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Streaming Tree entry reader with persistable resume cursors.
3
4use serde::{Deserialize, Serialize};
5
6use super::{
7    ContentHash, Tree, TreeEntry, TreeError,
8    tree_canonical::{
9        TREE_ENCODING_VERSION, TREE_HEADER_LEN, TreeHeader, decode_entry_frame, decode_header,
10    },
11    tree_source::{TreeBodyIntegrity, TreeByteSource},
12};
13
14/// Failure while streaming or range-resuming a canonical tree.
15#[derive(Debug, thiserror::Error)]
16pub enum TreeStreamError {
17    #[error("invalid tree entry: {0}")]
18    Invalid(#[from] TreeError),
19    #[error(
20        "unsupported tree encoding version {found} (this binary writes {TREE_ENCODING_VERSION})"
21    )]
22    UnsupportedVersion { found: u8 },
23    #[error("tree resume cursor does not match this object: {0}")]
24    CursorMismatch(String),
25    #[error("truncated tree frame at byte {offset}")]
26    TruncatedFrame { offset: u64 },
27    #[error("tree payload has {extra} trailing byte(s) after declared end")]
28    TrailingBytes { extra: u64 },
29    #[error("tree ended after {decoded} of {expected} declared entries")]
30    UnexpectedEof { expected: u64, decoded: u64 },
31    #[error("tree entry exceeds page byte limit ({decoded_bytes} > {max_decoded_bytes})")]
32    OversizedEntry {
33        decoded_bytes: usize,
34        max_decoded_bytes: usize,
35    },
36    #[error("tree page limits must be nonzero")]
37    InvalidPageLimits,
38    #[error("ranged tree resume requires a verified-placement object source")]
39    UnverifiedRange,
40    #[error("malformed tree encoding: {0}")]
41    Malformed(String),
42    #[error("tree I/O error: {0}")]
43    Io(#[from] std::io::Error),
44    #[error("decoded tree hash {found} does not match {expected}")]
45    HashMismatch {
46        expected: ContentHash,
47        found: ContentHash,
48    },
49}
50
51/// Caller-sized page budget. Zero limits fail closed.
52///
53/// Fields stay private so callers cannot bypass [`Self::new`].
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub struct TreePageLimits {
56    max_entries: usize,
57    max_decoded_bytes: usize,
58}
59
60impl TreePageLimits {
61    pub fn new(max_entries: usize, max_decoded_bytes: usize) -> Result<Self, TreeStreamError> {
62        if max_entries == 0 || max_decoded_bytes == 0 {
63            return Err(TreeStreamError::InvalidPageLimits);
64        }
65        Ok(Self {
66            max_entries,
67            max_decoded_bytes,
68        })
69    }
70
71    pub fn max_entries(&self) -> usize {
72        self.max_entries
73    }
74
75    pub fn max_decoded_bytes(&self) -> usize {
76        self.max_decoded_bytes
77    }
78}
79
80/// Persistable entry-boundary cursor bound to a tree id and encoding version.
81#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
82pub struct TreeResumeCursor {
83    pub(crate) tree_id: ContentHash,
84    pub(crate) encoding_version: u8,
85    pub(crate) ordinal: u64,
86    pub(crate) byte_offset: u64,
87    pub(crate) prev_name: Option<String>,
88}
89
90impl TreeResumeCursor {
91    pub fn start(tree_id: ContentHash) -> Self {
92        Self {
93            tree_id,
94            encoding_version: TREE_ENCODING_VERSION,
95            ordinal: 0,
96            byte_offset: TREE_HEADER_LEN as u64,
97            prev_name: None,
98        }
99    }
100
101    pub fn tree_id(&self) -> ContentHash {
102        self.tree_id
103    }
104
105    pub fn encoding_version(&self) -> u8 {
106        self.encoding_version
107    }
108
109    pub fn ordinal(&self) -> u64 {
110        self.ordinal
111    }
112
113    pub fn byte_offset(&self) -> u64 {
114        self.byte_offset
115    }
116
117    pub fn prev_name(&self) -> Option<&str> {
118        self.prev_name.as_deref()
119    }
120}
121
122/// One bounded page of decoded entries plus the cursor after the last entry.
123#[derive(Clone, Debug, PartialEq, Eq)]
124pub struct TreePage {
125    pub entries: Vec<TreeEntry>,
126    pub resume_cursor: TreeResumeCursor,
127}
128
129/// Incremental HTR4 reader. Does not materialize a full `Vec<TreeEntry>`.
130#[derive(Debug)]
131pub struct TreeEntryReader<S: TreeByteSource> {
132    source: S,
133    header: TreeHeader,
134    cursor: TreeResumeCursor,
135    hasher: Option<blake3::Hasher>,
136    decoded_logical_len: u64,
137    pending: Option<(TreeEntry, usize)>,
138    finished: bool,
139}
140
141impl<S: TreeByteSource> TreeEntryReader<S> {
142    pub fn open(
143        mut source: S,
144        expected_id: ContentHash,
145        resume: Option<&TreeResumeCursor>,
146    ) -> Result<Self, TreeStreamError> {
147        let mut header_buf = [0u8; TREE_HEADER_LEN];
148        source.read_exact_at(0, &mut header_buf)?;
149        let header = decode_header(&header_buf)?;
150        if header.tree_id != expected_id {
151            return Err(TreeStreamError::HashMismatch {
152                expected: expected_id,
153                found: header.tree_id,
154            });
155        }
156        let expected_len = TREE_HEADER_LEN as u64 + header.payload_len;
157        if source.len() < expected_len {
158            return Err(TreeStreamError::TruncatedFrame {
159                offset: source.len(),
160            });
161        }
162        if source.len() > expected_len {
163            return Err(TreeStreamError::TrailingBytes {
164                extra: source.len() - expected_len,
165            });
166        }
167        let cursor = resume
168            .cloned()
169            .unwrap_or_else(|| TreeResumeCursor::start(expected_id));
170        validate_cursor(&header, &cursor)?;
171        if cursor.ordinal > 0 && source.integrity() != TreeBodyIntegrity::VerifiedPlacement {
172            return Err(TreeStreamError::UnverifiedRange);
173        }
174        let hasher =
175            (cursor.ordinal == 0).then(|| ContentHash::typed_hasher("tree", header.logical_len));
176        let mut reader = Self {
177            source,
178            header,
179            cursor,
180            hasher,
181            decoded_logical_len: 0,
182            pending: None,
183            finished: false,
184        };
185        reader.arm_pending_at_cursor()?;
186        Ok(reader)
187    }
188
189    pub fn header(&self) -> &TreeHeader {
190        &self.header
191    }
192
193    pub fn bytes_read(&self) -> u64 {
194        self.source.bytes_read()
195    }
196
197    pub fn next_page(
198        &mut self,
199        limits: TreePageLimits,
200    ) -> Result<Option<TreePage>, TreeStreamError> {
201        if limits.max_entries() == 0 || limits.max_decoded_bytes() == 0 {
202            return Err(TreeStreamError::InvalidPageLimits);
203        }
204        if self.cursor.ordinal == self.header.entry_count {
205            return Ok(None);
206        }
207        let mut entries = Vec::new();
208        let mut decoded_bytes = 0usize;
209        while entries.len() < limits.max_entries() && self.cursor.ordinal < self.header.entry_count
210        {
211            let (entry, consumed) = self.take_next_entry()?;
212            let size = entry.decoded_size();
213            if size > limits.max_decoded_bytes() {
214                return Err(TreeStreamError::OversizedEntry {
215                    decoded_bytes: size,
216                    max_decoded_bytes: limits.max_decoded_bytes(),
217                });
218            }
219            if !entries.is_empty()
220                && decoded_bytes.saturating_add(size) > limits.max_decoded_bytes()
221            {
222                self.pending = Some((entry, consumed));
223                break;
224            }
225            self.commit_entry(&entry, consumed)?;
226            decoded_bytes += size;
227            entries.push(entry);
228        }
229        Ok(Some(TreePage {
230            entries,
231            resume_cursor: self.cursor.clone(),
232        }))
233    }
234
235    /// Yield one decoded entry. Used by full-object collect without a page ceiling.
236    pub fn next_entry(&mut self) -> Result<Option<TreeEntry>, TreeStreamError> {
237        if self.cursor.ordinal == self.header.entry_count {
238            return Ok(None);
239        }
240        let (entry, consumed) = self.take_next_entry()?;
241        self.commit_entry(&entry, consumed)?;
242        Ok(Some(entry))
243    }
244
245    pub fn finish_and_verify(&mut self) -> Result<(), TreeStreamError> {
246        if self.cursor.ordinal != self.header.entry_count {
247            return Err(TreeStreamError::UnexpectedEof {
248                expected: self.header.entry_count,
249                decoded: self.cursor.ordinal,
250            });
251        }
252        let payload_end = TREE_HEADER_LEN as u64 + self.header.payload_len;
253        if self.cursor.byte_offset != payload_end {
254            return Err(TreeStreamError::TrailingBytes {
255                extra: payload_end.abs_diff(self.cursor.byte_offset),
256            });
257        }
258        if let Some(hasher) = self.hasher.take() {
259            if self.decoded_logical_len != self.header.logical_len {
260                return Err(TreeStreamError::Malformed(
261                    "declared logical length does not match entries".into(),
262                ));
263            }
264            let found = ContentHash::from_bytes(hasher.finalize().into());
265            if found != self.header.tree_id {
266                return Err(TreeStreamError::HashMismatch {
267                    expected: self.header.tree_id,
268                    found,
269                });
270            }
271        } else if self.source.integrity() != TreeBodyIntegrity::VerifiedPlacement {
272            return Err(TreeStreamError::UnverifiedRange);
273        }
274        self.finished = true;
275        Ok(())
276    }
277
278    fn take_next_entry(&mut self) -> Result<(TreeEntry, usize), TreeStreamError> {
279        if let Some(pending) = self.pending.take() {
280            return Ok(pending);
281        }
282        self.read_entry_at(self.cursor.byte_offset)
283    }
284
285    fn read_entry_at(&mut self, offset: u64) -> Result<(TreeEntry, usize), TreeStreamError> {
286        let payload_end = TREE_HEADER_LEN as u64 + self.header.payload_len;
287        let mut len_buf = [0u8; 4];
288        self.source.read_exact_at(offset, &mut len_buf)?;
289        let frame_len = u64::from(u32::from_le_bytes(len_buf));
290        let frame_start = offset
291            .checked_add(4)
292            .ok_or(TreeStreamError::TruncatedFrame { offset })?;
293        let frame_end = frame_start
294            .checked_add(frame_len)
295            .ok_or(TreeStreamError::TruncatedFrame { offset })?;
296        if frame_end > payload_end || frame_end > self.source.len() {
297            return Err(TreeStreamError::TruncatedFrame { offset });
298        }
299        let frame_len =
300            usize::try_from(frame_len).map_err(|_| TreeStreamError::TruncatedFrame { offset })?;
301        let mut frame = vec![0u8; frame_len];
302        self.source.read_exact_at(frame_start, &mut frame)?;
303        let entry = decode_entry_frame(&frame)?;
304        Ok((entry, 4 + frame_len))
305    }
306
307    fn commit_entry(&mut self, entry: &TreeEntry, consumed: usize) -> Result<(), TreeStreamError> {
308        if let Some(previous) = self.cursor.prev_name.as_deref()
309            && previous >= entry.name()
310        {
311            return Err(TreeError::InvalidStructure(
312                "entries must be strictly sorted by name".into(),
313            )
314            .into());
315        }
316        if let Some(hasher) = &mut self.hasher {
317            entry.update_hasher(hasher);
318        }
319        self.decoded_logical_len = self
320            .decoded_logical_len
321            .checked_add(entry.encoded_len() as u64)
322            .ok_or_else(|| TreeStreamError::Malformed("logical length overflow".into()))?;
323        self.cursor.ordinal += 1;
324        self.cursor.byte_offset += consumed as u64;
325        self.cursor.prev_name = Some(entry.name().to_string());
326        Ok(())
327    }
328
329    fn arm_pending_at_cursor(&mut self) -> Result<(), TreeStreamError> {
330        if self.cursor.ordinal == 0 || self.cursor.ordinal == self.header.entry_count {
331            return Ok(());
332        }
333        let (entry, consumed) = self.read_entry_at(self.cursor.byte_offset)?;
334        if let Some(previous) = self.cursor.prev_name.as_deref()
335            && previous >= entry.name()
336        {
337            return Err(TreeStreamError::CursorMismatch(
338                "cursor previous name is not a valid predecessor".into(),
339            ));
340        }
341        self.pending = Some((entry, consumed));
342        Ok(())
343    }
344}
345
346#[cfg(test)]
347#[path = "tree_stream_proptests.rs"]
348mod tree_stream_proptests;
349#[cfg(test)]
350#[path = "tree_stream_tests.rs"]
351mod tree_stream_tests;
352
353fn validate_cursor(header: &TreeHeader, cursor: &TreeResumeCursor) -> Result<(), TreeStreamError> {
354    if cursor.encoding_version != TREE_ENCODING_VERSION {
355        return Err(TreeStreamError::CursorMismatch(format!(
356            "encoding version {} is not {TREE_ENCODING_VERSION}",
357            cursor.encoding_version
358        )));
359    }
360    if cursor.tree_id != header.tree_id {
361        return Err(TreeStreamError::CursorMismatch(
362            "cursor tree id does not match the opened object".into(),
363        ));
364    }
365    let payload_end = TREE_HEADER_LEN as u64 + header.payload_len;
366    if cursor.ordinal > header.entry_count {
367        return Err(TreeStreamError::CursorMismatch(
368            "cursor ordinal is past the declared entry count".into(),
369        ));
370    }
371    if cursor.ordinal == 0 {
372        if cursor.byte_offset != TREE_HEADER_LEN as u64 || cursor.prev_name.is_some() {
373            return Err(TreeStreamError::CursorMismatch(
374                "start cursor must be the first entry boundary".into(),
375            ));
376        }
377        return Ok(());
378    }
379    if cursor.ordinal == header.entry_count {
380        if cursor.byte_offset != payload_end {
381            return Err(TreeStreamError::CursorMismatch(
382                "end cursor is not the declared payload end".into(),
383            ));
384        }
385        return Ok(());
386    }
387    if cursor.byte_offset < TREE_HEADER_LEN as u64 || cursor.byte_offset >= payload_end {
388        return Err(TreeStreamError::CursorMismatch(
389            "cursor byte offset is not inside the payload".into(),
390        ));
391    }
392    Ok(())
393}
394
395impl Tree {
396    /// Decode HTR4 through the streaming reader and collect the eager `Tree`.
397    pub fn decode_canonical_streamed(data: &[u8]) -> Result<Self, TreeStreamError> {
398        let header = decode_header(data)?;
399        let mut reader = TreeEntryReader::open(
400            super::tree_source::BytesTreeSource::sequential_verify(bytes::Bytes::copy_from_slice(
401                data,
402            )),
403            header.tree_id,
404            None,
405        )?;
406        let mut entries = Vec::new();
407        while let Some(entry) = reader.next_entry()? {
408            entries.push(entry);
409        }
410        reader.finish_and_verify()?;
411        let tree = Tree::try_from_decoded_entries(entries).map_err(TreeStreamError::from)?;
412        let found = tree.hash();
413        if found != header.tree_id {
414            return Err(TreeStreamError::HashMismatch {
415                expected: header.tree_id,
416                found,
417            });
418        }
419        Ok(tree)
420    }
421}