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

sparse/
builder.rs
