1#[cfg(target_endian = "big")]
6assert!(false, "This library assumes little-endian!");
7
8pub mod builder;
9mod format;
10pub mod reader;
11
12use crate::format::{CHUNK_HEADER_SIZE, ChunkHeader, SPARSE_HEADER_SIZE, SparseHeader};
13use crate::reader::SparseReader;
14use anyhow::{Context, Result, bail, ensure};
15use core::fmt;
16use serde::de::DeserializeOwned;
17use std::fs::File;
18use std::io::{Cursor, Read, Seek, SeekFrom, Write};
19use std::path::Path;
20use tempfile::{NamedTempFile, TempPath};
21
22const BLK_SIZE: u32 = 0x1000;
25
26fn deserialize_from<'a, T: DeserializeOwned, R: Read + ?Sized>(source: &mut R) -> Result<T> {
27 let mut buf = vec![0u8; std::mem::size_of::<T>()];
28 source.read_exact(&mut buf[..]).context("Failed to read bytes")?;
29 Ok(bincode::deserialize(&buf[..])?)
30}
31
32pub trait Writer: Write + Seek {
34 fn set_len(&mut self, size: u64) -> Result<()>;
36}
37
38impl Writer for File {
39 fn set_len(&mut self, size: u64) -> Result<()> {
40 Ok(File::set_len(self, size)?)
41 }
42}
43
44impl Writer for Cursor<Vec<u8>> {
45 fn set_len(&mut self, size: u64) -> Result<()> {
46 Vec::resize(self.get_mut(), size as usize, 0u8);
47 Ok(())
48 }
49}
50
51struct LimitedReader<'a, R>(pub &'a mut R, pub usize);
57
58impl<'a, R: Read + Seek> Read for LimitedReader<'a, R> {
59 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
60 let offset = self.0.stream_position()?;
61 let avail = self.1.saturating_sub(offset as usize);
62 let to_read = std::cmp::min(avail, buf.len());
63 self.0.read(&mut buf[..to_read])
64 }
65}
66
67pub fn is_sparse_image<R: Read + Seek>(reader: &mut R) -> bool {
69 || -> Option<bool> {
70 let header: SparseHeader = deserialize_from(reader).ok()?;
71 let is_sparse = header.magic == format::SPARSE_HEADER_MAGIC;
72 reader.seek(SeekFrom::Start(0)).ok()?;
73 Some(is_sparse)
74 }()
75 .unwrap_or(false)
76}
77
78#[derive(Clone, PartialEq, Debug)]
79pub enum Chunk {
80 Raw { start: u64, size: u64 },
84 Fill { start: u64, size: u64, value: u32 },
88 DontCare { start: u64, size: u64 },
95 #[allow(dead_code)]
99 Crc32 { checksum: u32 },
100}
101
102impl Chunk {
103 pub fn read_metadata<R: Read>(reader: &mut R, offset: u64, block_size: u32) -> Result<Self> {
108 let header: ChunkHeader =
109 deserialize_from(reader).context("Failed to read chunk header")?;
110 ensure!(header.valid(), "Invalid chunk header");
111
112 let size = header.chunk_sz as u64 * block_size as u64;
113 match header.chunk_type {
114 format::CHUNK_TYPE_RAW => Ok(Self::Raw { start: offset, size }),
115 format::CHUNK_TYPE_FILL => {
116 let value: u32 =
117 deserialize_from(reader).context("Failed to deserialize fill value")?;
118 Ok(Self::Fill { start: offset, size, value })
119 }
120 format::CHUNK_TYPE_DONT_CARE => Ok(Self::DontCare { start: offset, size }),
121 format::CHUNK_TYPE_CRC32 => {
122 let checksum: u32 =
123 deserialize_from(reader).context("Failed to deserialize checksum")?;
124 Ok(Self::Crc32 { checksum })
125 }
126 _ => unreachable!(),
128 }
129 }
130
131 fn valid(&self, block_size: u32) -> bool {
132 self.output_size() % (block_size as u64) == 0
133 }
134
135 fn output_offset(&self) -> Option<u64> {
138 match self {
139 Self::Raw { start, .. } => Some(*start),
140 Self::Fill { start, .. } => Some(*start),
141 Self::DontCare { start, .. } => Some(*start),
142 Self::Crc32 { .. } => None,
143 }
144 }
145
146 fn output_size(&self) -> u64 {
148 match self {
149 Self::Raw { size, .. } => *size,
150 Self::Fill { size, .. } => *size,
151 Self::DontCare { size, .. } => *size,
152 Self::Crc32 { .. } => 0,
153 }
154 }
155
156 fn output_blocks(&self, block_size: u32) -> u32 {
158 self.output_size().div_ceil(block_size as u64) as u32
159 }
160
161 fn chunk_type(&self) -> u16 {
164 match self {
165 Self::Raw { .. } => format::CHUNK_TYPE_RAW,
166 Self::Fill { .. } => format::CHUNK_TYPE_FILL,
167 Self::DontCare { .. } => format::CHUNK_TYPE_DONT_CARE,
168 Self::Crc32 { .. } => format::CHUNK_TYPE_CRC32,
169 }
170 }
171
172 fn chunk_data_len(&self) -> u32 {
179 let header_size = format::CHUNK_HEADER_SIZE;
180 let data_size = match self {
181 Self::Raw { size, .. } => *size as u32,
182 Self::Fill { .. } => std::mem::size_of::<u32>() as u32,
183 Self::DontCare { .. } => 0,
184 Self::Crc32 { .. } => std::mem::size_of::<u32>() as u32,
185 };
186 header_size.checked_add(data_size).unwrap()
187 }
188
189 fn write<W: Write, R: Read>(
193 &self,
194 source: Option<&mut R>,
195 dest: &mut W,
196 block_size: u32,
197 ) -> Result<()> {
198 ensure!(self.valid(block_size), "Not writing invalid chunk",);
199 let header = ChunkHeader::new(
200 self.chunk_type(),
201 0x0,
202 self.output_blocks(block_size),
203 self.chunk_data_len(),
204 );
205
206 bincode::serialize_into(&mut *dest, &header)?;
207
208 match self {
209 Self::Raw { size, .. } => {
210 ensure!(source.is_some(), "No source for Raw chunk");
211 let n = std::io::copy(source.unwrap(), dest)?;
212 let size = *size as u64;
213 if n < size {
214 let zeroes = vec![0u8; (size - n) as usize];
215 dest.write_all(&zeroes)?;
216 }
217 }
218 Self::Fill { value, .. } => {
219 bincode::serialize_into(dest, value)?;
221 }
222 Self::DontCare { .. } => {
223 }
225 Self::Crc32 { checksum } => {
226 bincode::serialize_into(dest, checksum)?;
227 }
228 }
229 Ok(())
230 }
231}
232
233impl fmt::Display for Chunk {
234 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235 let message = match self {
236 Self::Raw { start, size } => {
237 format!("RawChunk: start: {}, total bytes: {}", start, size)
238 }
239 Self::Fill { start, size, value } => {
240 format!("FillChunk: start: {}, value: {}, n_blocks: {}", start, value, size)
241 }
242 Self::DontCare { start, size } => {
243 format!("DontCareChunk: start: {}, bytes: {}", start, size)
244 }
245 Self::Crc32 { checksum } => format!("Crc32Chunk: checksum: {:?}", checksum),
246 };
247 write!(f, "{}", message)
248 }
249}
250
251pub const NO_SOURCE: Option<&mut Cursor<&[u8]>> = None;
255
256#[derive(Clone, Debug, PartialEq)]
257struct SparseFileWriter {
258 chunks: Vec<Chunk>,
259}
260
261impl SparseFileWriter {
262 fn new(chunks: Vec<Chunk>) -> SparseFileWriter {
263 SparseFileWriter { chunks }
264 }
265
266 fn total_blocks(&self) -> u32 {
267 self.chunks.iter().map(|c| c.output_blocks(BLK_SIZE)).sum()
268 }
269
270 fn total_bytes(&self) -> u64 {
271 self.chunks.iter().map(|c| c.output_size() as u64).sum()
272 }
273
274 fn write<W: Write + Seek, R: Read + Seek>(&self, reader: &mut R, writer: &mut W) -> Result<()> {
275 let header = SparseHeader::new(
276 BLK_SIZE.try_into().unwrap(), self.total_blocks(), self.chunks.len().try_into().unwrap(), );
280
281 bincode::serialize_into(&mut *writer, &header)?;
282
283 for chunk in &self.chunks {
284 let mut reader = if let &Chunk::Raw { start, size } = chunk {
285 reader.seek(SeekFrom::Start(start))?;
286 Some(LimitedReader(reader, start as usize + size as usize))
287 } else {
288 None
289 };
290 chunk.write(reader.as_mut(), writer, BLK_SIZE)?;
291 }
292
293 Ok(())
294 }
295}
296
297impl fmt::Display for SparseFileWriter {
298 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
299 write!(f, r"SparseFileWriter: {} Chunks:", self.chunks.len())
300 }
301}
302
303fn add_sparse_chunk(r: &mut Vec<Chunk>, chunk: Chunk) -> Result<()> {
312 match r.last_mut() {
313 Some(last) => match (&last, &chunk) {
316 (Chunk::Raw { start, size }, Chunk::Raw { size: new_length, .. })
317 if size.checked_add(*new_length).is_some() =>
318 {
319 *last = Chunk::Raw { start: *start, size: size + new_length };
320 return Ok(());
321 }
322 (
323 Chunk::Fill { start, size, value },
324 Chunk::Fill { size: new_size, value: new_value, .. },
325 ) if value == new_value && size.checked_add(*new_size).is_some() => {
326 *last = Chunk::Fill { start: *start, size: size + new_size, value: *value };
327 return Ok(());
328 }
329 (Chunk::DontCare { start, size }, Chunk::DontCare { size: new_size, .. })
330 if size.checked_add(*new_size).is_some() =>
331 {
332 *last = Chunk::DontCare { start: *start, size: size + new_size };
333 return Ok(());
334 }
335 _ => {}
336 },
337 None => {}
338 }
339
340 r.push(chunk);
345 Ok(())
346}
347
348pub fn unsparse<W: Writer, R: Read + Seek>(source: &mut R, dest: &mut W) -> Result<()> {
350 let header: SparseHeader = deserialize_from(source).context("Failed to read header")?;
351 ensure!(header.valid(), "Invalid sparse image header {:?}", header);
352
353 for _ in 0..header.total_chunks {
354 expand_chunk(source, dest, header.blk_sz).context("Failed to expand chunk")?;
355 }
356 let offset = dest.stream_position()?;
358 dest.set_len(offset).context("Failed to truncate output")?;
359 dest.flush()?;
360 Ok(())
361}
362
363fn expand_chunk<R: Read + Seek, W: Write + Seek>(
365 source: &mut R,
366 dest: &mut W,
367 block_size: u32,
368) -> Result<()> {
369 let header: ChunkHeader =
370 deserialize_from(source).context("Failed to deserialize chunk header")?;
371 ensure!(header.valid(), "Invalid chunk header {:x?}", header);
372 let size = (header.chunk_sz * block_size) as usize;
373 match header.chunk_type {
374 format::CHUNK_TYPE_RAW => {
375 let limit = source.stream_position()? as usize + size;
376 std::io::copy(&mut LimitedReader(source, limit), dest)
377 .context("Failed to copy contents")?;
378 }
379 format::CHUNK_TYPE_FILL => {
380 let value: [u8; 4] =
381 deserialize_from(source).context("Failed to deserialize fill value")?;
382 assert!(size % 4 == 0);
383 let repeated = value.repeat(size / 4);
384 dest.write_all(&repeated).context("Failed to fill contents")?;
385 }
386 format::CHUNK_TYPE_DONT_CARE => {
387 dest.seek(SeekFrom::Current(size as i64)).context("Failed to skip contents")?;
388 }
389 format::CHUNK_TYPE_CRC32 => {
390 let _: u32 = deserialize_from(source).context("Failed to deserialize fill value")?;
391 }
392 _ => bail!("Invalid type {}", header.chunk_type),
393 };
394 Ok(())
395}
396
397fn resparse(
403 sparse_file: SparseFileWriter,
404 max_download_size: u64,
405) -> Result<Vec<SparseFileWriter>> {
406 if max_download_size <= BLK_SIZE as u64 {
407 anyhow::bail!(
408 "Given maximum download size ({}) is less than the block size ({})",
409 max_download_size,
410 BLK_SIZE
411 );
412 }
413 let mut ret = Vec::<SparseFileWriter>::new();
414
415 let sunk_file_length = format::SPARSE_HEADER_SIZE as u64
418 + Chunk::DontCare { start: 0, size: BLK_SIZE.into() }.chunk_data_len() as u64
419 + Chunk::Crc32 { checksum: 2345 }.chunk_data_len() as u64;
420
421 let mut chunk_pos = 0;
422 let mut output_offset = 0;
423 while chunk_pos < sparse_file.chunks.len() {
424 log::trace!("Starting a new file at chunk position: {}", chunk_pos);
425
426 let mut file_len = 0;
427 file_len += sunk_file_length;
428
429 let mut chunks = Vec::<Chunk>::new();
430 if chunk_pos > 0 {
431 log::trace!("Adding a DontCare chunk offset: {}", chunk_pos);
434 let dont_care = Chunk::DontCare { start: 0, size: output_offset };
435 chunks.push(dont_care);
436 }
437
438 loop {
439 match sparse_file.chunks.get(chunk_pos) {
440 Some(chunk) => {
441 let curr_chunk_data_len = chunk.chunk_data_len() as u64;
442 if (file_len + curr_chunk_data_len) > max_download_size {
443 log::trace!(
444 "Current file size is: {} and adding another chunk of len: {} would \
445 put us over our max: {}",
446 file_len,
447 curr_chunk_data_len,
448 max_download_size
449 );
450
451 let remainder_size = sparse_file.total_bytes() - output_offset;
457 let dont_care =
458 Chunk::DontCare { start: output_offset, size: remainder_size };
459 chunks.push(dont_care);
460 break;
461 }
462 log::trace!(
463 "chunk: {} curr_chunk_data_len: {} current file size: {} \
464 max_download_size: {} diff: {}",
465 chunk_pos,
466 curr_chunk_data_len,
467 file_len,
468 max_download_size,
469 (max_download_size - file_len - curr_chunk_data_len)
470 );
471 add_sparse_chunk(&mut chunks, chunk.clone())?;
472 file_len += curr_chunk_data_len;
473 chunk_pos = chunk_pos + 1;
474 output_offset += chunk.output_size() as u64;
475 }
476 None => {
477 log::trace!("Finished iterating chunks");
478 break;
479 }
480 }
481 }
482 let resparsed = SparseFileWriter::new(chunks);
483 log::trace!("resparse: Adding new SparseFile: {}", resparsed);
484 ret.push(resparsed);
485 }
486
487 Ok(ret)
488}
489
490pub fn resparse_sparse_img<R: Read + std::io::Seek>(
500 reader: &mut SparseReader<R>,
501 dir: &Path,
502 max_download_size: u64,
503) -> Result<Vec<TempPath>> {
504 log::debug!("Building writer from Reader");
505 let mut chunks = vec![];
506 let header_sunk = SPARSE_HEADER_SIZE as u64;
509 let mut raw_chunks_encountered = 0;
510 for (chunk, offset) in reader.chunks() {
511 log::info!("resparse_sparse_img. Processing chunk: {} with offset: {:#?}", chunk, offset);
512 if chunk.chunk_type() == format::CHUNK_TYPE_RAW {
513 raw_chunks_encountered += 1;
514 let blks = chunk.output_blocks(BLK_SIZE);
516 log::trace!("resparse_sparse_image: splitting RAW chunk into {} chunks", blks);
517 let sunk: u64 = header_sunk + (raw_chunks_encountered * CHUNK_HEADER_SIZE) as u64;
518 for i in 0..blks {
519 let start: u64 = offset.unwrap_or(0) + (BLK_SIZE * i) as u64 - sunk;
520 log::debug!("resparse_sparse_img: adding RAW chunk at start: {}", start);
521 chunks.push(Chunk::Raw { start, size: BLK_SIZE.into() })
522 }
523 } else {
524 chunks.push(chunk.clone());
525 }
526 }
527 let sparse_file = SparseFileWriter::new(chunks);
528
529 let mut ret = Vec::<TempPath>::new();
530 log::debug!("Resparsing sparse file");
531 for re_sparsed_file in resparse(sparse_file, max_download_size)? {
532 let (file, temp_path) = NamedTempFile::new_in(dir)?.into_parts();
533 let mut file_create = File::from(file);
534
535 log::debug!("Writing resparsed {} to disk", re_sparsed_file);
536 re_sparsed_file.write(reader, &mut file_create)?;
537
538 ret.push(temp_path);
539 }
540
541 log::debug!("Finished building sparse files");
542
543 Ok(ret)
544}
545
546pub fn build_sparse_files(
557 name: &str,
558 file_to_upload: &str,
559 dir: &Path,
560 max_download_size: u64,
561) -> Result<Vec<TempPath>> {
562 if max_download_size <= BLK_SIZE as u64 {
563 anyhow::bail!(
564 "Given maximum download size ({}) is less than the block size ({})",
565 max_download_size,
566 BLK_SIZE
567 );
568 }
569 log::debug!("Building sparse files for: {}. File: {}", name, file_to_upload);
570 let mut in_file = File::open(file_to_upload)?;
571
572 let mut total_read: usize = 0;
573 let mut chunks =
575 Vec::<Chunk>::with_capacity((in_file.metadata()?.len() as usize / BLK_SIZE as usize) + 1);
576
577 let mut buf = [0u8; BLK_SIZE as usize];
578 loop {
579 let read = in_file.read(&mut buf)?;
580 if read == 0 {
581 break;
582 }
583
584 let is_fill = buf.chunks(4).collect::<Vec<&[u8]>>().windows(2).all(|w| w[0] == w[1]);
585
586 if is_fill {
587 let value: u32 = bincode::deserialize(&buf[0..4])?;
593 let fill = Chunk::Fill {
595 start: total_read as u64,
596 size: buf.len().try_into().unwrap(),
597 value,
598 };
599 log::trace!("Sparsing file: {}. Created: {}", file_to_upload, fill);
600 chunks.push(fill);
601 } else {
602 let raw = Chunk::Raw { start: total_read as u64, size: buf.len().try_into().unwrap() };
604 log::trace!("Sparsing file: {}. Created: {}", file_to_upload, raw);
605 chunks.push(raw);
606 if read < buf.len() {
607 let skip_end =
611 Chunk::DontCare { start: (total_read + read) as u64, size: BLK_SIZE.into() };
612 chunks.push(skip_end);
613 }
614 }
615 total_read += read;
616 }
617
618 log::trace!("Creating sparse file from: {} chunks", chunks.len());
619
620 let sparse_file = SparseFileWriter::new(chunks);
630 log::trace!("Created sparse file: {}", sparse_file);
631
632 let mut ret = Vec::<TempPath>::new();
633 log::trace!("Resparsing sparse file");
634 for re_sparsed_file in resparse(sparse_file, max_download_size)? {
635 let (file, temp_path) = NamedTempFile::new_in(dir)?.into_parts();
636 let mut file_create = File::from(file);
637
638 log::trace!("Writing resparsed {} to disk", re_sparsed_file);
639 re_sparsed_file.write(&mut in_file, &mut file_create)?;
640
641 ret.push(temp_path);
642 }
643
644 log::debug!("Finished building sparse files");
645
646 Ok(ret)
647}
648
649#[cfg(test)]
653mod test {
654 #[cfg(target_os = "linux")]
655 use crate::build_sparse_files;
656
657 use super::builder::{DataSource, SparseImageBuilder};
658 use super::{
659 BLK_SIZE, Chunk, NO_SOURCE, SparseFileWriter, add_sparse_chunk, resparse, unsparse,
660 };
661 use rand::rngs::SmallRng;
662 use rand::{RngCore, SeedableRng};
663 use std::io::{Cursor, Read as _, Seek as _, SeekFrom, Write as _};
664 #[cfg(target_os = "linux")]
665 use std::path::Path;
666 #[cfg(target_os = "linux")]
667 use std::process::{Command, Stdio};
668 use tempfile::{NamedTempFile, TempDir};
669
670 #[test]
671 fn test_fill_into_bytes() {
672 let mut dest = Cursor::new(Vec::<u8>::new());
673
674 let fill_chunk = Chunk::Fill { start: 0, size: (5 * BLK_SIZE).into(), value: 365 };
675 fill_chunk.write(NO_SOURCE, &mut dest, BLK_SIZE).unwrap();
676 assert_eq!(dest.into_inner(), [194, 202, 0, 0, 5, 0, 0, 0, 16, 0, 0, 0, 109, 1, 0, 0]);
677 }
678
679 #[test]
680 fn test_raw_into_bytes() {
681 const EXPECTED_RAW_BYTES: [u8; 22] =
682 [193, 202, 0, 0, 1, 0, 0, 0, 12, 16, 0, 0, 49, 50, 51, 52, 53, 0, 0, 0, 0, 0];
683
684 let mut source = Cursor::new(Vec::<u8>::from(&b"12345"[..]));
685 let mut sparse = Cursor::new(Vec::<u8>::new());
686 let chunk = Chunk::Raw { start: 0, size: BLK_SIZE.into() };
687
688 chunk.write(Some(&mut source), &mut sparse, BLK_SIZE).unwrap();
689 let buf = sparse.into_inner();
690 assert_eq!(buf.len(), 4108);
691 assert_eq!(&buf[..EXPECTED_RAW_BYTES.len()], EXPECTED_RAW_BYTES);
692 assert_eq!(&buf[EXPECTED_RAW_BYTES.len()..], &[0u8; 4108 - EXPECTED_RAW_BYTES.len()]);
693 }
694
695 #[test]
696 fn test_dont_care_into_bytes() {
697 let mut dest = Cursor::new(Vec::<u8>::new());
698 let chunk = Chunk::DontCare { start: 0, size: (5 * BLK_SIZE).into() };
699
700 chunk.write(NO_SOURCE, &mut dest, BLK_SIZE).unwrap();
701 assert_eq!(dest.into_inner(), [195, 202, 0, 0, 5, 0, 0, 0, 12, 0, 0, 0]);
702 }
703
704 #[test]
705 fn test_sparse_file_into_bytes() {
706 let mut source = Cursor::new(Vec::<u8>::from(&b"123"[..]));
707 let mut sparse = Cursor::new(Vec::<u8>::new());
708 let mut chunks = Vec::<Chunk>::new();
709 let fill = Chunk::Fill { start: 0, size: 4096, value: 5 };
711 chunks.push(fill);
712 let raw = Chunk::Raw { start: 0, size: 12288 };
714 chunks.push(raw);
715 let dontcare = Chunk::DontCare { start: 0, size: 4096 };
717 chunks.push(dontcare);
718
719 let sparsefile = SparseFileWriter::new(chunks);
720 sparsefile.write(&mut source, &mut sparse).unwrap();
721
722 sparse.seek(SeekFrom::Start(0)).unwrap();
723 let mut unsparsed = Cursor::new(Vec::<u8>::new());
724 unsparse(&mut sparse, &mut unsparsed).unwrap();
725 let buf = unsparsed.into_inner();
726 assert_eq!(buf.len(), 4096 + 12288 + 4096);
727 {
728 let chunks = buf[..4096].chunks(4);
729 for chunk in chunks {
730 assert_eq!(chunk, &[5u8, 0, 0, 0]);
731 }
732 }
733 assert_eq!(&buf[4096..4099], b"123");
734 assert_eq!(&buf[4099..16384], &[0u8; 12285]);
735 assert_eq!(&buf[16384..], &[0u8; 4096]);
736 }
737
738 #[test]
742 fn test_resparse_bails_on_too_small_size() {
743 let sparse = SparseFileWriter::new(Vec::<Chunk>::new());
744 assert!(resparse(sparse, 4095).is_err());
745 }
746
747 #[test]
748 fn test_resparse_splits() {
749 let max_download_size = 4096 * 2;
750
751 let mut chunks = Vec::<Chunk>::new();
752 chunks.push(Chunk::Raw { start: 0, size: 4096 });
753 chunks.push(Chunk::Fill { start: 4096, size: 4096, value: 2 });
754 chunks.push(Chunk::Raw { start: 8192, size: 4096 });
757
758 let input_sparse_file = SparseFileWriter::new(chunks);
759 let resparsed_files = resparse(input_sparse_file, max_download_size).unwrap();
760 assert_eq!(2, resparsed_files.len());
761
762 assert_eq!(3, resparsed_files[0].chunks.len());
763 assert_eq!(Chunk::Raw { start: 0, size: 4096 }, resparsed_files[0].chunks[0]);
764 assert_eq!(Chunk::Fill { start: 4096, size: 4096, value: 2 }, resparsed_files[0].chunks[1]);
765 assert_eq!(Chunk::DontCare { start: 8192, size: 4096 }, resparsed_files[0].chunks[2]);
766
767 assert_eq!(2, resparsed_files[1].chunks.len());
768 assert_eq!(Chunk::DontCare { start: 0, size: 8192 }, resparsed_files[1].chunks[0]);
769 assert_eq!(Chunk::Raw { start: 8192, size: 4096 }, resparsed_files[1].chunks[1]);
770 }
771
772 #[test]
776 fn test_add_sparse_chunk_adds_empty() {
777 let init_vec = Vec::<Chunk>::new();
778 let mut res = init_vec.clone();
779 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
780 assert_eq!(0, init_vec.len());
781 assert_ne!(init_vec, res);
782 assert_eq!(Chunk::Fill { start: 0, size: 4096, value: 1 }, res[0]);
783 }
784
785 #[test]
786 fn test_add_sparse_chunk_fill() {
787 {
789 let mut init_vec = Vec::<Chunk>::new();
790 init_vec.push(Chunk::Fill { start: 0, size: 8192, value: 1 });
791 let mut res = init_vec.clone();
792 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 8192, value: 1 }).unwrap();
793 assert_eq!(1, res.len());
794 assert_eq!(Chunk::Fill { start: 0, size: 16384, value: 1 }, res[0]);
795 }
796
797 {
799 let mut init_vec = Vec::<Chunk>::new();
800 init_vec.push(Chunk::Fill { start: 0, size: 4096, value: 1 });
801 let mut res = init_vec.clone();
802 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 2 }).unwrap();
803 assert_ne!(res, init_vec);
804 assert_eq!(2, res.len());
805 assert_eq!(
806 res,
807 [
808 Chunk::Fill { start: 0, size: 4096, value: 1 },
809 Chunk::Fill { start: 0, size: 4096, value: 2 }
810 ]
811 );
812 }
813
814 {
816 let mut init_vec = Vec::<Chunk>::new();
817 init_vec.push(Chunk::Fill { start: 0, size: 4096, value: 2 });
818 let mut res = init_vec.clone();
819 add_sparse_chunk(&mut res, Chunk::DontCare { start: 0, size: 4096 }).unwrap();
820 assert_ne!(res, init_vec);
821 assert_eq!(2, res.len());
822 assert_eq!(
823 res,
824 [
825 Chunk::Fill { start: 0, size: 4096, value: 2 },
826 Chunk::DontCare { start: 0, size: 4096 }
827 ]
828 );
829 }
830
831 {
833 let mut init_vec = Vec::<Chunk>::new();
834 init_vec.push(Chunk::Fill { start: 0, size: 4096, value: 1 });
835 let mut res = init_vec.clone();
836 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: u64::MAX - 4095, value: 1 })
837 .unwrap();
838 assert_ne!(res, init_vec);
839 assert_eq!(2, res.len());
840 assert_eq!(
841 res,
842 [
843 Chunk::Fill { start: 0, size: 4096, value: 1 },
844 Chunk::Fill { start: 0, size: u64::MAX - 4095, value: 1 }
845 ]
846 );
847 }
848 }
849
850 #[test]
851 fn test_add_sparse_chunk_dont_care() {
852 {
854 let mut init_vec = Vec::<Chunk>::new();
855 init_vec.push(Chunk::DontCare { start: 0, size: 4096 });
856 let mut res = init_vec.clone();
857 add_sparse_chunk(&mut res, Chunk::DontCare { start: 0, size: 4096 }).unwrap();
858 assert_eq!(1, res.len());
859 assert_eq!(Chunk::DontCare { start: 0, size: 8192 }, res[0]);
860 }
861
862 {
864 let mut init_vec = Vec::<Chunk>::new();
865 init_vec.push(Chunk::DontCare { start: 0, size: 4096 });
866 let mut res = init_vec.clone();
867 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
868 assert_eq!(2, res.len());
869 assert_eq!(
870 res,
871 [
872 Chunk::DontCare { start: 0, size: 4096 },
873 Chunk::Fill { start: 0, size: 4096, value: 1 }
874 ]
875 );
876 }
877
878 {
880 let mut init_vec = Vec::<Chunk>::new();
881 init_vec.push(Chunk::DontCare { start: 0, size: 4096 });
882 let mut res = init_vec.clone();
883 add_sparse_chunk(&mut res, Chunk::DontCare { start: 0, size: u64::MAX - 4095 })
884 .unwrap();
885 assert_eq!(2, res.len());
886 assert_eq!(
887 res,
888 [
889 Chunk::DontCare { start: 0, size: 4096 },
890 Chunk::DontCare { start: 0, size: u64::MAX - 4095 }
891 ]
892 );
893 }
894 }
895
896 #[test]
897 fn test_add_sparse_chunk_raw() {
898 {
900 let mut init_vec = Vec::<Chunk>::new();
901 init_vec.push(Chunk::Raw { start: 0, size: 12288 });
902 let mut res = init_vec.clone();
903 add_sparse_chunk(&mut res, Chunk::Raw { start: 0, size: 16384 }).unwrap();
904 assert_eq!(1, res.len());
905 assert_eq!(Chunk::Raw { start: 0, size: 28672 }, res[0]);
906 }
907
908 {
910 let mut init_vec = Vec::<Chunk>::new();
911 init_vec.push(Chunk::Raw { start: 0, size: 12288 });
912 let mut res = init_vec.clone();
913 add_sparse_chunk(&mut res, Chunk::Fill { start: 3, size: 8192, value: 1 }).unwrap();
914 assert_eq!(2, res.len());
915 assert_eq!(
916 res,
917 [
918 Chunk::Raw { start: 0, size: 12288 },
919 Chunk::Fill { start: 3, size: 8192, value: 1 }
920 ]
921 );
922 }
923
924 {
926 let mut init_vec = Vec::<Chunk>::new();
927 init_vec.push(Chunk::Raw { start: 0, size: 4096 });
928 let mut res = init_vec.clone();
929 add_sparse_chunk(&mut res, Chunk::Raw { start: 0, size: u64::MAX - 4095 }).unwrap();
930 assert_eq!(2, res.len());
931 assert_eq!(
932 res,
933 [
934 Chunk::Raw { start: 0, size: 4096 },
935 Chunk::Raw { start: 0, size: u64::MAX - 4095 }
936 ]
937 );
938 }
939 }
940
941 #[test]
942 fn test_add_sparse_chunk_crc32() {
943 {
945 let mut init_vec = Vec::<Chunk>::new();
946 init_vec.push(Chunk::Crc32 { checksum: 1234 });
947 let mut res = init_vec.clone();
948 add_sparse_chunk(&mut res, Chunk::Crc32 { checksum: 2345 }).unwrap();
949 assert_eq!(2, res.len());
950 assert_eq!(res, [Chunk::Crc32 { checksum: 1234 }, Chunk::Crc32 { checksum: 2345 }]);
951 }
952
953 {
955 let mut init_vec = Vec::<Chunk>::new();
956 init_vec.push(Chunk::Crc32 { checksum: 1234 });
957 let mut res = init_vec.clone();
958 add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
959 assert_eq!(2, res.len());
960 assert_eq!(
961 res,
962 [Chunk::Crc32 { checksum: 1234 }, Chunk::Fill { start: 0, size: 4096, value: 1 }]
963 );
964 }
965 }
966
967 #[test]
972 fn test_roundtrip() {
973 let tmpdir = TempDir::new().unwrap();
974
975 let (mut file, _temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
977 let mut rng = SmallRng::from_os_rng();
978 let mut buf = Vec::<u8>::new();
979 buf.resize(1 * 4096, 0);
980 rng.fill_bytes(&mut buf);
981 file.write_all(&buf).unwrap();
982 file.flush().unwrap();
983 file.seek(SeekFrom::Start(0)).unwrap();
984 let content_size = buf.len();
985
986 let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
988 SparseImageBuilder::new()
989 .add_source(DataSource::Buffer(Box::new([0xffu8; 8192])))
990 .add_source(DataSource::Reader { reader: Box::new(file), size: content_size as u64 })
991 .add_source(DataSource::Fill(0xaaaa_aaaau32, 1024))
992 .add_source(DataSource::Skip(16384))
993 .build(&mut sparse_file)
994 .expect("Build sparse image failed");
995 sparse_file.seek(SeekFrom::Start(0)).unwrap();
996
997 let mut orig_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
998 unsparse(&mut sparse_file, &mut orig_file).expect("unsparse failed");
999 orig_file.seek(SeekFrom::Start(0)).unwrap();
1000
1001 let mut unsparsed_bytes = vec![];
1002 orig_file.read_to_end(&mut unsparsed_bytes).expect("Failed to read unsparsed image");
1003 assert_eq!(unsparsed_bytes.len(), 8192 + 20480 + content_size);
1004 assert_eq!(&unsparsed_bytes[..8192], &[0xffu8; 8192]);
1005 assert_eq!(&unsparsed_bytes[8192..8192 + content_size], &buf[..]);
1006 assert_eq!(&unsparsed_bytes[8192 + content_size..12288 + content_size], &[0xaau8; 4096]);
1007 assert_eq!(&unsparsed_bytes[12288 + content_size..], &[0u8; 16384]);
1008 }
1009
1010 #[test]
1011 #[cfg(target_os = "linux")]
1023 fn test_with_simg2img() {
1024 let simg2img_path = Path::new("./host_x64/test_data/storage/sparse/simg2img");
1025 assert!(
1026 Path::exists(simg2img_path),
1027 "simg2img binary must exist at {}",
1028 simg2img_path.display()
1029 );
1030
1031 let tmpdir = TempDir::new().unwrap();
1032
1033 let (mut file, temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
1035 let mut rng = SmallRng::from_os_rng();
1036 let mut buf = Vec::<u8>::new();
1037 buf.resize(50 * 4096 + 1244, 0);
1039 rng.fill_bytes(&mut buf);
1040 file.write_all(&buf).unwrap();
1041 file.flush().unwrap();
1042 file.seek(SeekFrom::Start(0)).unwrap();
1043
1044 let files = build_sparse_files(
1046 "test",
1047 temp_path.to_path_buf().to_str().expect("Should succeed"),
1048 tmpdir.path(),
1049 4096 * 2,
1050 )
1051 .unwrap();
1052
1053 let mut simg2img_output = tmpdir.path().to_path_buf();
1054 simg2img_output.push("output");
1055
1056 let mut simg2img = Command::new(simg2img_path)
1057 .args(&files[..])
1058 .arg(&simg2img_output)
1059 .stdout(Stdio::piped())
1060 .stderr(Stdio::piped())
1061 .spawn()
1062 .expect("Failed to spawn simg2img");
1063 let res = simg2img.wait().expect("simg2img did was not running");
1064 assert!(res.success(), "simg2img did not succeed");
1065 let mut simg2img_stdout = simg2img.stdout.take().expect("Get stdout from simg2img");
1066 let mut simg2img_stderr = simg2img.stderr.take().expect("Get stderr from simg2img");
1067
1068 let mut stdout = String::new();
1069 simg2img_stdout.read_to_string(&mut stdout).expect("Reading simg2img stdout");
1070 assert_eq!(stdout, "");
1071
1072 let mut stderr = String::new();
1073 simg2img_stderr.read_to_string(&mut stderr).expect("Reading simg2img stderr");
1074 assert_eq!(stderr, "");
1075
1076 let simg2img_output_bytes =
1077 std::fs::read(simg2img_output).expect("Failed to read simg2img output");
1078
1079 assert_eq!(
1080 buf,
1081 simg2img_output_bytes[0..buf.len()],
1082 "Output from simg2img should match our generated file"
1083 );
1084
1085 assert_eq!(
1086 simg2img_output_bytes[buf.len()..],
1087 vec![0u8; simg2img_output_bytes.len() - buf.len()],
1088 "The remainder of our simg2img_output_bytes should be 0"
1089 );
1090 }
1091
1092 #[test]
1093 #[cfg(target_os = "linux")]
1094 fn test_resparse_from_sparse() {
1095 use crate::reader::SparseReader;
1096 use crate::resparse_sparse_img;
1097
1098 let simg2img_path = Path::new("./host_x64/test_data/storage/sparse/simg2img");
1099 assert!(
1100 Path::exists(simg2img_path),
1101 "simg2img binary must exist at {}",
1102 simg2img_path.display()
1103 );
1104
1105 let tmpdir = TempDir::new().unwrap();
1106
1107 let (mut file, _temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
1109 let mut rng = SmallRng::from_os_rng();
1110 let mut buf = Vec::<u8>::new();
1111 buf.resize(10 * 4096, 0);
1112 rng.fill_bytes(&mut buf);
1113 file.write_all(&buf).unwrap();
1114 file.flush().unwrap();
1115 file.seek(SeekFrom::Start(0)).unwrap();
1116 let content_size = buf.len();
1117
1118 let sparse_file_tmp = NamedTempFile::new_in(&tmpdir).unwrap();
1120 let mut sparse_file = sparse_file_tmp.into_file();
1121 SparseImageBuilder::new()
1122 .add_source(DataSource::Buffer(Box::new([0xffu8; 4096 * 2])))
1123 .add_source(DataSource::Reader { reader: Box::new(file), size: content_size as u64 })
1124 .add_source(DataSource::Fill(0xaaaa_aaaau32, 1024))
1125 .add_source(DataSource::Skip(16384))
1126 .build(&mut sparse_file)
1127 .expect("Build sparse image failed");
1128 sparse_file.seek(SeekFrom::Start(0)).unwrap();
1129
1130 let mut reader = SparseReader::new(sparse_file).expect("create reader");
1131
1132 let files = resparse_sparse_img(&mut reader, tmpdir.path(), 4096 * 3).unwrap();
1133
1134 let mut simg2img_output_sparsed = tmpdir.path().to_path_buf();
1136 simg2img_output_sparsed.push("output_sparsed");
1137
1138 let mut simg2img_sparsed = Command::new(simg2img_path)
1139 .args(&files[..])
1140 .arg(&simg2img_output_sparsed)
1141 .stdout(Stdio::piped())
1142 .stderr(Stdio::piped())
1143 .spawn()
1144 .expect("Failed to spawn simg2img");
1145 let res = simg2img_sparsed.wait().expect("simg2img did was not running");
1146 assert!(res.success(), "simg2img did not succeed");
1147 let mut simg2img_stdout = simg2img_sparsed.stdout.take().expect("Get stdout from simg2img");
1148 let mut simg2img_stderr = simg2img_sparsed.stderr.take().expect("Get stderr from simg2img");
1149
1150 let mut stdout = String::new();
1151 simg2img_stdout.read_to_string(&mut stdout).expect("Reading simg2img stdout");
1152 assert_eq!(stdout, "");
1153
1154 let mut stderr = String::new();
1155 simg2img_stderr.read_to_string(&mut stderr).expect("Reading simg2img stderr");
1156 assert_eq!(stderr, "");
1157 }
1158}