1use alloc::vec::Vec;
55
56use xxhash_rust::xxh3::{Xxh3, xxh3_64};
57
58use crate::error::Error;
59
60pub const MAGIC: u32 = 0x504C_474D;
62
63pub const FORMAT_VERSION: u16 = 1;
65
66pub const FLAG_VECTORS: u16 = 1;
68
69const HEADER: usize = 64;
71
72const ENTRY: usize = 32;
74
75const ALIGN: usize = 64;
77
78const FILE_HASH_AT: usize = 20;
80
81fn align_up(v: u64) -> u64 {
83 v.div_ceil(ALIGN as u64) * ALIGN as u64
84}
85
86fn file_hash(bytes: &[u8]) -> u64 {
88 let mut h = Xxh3::new();
89 h.update(&bytes[..FILE_HASH_AT]);
90 h.update(&[0u8; 8]);
91 h.update(&bytes[FILE_HASH_AT + 8..]);
92 h.digest()
93}
94
95#[derive(Debug, Default)]
109pub struct SnapshotWriter {
110 sections: Vec<(u16, Vec<u8>)>,
111}
112
113impl SnapshotWriter {
114 pub fn new() -> Self {
116 Self::default()
117 }
118
119 pub fn section(&mut self, kind: u16, bytes: Vec<u8>) -> Result<(), Error> {
126 if self.sections.iter().any(|&(k, _)| k == kind) {
127 return Err(Error::Corrupt("duplicate section kind"));
128 }
129 self.sections.push((kind, bytes));
130 Ok(())
131 }
132
133 pub fn finish(self, config: &[u8], flags: u16, created_at: u64, engine_ver: &str) -> Vec<u8> {
140 let config_end = HEADER as u64 + config.len() as u64;
141 let table_start = align_up(config_end);
142 let table_end = table_start + (self.sections.len() * ENTRY) as u64;
143
144 let mut offsets = Vec::with_capacity(self.sections.len());
146 let mut cursor = align_up(table_end);
147 for (_, bytes) in &self.sections {
148 offsets.push(cursor);
149 cursor = align_up(cursor + bytes.len() as u64);
150 }
151 let file_len = cursor as usize;
152
153 let mut out = alloc::vec![0u8; file_len];
154 out[0..4].copy_from_slice(&MAGIC.to_le_bytes());
155 out[4..6].copy_from_slice(&FORMAT_VERSION.to_le_bytes());
156 out[6..8].copy_from_slice(&flags.to_le_bytes());
157 out[8..10].copy_from_slice(&(self.sections.len() as u16).to_le_bytes());
158 out[16..20].copy_from_slice(&(config.len() as u32).to_le_bytes());
159 out[28..36].copy_from_slice(&created_at.to_le_bytes());
160 let ver = engine_ver.as_bytes();
161 let ver_len = ver.len().min(24);
162 out[36..36 + ver_len].copy_from_slice(&ver[..ver_len]);
163 out[HEADER..HEADER + config.len()].copy_from_slice(config);
164
165 for (i, (kind, bytes)) in self.sections.iter().enumerate() {
166 let at = table_start as usize + i * ENTRY;
167 out[at..at + 2].copy_from_slice(&kind.to_le_bytes());
168 out[at + 2..at + 4].copy_from_slice(&(ALIGN as u16).to_le_bytes());
169 out[at + 8..at + 16].copy_from_slice(&offsets[i].to_le_bytes());
170 out[at + 16..at + 24].copy_from_slice(&(bytes.len() as u64).to_le_bytes());
171 out[at + 24..at + 32].copy_from_slice(&xxh3_64(bytes).to_le_bytes());
172 let start = offsets[i] as usize;
173 out[start..start + bytes.len()].copy_from_slice(bytes);
174 }
175
176 let hash = file_hash(&out);
177 out[FILE_HASH_AT..FILE_HASH_AT + 8].copy_from_slice(&hash.to_le_bytes());
178 out
179 }
180}
181
182#[derive(Debug, Clone, Copy)]
185#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
186pub struct SectionMeta {
187 pub kind: u16,
189 pub len: u64,
191 pub hash: u64,
193}
194
195#[derive(Debug)]
199#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
200pub struct Prefix {
201 pub bytes: Vec<u8>,
205 pub offsets: Vec<u64>,
207 pub file_len: u64,
209}
210
211pub trait SnapshotSink {
216 fn write(&mut self, bytes: &[u8]) -> Result<(), Error>;
221
222 fn patch(&mut self, at: u64, bytes: &[u8]) -> Result<(), Error>;
227}
228
229impl SnapshotSink for &mut Vec<u8> {
230 fn write(&mut self, bytes: &[u8]) -> Result<(), Error> {
231 self.extend_from_slice(bytes);
232 Ok(())
233 }
234
235 fn patch(&mut self, at: u64, bytes: &[u8]) -> Result<(), Error> {
236 let at = at as usize;
237 self[at..at + bytes.len()].copy_from_slice(bytes);
238 Ok(())
239 }
240}
241
242pub const FILE_HASH_OFFSET: u64 = FILE_HASH_AT as u64;
245
246pub fn build_prefix(
253 config: &[u8],
254 flags: u16,
255 created_at: u64,
256 engine_ver: &str,
257 metas: &[SectionMeta],
258) -> Prefix {
259 let count = metas.len();
260 let config_end = HEADER as u64 + config.len() as u64;
261 let table_start = align_up(config_end);
262 let table_end = table_start + (count * ENTRY) as u64;
263
264 let mut offsets = Vec::with_capacity(count);
265 let mut cursor = align_up(table_end);
266 for m in metas {
267 offsets.push(cursor);
268 cursor = align_up(cursor + m.len);
269 }
270 let file_len = cursor;
271
272 let prefix_len = align_up(table_end) as usize;
273 let mut out = alloc::vec![0u8; prefix_len];
274 out[0..4].copy_from_slice(&MAGIC.to_le_bytes());
275 out[4..6].copy_from_slice(&FORMAT_VERSION.to_le_bytes());
276 out[6..8].copy_from_slice(&flags.to_le_bytes());
277 out[8..10].copy_from_slice(&(count as u16).to_le_bytes());
278 out[16..20].copy_from_slice(&(config.len() as u32).to_le_bytes());
279 out[28..36].copy_from_slice(&created_at.to_le_bytes());
280 let ver = engine_ver.as_bytes();
281 let ver_len = ver.len().min(24);
282 out[36..36 + ver_len].copy_from_slice(&ver[..ver_len]);
283 out[HEADER..HEADER + config.len()].copy_from_slice(config);
284
285 for (i, m) in metas.iter().enumerate() {
286 let at = table_start as usize + i * ENTRY;
287 out[at..at + 2].copy_from_slice(&m.kind.to_le_bytes());
288 out[at + 2..at + 4].copy_from_slice(&(ALIGN as u16).to_le_bytes());
289 out[at + 8..at + 16].copy_from_slice(&offsets[i].to_le_bytes());
290 out[at + 16..at + 24].copy_from_slice(&m.len.to_le_bytes());
291 out[at + 24..at + 32].copy_from_slice(&m.hash.to_le_bytes());
292 }
293 Prefix {
295 bytes: out,
296 offsets,
297 file_len,
298 }
299}
300
301pub fn pad_len(offset: u64, len: u64) -> usize {
304 (align_up(offset + len) - (offset + len)) as usize
305}
306
307#[derive(Debug)]
309pub struct Snapshot<'a> {
310 bytes: &'a [u8],
311 pub flags: u16,
313 pub created_at: u64,
315 config_len: usize,
316 sections: Vec<(u16, usize, usize, u64)>,
320 engine_ver_len: usize,
321}
322
323impl<'a> Snapshot<'a> {
324 pub fn parse(bytes: &'a [u8]) -> Result<Self, Error> {
335 if bytes.len() < HEADER {
336 return Err(Error::Corrupt("snapshot shorter than its header"));
337 }
338 if !bytes.len().is_multiple_of(ALIGN) {
339 return Err(Error::Corrupt("snapshot length is not 64-byte aligned"));
340 }
341 if u32::from_le_bytes(bytes[0..4].try_into().unwrap()) != MAGIC {
342 return Err(Error::Corrupt("bad magic"));
343 }
344 let version = u16::from_le_bytes(bytes[4..6].try_into().unwrap());
345 if version != FORMAT_VERSION {
346 return Err(Error::UnsupportedVersion(version));
347 }
348 let flags = u16::from_le_bytes(bytes[6..8].try_into().unwrap());
349 if flags & !FLAG_VECTORS != 0 {
350 return Err(Error::Corrupt("unknown flag bits set"));
351 }
352 let section_cnt = u16::from_le_bytes(bytes[8..10].try_into().unwrap()) as usize;
353 if bytes[10..16] != [0u8; 6] || bytes[60..64] != [0u8; 4] {
354 return Err(Error::Corrupt("reserved header bytes must be zero"));
355 }
356 let config_len = u32::from_le_bytes(bytes[16..20].try_into().unwrap()) as usize;
357 let created_at = u64::from_le_bytes(bytes[28..36].try_into().unwrap());
358 let ver_bytes = &bytes[36..60];
359 let engine_ver_len = ver_bytes.iter().position(|&b| b == 0).unwrap_or(24);
360 if ver_bytes[engine_ver_len..].iter().any(|&b| b != 0) {
361 return Err(Error::Corrupt("engine version is not zero-terminated"));
362 }
363 if core::str::from_utf8(&ver_bytes[..engine_ver_len]).is_err() {
364 return Err(Error::Corrupt("engine version is not UTF-8"));
365 }
366
367 let file_len = bytes.len() as u64;
368 let config_end = HEADER as u64 + config_len as u64;
369 let table_start = align_up(config_end);
370 let table_end = table_start + (section_cnt * ENTRY) as u64;
371 if table_end > file_len {
372 return Err(Error::Corrupt("section table out of bounds"));
373 }
374 if bytes[config_end as usize..table_start as usize]
375 .iter()
376 .any(|&b| b != 0)
377 {
378 return Err(Error::Corrupt("nonzero padding after the config block"));
379 }
380
381 let mut sections = Vec::with_capacity(section_cnt);
383 let mut expected = align_up(table_end);
384 if bytes[table_end as usize..expected as usize]
385 .iter()
386 .any(|&b| b != 0)
387 {
388 return Err(Error::Corrupt("nonzero padding after the section table"));
389 }
390 for i in 0..section_cnt {
391 let at = table_start as usize + i * ENTRY;
392 let kind = u16::from_le_bytes(bytes[at..at + 2].try_into().unwrap());
393 let align = u16::from_le_bytes(bytes[at + 2..at + 4].try_into().unwrap());
394 if align as usize != ALIGN {
395 return Err(Error::Corrupt("section alignment must be 64"));
396 }
397 if bytes[at + 4..at + 8] != [0u8; 4] {
398 return Err(Error::Corrupt("reserved section bytes must be zero"));
399 }
400 let offset = u64::from_le_bytes(bytes[at + 8..at + 16].try_into().unwrap());
401 let len = u64::from_le_bytes(bytes[at + 16..at + 24].try_into().unwrap());
402 let want = u64::from_le_bytes(bytes[at + 24..at + 32].try_into().unwrap());
403 if offset != expected {
404 return Err(Error::Corrupt("sections must be contiguous in table order"));
405 }
406 let end = offset
407 .checked_add(len)
408 .ok_or(Error::Corrupt("section length overflow"))?;
409 if end > file_len {
410 return Err(Error::Corrupt("section out of bounds"));
411 }
412 if sections.iter().any(|&(k, _, _, _)| k == kind) {
413 return Err(Error::Corrupt("duplicate section kind"));
414 }
415 expected = align_up(end);
416 if bytes[end as usize..expected.min(file_len) as usize]
417 .iter()
418 .any(|&b| b != 0)
419 {
420 return Err(Error::Corrupt("nonzero padding after a section"));
421 }
422 sections.push((kind, offset as usize, len as usize, want));
423 }
424 if expected != file_len {
425 return Err(Error::Corrupt("trailing bytes after the last section"));
426 }
427
428 Ok(Self {
429 bytes,
430 flags,
431 created_at,
432 config_len,
433 sections,
434 engine_ver_len,
435 })
436 }
437
438 pub fn config(&self) -> &'a [u8] {
441 &self.bytes[HEADER..HEADER + self.config_len]
442 }
443
444 pub fn section(&self, kind: u16) -> Option<&'a [u8]> {
447 self.sections
448 .iter()
449 .find(|&&(k, _, _, _)| k == kind)
450 .map(|&(_, start, len, _)| &self.bytes[start..start + len])
451 }
452
453 pub fn engine_ver(&self) -> &'a str {
455 core::str::from_utf8(&self.bytes[36..36 + self.engine_ver_len])
456 .expect("validated during parse")
457 }
458
459 pub fn scrub(&self) -> ScrubCursor<'a> {
462 self.scrub_with_budget(DEFAULT_SCRUB_BUDGET)
463 }
464
465 pub fn scrub_with_budget(&self, budget: usize) -> ScrubCursor<'a> {
468 ScrubCursor {
469 bytes: self.bytes,
470 sections: self
471 .sections
472 .iter()
473 .map(|&(_, s, l, w)| (s, l, w))
474 .collect(),
475 budget: budget.max(1),
476 pos: 0,
477 sec: 0,
478 sec_hash: Xxh3::new(),
479 file: Xxh3::new(),
480 done: false,
481 }
482 }
483}
484
485pub const DEFAULT_SCRUB_BUDGET: usize = 1 << 20;
492
493#[derive(Debug, Clone, Copy, PartialEq, Eq)]
495#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
496pub struct ScrubProgress {
497 pub done_bytes: u64,
499 pub total_bytes: u64,
501}
502
503pub struct ScrubCursor<'a> {
514 bytes: &'a [u8],
515 sections: Vec<(usize, usize, u64)>,
517 budget: usize,
518 pos: usize,
519 sec: usize,
521 sec_hash: Xxh3,
523 file: Xxh3,
525 done: bool,
526}
527
528impl ScrubCursor<'_> {
529 fn feed_file(&mut self, from: usize, to: usize) {
532 let before_end = to.min(FILE_HASH_AT);
533 if from < before_end {
534 self.file.update(&self.bytes[from..before_end]);
535 }
536 let z_start = from.max(FILE_HASH_AT);
537 let z_end = to.min(FILE_HASH_AT + 8);
538 if z_start < z_end {
539 self.file.update(&[0u8; 8][..z_end - z_start]);
540 }
541 let after_start = from.max(FILE_HASH_AT + 8);
542 if after_start < to {
543 self.file.update(&self.bytes[after_start..to]);
544 }
545 }
546
547 fn complete_section(&mut self, want: u64) -> Result<(), Error> {
550 if self.sec_hash.digest() != want {
551 self.done = true;
552 return Err(Error::Corrupt("section checksum mismatch"));
553 }
554 self.sec += 1;
555 self.sec_hash = Xxh3::new();
556 Ok(())
557 }
558}
559
560impl Iterator for ScrubCursor<'_> {
561 type Item = Result<ScrubProgress, Error>;
562
563 fn next(&mut self) -> Option<Self::Item> {
564 if self.done {
565 return None;
566 }
567 let n = self.sections.len();
568 let file_len = self.bytes.len();
569 let mut budget_left = self.budget;
570
571 while budget_left > 0 && self.pos < file_len {
572 while self.sec < n {
574 let (start, len, want) = self.sections[self.sec];
575 if self.pos == start && len == 0 {
576 if let Err(e) = self.complete_section(want) {
577 return Some(Err(e));
578 }
579 } else {
580 break;
581 }
582 }
583
584 let (in_body, boundary) = match self.sections.get(self.sec) {
585 Some(&(start, len, _)) if self.pos >= start => (true, start + len),
586 Some(&(start, _, _)) => (false, start),
587 None => (false, file_len),
588 };
589 let end = (self.pos + budget_left).min(boundary);
590 self.feed_file(self.pos, end);
591 if in_body {
592 self.sec_hash.update(&self.bytes[self.pos..end]);
593 }
594 budget_left -= end - self.pos;
595 self.pos = end;
596
597 if in_body && self.pos == boundary {
598 let want = self.sections[self.sec].2;
599 if let Err(e) = self.complete_section(want) {
600 return Some(Err(e));
601 }
602 }
603 }
604
605 if self.pos == file_len {
606 while self.sec < n {
608 let want = self.sections[self.sec].2;
609 if let Err(e) = self.complete_section(want) {
610 return Some(Err(e));
611 }
612 }
613 self.done = true;
614 let stored = u64::from_le_bytes(
615 self.bytes[FILE_HASH_AT..FILE_HASH_AT + 8]
616 .try_into()
617 .unwrap(),
618 );
619 if stored != 0 && self.file.digest() != stored {
621 return Some(Err(Error::Corrupt("file checksum mismatch")));
622 }
623 }
624
625 Some(Ok(ScrubProgress {
626 done_bytes: self.pos as u64,
627 total_bytes: file_len as u64,
628 }))
629 }
630}