Skip to main content

coven_format/
chunks.rs

1//! Fixed-size plaintext chunks and bounded reconstruction of individual frames.
2
3use crate::error::{require, Error, Rule};
4use crate::{frame_length, FRAME_PREFIX_LEN};
5
6/// Plaintext bytes in a full file, write-part or snapshot chunk.
7pub const CHUNK_SIZE: usize = 64 * 1024;
8
9/// How many chunks `length` plaintext bytes cut into; none when empty.
10pub fn chunk_count(length: u64) -> u64 {
11    length.div_ceil(CHUNK_SIZE as u64)
12}
13
14/// Where chunk `index` begins in the plaintext.
15pub fn chunk_start(index: u64) -> u64 {
16    index * CHUNK_SIZE as u64
17}
18
19/// The plaintext length of chunk `index` of `length` bytes: full, except a
20/// shorter last chunk. The caller keeps `index` below the chunk count.
21pub fn chunk_length(length: u64, index: u64) -> usize {
22    (length - chunk_start(index)).min(CHUNK_SIZE as u64) as usize
23}
24
25/// Cut a sequence of bounded byte pieces into fixed-size chunks, retaining at most one frame
26/// and one output chunk. Only the last chunk may be shorter than [`CHUNK_SIZE`].
27pub struct PlaintextChunks<I> {
28    frames: I,
29    frame: Vec<u8>,
30    offset: usize,
31    ended: bool,
32}
33
34impl<I: Iterator<Item = Result<Vec<u8>, Error>>> PlaintextChunks<I> {
35    /// The pieces come from a write, snapshot or slot encoder, in stream order.
36    pub fn new(frames: I) -> Self {
37        Self { frames, frame: Vec::new(), offset: 0, ended: false }
38    }
39
40    fn chunk(&mut self) -> Result<Vec<u8>, Error> {
41        let mut chunk = Vec::new();
42        chunk.try_reserve_exact(CHUNK_SIZE)?;
43        while chunk.len() < CHUNK_SIZE {
44            if self.offset == self.frame.len() {
45                match self.frames.next() {
46                    Some(frame) => {
47                        self.frame = frame?;
48                        self.offset = 0;
49                    }
50                    None => {
51                        self.ended = true;
52                        break;
53                    }
54                }
55            }
56            let count = (CHUNK_SIZE - chunk.len()).min(self.frame.len() - self.offset);
57            chunk.extend_from_slice(&self.frame[self.offset..self.offset + count]);
58            self.offset += count;
59        }
60        Ok(chunk)
61    }
62}
63
64impl<I: Iterator<Item = Result<Vec<u8>, Error>>> Iterator for PlaintextChunks<I> {
65    type Item = Result<Vec<u8>, Error>;
66
67    fn next(&mut self) -> Option<Self::Item> {
68        if self.ended {
69            return None;
70        }
71        match self.chunk() {
72            Ok(chunk) if chunk.is_empty() => None,
73            Ok(chunk) => Some(Ok(chunk)),
74            Err(error) => {
75                self.ended = true;
76                Some(Err(error))
77            }
78        }
79    }
80}
81
82/// Holds only one unfinished frame. A known remaining length lets the write
83/// decoder refuse a frame that crosses its part's end before reserving memory.
84pub struct FrameDecoder {
85    bytes: Vec<u8>,
86    remaining: Option<u64>,
87}
88
89impl FrameDecoder {
90    /// Bound the complete stream when its length is known.
91    pub fn new(remaining: Option<u64>) -> Self {
92        Self { bytes: Vec::new(), remaining }
93    }
94
95    /// Consume bytes through one bounded frame, retaining an unfinished frame.
96    pub fn next(&mut self, input: &mut &[u8]) -> Result<Option<Vec<u8>>, Error> {
97        loop {
98            let target = if self.bytes.len() < FRAME_PREFIX_LEN {
99                FRAME_PREFIX_LEN
100            } else {
101                frame_length(&self.bytes)?
102            };
103            if let Some(remaining) = self.remaining {
104                require(
105                    (target - self.bytes.len()) as u64 <= remaining,
106                    "frame crosses stream end",
107                    Rule::StreamLength,
108                )?;
109            }
110            if self.bytes.len() == target {
111                // A complete prefix may also be an empty frame.
112                let length = frame_length(&self.bytes)?;
113                if length == self.bytes.len() {
114                    return Ok(Some(std::mem::take(&mut self.bytes)));
115                }
116                continue;
117            }
118            if input.is_empty() {
119                return Ok(None);
120            }
121            let count = (target - self.bytes.len()).min(input.len());
122            self.bytes.try_reserve_exact(count)?;
123            self.bytes.extend_from_slice(&input[..count]);
124            *input = &input[count..];
125            if let Some(remaining) = &mut self.remaining {
126                *remaining -= count as u64;
127            }
128        }
129    }
130
131    /// Reject an incomplete frame or a stream shorter than its declared length.
132    pub fn finish(&self) -> Result<(), Error> {
133        if self.bytes.is_empty() && self.remaining.is_none_or(|n| n == 0) {
134            Ok(())
135        } else {
136            Err(Error::Truncated)
137        }
138    }
139}
140
141#[cfg(test)]
142#[path = "chunks_tests.rs"]
143mod tests;