sparse/
builder.rs

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::{CHUNK_HEADER_SIZE, SPARSE_HEADER_SIZE};
6use crate::{Chunk, SparseHeader, BLK_SIZE, NO_SOURCE};
7use anyhow::{ensure, Context, Result};
8use std::io::{Cursor, Read, Seek, SeekFrom, Write};
9use std::ops::Range;
10
11/// Input data for a SparseImageBuilder.
12pub enum DataSource {
13    Buffer(Box<[u8]>),
14    /// Read `size` bytes from `reader`.
15    Reader {
16        reader: Box<dyn Read>,
17        size: u64,
18    },
19    /// Skips this many bytes.
20    Skip(u64),
21    /// Repeats the given u32, this many times.
22    Fill(u32, u64),
23}
24
25/// Builds sparse image files from a set of input DataSources.
26pub struct SparseImageBuilder {
27    block_size: u32,
28    chunks: Vec<DataSource>,
29
30    // The `total_sz` field in the chunk's header is 32 bits and includes the 12 bytes for the
31    // header itself. The chunk size must be a multiple of the block size so the maximum chunk size
32    // is 4GiB minus the block size provided the block size is greater than or equal to 12. This
33    // field could be derived from `block_size` but is stored separately so it can be set in tests
34    // to avoid creating 4GiB data sources.
35    max_chunk_size: u32,
36}
37
38impl SparseImageBuilder {
39    pub fn new() -> Self {
40        Self { block_size: BLK_SIZE, chunks: vec![], max_chunk_size: u32::MAX - BLK_SIZE + 1 }
41    }
42
43    pub fn set_block_size(mut self, block_size: u32) -> Self {
44        assert!(
45            block_size >= CHUNK_HEADER_SIZE,
46            "The block size must be greater than {}",
47            CHUNK_HEADER_SIZE
48        );
49        self.max_chunk_size = u32::MAX - block_size + 1;
50        self.block_size = block_size;
51        self
52    }
53
54    pub fn add_chunk(mut self, source: DataSource) -> Self {
55        self.chunks.push(source);
56        self
57    }
58
59    pub fn build<W: Write + Seek>(self, output: &mut W) -> Result<()> {
60        // We'll fill the header in later.
61        output.seek(SeekFrom::Start(SPARSE_HEADER_SIZE as u64))?;
62        let mut chunk_writer = ChunkWriter::new(self.block_size, output);
63        for input_chunk in self.chunks {
64            match input_chunk {
65                DataSource::Buffer(buf) => {
66                    ensure!(
67                        buf.len() % self.block_size as usize == 0,
68                        "Invalid buffer length {}",
69                        buf.len()
70                    );
71                    for slice in buf.chunks(self.max_chunk_size as usize) {
72                        chunk_writer
73                            .write_raw_chunk(slice.len().try_into().unwrap(), Cursor::new(slice))?;
74                    }
75                }
76                DataSource::Reader { mut reader, size } => {
77                    ensure!(size % self.block_size as u64 == 0, "Invalid Reader length {}", size);
78                    for size in ChunkedRange::new(0..size, self.max_chunk_size) {
79                        chunk_writer.write_raw_chunk(size, (&mut reader).take(size as u64))?;
80                    }
81                }
82                DataSource::Skip(size) => {
83                    ensure!(size % self.block_size as u64 == 0, "Invalid Skip length {}", size);
84                    for size in ChunkedRange::new(0..size, self.max_chunk_size) {
85                        chunk_writer.write_dont_care_chunk(size)?;
86                    }
87                }
88                DataSource::Fill(value, count) => {
89                    let size = count * std::mem::size_of::<u32>() as u64;
90                    ensure!(size % self.block_size as u64 == 0, "Invalid Fill length {}", size);
91                    for size in ChunkedRange::new(0..size, self.max_chunk_size) {
92                        chunk_writer.write_fill_chunk(size, value)?;
93                    }
94                }
95            };
96        }
97
98        let ChunkWriter { num_blocks, num_chunks, .. } = chunk_writer;
99        output.seek(SeekFrom::Start(0))?;
100        let header = SparseHeader::new(self.block_size, num_blocks, num_chunks);
101        bincode::serialize_into(&mut *output, &header)?;
102
103        output.flush()?;
104        Ok(())
105    }
106}
107
108struct ChunkWriter<'a, W> {
109    block_size: u32,
110    current_offset: u64,
111    num_chunks: u32,
112    num_blocks: u32,
113    writer: &'a mut W,
114}
115
116impl<'a, W: Write> ChunkWriter<'a, W> {
117    fn new(block_size: u32, writer: &'a mut W) -> Self {
118        Self { block_size, current_offset: 0, num_chunks: 0, num_blocks: 0, writer }
119    }
120
121    fn write_chunk_impl<R: Read>(&mut self, chunk: Chunk, source: Option<&mut R>) -> Result<()> {
122        chunk.write(source, &mut self.writer, self.block_size)?;
123        self.num_blocks = self
124            .num_blocks
125            .checked_add(chunk.output_blocks(self.block_size))
126            .context("Sparse image would contain too many blocks")?;
127        // The number of blocks and chunks are both a u32. Each chunk contains at least 1 block so
128        // the number of blocks will overflow above before the number of chunks.
129        self.num_chunks += 1;
130        self.current_offset += chunk.output_size() as u64;
131        Ok(())
132    }
133
134    fn write_raw_chunk<R: Read>(&mut self, size: u32, mut source: R) -> Result<()> {
135        self.write_chunk_impl(Chunk::Raw { start: self.current_offset, size }, Some(&mut source))
136    }
137
138    fn write_dont_care_chunk(&mut self, size: u32) -> Result<()> {
139        self.write_chunk_impl(Chunk::DontCare { start: self.current_offset, size }, NO_SOURCE)
140    }
141
142    fn write_fill_chunk(&mut self, size: u32, value: u32) -> Result<()> {
143        self.write_chunk_impl(Chunk::Fill { start: self.current_offset, size, value }, NO_SOURCE)
144    }
145}
146
147/// An iterator that yields `max_chunk_size` `(range.end - range.start) / max_chunk_size` times
148/// followed by `(range.end - range.start) % max_chunk_size` if it's none zero.
149///
150/// # Examples
151/// ```
152/// assert_eq!(ChunkedRange::new(0..10, 5).collect::<Vec<_>>(), vec![5, 5]);
153/// assert_eq!(ChunkedRange::new(0..13, 5).collect::<Vec<_>>(), vec![5, 5, 3]);
154/// ```
155struct ChunkedRange {
156    range: Range<u64>,
157    max_chunk_size: u32,
158}
159
160impl ChunkedRange {
161    fn new(range: Range<u64>, max_chunk_size: u32) -> Self {
162        Self { range, max_chunk_size }
163    }
164}
165
166impl Iterator for ChunkedRange {
167    type Item = u32;
168
169    fn next(&mut self) -> Option<Self::Item> {
170        let size = self.range.end - self.range.start;
171        if size == 0 {
172            None
173        } else if size >= self.max_chunk_size as u64 {
174            self.range.start += self.max_chunk_size as u64;
175            Some(self.max_chunk_size)
176        } else {
177            self.range.start = self.range.end;
178            Some(size as u32)
179        }
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use crate::format::CHUNK_HEADER_SIZE;
187    use crate::reader::SparseReader;
188
189    #[test]
190    fn test_chunked_range() {
191        assert_eq!(&ChunkedRange::new(0..0, 32).collect::<Vec<_>>(), &[]);
192        assert_eq!(&ChunkedRange::new(0..10, 32).collect::<Vec<_>>(), &[10]);
193        assert_eq!(&ChunkedRange::new(100..101, 32).collect::<Vec<_>>(), &[1]);
194        assert_eq!(&ChunkedRange::new(0..100, 32).collect::<Vec<_>>(), &[32, 32, 32, 4]);
195        assert_eq!(&ChunkedRange::new(10..100, 32).collect::<Vec<_>>(), &[32, 32, 26]);
196        assert_eq!(
197            &ChunkedRange::new((u32::MAX as u64)..(u32::MAX as u64 + 80), 32).collect::<Vec<_>>(),
198            &[32, 32, 16]
199        );
200        assert_eq!(
201            &ChunkedRange::new((u64::MAX - 50)..u64::MAX, 32).collect::<Vec<_>>(),
202            &[32, 18]
203        );
204    }
205
206    #[test]
207    fn test_build_with_buffer() {
208        let mut builder = SparseImageBuilder::new();
209        builder.max_chunk_size = BLK_SIZE;
210        let mut buf = Vec::with_capacity((BLK_SIZE * 2) as usize);
211        let part1 = vec![0xABu8; BLK_SIZE as usize];
212        let part2 = vec![0xCDu8; BLK_SIZE as usize];
213        buf.extend_from_slice(&part1);
214        buf.extend_from_slice(&part2);
215        let mut output = vec![];
216        builder
217            .add_chunk(DataSource::Buffer(buf.into_boxed_slice()))
218            .build(&mut Cursor::new(&mut output))
219            .unwrap();
220
221        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
222        assert_eq!(
223            &reader.chunks(),
224            &[
225                (
226                    Chunk::Raw { start: 0, size: BLK_SIZE },
227                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as u64)
228                ),
229                (
230                    Chunk::Raw { start: BLK_SIZE as u64, size: BLK_SIZE },
231                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as u64)
232                )
233            ]
234        );
235        assert_eq!(
236            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as usize
237                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE + BLK_SIZE) as usize],
238            &part1
239        );
240        assert_eq!(
241            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as usize
242                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE * 2) as usize],
243            &part2
244        );
245    }
246
247    #[test]
248    fn test_build_with_reader() {
249        let part1 = vec![0xABu8; BLK_SIZE as usize];
250        let part2 = vec![0xCDu8; BLK_SIZE as usize];
251        let mut buf = Vec::with_capacity(BLK_SIZE as usize * 2);
252        buf.extend_from_slice(&part1);
253        buf.extend_from_slice(&part2);
254
255        let mut builder = SparseImageBuilder::new();
256        builder.max_chunk_size = BLK_SIZE;
257        let mut output = vec![];
258
259        let reader1 = Cursor::new(buf.clone());
260        let mut reader2 = Cursor::new(buf);
261        reader2.seek(SeekFrom::Start(BLK_SIZE as u64)).unwrap();
262
263        builder
264            .add_chunk(DataSource::Reader {
265                reader: Box::new(reader1),
266                size: (BLK_SIZE * 2) as u64,
267            })
268            .add_chunk(DataSource::Reader { reader: Box::new(reader2), size: BLK_SIZE as u64 })
269            .build(&mut Cursor::new(&mut output))
270            .unwrap();
271
272        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
273        assert_eq!(
274            &reader.chunks(),
275            &[
276                (
277                    Chunk::Raw { start: 0, size: BLK_SIZE },
278                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as u64)
279                ),
280                (
281                    Chunk::Raw { start: BLK_SIZE as u64, size: BLK_SIZE },
282                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as u64)
283                ),
284                (
285                    Chunk::Raw { start: (BLK_SIZE * 2) as u64, size: BLK_SIZE },
286                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 2) as u64)
287                ),
288            ]
289        );
290        assert_eq!(
291            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as usize
292                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE + BLK_SIZE) as usize],
293            &part1
294        );
295        assert_eq!(
296            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as usize
297                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE * 2) as usize],
298            &part2
299        );
300        assert_eq!(
301            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 2) as usize
302                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 3) as usize],
303            &part2
304        );
305    }
306
307    #[test]
308    fn test_build_with_skip() {
309        let mut builder = SparseImageBuilder::new();
310        builder.max_chunk_size = BLK_SIZE;
311        let mut output = vec![];
312        builder
313            .add_chunk(DataSource::Skip((BLK_SIZE * 2) as u64))
314            .build(&mut Cursor::new(&mut output))
315            .unwrap();
316
317        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
318        assert_eq!(
319            &reader.chunks(),
320            &[
321                (Chunk::DontCare { start: 0, size: BLK_SIZE }, None),
322                (Chunk::DontCare { start: BLK_SIZE as u64, size: BLK_SIZE }, None)
323            ]
324        );
325    }
326
327    #[test]
328    fn test_build_with_fill() {
329        let mut builder = SparseImageBuilder::new();
330        builder.max_chunk_size = BLK_SIZE;
331        let mut output = vec![];
332        builder
333            .add_chunk(DataSource::Fill(0xAB, (BLK_SIZE / 2) as u64))
334            .build(&mut Cursor::new(&mut output))
335            .unwrap();
336
337        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
338        assert_eq!(
339            &reader.chunks(),
340            &[
341                (Chunk::Fill { start: 0, size: BLK_SIZE, value: 0xAB }, None),
342                (Chunk::Fill { start: BLK_SIZE as u64, size: BLK_SIZE, value: 0xAB }, None)
343            ]
344        );
345    }
346
347    #[test]
348    fn test_overflow_block_count() {
349        struct Sink;
350
351        impl Write for Sink {
352            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
353                Ok(buf.len())
354            }
355
356            fn flush(&mut self) -> std::io::Result<()> {
357                Ok(())
358            }
359        }
360
361        impl Seek for Sink {
362            fn seek(&mut self, _pos: SeekFrom) -> std::io::Result<u64> {
363                Ok(0)
364            }
365        }
366
367        let result = SparseImageBuilder::new()
368            .set_block_size(16)
369            .add_chunk(DataSource::Skip(u64::MAX - 15))
370            .build(&mut Sink);
371        assert!(result.is_err());
372    }
373}
sparse/builder.rs

sparse/
builder.rs
