sparse/
lib.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#[cfg(target_endian = "big")]
6assert!(false, "This library assumes little-endian!");
7
8pub mod builder;
9mod format;
10pub mod reader;
11
12use crate::format::{ChunkHeader, SparseHeader};
13use anyhow::{bail, ensure, Context, Result};
14use core::fmt;
15use serde::de::DeserializeOwned;
16use std::fs::File;
17use std::io::{Cursor, Read, Seek, SeekFrom, Write};
18use std::path::Path;
19use tempfile::{NamedTempFile, TempPath};
20
21// Size of blocks to write.  Note that the format supports varied block sizes; this is the preferred
22// size by this library.
23const BLK_SIZE: usize = 0x1000;
24
25fn deserialize_from<'a, T: DeserializeOwned, R: Read + ?Sized>(source: &mut R) -> Result<T> {
26    let mut buf = vec![0u8; std::mem::size_of::<T>()];
27    source.read_exact(&mut buf[..]).context("Failed to read bytes")?;
28    Ok(bincode::deserialize(&buf[..])?)
29}
30
31/// A union trait for `Read` and `Seek`.
32pub trait Reader: Read + Seek {}
33
34impl<T: Read + Seek> Reader for T {}
35
36/// A union trait for `Write` and `Seek` that also allows truncation.
37pub trait Writer: Write + Seek {
38    /// Sets the length of the output stream.
39    fn set_len(&mut self, size: u64) -> Result<()>;
40}
41
42impl Writer for File {
43    fn set_len(&mut self, size: u64) -> Result<()> {
44        Ok(File::set_len(self, size)?)
45    }
46}
47
48impl Writer for Cursor<Vec<u8>> {
49    fn set_len(&mut self, size: u64) -> Result<()> {
50        Vec::resize(self.get_mut(), size as usize, 0u8);
51        Ok(())
52    }
53}
54
55// A wrapper around a Reader, which makes it seem like the underlying stream is only self.1 bytes
56// long.  The underlying reader is still advanced upon reading.
57// This is distinct from `std::io::Take` in that it does not modify the seek offset of the
58// underlying reader.  In other words, `LimitedReader` can be used to read a window within the
59// reader (by setting seek offset to the start, and the size limit to the end).
60struct LimitedReader<'a>(pub &'a mut dyn Reader, pub usize);
61
62impl<'a> Read for LimitedReader<'a> {
63    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
64        let offset = self.0.stream_position()?;
65        let avail = self.1.saturating_sub(offset as usize);
66        let to_read = std::cmp::min(avail, buf.len());
67        self.0.read(&mut buf[..to_read])
68    }
69}
70
71impl<'a> Seek for LimitedReader<'a> {
72    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
73        self.0.seek(pos)
74    }
75}
76
77/// Returns whether the image in `reader` appears to be in the sparse format.
78pub fn is_sparse_image<R: Reader>(reader: &mut R) -> bool {
79    || -> Option<bool> {
80        let header: SparseHeader = deserialize_from(reader).ok()?;
81        let is_sparse = header.magic == format::SPARSE_HEADER_MAGIC;
82        reader.seek(SeekFrom::Start(0)).ok()?;
83        Some(is_sparse)
84    }()
85    .unwrap_or(false)
86}
87
88#[derive(Clone, PartialEq, Debug)]
89enum Chunk {
90    /// `Raw` represents a set of blocks to be written to disk as-is.
91    /// `start` is the offset in the expanded image at which the Raw section starts.
92    /// `start` and `size` are in bytes, but must be block-aligned.
93    Raw { start: u64, size: usize },
94    /// `Fill` represents a Chunk that has the `value` repeated enough to fill `size` bytes.
95    /// `start` is the offset in the expanded image at which the Fill section starts.
96    /// `start` and `size` are in bytes, but must be block-aligned.
97    Fill { start: u64, size: usize, value: u32 },
98    /// `DontCare` represents a set of blocks that need to be "offset" by the
99    /// image recipient.  If an image needs to be broken up into two sparse images, and we flash n
100    /// bytes for Sparse Image 1, Sparse Image 2 needs to start with a DontCareChunk with
101    /// (n/blocksize) blocks as its "size" property.
102    /// `start` is the offset in the expanded image at which the DontCare section starts.
103    /// `start` and `size` are in bytes, but must be block-aligned.
104    DontCare { start: u64, size: usize },
105    /// `Crc32Chunk` is used as a checksum of a given set of Chunks for a SparseImage.  This is not
106    /// required and unused in most implementations of the Sparse Image format. The type is included
107    /// for completeness. It has 4 bytes of CRC32 checksum as describable in a u32.
108    #[allow(dead_code)]
109    Crc32 { checksum: u32 },
110}
111
112impl Chunk {
113    /// Attempts to read a `Chunk` from `reader`.  The reader will be positioned at the first byte
114    /// following the chunk header and any extra data; for a Raw chunk this means it will point at
115    /// the data payload, and for other chunks it will point at the next chunk header (or EOF).
116    /// `offset` is the current offset in the logical volume.
117    pub fn read_metadata<R: Reader>(reader: &mut R, offset: u64, block_size: u32) -> Result<Self> {
118        let header: ChunkHeader =
119            deserialize_from(reader).context("Failed to read chunk header")?;
120        ensure!(header.valid(), "Invalid chunk header");
121
122        let size = (header.chunk_sz * block_size) as usize;
123        match header.chunk_type {
124            format::CHUNK_TYPE_RAW => Ok(Self::Raw { start: offset, size }),
125            format::CHUNK_TYPE_FILL => {
126                let value: u32 =
127                    deserialize_from(reader).context("Failed to deserialize fill value")?;
128                Ok(Self::Fill { start: offset, size, value })
129            }
130            format::CHUNK_TYPE_DONT_CARE => Ok(Self::DontCare { start: offset, size }),
131            format::CHUNK_TYPE_CRC32 => {
132                let checksum: u32 =
133                    deserialize_from(reader).context("Failed to deserialize checksum")?;
134                Ok(Self::Crc32 { checksum })
135            }
136            // We already validated the chunk_type in `ChunkHeader::is_valid`.
137            _ => unreachable!(),
138        }
139    }
140
141    fn valid(&self, block_size: usize) -> bool {
142        self.output_size() % block_size == 0
143    }
144
145    /// Returns the offset into the logical image the chunk refers to, or None if the chunk has no
146    /// output data.
147    fn output_offset(&self) -> Option<u64> {
148        match self {
149            Self::Raw { start, .. } => Some(*start),
150            Self::Fill { start, .. } => Some(*start),
151            Self::DontCare { start, .. } => Some(*start),
152            Self::Crc32 { .. } => None,
153        }
154    }
155
156    /// Return number of bytes the chunk expands to when written to the partition.
157    fn output_size(&self) -> usize {
158        match self {
159            Self::Raw { size, .. } => *size,
160            Self::Fill { size, .. } => *size,
161            Self::DontCare { size, .. } => *size,
162            Self::Crc32 { .. } => 0,
163        }
164    }
165
166    /// Return number of blocks the chunk expands to when written to the partition.
167    fn output_blocks(&self, block_size: usize) -> u32 {
168        let size_bytes = self.output_size();
169        ((size_bytes + block_size - 1) / block_size) as u32
170    }
171
172    /// `chunk_type` returns the integer flag to represent the type of chunk
173    /// to use in the ChunkHeader
174    fn chunk_type(&self) -> u16 {
175        match self {
176            Self::Raw { .. } => format::CHUNK_TYPE_RAW,
177            Self::Fill { .. } => format::CHUNK_TYPE_FILL,
178            Self::DontCare { .. } => format::CHUNK_TYPE_DONT_CARE,
179            Self::Crc32 { .. } => format::CHUNK_TYPE_CRC32,
180        }
181    }
182
183    /// `chunk_data_len` returns the length of the chunk's header plus the
184    /// length of the data when serialized
185    fn chunk_data_len(&self) -> usize {
186        let header_size = format::CHUNK_HEADER_SIZE;
187        let data_size = match self {
188            Self::Raw { size, .. } => *size,
189            Self::Fill { .. } => std::mem::size_of::<u32>(),
190            Self::DontCare { .. } => 0,
191            Self::Crc32 { .. } => std::mem::size_of::<u32>(),
192        };
193        header_size + data_size
194    }
195
196    /// Writes the chunk to the given Writer.  `source` is a Reader containing the data payload for
197    /// a Raw type chunk, with the seek offset pointing to the first byte of the data payload, and
198    /// with exactly enough bytes available for the rest of the data payload.
199    fn write<W: Write + Seek, R: Read + Seek>(
200        &self,
201        source: Option<&mut R>,
202        dest: &mut W,
203    ) -> Result<()> {
204        ensure!(self.valid(BLK_SIZE), "Not writing invalid chunk");
205        let header = ChunkHeader::new(
206            self.chunk_type(),
207            0x0,
208            self.output_blocks(BLK_SIZE),
209            self.chunk_data_len() as u32,
210        );
211
212        let header_bytes: Vec<u8> = bincode::serialize(&header)?;
213        std::io::copy(&mut Cursor::new(header_bytes), dest)?;
214
215        match self {
216            Self::Raw { size, .. } => {
217                ensure!(source.is_some(), "No source for Raw chunk");
218                let n = std::io::copy(source.unwrap(), dest)? as usize;
219                if n < *size {
220                    let zeroes = vec![0u8; *size - n];
221                    std::io::copy(&mut Cursor::new(zeroes), dest)?;
222                }
223            }
224            Self::Fill { value, .. } => {
225                // Serliaze the value,
226                bincode::serialize_into(dest, value)?;
227            }
228            Self::DontCare { .. } => {
229                // DontCare has no data to write
230            }
231            Self::Crc32 { checksum } => {
232                bincode::serialize_into(dest, checksum)?;
233            }
234        }
235        Ok(())
236    }
237}
238
239impl fmt::Display for Chunk {
240    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
241        let message = match self {
242            Self::Raw { start, size } => {
243                format!("RawChunk: start: {}, total bytes: {}", start, size)
244            }
245            Self::Fill { start, size, value } => {
246                format!("FillChunk: start: {}, value: {}, n_blocks: {}", start, value, size)
247            }
248            Self::DontCare { start, size } => {
249                format!("DontCareChunk: start: {}, bytes: {}", start, size)
250            }
251            Self::Crc32 { checksum } => format!("Crc32Chunk: checksum: {:?}", checksum),
252        };
253        write!(f, "{}", message)
254    }
255}
256
257#[derive(Clone, Debug, PartialEq)]
258struct SparseFileWriter {
259    chunks: Vec<Chunk>,
260}
261
262impl SparseFileWriter {
263    fn new(chunks: Vec<Chunk>) -> SparseFileWriter {
264        SparseFileWriter { chunks }
265    }
266
267    fn total_blocks(&self) -> u32 {
268        self.chunks.iter().map(|c| c.output_blocks(BLK_SIZE)).sum()
269    }
270
271    fn total_bytes(&self) -> usize {
272        self.chunks.iter().map(|c| c.output_size()).sum()
273    }
274
275    fn write<W: Write + Seek, R: Read + Seek>(&self, reader: &mut R, writer: &mut W) -> Result<()> {
276        let header = SparseHeader::new(
277            BLK_SIZE.try_into().unwrap(),          // Size of the blocks
278            self.total_blocks(),                   // Total blocks in this image
279            self.chunks.len().try_into().unwrap(), // Total chunks in this image
280        );
281
282        let header_bytes: Vec<u8> = bincode::serialize(&header)?;
283        std::io::copy(&mut Cursor::new(header_bytes), writer)?;
284
285        for chunk in &self.chunks {
286            let mut reader = if let &Chunk::Raw { start, size } = chunk {
287                reader.seek(SeekFrom::Start(start))?;
288                Some(LimitedReader(reader, start as usize + size))
289            } else {
290                None
291            };
292            chunk.write(reader.as_mut(), writer)?;
293        }
294
295        Ok(())
296    }
297}
298
299impl fmt::Display for SparseFileWriter {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        write!(f, r"SparseFileWriter: {} Chunks:", self.chunks.len())
302    }
303}
304
305/// `add_sparse_chunk` takes the input vec, v and given `Chunk`, chunk, and
306/// attempts to add the chunk to the end of the vec. If the current last chunk
307/// is the same kind of Chunk as the `chunk`, then it will merge the two chunks
308/// into one chunk.
309///
310/// Example: A `FillChunk` with value 0 and size 1 is the last chunk
311/// in `v`, and `chunk` is a FillChunk with value 0 and size 1, after this,
312/// `v`'s last element will be a FillChunk with value 0 and size 2.
313fn add_sparse_chunk(r: &mut Vec<Chunk>, chunk: Chunk) -> Result<()> {
314    match r.last_mut() {
315        // We've got something in the Vec... if they are both the same type,
316        // merge them, otherwise, just push the new one
317        Some(last) => match (&last, &chunk) {
318            (Chunk::Raw { start, size }, Chunk::Raw { size: new_length, .. }) => {
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 => {
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                *last = Chunk::DontCare { start: *start, size: size + new_size };
331                return Ok(());
332            }
333            _ => {}
334        },
335        None => {}
336    }
337
338    // If the chunk types differ they cannot be merged.
339    // If they are both Fill but have different values, they cannot be merged.
340    // Crc32 cannot be merged.
341    // If we dont have any chunks then we add it
342    r.push(chunk);
343    Ok(())
344}
345
346/// Reads a sparse image from `source` and expands it to its unsparsed representation in `dest`.
347pub fn unsparse<W: Writer, R: Reader>(source: &mut R, dest: &mut W) -> Result<()> {
348    let header: SparseHeader = deserialize_from(source).context("Failed to read header")?;
349    ensure!(header.valid(), "Invalid sparse image header {:?}", header);
350    let num_chunks = header.total_chunks as usize;
351
352    for _ in 0..num_chunks {
353        expand_chunk(source, dest, header.blk_sz).context("Failed to expand chunk")?;
354    }
355    // Truncate output to its current seek offset, in case the last chunk we wrote was DontNeed.
356    let offset = dest.stream_position()?;
357    dest.set_len(offset).context("Failed to truncate output")?;
358    dest.flush()?;
359    Ok(())
360}
361
362/// Reads a chunk from `source`, and expands it, writing the result to `dest`.
363fn expand_chunk<R: Read + Seek, W: Write + Seek>(
364    source: &mut R,
365    dest: &mut W,
366    block_size: u32,
367) -> Result<()> {
368    let header: ChunkHeader =
369        deserialize_from(source).context("Failed to deserialize chunk header")?;
370    ensure!(header.valid(), "Invalid chunk header {:x?}", header);
371    let size = (header.chunk_sz * block_size) as usize;
372    match header.chunk_type {
373        format::CHUNK_TYPE_RAW => {
374            let limit = source.stream_position()? as usize + size;
375            std::io::copy(&mut LimitedReader(source, limit), dest)
376                .context("Failed to copy contents")?;
377        }
378        format::CHUNK_TYPE_FILL => {
379            let value: [u8; 4] =
380                deserialize_from(source).context("Failed to deserialize fill value")?;
381            assert!(size % 4 == 0);
382            let repeated = value.repeat(size / 4);
383            std::io::copy(&mut Cursor::new(repeated), dest).context("Failed to fill contents")?;
384        }
385        format::CHUNK_TYPE_DONT_CARE => {
386            dest.seek(SeekFrom::Current(size as i64)).context("Failed to skip contents")?;
387        }
388        format::CHUNK_TYPE_CRC32 => {
389            let _: u32 = deserialize_from(source).context("Failed to deserialize fill value")?;
390        }
391        _ => bail!("Invalid type {}", header.chunk_type),
392    };
393    Ok(())
394}
395
396/// `resparse` takes a SparseFile and a maximum size and will
397/// break the single SparseFile into multiple SparseFiles whose
398/// size will not exceed the maximum_download_size.
399///
400/// This will return an error if max_download_size is <= BLK_SIZE
401fn resparse(
402    sparse_file: SparseFileWriter,
403    max_download_size: u64,
404) -> Result<Vec<SparseFileWriter>> {
405    if max_download_size as usize <= BLK_SIZE {
406        anyhow::bail!(
407            "Given maximum download size ({}) is less than the block size ({})",
408            max_download_size,
409            BLK_SIZE
410        );
411    }
412    let mut ret = Vec::<SparseFileWriter>::new();
413
414    // File length already starts with a header for the SparseFile as
415    // well as the size of a potential DontCare and Crc32 Chunk
416    let sunk_file_length = format::SPARSE_HEADER_SIZE
417        + (Chunk::DontCare { start: 0, size: BLK_SIZE }.chunk_data_len()
418            + Chunk::Crc32 { checksum: 2345 }.chunk_data_len());
419
420    let mut chunk_pos = 0;
421    let mut output_offset = 0;
422    while chunk_pos < sparse_file.chunks.len() {
423        log::trace!("Starting a new file at chunk position: {}", chunk_pos);
424
425        let mut file_len = 0;
426        file_len += sunk_file_length;
427
428        let mut chunks = Vec::<Chunk>::new();
429        if chunk_pos > 0 {
430            // If we already have some chunks... add a DontCare block to
431            // move the pointer
432            log::trace!("Adding a DontCare chunk offset: {}", chunk_pos);
433            let dont_care = Chunk::DontCare { start: 0, size: output_offset };
434            chunks.push(dont_care);
435        }
436
437        loop {
438            match sparse_file.chunks.get(chunk_pos) {
439                Some(chunk) => {
440                    let curr_chunk_data_len = chunk.chunk_data_len();
441                    if (file_len + curr_chunk_data_len) as u64 > max_download_size {
442                        log::trace!("Current file size is: {} and adding another chunk of len: {} would put us over our max: {}", file_len, curr_chunk_data_len, max_download_size);
443
444                        // Add a dont care chunk to cover everything to the end of the image.
445                        // While this is not strictly speaking needed, other tools
446                        // (simg2simg) produce this chunk, and the Sparse image inspection tool
447                        // simg_dump will produce a warning if a sparse file does not have the same
448                        // number of output blocks as declared in the header.
449                        let remainder_size = sparse_file.total_bytes() - output_offset;
450                        let dont_care =
451                            Chunk::DontCare { start: output_offset as u64, size: remainder_size };
452                        chunks.push(dont_care);
453                        break;
454                    }
455                    log::trace!("chunk: {} curr_chunk_data_len: {} current file size: {} max_download_size: {} diff: {}", chunk_pos, curr_chunk_data_len, file_len, max_download_size, (max_download_size as usize - file_len - curr_chunk_data_len) );
456                    add_sparse_chunk(&mut chunks, chunk.clone())?;
457                    file_len += curr_chunk_data_len;
458                    chunk_pos = chunk_pos + 1;
459                    output_offset += chunk.output_size();
460                }
461                None => {
462                    log::trace!("Finished iterating chunks");
463                    break;
464                }
465            }
466        }
467        let resparsed = SparseFileWriter::new(chunks);
468        log::trace!("resparse: Adding new SparseFile: {}", resparsed);
469        ret.push(resparsed);
470    }
471
472    Ok(ret)
473}
474
475/// Takes the given `file_to_upload` for the `named` partition and creates a
476/// set of temporary files in the given `dir` in Sparse Image Format. With the
477/// provided `max_download_size` constraining file size.
478///
479/// # Arguments
480///
481/// * `name` - Name of the partition the image. Used for logs only.
482/// * `file_to_upload` - Path to the file to translate to sparse image format.
483/// * `dir` - Path to write the Sparse file(s).
484/// * `max_download_size` - Maximum size that can be downloaded by the device.
485pub fn build_sparse_files(
486    name: &str,
487    file_to_upload: &str,
488    dir: &Path,
489    max_download_size: u64,
490) -> Result<Vec<TempPath>> {
491    if max_download_size as usize <= BLK_SIZE {
492        anyhow::bail!(
493            "Given maximum download size ({}) is less than the block size ({})",
494            max_download_size,
495            BLK_SIZE
496        );
497    }
498    log::debug!("Building sparse files for: {}. File: {}", name, file_to_upload);
499    let mut in_file = File::open(file_to_upload)?;
500
501    let mut total_read: usize = 0;
502    // Preallocate vector to avoid reallocations as it grows.
503    let mut chunks =
504        Vec::<Chunk>::with_capacity((in_file.metadata()?.len() as usize / BLK_SIZE) + 1);
505    let mut buf = [0u8; BLK_SIZE];
506    loop {
507        let read = in_file.read(&mut buf)?;
508        if read == 0 {
509            break;
510        }
511
512        let is_fill = buf.chunks(4).collect::<Vec<&[u8]>>().windows(2).all(|w| w[0] == w[1]);
513        if is_fill {
514            // The Android Sparse Image Format specifies that a fill block
515            // is a four-byte u32 repeated to fill BLK_SIZE. Here we use
516            // bincode::deserialize to get the repeated four byte pattern from
517            // the buffer so that it can be serialized later when we write
518            // the sparse file with bincode::serialize.
519            let value: u32 = bincode::deserialize(&buf[0..4])?;
520            // Add a fill chunk
521            let fill = Chunk::Fill { start: total_read as u64, size: buf.len(), value };
522            log::trace!("Sparsing file: {}. Created: {}", file_to_upload, fill);
523            chunks.push(fill);
524        } else {
525            // Add a raw chunk
526            let raw = Chunk::Raw { start: total_read as u64, size: buf.len() };
527            log::trace!("Sparsing file: {}. Created: {}", file_to_upload, raw);
528            chunks.push(raw);
529        }
530        total_read += read;
531    }
532
533    log::trace!("Creating sparse file from: {} chunks", chunks.len());
534
535    // At this point we are making a new sparse file fom an unoptomied set of
536    // Chunks. This primarliy means that adjacent Fill chunks of same value are
537    // not collapsed into a single Fill chunk (with a larger size). The advantage
538    // to this two pass approach is that (with some future work), we can create
539    // the "unoptomized" sparse file from a given image, and then "resparse" it
540    // as many times as desired with different `max_download_size` parameters.
541    // This would simplify the scenario where we want to flash the same image
542    // to multiple physical devices which may have slight differences in their
543    // hardware (and therefore different `max_download_size`es)
544    let sparse_file = SparseFileWriter::new(chunks);
545    log::trace!("Created sparse file: {}", sparse_file);
546
547    let mut ret = Vec::<TempPath>::new();
548    log::trace!("Resparsing sparse file");
549    for re_sparsed_file in resparse(sparse_file, max_download_size)? {
550        let (file, temp_path) = NamedTempFile::new_in(dir)?.into_parts();
551        let mut file_create = File::from(file);
552
553        log::trace!("Writing resparsed {} to disk", re_sparsed_file);
554        re_sparsed_file.write(&mut in_file, &mut file_create)?;
555
556        ret.push(temp_path);
557    }
558
559    log::debug!("Finished building sparse files");
560
561    Ok(ret)
562}
563
564////////////////////////////////////////////////////////////////////////////////
565// tests
566
567#[cfg(test)]
568mod test {
569    #[cfg(target_os = "linux")]
570    use crate::build_sparse_files;
571
572    use super::builder::{DataSource, SparseImageBuilder};
573    use super::{add_sparse_chunk, resparse, unsparse, Chunk, SparseFileWriter, BLK_SIZE};
574    use rand::rngs::SmallRng;
575    use rand::{RngCore, SeedableRng};
576    use std::io::{Cursor, Read as _, Seek as _, SeekFrom, Write as _};
577    #[cfg(target_os = "linux")]
578    use std::path::Path;
579    #[cfg(target_os = "linux")]
580    use std::process::{Command, Stdio};
581    use tempfile::{NamedTempFile, TempDir};
582
583    #[test]
584    fn test_fill_into_bytes() {
585        let mut dest = Cursor::new(Vec::<u8>::new());
586
587        let fill_chunk = Chunk::Fill { start: 0, size: 5 * BLK_SIZE, value: 365 };
588        // We have to convince the compiler that there's a specific type here.
589        fill_chunk.write(None::<&mut Cursor<Vec<u8>>>, &mut dest).unwrap();
590        assert_eq!(dest.into_inner(), [194, 202, 0, 0, 5, 0, 0, 0, 16, 0, 0, 0, 109, 1, 0, 0]);
591    }
592
593    #[test]
594    fn test_raw_into_bytes() {
595        const EXPECTED_RAW_BYTES: [u8; 22] =
596            [193, 202, 0, 0, 1, 0, 0, 0, 12, 16, 0, 0, 49, 50, 51, 52, 53, 0, 0, 0, 0, 0];
597
598        let mut source = Cursor::new(Vec::<u8>::from(&b"12345"[..]));
599        let mut sparse = Cursor::new(Vec::<u8>::new());
600        let chunk = Chunk::Raw { start: 0, size: BLK_SIZE };
601
602        chunk.write(Some(&mut source), &mut sparse).unwrap();
603        let buf = sparse.into_inner();
604        assert_eq!(buf.len(), 4108);
605        assert_eq!(&buf[..EXPECTED_RAW_BYTES.len()], EXPECTED_RAW_BYTES);
606        assert_eq!(&buf[EXPECTED_RAW_BYTES.len()..], &[0u8; 4108 - EXPECTED_RAW_BYTES.len()]);
607    }
608
609    #[test]
610    fn test_dont_care_into_bytes() {
611        let mut dest = Cursor::new(Vec::<u8>::new());
612        let chunk = Chunk::DontCare { start: 0, size: 5 * BLK_SIZE };
613
614        // We have to convince the compiler that there's a specific type here.
615        chunk.write(None::<&mut Cursor<Vec<u8>>>, &mut dest).unwrap();
616        assert_eq!(dest.into_inner(), [195, 202, 0, 0, 5, 0, 0, 0, 12, 0, 0, 0]);
617    }
618
619    #[test]
620    fn test_sparse_file_into_bytes() {
621        let mut source = Cursor::new(Vec::<u8>::from(&b"123"[..]));
622        let mut sparse = Cursor::new(Vec::<u8>::new());
623        let mut chunks = Vec::<Chunk>::new();
624        // Add a fill chunk
625        let fill = Chunk::Fill { start: 0, size: 4096, value: 5 };
626        chunks.push(fill);
627        // Add a raw chunk
628        let raw = Chunk::Raw { start: 0, size: 12288 };
629        chunks.push(raw);
630        // Add a dontcare chunk
631        let dontcare = Chunk::DontCare { start: 0, size: 4096 };
632        chunks.push(dontcare);
633
634        let sparsefile = SparseFileWriter::new(chunks);
635        sparsefile.write(&mut source, &mut sparse).unwrap();
636
637        sparse.seek(SeekFrom::Start(0)).unwrap();
638        let mut unsparsed = Cursor::new(Vec::<u8>::new());
639        unsparse(&mut sparse, &mut unsparsed).unwrap();
640        let buf = unsparsed.into_inner();
641        assert_eq!(buf.len(), 4096 + 12288 + 4096);
642        {
643            let chunks = buf[..4096].chunks(4);
644            for chunk in chunks {
645                assert_eq!(chunk, &[5u8, 0, 0, 0]);
646            }
647        }
648        assert_eq!(&buf[4096..4099], b"123");
649        assert_eq!(&buf[4099..16384], &[0u8; 12285]);
650        assert_eq!(&buf[16384..], &[0u8; 4096]);
651    }
652
653    ////////////////////////////////////////////////////////////////////////////
654    // Tests for resparse
655
656    #[test]
657    fn test_resparse_bails_on_too_small_size() {
658        let sparse = SparseFileWriter::new(Vec::<Chunk>::new());
659        assert!(resparse(sparse, 4095).is_err());
660    }
661
662    #[test]
663    fn test_resparse_splits() {
664        let max_download_size = 4096 * 2;
665
666        let mut chunks = Vec::<Chunk>::new();
667        chunks.push(Chunk::Raw { start: 0, size: 4096 });
668        chunks.push(Chunk::Fill { start: 4096, size: 4096, value: 2 });
669        // We want 2 sparse files with the second sparse file having a
670        // DontCare chunk and then this chunk
671        chunks.push(Chunk::Raw { start: 8192, size: 4096 });
672
673        let input_sparse_file = SparseFileWriter::new(chunks);
674        let resparsed_files = resparse(input_sparse_file, max_download_size).unwrap();
675        assert_eq!(2, resparsed_files.len());
676
677        assert_eq!(3, resparsed_files[0].chunks.len());
678        assert_eq!(Chunk::Raw { start: 0, size: 4096 }, resparsed_files[0].chunks[0]);
679        assert_eq!(Chunk::Fill { start: 4096, size: 4096, value: 2 }, resparsed_files[0].chunks[1]);
680        assert_eq!(Chunk::DontCare { start: 8192, size: 4096 }, resparsed_files[0].chunks[2]);
681
682        assert_eq!(2, resparsed_files[1].chunks.len());
683        assert_eq!(Chunk::DontCare { start: 0, size: 8192 }, resparsed_files[1].chunks[0]);
684        assert_eq!(Chunk::Raw { start: 8192, size: 4096 }, resparsed_files[1].chunks[1]);
685    }
686
687    ////////////////////////////////////////////////////////////////////////////
688    // Tests for add_sparse_chunk
689
690    #[test]
691    fn test_add_sparse_chunk_adds_empty() {
692        let init_vec = Vec::<Chunk>::new();
693        let mut res = init_vec.clone();
694        add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
695        assert_eq!(0, init_vec.len());
696        assert_ne!(init_vec, res);
697        assert_eq!(Chunk::Fill { start: 0, size: 4096, value: 1 }, res[0]);
698    }
699
700    #[test]
701    fn test_add_sparse_chunk_fill() {
702        // Test merge
703        {
704            let mut init_vec = Vec::<Chunk>::new();
705            init_vec.push(Chunk::Fill { start: 0, size: 8192, value: 1 });
706            let mut res = init_vec.clone();
707            add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 8192, value: 1 }).unwrap();
708            assert_eq!(1, res.len());
709            assert_eq!(Chunk::Fill { start: 0, size: 16384, value: 1 }, res[0]);
710        }
711
712        // Test dont merge on different value
713        {
714            let mut init_vec = Vec::<Chunk>::new();
715            init_vec.push(Chunk::Fill { start: 0, size: 4096, value: 1 });
716            let mut res = init_vec.clone();
717            add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 2 }).unwrap();
718            assert_ne!(res, init_vec);
719            assert_eq!(2, res.len());
720            assert_eq!(
721                res,
722                [
723                    Chunk::Fill { start: 0, size: 4096, value: 1 },
724                    Chunk::Fill { start: 0, size: 4096, value: 2 }
725                ]
726            );
727        }
728
729        // Test dont merge on different type
730        {
731            let mut init_vec = Vec::<Chunk>::new();
732            init_vec.push(Chunk::Fill { start: 0, size: 4096, value: 2 });
733            let mut res = init_vec.clone();
734            add_sparse_chunk(&mut res, Chunk::DontCare { start: 0, size: 4096 }).unwrap();
735            assert_ne!(res, init_vec);
736            assert_eq!(2, res.len());
737            assert_eq!(
738                res,
739                [
740                    Chunk::Fill { start: 0, size: 4096, value: 2 },
741                    Chunk::DontCare { start: 0, size: 4096 }
742                ]
743            );
744        }
745    }
746
747    #[test]
748    fn test_add_sparse_chunk_dont_care() {
749        // Test they merge
750        {
751            let mut init_vec = Vec::<Chunk>::new();
752            init_vec.push(Chunk::DontCare { start: 0, size: 4096 });
753            let mut res = init_vec.clone();
754            add_sparse_chunk(&mut res, Chunk::DontCare { start: 0, size: 4096 }).unwrap();
755            assert_eq!(1, res.len());
756            assert_eq!(Chunk::DontCare { start: 0, size: 8192 }, res[0]);
757        }
758
759        // Test they dont merge on different type
760        {
761            let mut init_vec = Vec::<Chunk>::new();
762            init_vec.push(Chunk::DontCare { start: 0, size: 4096 });
763            let mut res = init_vec.clone();
764            add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
765            assert_eq!(2, res.len());
766            assert_eq!(
767                res,
768                [
769                    Chunk::DontCare { start: 0, size: 4096 },
770                    Chunk::Fill { start: 0, size: 4096, value: 1 }
771                ]
772            );
773        }
774    }
775
776    #[test]
777    fn test_add_sparse_chunk_raw() {
778        // Test they merge
779        {
780            let mut init_vec = Vec::<Chunk>::new();
781            init_vec.push(Chunk::Raw { start: 0, size: 12288 });
782            let mut res = init_vec.clone();
783            add_sparse_chunk(&mut res, Chunk::Raw { start: 0, size: 16384 }).unwrap();
784            assert_eq!(1, res.len());
785            assert_eq!(Chunk::Raw { start: 0, size: 28672 }, res[0]);
786        }
787
788        // Test they dont merge on different type
789        {
790            let mut init_vec = Vec::<Chunk>::new();
791            init_vec.push(Chunk::Raw { start: 0, size: 12288 });
792            let mut res = init_vec.clone();
793            add_sparse_chunk(&mut res, Chunk::Fill { start: 3, size: 8192, value: 1 }).unwrap();
794            assert_eq!(2, res.len());
795            assert_eq!(
796                res,
797                [
798                    Chunk::Raw { start: 0, size: 12288 },
799                    Chunk::Fill { start: 3, size: 8192, value: 1 }
800                ]
801            );
802        }
803    }
804
805    #[test]
806    fn test_add_sparse_chunk_crc32() {
807        // Test they dont merge on same type (Crc32 is special)
808        {
809            let mut init_vec = Vec::<Chunk>::new();
810            init_vec.push(Chunk::Crc32 { checksum: 1234 });
811            let mut res = init_vec.clone();
812            add_sparse_chunk(&mut res, Chunk::Crc32 { checksum: 2345 }).unwrap();
813            assert_eq!(2, res.len());
814            assert_eq!(res, [Chunk::Crc32 { checksum: 1234 }, Chunk::Crc32 { checksum: 2345 }]);
815        }
816
817        // Test they dont merge on different type
818        {
819            let mut init_vec = Vec::<Chunk>::new();
820            init_vec.push(Chunk::Crc32 { checksum: 1234 });
821            let mut res = init_vec.clone();
822            add_sparse_chunk(&mut res, Chunk::Fill { start: 0, size: 4096, value: 1 }).unwrap();
823            assert_eq!(2, res.len());
824            assert_eq!(
825                res,
826                [Chunk::Crc32 { checksum: 1234 }, Chunk::Fill { start: 0, size: 4096, value: 1 }]
827            );
828        }
829    }
830
831    ////////////////////////////////////////////////////////////////////////////
832    // Integration
833    //
834
835    #[test]
836    fn test_roundtrip() {
837        let tmpdir = TempDir::new().unwrap();
838
839        // Generate a large temporary file
840        let (mut file, _temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
841        let mut rng = SmallRng::from_entropy();
842        let mut buf = Vec::<u8>::new();
843        buf.resize(1 * 4096, 0);
844        rng.fill_bytes(&mut buf);
845        file.write_all(&buf).unwrap();
846        file.flush().unwrap();
847        file.seek(SeekFrom::Start(0)).unwrap();
848        let content_size = buf.len();
849
850        // build a sparse file
851        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
852        SparseImageBuilder::new()
853            .add_chunk(DataSource::Buffer(Box::new([0xffu8; 8192])))
854            .add_chunk(DataSource::Reader(Box::new(file)))
855            .add_chunk(DataSource::Fill(0xaaaa_aaaau32, 1024))
856            .add_chunk(DataSource::Skip(16384))
857            .build(&mut sparse_file)
858            .expect("Build sparse image failed");
859        sparse_file.seek(SeekFrom::Start(0)).unwrap();
860
861        let mut orig_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
862        unsparse(&mut sparse_file, &mut orig_file).expect("unsparse failed");
863        orig_file.seek(SeekFrom::Start(0)).unwrap();
864
865        let mut unsparsed_bytes = vec![];
866        orig_file.read_to_end(&mut unsparsed_bytes).expect("Failed to read unsparsed image");
867        assert_eq!(unsparsed_bytes.len(), 8192 + 20480 + content_size);
868        assert_eq!(&unsparsed_bytes[..8192], &[0xffu8; 8192]);
869        assert_eq!(&unsparsed_bytes[8192..8192 + content_size], &buf[..]);
870        assert_eq!(&unsparsed_bytes[8192 + content_size..12288 + content_size], &[0xaau8; 4096]);
871        assert_eq!(&unsparsed_bytes[12288 + content_size..], &[0u8; 16384]);
872    }
873
874    #[test]
875    /// test_with_simg2img is a "round trip" test that does the following
876    ///
877    /// 1. Generates a pseudorandom temporary file
878    /// 2. Builds sparse files out of it
879    /// 3. Uses the android tool simg2img to take the sparse files and generate
880    ///    the "original" image file out of them.
881    /// 4. Asserts the originally created file and the one created by simg2img
882    ///    have binary equivalent contents.
883    ///
884    /// This gives us a reasonable expectation of correctness given that the
885    /// Android-provided sparse tools are able to interpret our sparse images.
886    #[cfg(target_os = "linux")]
887    fn test_with_simg2img() {
888        let simg2img_path = Path::new("./host_x64/test_data/storage/sparse/simg2img");
889        assert!(
890            Path::exists(simg2img_path),
891            "simg2img binary must exist at {}",
892            simg2img_path.display()
893        );
894
895        let tmpdir = TempDir::new().unwrap();
896
897        // Generate a large temporary file
898        let (mut file, temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
899        let mut rng = SmallRng::from_entropy();
900        let mut buf = Vec::<u8>::new();
901        buf.resize(50 * 4096, 0);
902        rng.fill_bytes(&mut buf);
903        file.write_all(&buf).unwrap();
904        file.flush().unwrap();
905        file.seek(SeekFrom::Start(0)).unwrap();
906
907        // build a sparse file
908        let files = build_sparse_files(
909            "test",
910            temp_path.to_path_buf().to_str().expect("Should succeed"),
911            tmpdir.path(),
912            4096 * 2,
913        )
914        .unwrap();
915
916        let mut simg2img_output = tmpdir.path().to_path_buf();
917        simg2img_output.push("output");
918
919        let mut simg2img = Command::new(simg2img_path)
920            .args(&files[..])
921            .arg(&simg2img_output)
922            .stdout(Stdio::piped())
923            .stderr(Stdio::piped())
924            .spawn()
925            .expect("Failed to spawn simg2img");
926        let res = simg2img.wait().expect("simg2img did was not running");
927        assert!(res.success(), "simg2img did not succeed");
928        let mut simg2img_stdout = simg2img.stdout.take().expect("Get stdout from simg2img");
929        let mut simg2img_stderr = simg2img.stderr.take().expect("Get stderr from simg2img");
930
931        let mut stdout = String::new();
932        simg2img_stdout.read_to_string(&mut stdout).expect("Reading simg2img stdout");
933        assert_eq!(stdout, "");
934
935        let mut stderr = String::new();
936        simg2img_stderr.read_to_string(&mut stderr).expect("Reading simg2img stderr");
937        assert_eq!(stderr, "");
938
939        let simg2img_output_bytes =
940            std::fs::read(simg2img_output).expect("Failed to read simg2img output");
941
942        assert_eq!(
943            buf, simg2img_output_bytes,
944            "Output from simg2img should match our generated file"
945        );
946    }
947}
sparse/lib.rs

sparse/
lib.rs
