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(
136            Chunk::Raw { start: self.current_offset, size: size.into() },
137            Some(&mut source),
138        )
139    }
140
141    fn write_dont_care_chunk(&mut self, size: u32) -> Result<()> {
142        self.write_chunk_impl(
143            Chunk::DontCare { start: self.current_offset, size: size.into() },
144            NO_SOURCE,
145        )
146    }
147
148    fn write_fill_chunk(&mut self, size: u32, value: u32) -> Result<()> {
149        self.write_chunk_impl(
150            Chunk::Fill { start: self.current_offset, size: size.into(), value },
151            NO_SOURCE,
152        )
153    }
154}
155
156/// An iterator that yields `max_chunk_size` `(range.end - range.start) / max_chunk_size` times
157/// followed by `(range.end - range.start) % max_chunk_size` if it's none zero.
158///
159/// # Examples
160/// ```
161/// assert_eq!(ChunkedRange::new(0..10, 5).collect::<Vec<_>>(), vec![5, 5]);
162/// assert_eq!(ChunkedRange::new(0..13, 5).collect::<Vec<_>>(), vec![5, 5, 3]);
163/// ```
164struct ChunkedRange {
165    range: Range<u64>,
166    max_chunk_size: u32,
167}
168
169impl ChunkedRange {
170    fn new(range: Range<u64>, max_chunk_size: u32) -> Self {
171        Self { range, max_chunk_size }
172    }
173}
174
175impl Iterator for ChunkedRange {
176    type Item = u32;
177
178    fn next(&mut self) -> Option<Self::Item> {
179        let size = self.range.end - self.range.start;
180        if size == 0 {
181            None
182        } else if size >= self.max_chunk_size as u64 {
183            self.range.start += self.max_chunk_size as u64;
184            Some(self.max_chunk_size)
185        } else {
186            self.range.start = self.range.end;
187            Some(size as u32)
188        }
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::format::CHUNK_HEADER_SIZE;
196    use crate::reader::SparseReader;
197
198    #[test]
199    fn test_chunked_range() {
200        assert_eq!(&ChunkedRange::new(0..0, 32).collect::<Vec<_>>(), &[]);
201        assert_eq!(&ChunkedRange::new(0..10, 32).collect::<Vec<_>>(), &[10]);
202        assert_eq!(&ChunkedRange::new(100..101, 32).collect::<Vec<_>>(), &[1]);
203        assert_eq!(&ChunkedRange::new(0..100, 32).collect::<Vec<_>>(), &[32, 32, 32, 4]);
204        assert_eq!(&ChunkedRange::new(10..100, 32).collect::<Vec<_>>(), &[32, 32, 26]);
205        assert_eq!(
206            &ChunkedRange::new((u32::MAX as u64)..(u32::MAX as u64 + 80), 32).collect::<Vec<_>>(),
207            &[32, 32, 16]
208        );
209        assert_eq!(
210            &ChunkedRange::new((u64::MAX - 50)..u64::MAX, 32).collect::<Vec<_>>(),
211            &[32, 18]
212        );
213    }
214
215    #[test]
216    fn test_build_with_buffer() {
217        let mut builder = SparseImageBuilder::new();
218        builder.max_chunk_size = BLK_SIZE;
219        let mut buf = Vec::with_capacity((BLK_SIZE * 2) as usize);
220        let part1 = vec![0xABu8; BLK_SIZE as usize];
221        let part2 = vec![0xCDu8; BLK_SIZE as usize];
222        buf.extend_from_slice(&part1);
223        buf.extend_from_slice(&part2);
224        let mut output = vec![];
225        builder
226            .add_chunk(DataSource::Buffer(buf.into_boxed_slice()))
227            .build(&mut Cursor::new(&mut output))
228            .unwrap();
229
230        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
231        assert_eq!(
232            reader.chunks(),
233            &[
234                (
235                    Chunk::Raw { start: 0, size: BLK_SIZE.into() },
236                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as u64)
237                ),
238                (
239                    Chunk::Raw { start: BLK_SIZE as u64, size: BLK_SIZE.into() },
240                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as u64)
241                )
242            ]
243        );
244        assert_eq!(
245            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as usize
246                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE + BLK_SIZE) as usize],
247            &part1
248        );
249        assert_eq!(
250            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as usize
251                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE * 2) as usize],
252            &part2
253        );
254    }
255
256    #[test]
257    fn test_build_with_reader() {
258        let part1 = vec![0xABu8; BLK_SIZE as usize];
259        let part2 = vec![0xCDu8; BLK_SIZE as usize];
260        let mut buf = Vec::with_capacity(BLK_SIZE as usize * 2);
261        buf.extend_from_slice(&part1);
262        buf.extend_from_slice(&part2);
263
264        let mut builder = SparseImageBuilder::new();
265        builder.max_chunk_size = BLK_SIZE;
266        let mut output = vec![];
267
268        let reader1 = Cursor::new(buf.clone());
269        let mut reader2 = Cursor::new(buf);
270        reader2.seek(SeekFrom::Start(BLK_SIZE as u64)).unwrap();
271
272        builder
273            .add_chunk(DataSource::Reader {
274                reader: Box::new(reader1),
275                size: (BLK_SIZE * 2) as u64,
276            })
277            .add_chunk(DataSource::Reader { reader: Box::new(reader2), size: BLK_SIZE as u64 })
278            .build(&mut Cursor::new(&mut output))
279            .unwrap();
280
281        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
282        assert_eq!(
283            reader.chunks(),
284            &[
285                (
286                    Chunk::Raw { start: 0, size: BLK_SIZE.into() },
287                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as u64)
288                ),
289                (
290                    Chunk::Raw { start: BLK_SIZE as u64, size: BLK_SIZE.into() },
291                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as u64)
292                ),
293                (
294                    Chunk::Raw { start: (BLK_SIZE * 2) as u64, size: BLK_SIZE.into() },
295                    Some((SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 2) as u64)
296                ),
297            ]
298        );
299        assert_eq!(
300            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE) as usize
301                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE + BLK_SIZE) as usize],
302            &part1
303        );
304        assert_eq!(
305            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE) as usize
306                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 2 + BLK_SIZE * 2) as usize],
307            &part2
308        );
309        assert_eq!(
310            &output[(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 2) as usize
311                ..(SPARSE_HEADER_SIZE + CHUNK_HEADER_SIZE * 3 + BLK_SIZE * 3) as usize],
312            &part2
313        );
314    }
315
316    #[test]
317    fn test_build_with_skip() {
318        let mut builder = SparseImageBuilder::new();
319        builder.max_chunk_size = BLK_SIZE;
320        let mut output = vec![];
321        builder
322            .add_chunk(DataSource::Skip((BLK_SIZE * 2) as u64))
323            .build(&mut Cursor::new(&mut output))
324            .unwrap();
325
326        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
327        assert_eq!(
328            reader.chunks(),
329            &[
330                (Chunk::DontCare { start: 0, size: BLK_SIZE.into() }, None),
331                (Chunk::DontCare { start: BLK_SIZE as u64, size: BLK_SIZE.into() }, None)
332            ]
333        );
334    }
335
336    #[test]
337    fn test_build_with_fill() {
338        let mut builder = SparseImageBuilder::new();
339        builder.max_chunk_size = BLK_SIZE;
340        let mut output = vec![];
341        builder
342            .add_chunk(DataSource::Fill(0xAB, (BLK_SIZE / 2) as u64))
343            .build(&mut Cursor::new(&mut output))
344            .unwrap();
345
346        let reader = SparseReader::new(Cursor::new(&output)).unwrap();
347        assert_eq!(
348            reader.chunks(),
349            &[
350                (Chunk::Fill { start: 0, size: BLK_SIZE.into(), value: 0xAB }, None),
351                (Chunk::Fill { start: BLK_SIZE as u64, size: BLK_SIZE.into(), value: 0xAB }, None)
352            ]
353        );
354    }
355
356    #[test]
357    fn test_overflow_block_count() {
358        struct Sink;
359
360        impl Write for Sink {
361            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
362                Ok(buf.len())
363            }
364
365            fn flush(&mut self) -> std::io::Result<()> {
366                Ok(())
367            }
368        }
369
370        impl Seek for Sink {
371            fn seek(&mut self, _pos: SeekFrom) -> std::io::Result<u64> {
372                Ok(0)
373            }
374        }
375
376        let result = SparseImageBuilder::new()
377            .set_block_size(16)
378            .add_chunk(DataSource::Skip(u64::MAX - 15))
379            .build(&mut Sink);
380        assert!(result.is_err());
381    }
382}
sparse/builder.rs

sparse/
builder.rs
