disk/
zstd.rs

1// Copyright 2024 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Use seekable zstd archive of raw disk image as read only disk
6
7use std::cmp::min;
8use std::fs::File;
9use std::io;
10use std::io::ErrorKind;
11use std::io::Read;
12use std::io::Seek;
13use std::sync::Arc;
14use std::sync::RwLock;
15
16use anyhow::bail;
17use anyhow::Context;
18use async_trait::async_trait;
19use base::AsRawDescriptor;
20use base::FileAllocate;
21use base::FileReadWriteAtVolatile;
22use base::FileSetLen;
23use base::RawDescriptor;
24use base::VolatileSlice;
25use cros_async::BackingMemory;
26use cros_async::Executor;
27use cros_async::IoOptions;
28use cros_async::IoSource;
29
30use crate::AsyncDisk;
31use crate::DiskFile;
32use crate::DiskGetLen;
33use crate::Error as DiskError;
34use crate::Result as DiskResult;
35use crate::ToAsyncDisk;
36
37// Zstandard frame magic
38pub const ZSTD_FRAME_MAGIC: u32 = 0xFD2FB528;
39
40// Skippable frame magic can be anything between [0x184D2A50, 0x184D2A5F]
41pub const ZSTD_SKIPPABLE_MAGIC_LOW: u32 = 0x184D2A50;
42pub const ZSTD_SKIPPABLE_MAGIC_HIGH: u32 = 0x184D2A5F;
43pub const ZSTD_SEEK_TABLE_MAGIC: u32 = 0x8F92EAB1;
44
45pub const ZSTD_DEFAULT_FRAME_SIZE: usize = 128 << 10; // 128KB
46
47#[derive(Clone, Debug)]
48pub struct ZstdSeekTable {
49    // Cumulative sum of decompressed sizes of all frames before the indexed frame.
50    // The last element is the total decompressed size of the zstd archive.
51    cumulative_decompressed_sizes: Vec<u64>,
52    // Cumulative sum of compressed sizes of all frames before the indexed frame.
53    // The last element is the total compressed size of the zstd archive.
54    cumulative_compressed_sizes: Vec<u64>,
55}
56
57impl ZstdSeekTable {
58    /// Read seek table entries from seek_table_entries
59    pub fn from_footer(
60        seek_table_entries: &[u8],
61        num_frames: u32,
62        checksum_flag: bool,
63    ) -> anyhow::Result<ZstdSeekTable> {
64        let mut cumulative_decompressed_size: u64 = 0;
65        let mut cumulative_compressed_size: u64 = 0;
66        let mut cumulative_decompressed_sizes = Vec::with_capacity(num_frames as usize + 1);
67        let mut cumulative_compressed_sizes = Vec::with_capacity(num_frames as usize + 1);
68        let mut offset = 0;
69        cumulative_decompressed_sizes.push(0);
70        cumulative_compressed_sizes.push(0);
71        for _ in 0..num_frames {
72            let compressed_size = u32::from_le_bytes(
73                seek_table_entries
74                    .get(offset..offset + 4)
75                    .context("failed to parse seektable entry")?
76                    .try_into()?,
77            );
78            let decompressed_size = u32::from_le_bytes(
79                seek_table_entries
80                    .get(offset + 4..offset + 8)
81                    .context("failed to parse seektable entry")?
82                    .try_into()?,
83            );
84            cumulative_decompressed_size += decompressed_size as u64;
85            cumulative_compressed_size += compressed_size as u64;
86            cumulative_decompressed_sizes.push(cumulative_decompressed_size);
87            cumulative_compressed_sizes.push(cumulative_compressed_size);
88            offset += 8 + (checksum_flag as usize * 4);
89        }
90        cumulative_decompressed_sizes.push(cumulative_decompressed_size);
91        cumulative_compressed_sizes.push(cumulative_compressed_size);
92
93        Ok(ZstdSeekTable {
94            cumulative_decompressed_sizes,
95            cumulative_compressed_sizes,
96        })
97    }
98
99    /// Returns the index of the frame that contains the given decompressed offset.
100    pub fn find_frame_index(&self, decompressed_offset: u64) -> Option<usize> {
101        if self.cumulative_decompressed_sizes.is_empty()
102            || decompressed_offset >= *self.cumulative_decompressed_sizes.last().unwrap()
103        {
104            return None;
105        }
106        self.cumulative_decompressed_sizes
107            .partition_point(|&size| size <= decompressed_offset)
108            .checked_sub(1)
109    }
110}
111
112#[derive(Debug)]
113pub struct ZstdDisk {
114    file: File,
115    seek_table: ZstdSeekTable,
116    cache: RwLock<Option<ZstdFrameCache>>,
117}
118
119#[derive(Debug)]
120struct ZstdFrameCache {
121    frame_index: usize,
122    data: Vec<u8>,
123}
124
125impl ZstdDisk {
126    pub fn from_file(mut file: File) -> anyhow::Result<ZstdDisk> {
127        // Verify file is large enough to contain a seek table (17 bytes)
128        if file.metadata()?.len() < 17 {
129            return Err(anyhow::anyhow!("File too small to contain zstd seek table"));
130        }
131
132        // Read last 9 bytes as seek table footer
133        let mut seektable_footer = [0u8; 9];
134        file.seek(std::io::SeekFrom::End(-9))?;
135        file.read_exact(&mut seektable_footer)?;
136
137        // Verify last 4 bytes of footer is seek table magic
138        if u32::from_le_bytes(seektable_footer[5..9].try_into()?) != ZSTD_SEEK_TABLE_MAGIC {
139            return Err(anyhow::anyhow!("Invalid zstd seek table magic"));
140        }
141
142        // Get number of frame from seek table
143        let num_frames = u32::from_le_bytes(seektable_footer[0..4].try_into()?);
144
145        // Read flags from seek table descriptor
146        let checksum_flag = (seektable_footer[4] >> 7) & 1 != 0;
147        if (seektable_footer[4] & 0x7C) != 0 {
148            bail!(
149                "This zstd seekable decoder cannot parse seek table with non-zero reserved flags"
150            );
151        }
152
153        let seek_table_entries_size = num_frames * (8 + (checksum_flag as u32 * 4));
154
155        // Seek to the beginning of the seek table
156        file.seek(std::io::SeekFrom::End(
157            -(9 + seek_table_entries_size as i64),
158        ))?;
159
160        // Return new ZstdDisk
161        let mut seek_table_entries: Vec<u8> = vec![0u8; seek_table_entries_size as usize];
162        file.read_exact(&mut seek_table_entries)?;
163
164        let seek_table =
165            ZstdSeekTable::from_footer(&seek_table_entries, num_frames, checksum_flag)?;
166
167        Ok(ZstdDisk {
168            file,
169            seek_table,
170            cache: RwLock::new(None),
171        })
172    }
173}
174
175impl DiskGetLen for ZstdDisk {
176    fn get_len(&self) -> std::io::Result<u64> {
177        self.seek_table
178            .cumulative_decompressed_sizes
179            .last()
180            .copied()
181            .ok_or(io::ErrorKind::InvalidData.into())
182    }
183}
184
185impl FileSetLen for ZstdDisk {
186    fn set_len(&self, _len: u64) -> std::io::Result<()> {
187        Err(io::Error::new(
188            io::ErrorKind::PermissionDenied,
189            "unsupported operation",
190        ))
191    }
192}
193
194impl AsRawDescriptor for ZstdDisk {
195    fn as_raw_descriptor(&self) -> RawDescriptor {
196        self.file.as_raw_descriptor()
197    }
198}
199
200struct CompressedReadInstruction {
201    frame_index: usize,
202    // byte offset of the entire compressed file to start read from
203    read_offset: u64,
204    // number of bytes to read from the compressed file
205    read_size: u64,
206}
207
208fn compresed_frame_read_instruction(
209    seek_table: &ZstdSeekTable,
210    offset: u64,
211) -> anyhow::Result<CompressedReadInstruction> {
212    let frame_index = seek_table
213        .find_frame_index(offset)
214        .with_context(|| format!("no frame for offset {offset}"))?;
215    let compressed_offset = seek_table.cumulative_compressed_sizes[frame_index];
216    let next_compressed_offset = seek_table
217        .cumulative_compressed_sizes
218        .get(frame_index + 1)
219        .context("Offset out of range (next_compressed_offset overflow)")?;
220    let compressed_size = next_compressed_offset - compressed_offset;
221    Ok(CompressedReadInstruction {
222        frame_index,
223        read_offset: compressed_offset,
224        read_size: compressed_size,
225    })
226}
227
228fn copy_to_volatile_slice(src: &[u8], dst: VolatileSlice) -> io::Result<usize> {
229    let read_len = min(dst.size(), src.len());
230    let data_to_copy = &src[..read_len];
231    dst.sub_slice(0, read_len)
232        .map_err(io::Error::other)?
233        .copy_from(data_to_copy);
234    Ok(data_to_copy.len())
235}
236
237impl FileReadWriteAtVolatile for ZstdDisk {
238    fn read_at_volatile(&self, slice: VolatileSlice, offset: u64) -> io::Result<usize> {
239        let read_instruction = compresed_frame_read_instruction(&self.seek_table, offset)
240            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
241
242        // Try obtain read lock of cache
243        if let Some(cache) = self.cache.try_read().ok().as_ref().and_then(|g| g.as_ref()) {
244            if cache.frame_index == read_instruction.frame_index {
245                // Cache hit
246                let decompressed_offset_in_frame = offset
247                    - self.seek_table.cumulative_decompressed_sizes[read_instruction.frame_index];
248                return copy_to_volatile_slice(
249                    &cache.data[decompressed_offset_in_frame as usize..],
250                    slice,
251                );
252            }
253        }
254
255        let mut compressed_data = vec![0u8; read_instruction.read_size as usize];
256
257        let compressed_frame_slice = VolatileSlice::new(compressed_data.as_mut_slice());
258
259        self.file
260            .read_at_volatile(compressed_frame_slice, read_instruction.read_offset)
261            .map_err(io::Error::other)?;
262
263        let mut decompressor: zstd::bulk::Decompressor<'_> = zstd::bulk::Decompressor::new()?;
264        let mut decompressed_data = Vec::with_capacity(ZSTD_DEFAULT_FRAME_SIZE);
265        let decoded_size =
266            decompressor.decompress_to_buffer(&compressed_data, &mut decompressed_data)?;
267
268        let decompressed_offset_in_frame =
269            offset - self.seek_table.cumulative_decompressed_sizes[read_instruction.frame_index];
270
271        if decompressed_offset_in_frame >= decoded_size as u64 {
272            return Err(io::Error::new(
273                io::ErrorKind::InvalidData,
274                "BUG: Frame offset larger than decoded size",
275            ));
276        }
277
278        let updated_cache = ZstdFrameCache {
279            frame_index: read_instruction.frame_index,
280            data: decompressed_data,
281        };
282
283        let result = copy_to_volatile_slice(
284            &updated_cache.data[decompressed_offset_in_frame as usize..],
285            slice,
286        );
287
288        if let Ok(mut cache) = self.cache.try_write() {
289            *cache = Some(updated_cache);
290        };
291        result
292    }
293
294    fn write_at_volatile(&self, _slice: VolatileSlice, _offset: u64) -> io::Result<usize> {
295        Err(io::Error::new(
296            io::ErrorKind::PermissionDenied,
297            "unsupported operation",
298        ))
299    }
300}
301
302pub struct AsyncZstdDisk {
303    inner: IoSource<File>,
304    seek_table: ZstdSeekTable,
305    cache: RwLock<Option<ZstdFrameCache>>,
306}
307
308impl ToAsyncDisk for ZstdDisk {
309    fn to_async_disk(self: Box<Self>, ex: &Executor) -> DiskResult<Box<dyn AsyncDisk>> {
310        Ok(Box::new(AsyncZstdDisk {
311            inner: ex.async_from(self.file).map_err(DiskError::ToAsync)?,
312            seek_table: self.seek_table,
313            cache: RwLock::new(None),
314        }))
315    }
316}
317
318impl DiskGetLen for AsyncZstdDisk {
319    fn get_len(&self) -> io::Result<u64> {
320        self.seek_table
321            .cumulative_decompressed_sizes
322            .last()
323            .copied()
324            .ok_or(io::ErrorKind::InvalidData.into())
325    }
326}
327
328impl FileSetLen for AsyncZstdDisk {
329    fn set_len(&self, _len: u64) -> io::Result<()> {
330        Err(io::Error::new(
331            io::ErrorKind::PermissionDenied,
332            "unsupported operation",
333        ))
334    }
335}
336
337impl FileAllocate for AsyncZstdDisk {
338    fn allocate(&self, _offset: u64, _length: u64) -> io::Result<()> {
339        Err(io::Error::new(
340            io::ErrorKind::PermissionDenied,
341            "unsupported operation",
342        ))
343    }
344}
345
346fn copy_to_mem(
347    decompressed_data: &[u8],
348    mem: Arc<dyn BackingMemory + Send + Sync>,
349    mem_offsets: cros_async::MemRegionIter,
350) -> DiskResult<usize> {
351    // Copy the decompressed data to the provided memory regions.
352    let mut total_copied = 0;
353    for mem_region in mem_offsets {
354        let src_slice = &decompressed_data[total_copied..];
355        let dst_slice = mem
356            .get_volatile_slice(mem_region)
357            .map_err(DiskError::GuestMemory)?;
358
359        let to_copy = min(src_slice.len(), dst_slice.size());
360
361        if to_copy > 0 {
362            dst_slice
363                .sub_slice(0, to_copy)
364                .map_err(|e| DiskError::ReadingData(io::Error::other(e)))?
365                .copy_from(&src_slice[..to_copy]);
366
367            total_copied += to_copy;
368
369            // if fully copied destination buffers, break the loop.
370            if total_copied == dst_slice.size() {
371                break;
372            }
373        }
374    }
375
376    Ok(total_copied)
377}
378
379#[async_trait(?Send)]
380impl AsyncDisk for AsyncZstdDisk {
381    async fn flush(&self) -> DiskResult<()> {
382        // zstd is read-only, nothing to flush.
383        Ok(())
384    }
385
386    async fn fsync(&self) -> DiskResult<()> {
387        // Do nothing because it's read-only.
388        Ok(())
389    }
390
391    async fn fdatasync(&self) -> DiskResult<()> {
392        // Do nothing because it's read-only.
393        Ok(())
394    }
395
396    /// Reads data from `file_offset` of decompressed disk image till the end of current
397    /// zstd frame and write them into memory `mem` at `mem_offsets`. This function should
398    /// function the same as running `preadv()` on decompressed zstd image and reading into
399    /// the array of `iovec`s specified with `mem` and `mem_offsets`.
400    async fn read_to_mem<'a>(
401        &'a self,
402        file_offset: u64,
403        mem: Arc<dyn BackingMemory + Send + Sync>,
404        mem_offsets: cros_async::MemRegionIter<'a>,
405        options: IoOptions,
406    ) -> DiskResult<usize> {
407        let read_instruction = compresed_frame_read_instruction(&self.seek_table, file_offset)
408            .map_err(|e| DiskError::ReadingData(io::Error::new(io::ErrorKind::InvalidData, e)))?;
409
410        // Try obtain read lock of cache
411        if let Some(cache) = self.cache.try_read().ok().as_ref().and_then(|g| g.as_ref()) {
412            if cache.frame_index == read_instruction.frame_index {
413                // Cache hit
414                let decompressed_offset_in_frame = file_offset
415                    - self.seek_table.cumulative_decompressed_sizes[read_instruction.frame_index];
416                return copy_to_mem(
417                    &cache.data[decompressed_offset_in_frame as usize..],
418                    mem,
419                    mem_offsets,
420                );
421            }
422        }
423
424        let compressed_data = vec![0u8; read_instruction.read_size as usize];
425
426        let (compressed_read_size, compressed_data) = self
427            .inner
428            .read_to_vec(Some(read_instruction.read_offset), compressed_data, options)
429            .await
430            .map_err(|e| DiskError::ReadingData(io::Error::other(e)))?;
431
432        if compressed_read_size != read_instruction.read_size as usize {
433            return Err(DiskError::ReadingData(io::Error::new(
434                ErrorKind::UnexpectedEof,
435                "Read from compressed data result in wrong length",
436            )));
437        }
438
439        let mut decompressor: zstd::bulk::Decompressor<'_> =
440            zstd::bulk::Decompressor::new().map_err(DiskError::ReadingData)?;
441        let mut decompressed_data = Vec::with_capacity(ZSTD_DEFAULT_FRAME_SIZE);
442        let decoded_size = decompressor
443            .decompress_to_buffer(&compressed_data, &mut decompressed_data)
444            .map_err(DiskError::ReadingData)?;
445
446        let decompressed_offset_in_frame = file_offset
447            - self.seek_table.cumulative_decompressed_sizes[read_instruction.frame_index];
448
449        if decompressed_offset_in_frame as usize > decoded_size {
450            return Err(DiskError::ReadingData(io::Error::new(
451                ErrorKind::InvalidData,
452                "BUG: Frame offset larger than decoded size",
453            )));
454        }
455
456        // Copy the decompressed data to the provided memory regions.
457        let result = copy_to_mem(
458            &decompressed_data[decompressed_offset_in_frame as usize..],
459            mem,
460            mem_offsets,
461        );
462
463        let updated_cache = ZstdFrameCache {
464            frame_index: read_instruction.frame_index,
465            data: decompressed_data,
466        };
467
468        if let Ok(mut cache) = self.cache.try_write() {
469            *cache = Some(updated_cache);
470        };
471        result
472    }
473
474    async fn write_from_mem<'a>(
475        &'a self,
476        _file_offset: u64,
477        _mem: Arc<dyn BackingMemory + Send + Sync>,
478        _mem_offsets: cros_async::MemRegionIter<'a>,
479        _options: IoOptions,
480    ) -> DiskResult<usize> {
481        Err(DiskError::UnsupportedOperation)
482    }
483
484    async fn punch_hole(&self, _file_offset: u64, _length: u64) -> DiskResult<()> {
485        Err(DiskError::UnsupportedOperation)
486    }
487
488    async fn write_zeroes_at(&self, _file_offset: u64, _length: u64) -> DiskResult<()> {
489        Err(DiskError::UnsupportedOperation)
490    }
491}
492
493impl DiskFile for ZstdDisk {}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    #[test]
500    fn test_find_frame_index_empty() {
501        let seek_table = ZstdSeekTable {
502            cumulative_decompressed_sizes: vec![0],
503            cumulative_compressed_sizes: vec![0],
504        };
505        assert_eq!(seek_table.find_frame_index(0), None);
506        assert_eq!(seek_table.find_frame_index(5), None);
507    }
508
509    #[test]
510    fn test_find_frame_index_single_frame() {
511        let seek_table = ZstdSeekTable {
512            cumulative_decompressed_sizes: vec![0, 100],
513            cumulative_compressed_sizes: vec![0, 50],
514        };
515        assert_eq!(seek_table.find_frame_index(0), Some(0));
516        assert_eq!(seek_table.find_frame_index(50), Some(0));
517        assert_eq!(seek_table.find_frame_index(99), Some(0));
518        assert_eq!(seek_table.find_frame_index(100), None);
519    }
520
521    #[test]
522    fn test_find_frame_index_multiple_frames() {
523        let seek_table = ZstdSeekTable {
524            cumulative_decompressed_sizes: vec![0, 100, 300, 500],
525            cumulative_compressed_sizes: vec![0, 50, 120, 200],
526        };
527        assert_eq!(seek_table.find_frame_index(0), Some(0));
528        assert_eq!(seek_table.find_frame_index(99), Some(0));
529        assert_eq!(seek_table.find_frame_index(100), Some(1));
530        assert_eq!(seek_table.find_frame_index(299), Some(1));
531        assert_eq!(seek_table.find_frame_index(300), Some(2));
532        assert_eq!(seek_table.find_frame_index(499), Some(2));
533        assert_eq!(seek_table.find_frame_index(500), None);
534        assert_eq!(seek_table.find_frame_index(1000), None);
535    }
536
537    #[test]
538    fn test_find_frame_index_with_skippable_frames() {
539        let seek_table = ZstdSeekTable {
540            cumulative_decompressed_sizes: vec![0, 100, 100, 100, 300],
541            cumulative_compressed_sizes: vec![0, 50, 60, 70, 150],
542        };
543        assert_eq!(seek_table.find_frame_index(0), Some(0));
544        assert_eq!(seek_table.find_frame_index(99), Some(0));
545        // Correctly skips the skippable frames.
546        assert_eq!(seek_table.find_frame_index(100), Some(3));
547        assert_eq!(seek_table.find_frame_index(299), Some(3));
548        assert_eq!(seek_table.find_frame_index(300), None);
549    }
550
551    #[test]
552    fn test_find_frame_index_with_last_skippable_frame() {
553        let seek_table = ZstdSeekTable {
554            cumulative_decompressed_sizes: vec![0, 20, 40, 40, 60, 60, 80, 80],
555            cumulative_compressed_sizes: vec![0, 10, 20, 30, 40, 50, 60, 70],
556        };
557        assert_eq!(seek_table.find_frame_index(0), Some(0));
558        assert_eq!(seek_table.find_frame_index(20), Some(1));
559        assert_eq!(seek_table.find_frame_index(21), Some(1));
560        assert_eq!(seek_table.find_frame_index(79), Some(5));
561        assert_eq!(seek_table.find_frame_index(80), None);
562        assert_eq!(seek_table.find_frame_index(300), None);
563    }
564}
disk/zstd.rs

disk/
zstd.rs
