sparse/
reader.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::{Chunk, SparseDataType, SparseError, SparseHeader, deserialize_from};
6
7use byteorder::{ByteOrder as _, LE};
8use std::io::{Read, Seek, SeekFrom};
9
10/// SparseReader is an implementation of std::io::Read which transparently unpacks the underlying
11/// sparse image as it is read.
12/// If random access reads are not required, it is more performant to use `unsparse` to completely
13/// unpack a sparse image.
14pub struct SparseReader<R> {
15    reader: R,
16    // Offset into the logical (unsparsed) image.
17    offset: u64,
18    // Size of the logical (unsparsed) image.
19    size: u64,
20    // The second field is the offset into `reader` at which the payload of the chunk appears, for
21    // Raw chunks.
22    chunks: Vec<(Chunk, Option<u64>)>,
23    // The block size of each chunk.
24    block_size: u32,
25}
26
27/// Helper utility class that rewinds a seekable reader on drop.
28///
29/// This is useful if, for example, a factory function reads a file that may or may not
30/// have a particular header and corresponding file structure. It is convenient for
31/// fallback code paths if the factory rewinds the file to the start.
32struct Rewinder<'a, R: Read + Seek> {
33    pub reader: &'a mut R,
34}
35
36impl<'a, R: Read + Seek> Drop for Rewinder<'a, R> {
37    fn drop(&mut self) {
38        if let Err(e) = self.reader.seek(SeekFrom::Start(0)) {
39            log::error!("File seek error: {e}");
40        }
41    }
42}
43
44impl<R: Read + Seek> SparseReader<R> {
45    /// Attempts to create a SparseReader from the given image.  Returns failure if the image is
46    /// malformed.
47    pub fn new(mut reader: R) -> std::result::Result<Self, SparseError> {
48        let rewinder = Rewinder { reader: &mut reader };
49        let header: SparseHeader = deserialize_from(rewinder.reader)
50            .map_err(|e| SparseError::Deserialize { ty: SparseDataType::Header, source: e })?;
51        if !header.valid() {
52            return Err(SparseError::InvalidHeader);
53        }
54        let num_chunks = header.total_chunks as usize;
55
56        let mut chunks = vec![];
57        let mut offset = 0;
58        for _ in 0..num_chunks {
59            let chunk = Chunk::read_metadata(rewinder.reader, offset, header.blk_sz)?;
60            let data_offset = if chunk.chunk_type() == crate::format::CHUNK_TYPE_RAW {
61                let data_offset = rewinder.reader.stream_position()?;
62                // Skip past the data payload
63                rewinder.reader.seek(SeekFrom::Current(chunk.output_size() as i64))?;
64                Some(data_offset)
65            } else {
66                None
67            };
68            offset += chunk.output_size() as u64;
69            chunks.push((chunk, data_offset));
70        }
71
72        drop(rewinder);
73        Ok(Self { reader, offset: 0, size: offset, chunks, block_size: header.blk_sz })
74    }
75
76    /// Returns the index of the current chunk in `self.chunks`.
77    fn current_chunk(&self) -> Option<usize> {
78        let mut off = 0;
79        let mut i = 0;
80        for (chunk, _) in &self.chunks {
81            let size = chunk.output_size() as u64;
82            if self.offset >= off && self.offset < off + size {
83                return Some(i);
84            }
85            off += size;
86            i += 1;
87        }
88        None
89    }
90
91    pub fn is_sparse_file(reader: &mut R) -> std::result::Result<bool, SparseError> {
92        let header: SparseHeader = deserialize_from(reader)
93            .map_err(|e| SparseError::Deserialize { ty: SparseDataType::Header, source: e })?;
94        let res = header.valid();
95        reader.seek(SeekFrom::Start(0)).map_err(|e| SparseError::Io(e))?;
96        Ok(res)
97    }
98
99    pub fn chunks(&self) -> &Vec<(Chunk, Option<u64>)> {
100        &self.chunks
101    }
102
103    pub fn unsparsed_size(&self) -> u64 {
104        self.size
105    }
106
107    pub fn block_size(&self) -> u32 {
108        self.block_size
109    }
110
111    /// Destroy the sparse reader and return all its fields.
112    pub fn destruct(self) -> (R, u64, u64, Vec<(Chunk, Option<u64>)>, u32) {
113        let Self { reader, offset, size, chunks, block_size } = self;
114        (reader, offset, size, chunks, block_size)
115    }
116}
117
118// It's assumed that `reader` already points at the right offset to read from the chunk, and `buf`
119// won't read past the end of the chunk.
120// `output_offset` is the logical position in the output stream.
121fn read_from_chunk<R: Read + Seek>(
122    reader: &mut R,
123    chunk: &Chunk,
124    output_offset: u64,
125    buf: &mut [u8],
126) -> std::io::Result<usize> {
127    match chunk {
128        Chunk::Raw { .. } => reader.read(buf),
129        Chunk::Fill { value, .. } => {
130            let mut value_bytes = value.to_le_bytes();
131            value_bytes.rotate_left(output_offset as usize % std::mem::size_of::<u32>());
132            let value_rotated = LE::read_u32(&value_bytes);
133            // Safety: `std::slice::align_to_mut` requires that everything in the dst slice is a
134            // valid type, which is true when going from [u8; 4] to [u32; 1].
135            let (prefix, wholes, suffix) = unsafe { buf.align_to_mut::<u32>() };
136            prefix.copy_from_slice(&value_bytes[value_bytes.len() - prefix.len()..]);
137            wholes.fill(value_rotated);
138            suffix.copy_from_slice(&value_bytes[..suffix.len()]);
139            Ok(buf.len())
140        }
141        Chunk::DontCare { .. } => {
142            buf.fill(0);
143            Ok(buf.len())
144        }
145        _ => unreachable!(),
146    }
147}
148
149impl<R: Read + Seek> Read for SparseReader<R> {
150    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
151        let mut bytes_read = 0;
152        while bytes_read < buf.len() {
153            let current_chunk_idx = match self.current_chunk() {
154                Some(i) => i,
155                None => return Ok(bytes_read),
156            };
157            let (current_chunk, chunk_start_offset) = &self.chunks[current_chunk_idx];
158            let offset_in_chunk = self.offset - current_chunk.output_offset().unwrap();
159            debug_assert!(offset_in_chunk < current_chunk.output_size() as u64);
160            let to_read = std::cmp::min(
161                buf.len() - bytes_read,
162                current_chunk.output_size() as usize - offset_in_chunk as usize,
163            );
164            if let Some(offset) = chunk_start_offset {
165                self.reader.seek(SeekFrom::Start(*offset + offset_in_chunk))?;
166            }
167            let bytes_read_from_chunk = read_from_chunk(
168                &mut self.reader,
169                current_chunk,
170                self.offset,
171                &mut buf[bytes_read..bytes_read + to_read],
172            )?;
173            bytes_read += bytes_read_from_chunk;
174            self.offset += bytes_read_from_chunk as u64;
175        }
176        Ok(bytes_read)
177    }
178}
179
180impl<R: Read + Seek> Seek for SparseReader<R> {
181    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
182        self.offset = match pos {
183            SeekFrom::Start(pos) => pos,
184            SeekFrom::Current(delta) => self
185                .offset
186                .checked_add_signed(delta)
187                .ok_or_else(|| std::io::Error::from(std::io::ErrorKind::InvalidInput))?,
188            SeekFrom::End(delta) => self
189                .size
190                .checked_add_signed(delta)
191                .ok_or_else(|| std::io::Error::from(std::io::ErrorKind::InvalidInput))?,
192        };
193        Ok(self.offset)
194    }
195}
196
197#[cfg(test)]
198mod test {
199    use crate::builder::{DataSource, SparseImageBuilder};
200    use crate::reader::SparseReader;
201    use rand::Rng as _;
202    use rand::rngs::SmallRng;
203    use std::io::{Read as _, Seek as _, SeekFrom, Write as _};
204    use tempfile::{NamedTempFile, TempDir};
205
206    #[test]
207    fn empty_reader() {
208        let tmpdir = TempDir::new().unwrap();
209
210        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
211        SparseImageBuilder::new().build(&mut sparse_file).expect("Build sparse image failed");
212        sparse_file.seek(SeekFrom::Start(0)).unwrap();
213
214        let mut reader =
215            SparseReader::new(Box::new(sparse_file)).expect("Failed to create SparseReader");
216
217        let mut unsparsed_bytes = vec![];
218        reader.read_to_end(&mut unsparsed_bytes).expect("Failed to read unsparsed image");
219        assert_eq!(unsparsed_bytes.len(), 0);
220    }
221
222    #[test]
223    fn is_sparse_file() {
224        let tmpdir = TempDir::new().unwrap();
225
226        let data = {
227            let mut data = Box::new([0u8; 8192]);
228            let mut i: u8 = 0;
229            for d in data.as_mut() {
230                *d = i;
231                i = i.wrapping_add(1);
232            }
233            data
234        };
235
236        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
237        SparseImageBuilder::new()
238            .add_source(DataSource::Buffer(data))
239            .build(&mut sparse_file)
240            .expect("Build sparse image failed");
241        sparse_file.seek(SeekFrom::Start(0)).unwrap();
242
243        assert!(SparseReader::is_sparse_file(&mut sparse_file).expect("Should be a sparse file"));
244
245        let mut garbage_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
246        let garbage_data = vec![1; 4096];
247        garbage_file.write_all(&garbage_data).expect("Writing garbage file");
248        garbage_file.seek(SeekFrom::Start(0)).unwrap();
249
250        assert!(!SparseReader::is_sparse_file(&mut garbage_file).unwrap());
251    }
252
253    #[test]
254    fn seek() {
255        let tmpdir = TempDir::new().unwrap();
256
257        let data = {
258            let mut data = Box::new([0u8; 8192]);
259            let mut i: u8 = 0;
260            for d in data.as_mut() {
261                *d = i;
262                i = i.wrapping_add(1);
263            }
264            data
265        };
266
267        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
268        SparseImageBuilder::new()
269            .add_source(DataSource::Buffer(data))
270            .build(&mut sparse_file)
271            .expect("Build sparse image failed");
272        sparse_file.seek(SeekFrom::Start(0)).unwrap();
273        let mut reader =
274            SparseReader::new(Box::new(sparse_file)).expect("Failed to create SparseReader");
275
276        let mut buf = [0u8; 1];
277        assert_eq!(0, reader.seek(SeekFrom::Start(0)).unwrap());
278        assert_eq!(1, reader.read(&mut buf).unwrap());
279        assert_eq!(buf[0], 0u8);
280
281        assert_eq!(100, reader.seek(SeekFrom::Start(100)).unwrap());
282        assert_eq!(1, reader.read(&mut buf).unwrap());
283        assert_eq!(buf[0], 100u8);
284
285        assert_eq!(99, reader.seek(SeekFrom::Current(-2)).unwrap());
286        assert_eq!(1, reader.read(&mut buf).unwrap());
287        assert_eq!(buf[0], 99u8);
288
289        assert_eq!(100, reader.seek(SeekFrom::Current(0)).unwrap());
290        assert_eq!(1, reader.read(&mut buf).unwrap());
291        assert_eq!(buf[0], 100u8);
292
293        assert_eq!(102, reader.seek(SeekFrom::Current(1)).unwrap());
294        assert_eq!(1, reader.read(&mut buf).unwrap());
295        assert_eq!(buf[0], 102u8);
296
297        assert_eq!(8191, reader.seek(SeekFrom::End(-1)).unwrap());
298        assert_eq!(1, reader.read(&mut buf).unwrap());
299        assert_eq!(buf[0], 255u8);
300
301        assert_eq!(8192, reader.seek(SeekFrom::End(0)).unwrap());
302        assert_eq!(0, reader.read(&mut buf).unwrap());
303
304        assert_eq!(8193, reader.seek(SeekFrom::End(1)).unwrap());
305        assert_eq!(0, reader.read(&mut buf).unwrap());
306    }
307
308    #[test]
309    fn read_past_eof() {
310        let tmpdir = TempDir::new().unwrap();
311
312        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
313        SparseImageBuilder::new()
314            .add_source(DataSource::Buffer(Box::new([0xffu8; 8192])))
315            .build(&mut sparse_file)
316            .expect("Build sparse image failed");
317        sparse_file.seek(SeekFrom::Start(0)).unwrap();
318
319        let mut reader =
320            SparseReader::new(Box::new(sparse_file)).expect("Failed to create SparseReader");
321
322        let mut buf = [0u8; 2];
323
324        reader.seek(SeekFrom::Start(8191)).expect("Seek failed");
325        assert_eq!(reader.read(&mut buf).expect("Failed to read"), 1);
326
327        reader.seek(SeekFrom::Start(8192)).expect("Seek failed");
328        assert_eq!(reader.read(&mut buf).expect("Failed to read"), 0);
329    }
330
331    #[test]
332    fn full_read() {
333        let tmpdir = TempDir::new().unwrap();
334
335        // Generate a large temporary file
336        let (mut file, _temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
337        let mut rng: SmallRng = rand::make_rng();
338        let mut data = Vec::<u8>::new();
339        data.resize(100 * 4096, 0);
340        rng.fill_bytes(&mut data);
341        file.write_all(&data).unwrap();
342        file.flush().unwrap();
343        file.seek(SeekFrom::Start(0)).unwrap();
344        let content_size = data.len();
345
346        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
347        SparseImageBuilder::new()
348            .add_source(DataSource::Buffer(Box::new([0xffu8; 8192])))
349            .add_source(DataSource::Reader { reader: Box::new(file), size: content_size as u64 })
350            .add_source(DataSource::Skip(16384))
351            .add_source(DataSource::Fill(0xaaaa_aaaau32, 1024))
352            .add_source(DataSource::Skip(4096))
353            .build(&mut sparse_file)
354            .expect("Build sparse image failed");
355        sparse_file.seek(SeekFrom::Start(0)).unwrap();
356
357        let mut reader =
358            SparseReader::new(Box::new(sparse_file)).expect("Failed to create SparseReader");
359
360        let mut unsparsed_bytes = vec![];
361        reader.read_to_end(&mut unsparsed_bytes).expect("Failed to read unsparsed image");
362        assert_eq!(unsparsed_bytes.len(), 8192 + content_size + 16384 + 4096 + 4096);
363        assert_eq!(&unsparsed_bytes[..8192], &[0xffu8; 8192]);
364        assert_eq!(&unsparsed_bytes[8192..8192 + content_size], &data[..]);
365        assert_eq!(
366            &unsparsed_bytes[8192 + content_size..8192 + content_size + 16384],
367            &[0u8; 16384]
368        );
369        assert_eq!(
370            &unsparsed_bytes[8192 + content_size + 16384..8192 + content_size + 16384 + 4096],
371            &[0xaau8; 4096]
372        );
373        assert_eq!(&unsparsed_bytes[8192 + content_size + 16384 + 4096..], &[0u8; 4096]);
374    }
375
376    #[test]
377    fn unaligned_reads() {
378        let tmpdir = TempDir::new().unwrap();
379
380        // Generate a large temporary file
381        let (mut file, _temp_path) = NamedTempFile::new_in(&tmpdir).unwrap().into_parts();
382        let mut rng: SmallRng = rand::make_rng();
383        let mut data = Vec::<u8>::new();
384        data.resize(100 * 4096, 0);
385        rng.fill_bytes(&mut data);
386        file.write_all(&data).unwrap();
387        file.flush().unwrap();
388        file.seek(SeekFrom::Start(0)).unwrap();
389        let content_size = data.len();
390
391        let mut sparse_file = NamedTempFile::new_in(&tmpdir).unwrap().into_file();
392        SparseImageBuilder::new()
393            .add_source(DataSource::Buffer(Box::new([0xffu8; 8192])))
394            .add_source(DataSource::Reader { reader: Box::new(file), size: content_size as u64 })
395            .add_source(DataSource::Skip(16384))
396            .add_source(DataSource::Fill(0x0102_0304u32, 1024))
397            .add_source(DataSource::Skip(4096))
398            .build(&mut sparse_file)
399            .expect("Build sparse image failed");
400        sparse_file.seek(SeekFrom::Start(0)).unwrap();
401
402        let mut reader =
403            SparseReader::new(Box::new(sparse_file)).expect("Failed to create SparseReader");
404
405        let mut buffer = [0u8; 4096];
406
407        // Do an unaligned read from each section
408
409        // DataSource::Buffer
410        reader.seek(SeekFrom::Start(10)).expect("Failed to seek");
411        let _ = reader.read(&mut buffer[..20]).expect("Failed to read");
412        assert_eq!(&buffer[..20], &[0xffu8; 20]);
413
414        // DataSource::File
415        reader.seek(SeekFrom::Start(8192 + 4095)).expect("Failed to seek");
416        let _ = reader.read(&mut buffer[..2]).expect("Failed to read");
417        assert_eq!(&buffer[..2], &data[4095..4097]);
418
419        // DataSource::Skip
420        reader.seek(SeekFrom::Start(8192 + content_size as u64 + 4090)).expect("Failed to seek");
421        let _ = reader.read(&mut buffer[..6]).expect("Failed to read");
422        assert_eq!(&buffer[..6], &[0u8; 6]);
423
424        // DataSource::Fill
425        reader
426            .seek(SeekFrom::Start(8192 + content_size as u64 + 16384 + 3))
427            .expect("Failed to seek");
428        let _ = reader.read(&mut buffer[..9]).expect("Failed to read");
429        // Bear in mind the byte ordering is LE, so 0x01020304 == [0x04, 0x03, 0x02, 0x01]
430        assert_eq!(&buffer[..9], &[0x01, 0x04, 0x03, 0x02, 0x01, 0x04, 0x03, 0x02, 0x01]);
431
432        // DataSource::Skip
433        reader
434            .seek(SeekFrom::Start(8192 + content_size as u64 + 16384 + 4096 + 1))
435            .expect("Failed to seek");
436        let _ = reader.read(&mut buffer[..4095]).expect("Failed to read");
437        assert_eq!(&buffer[..4095], &[0u8; 4095]);
438
439        // Do an unaligned read spanning two sections (the last Fill and Skip)
440        reader
441            .seek(SeekFrom::Start(8192 + content_size as u64 + 16384 + 4090))
442            .expect("Failed to seek");
443        let _ = reader.read(&mut buffer[..9]).expect("Failed to read");
444        assert_eq!(&buffer[..9], &[0x02, 0x01, 0x04, 0x03, 0x02, 0x01, 0x00, 0x00, 0x00]);
445    }
446}
sparse/reader.rs

sparse/
reader.rs
