1use crate::{hash::digest32, EmbeddedError, EmbeddedResult};
2use alloc::vec::Vec;
3
4#[derive(Debug, Clone, Eq, PartialEq)]
5pub struct ChunkCursor {
6 pub transfer_id: u32,
7 pub total_size: u32,
8 pub next_offset: u32,
9 pub expected_sequence: u16,
10 pub chunk_size: u16,
11}
12
13impl ChunkCursor {
14 pub fn validate(&self) -> EmbeddedResult<()> {
15 if self.transfer_id == 0 || self.chunk_size == 0 {
16 return Err(EmbeddedError::InvalidInput);
17 }
18 if self.next_offset > self.total_size {
19 return Err(EmbeddedError::InvalidCursor);
20 }
21 Ok(())
22 }
23}
24
25#[derive(Debug, Clone, Eq, PartialEq)]
26pub struct AttachmentChunk {
27 pub transfer_id: u32,
28 pub sequence: u16,
29 pub payload: Vec<u8>,
30}
31
32#[derive(Debug, Clone, Copy, Eq, PartialEq)]
33pub enum ChunkApply {
34 Appended,
35 DuplicateAccepted,
36}
37
38#[derive(Debug, Clone)]
39pub struct AttachmentReceiver {
40 transfer_id: u32,
41 total_size: u32,
42 chunk_size: u16,
43 expected_sequence: u16,
44 next_offset: u32,
45 bytes: Vec<u8>,
46}
47
48impl AttachmentReceiver {
49 pub fn start(transfer_id: u32, total_size: u32, chunk_size: u16) -> EmbeddedResult<Self> {
50 if transfer_id == 0 || chunk_size == 0 || total_size == 0 {
51 return Err(EmbeddedError::InvalidArgument);
52 }
53 let cap = usize::try_from(total_size).map_err(|_| EmbeddedError::InvalidArgument)?;
54 Ok(Self {
55 transfer_id,
56 total_size,
57 chunk_size,
58 expected_sequence: 0,
59 next_offset: 0,
60 bytes: Vec::with_capacity(cap.min(128 * 1024)),
61 })
62 }
63
64 pub fn resume(
65 transfer_id: u32,
66 total_size: u32,
67 chunk_size: u16,
68 existing_bytes: Vec<u8>,
69 ) -> EmbeddedResult<Self> {
70 if transfer_id == 0 || chunk_size == 0 || total_size == 0 {
71 return Err(EmbeddedError::InvalidArgument);
72 }
73 let existing_len =
74 u32::try_from(existing_bytes.len()).map_err(|_| EmbeddedError::InvalidArgument)?;
75 if existing_len > total_size {
76 return Err(EmbeddedError::InvalidCursor);
77 }
78 let chunk_size_u32 = u32::from(chunk_size);
79 let expected_sequence =
80 (existing_len / chunk_size_u32).try_into().map_err(|_| EmbeddedError::InvalidCursor)?;
81 Ok(Self {
82 transfer_id,
83 total_size,
84 chunk_size,
85 expected_sequence,
86 next_offset: existing_len,
87 bytes: existing_bytes,
88 })
89 }
90
91 pub fn cursor(&self) -> ChunkCursor {
92 ChunkCursor {
93 transfer_id: self.transfer_id,
94 total_size: self.total_size,
95 next_offset: self.next_offset,
96 expected_sequence: self.expected_sequence,
97 chunk_size: self.chunk_size,
98 }
99 }
100
101 pub fn apply_chunk(
102 &mut self,
103 offset: u32,
104 chunk: &AttachmentChunk,
105 ) -> EmbeddedResult<ChunkApply> {
106 if chunk.transfer_id != self.transfer_id {
107 return Err(EmbeddedError::NotFound);
108 }
109 if chunk.payload.is_empty() {
110 return Err(EmbeddedError::InvalidArgument);
111 }
112 if chunk.payload.len() > usize::from(self.chunk_size) {
113 return Err(EmbeddedError::InvalidArgument);
114 }
115
116 if offset != self.next_offset {
117 if offset < self.next_offset {
118 return self.handle_duplicate(offset, chunk);
119 }
120 return Err(EmbeddedError::InvalidCursor);
121 }
122
123 if chunk.sequence != self.expected_sequence {
124 return Err(EmbeddedError::SeqGap);
125 }
126
127 let payload_len_u32 =
128 u32::try_from(chunk.payload.len()).map_err(|_| EmbeddedError::InvalidArgument)?;
129 let new_offset =
130 self.next_offset.checked_add(payload_len_u32).ok_or(EmbeddedError::InvalidArgument)?;
131 if new_offset > self.total_size {
132 return Err(EmbeddedError::InvalidArgument);
133 }
134
135 self.bytes.extend_from_slice(&chunk.payload);
136 self.next_offset = new_offset;
137 self.expected_sequence = self.expected_sequence.saturating_add(1);
138 Ok(ChunkApply::Appended)
139 }
140
141 pub fn commit(self, expected_sha256: Option<[u8; 32]>) -> EmbeddedResult<Vec<u8>> {
142 if self.next_offset != self.total_size {
143 return Err(EmbeddedError::InvalidArgument);
144 }
145 if let Some(expected) = expected_sha256 {
146 let digest = digest32(&self.bytes);
147 if digest != expected {
148 return Err(EmbeddedError::ChecksumMismatch);
149 }
150 }
151 Ok(self.bytes)
152 }
153
154 fn handle_duplicate(&self, offset: u32, chunk: &AttachmentChunk) -> EmbeddedResult<ChunkApply> {
155 let start = usize::try_from(offset).map_err(|_| EmbeddedError::InvalidCursor)?;
156 let end = start.checked_add(chunk.payload.len()).ok_or(EmbeddedError::InvalidCursor)?;
157 if end > self.bytes.len() {
158 return Err(EmbeddedError::InvalidCursor);
159 }
160 if self.bytes[start..end] == chunk.payload {
161 return Ok(ChunkApply::DuplicateAccepted);
162 }
163 Err(EmbeddedError::IdempotencyConflict)
164 }
165}
166
167pub struct AttachmentChunker<'a> {
168 transfer_id: u32,
169 chunk_size: u16,
170 data: &'a [u8],
171 next_offset: usize,
172 next_sequence: u16,
173}
174
175impl<'a> AttachmentChunker<'a> {
176 pub fn new(transfer_id: u32, chunk_size: u16, data: &'a [u8]) -> EmbeddedResult<Self> {
177 if transfer_id == 0 || chunk_size == 0 || data.is_empty() {
178 return Err(EmbeddedError::InvalidArgument);
179 }
180 Ok(Self { transfer_id, chunk_size, data, next_offset: 0, next_sequence: 0 })
181 }
182
183 pub fn resume_from_offset(
184 transfer_id: u32,
185 chunk_size: u16,
186 data: &'a [u8],
187 offset: usize,
188 ) -> EmbeddedResult<Self> {
189 if offset > data.len() {
190 return Err(EmbeddedError::InvalidCursor);
191 }
192 let mut chunker = Self::new(transfer_id, chunk_size, data)?;
193 chunker.next_offset = offset;
194 chunker.next_sequence = u16::try_from(offset / usize::from(chunk_size))
195 .map_err(|_| EmbeddedError::InvalidCursor)?;
196 Ok(chunker)
197 }
198
199 pub fn next_chunk(&mut self) -> Option<(u32, AttachmentChunk)> {
200 if self.next_offset >= self.data.len() {
201 return None;
202 }
203 let end = (self.next_offset + usize::from(self.chunk_size)).min(self.data.len());
204 let payload = self.data[self.next_offset..end].to_vec();
205 let offset = u32::try_from(self.next_offset).ok()?;
206 let sequence = self.next_sequence;
207 self.next_offset = end;
208 self.next_sequence = self.next_sequence.saturating_add(1);
209 Some((offset, AttachmentChunk { transfer_id: self.transfer_id, sequence, payload }))
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::{AttachmentChunk, AttachmentChunker, AttachmentReceiver, ChunkApply};
216 use crate::{hash::digest32, EmbeddedError};
217
218 #[test]
219 fn roundtrip_append_duplicate_and_commit() {
220 let payload = b"abcdefghijklmnopqrstuvwxyz";
221 let mut chunker = AttachmentChunker::new(9, 8, payload).expect("chunker");
222 let mut receiver = AttachmentReceiver::start(9, payload.len() as u32, 8).expect("receiver");
223
224 let (offset0, chunk0) = chunker.next_chunk().expect("chunk0");
225 assert_eq!(
226 receiver.apply_chunk(offset0, &chunk0).expect("apply chunk0"),
227 ChunkApply::Appended
228 );
229 assert_eq!(
230 receiver.apply_chunk(offset0, &chunk0).expect("apply duplicate"),
231 ChunkApply::DuplicateAccepted
232 );
233
234 while let Some((offset, chunk)) = chunker.next_chunk() {
235 assert_eq!(receiver.apply_chunk(offset, &chunk).expect("append"), ChunkApply::Appended);
236 }
237
238 let checksum = digest32(payload);
239 let out = receiver.commit(Some(checksum)).expect("commit");
240 assert_eq!(out, payload);
241 }
242
243 #[test]
244 fn rejects_seq_gap_and_offset_mismatch() {
245 let mut receiver = AttachmentReceiver::start(77, 6, 3).expect("receiver");
246 let bad_seq = AttachmentChunk { transfer_id: 77, sequence: 1, payload: b"abc".to_vec() };
247 let err = receiver.apply_chunk(0, &bad_seq).expect_err("seq gap");
248 assert_eq!(err, EmbeddedError::SeqGap);
249
250 let good = AttachmentChunk { transfer_id: 77, sequence: 0, payload: b"abc".to_vec() };
251 receiver.apply_chunk(0, &good).expect("first chunk");
252
253 let next = AttachmentChunk { transfer_id: 77, sequence: 1, payload: b"def".to_vec() };
254 let err = receiver.apply_chunk(4, &next).expect_err("offset mismatch");
255 assert_eq!(err, EmbeddedError::InvalidCursor);
256 }
257
258 #[test]
259 fn duplicate_conflict_maps_to_idempotency_conflict() {
260 let mut receiver = AttachmentReceiver::start(21, 4, 2).expect("receiver");
261 let first = AttachmentChunk { transfer_id: 21, sequence: 0, payload: b"ab".to_vec() };
262 receiver.apply_chunk(0, &first).expect("first");
263 let conflicting = AttachmentChunk { transfer_id: 21, sequence: 0, payload: b"zz".to_vec() };
264 let err = receiver.apply_chunk(0, &conflicting).expect_err("conflict");
265 assert_eq!(err, EmbeddedError::IdempotencyConflict);
266 }
267
268 #[test]
269 fn commit_checks_completeness_and_checksum() {
270 let mut receiver = AttachmentReceiver::start(33, 3, 3).expect("receiver");
271 let chunk = AttachmentChunk { transfer_id: 33, sequence: 0, payload: b"abc".to_vec() };
272 receiver.apply_chunk(0, &chunk).expect("chunk");
273 let err = receiver.clone().commit(Some([0_u8; 32])).expect_err("bad checksum");
274 assert_eq!(err, EmbeddedError::ChecksumMismatch);
275
276 let mut incomplete = AttachmentReceiver::start(34, 4, 4).expect("incomplete");
277 let c = AttachmentChunk { transfer_id: 34, sequence: 0, payload: b"abc".to_vec() };
278 incomplete.apply_chunk(0, &c).expect("partial");
279 let err = incomplete.commit(None).expect_err("incomplete commit");
280 assert_eq!(err, EmbeddedError::InvalidArgument);
281 }
282
283 #[test]
284 fn resume_cursor_semantics() {
285 let existing = b"abcdef".to_vec();
286 let receiver = AttachmentReceiver::resume(55, 10, 3, existing).expect("resume");
287 let cursor = receiver.cursor();
288 assert_eq!(cursor.transfer_id, 55);
289 assert_eq!(cursor.next_offset, 6);
290 assert_eq!(cursor.expected_sequence, 2);
291 assert_eq!(cursor.chunk_size, 3);
292 }
293
294 #[test]
295 fn chunker_resume_from_offset() {
296 let payload = b"0123456789";
297 let mut chunker = AttachmentChunker::resume_from_offset(9, 4, payload, 4).expect("resume");
298 let (offset, chunk) = chunker.next_chunk().expect("next");
299 assert_eq!(offset, 4);
300 assert_eq!(chunk.sequence, 1);
301 assert_eq!(chunk.payload, b"4567");
302 }
303}