1use crate::error::{require, Error, Rule};
4use crate::{frame_length, FRAME_PREFIX_LEN};
5
6pub const CHUNK_SIZE: usize = 64 * 1024;
8
9pub fn chunk_count(length: u64) -> u64 {
11 length.div_ceil(CHUNK_SIZE as u64)
12}
13
14pub fn chunk_start(index: u64) -> u64 {
16 index * CHUNK_SIZE as u64
17}
18
19pub fn chunk_length(length: u64, index: u64) -> usize {
22 (length - chunk_start(index)).min(CHUNK_SIZE as u64) as usize
23}
24
25pub 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 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
82pub struct FrameDecoder {
85 bytes: Vec<u8>,
86 remaining: Option<u64>,
87}
88
89impl FrameDecoder {
90 pub fn new(remaining: Option<u64>) -> Self {
92 Self { bytes: Vec::new(), remaining }
93 }
94
95 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 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 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;