sparse/
builder.rs

Help
1// Copyright 2023 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
5use crate::format::SPARSE_HEADER_SIZE;
6use crate::{Chunk, Reader, SparseHeader, Writer, BLK_SIZE};
7use anyhow::{ensure, Result};
8use std::io::{Cursor, SeekFrom};
9
10/// Input data for a SparseImageBuilder.
11pub enum DataSource {
12    Buffer(Box<[u8]>),
13    /// Read everything from the reader.
14    Reader(Box<dyn Reader>),
15    /// Skips this many bytes.
16    Skip(u64),
17    /// Repeats the given u32, this many times.
18    Fill(u32, u64),
19}
20
21/// Builds sparse image files from a set of input DataSources.
22pub struct SparseImageBuilder {
23    block_size: u32,
24    chunks: Vec<DataSource>,
25}
26
27impl SparseImageBuilder {
28    pub fn new() -> Self {
29        Self { block_size: BLK_SIZE as u32, chunks: vec![] }
30    }
31
32    pub fn set_block_size(mut self, block_size: u32) -> Self {
33        self.block_size = block_size;
34        self
35    }
36
37    pub fn add_chunk(mut self, source: DataSource) -> Self {
38        self.chunks.push(source);
39        self
40    }
41
42    pub fn build<W: Writer>(self, output: &mut W) -> Result<()> {
43        // We'll fill the header in later.
44        output.seek(SeekFrom::Start(SPARSE_HEADER_SIZE as u64))?;
45
46        let mut num_blocks = 0;
47        let num_chunks = self.chunks.len() as u32;
48        let mut current_offset = 0;
49        for input_chunk in self.chunks {
50            let (chunk, mut source) = match input_chunk {
51                DataSource::Buffer(buf) => {
52                    ensure!(
53                        buf.len() % self.block_size as usize == 0,
54                        "Invalid buffer length {}",
55                        buf.len()
56                    );
57                    (
58                        Chunk::Raw { start: current_offset, size: buf.len() },
59                        Some(Box::new(Cursor::new(buf)) as Box<dyn Reader>),
60                    )
61                }
62                DataSource::Reader(mut reader) => {
63                    let size = reader.seek(SeekFrom::End(0))?;
64                    reader.seek(SeekFrom::Start(0))?;
65                    ensure!(size % self.block_size as u64 == 0, "Invalid Reader length {}", size);
66                    (Chunk::Raw { start: current_offset, size: size as usize }, Some(reader))
67                }
68                DataSource::Skip(size) => {
69                    ensure!(size % self.block_size as u64 == 0, "Invalid Skip length {}", size);
70                    (Chunk::DontCare { start: current_offset, size: size as usize }, None)
71                }
72                DataSource::Fill(value, count) => {
73                    let size = count as usize * std::mem::size_of::<u32>();
74                    ensure!(size % self.block_size as usize == 0, "Invalid Fill length {}", size);
75                    (Chunk::Fill { start: current_offset, size, value }, None)
76                }
77            };
78            chunk.write(source.as_mut(), output)?;
79            num_blocks += chunk.output_blocks(self.block_size as usize) as u32;
80            current_offset += chunk.output_size() as u64;
81        }
82
83        output.seek(SeekFrom::Start(0))?;
84        let header = SparseHeader::new(self.block_size, num_blocks, num_chunks);
85        let header_bytes: Vec<u8> = bincode::serialize(&header)?;
86        std::io::copy(&mut Cursor::new(header_bytes), output)?;
87
88        output.flush()?;
89        Ok(())
90    }
91}
sparse/builder.rs

sparse/
builder.rs
