qcow: Add qcow support

Extracted from crosvm (commit:f82d632), with clippy fixes.

Signed-off-by: Chao Peng <chao.p.peng@linux.intel.com>

qcow/src/qcow_raw_file.rs

@@ -0,0 +1,136 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::fs::File;
+use std::io::{self, BufWriter, Seek, SeekFrom};
+use std::mem::size_of;
+
+use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
+
+/// A qcow file. Allows reading/writing clusters and appending clusters.
+#[derive(Debug)]
+pub struct QcowRawFile {
+    file: File,
+    cluster_size: u64,
+    cluster_mask: u64,
+}
+
+impl QcowRawFile {
+    /// Creates a `QcowRawFile` from the given `File`, `None` is returned if `cluster_size` is not
+    /// a power of two.
+    pub fn from(file: File, cluster_size: u64) -> Option<Self> {
+        if cluster_size.count_ones() != 1 {
+            return None;
+        }
+        Some(QcowRawFile {
+            file,
+            cluster_size,
+            cluster_mask: cluster_size - 1,
+        })
+    }
+
+    /// Reads `count` 64 bit offsets and returns them as a vector.
+    /// `mask` optionally ands out some of the bits on the file.
+    pub fn read_pointer_table(
+        &mut self,
+        offset: u64,
+        count: u64,
+        mask: Option<u64>,
+    ) -> io::Result<Vec<u64>> {
+        let mut table = vec![0; count as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_u64_into::<BigEndian>(&mut table)?;
+        if let Some(m) = mask {
+            for ptr in &mut table {
+                *ptr &= m;
+            }
+        }
+        Ok(table)
+    }
+
+    /// Reads a cluster's worth of 64 bit offsets and returns them as a vector.
+    /// `mask` optionally ands out some of the bits on the file.
+    pub fn read_pointer_cluster(&mut self, offset: u64, mask: Option<u64>) -> io::Result<Vec<u64>> {
+        let count = self.cluster_size / size_of::<u64>() as u64;
+        self.read_pointer_table(offset, count, mask)
+    }
+
+    /// Writes `table` of u64 pointers to `offset` in the file.
+    /// `non_zero_flags` will be ORed with all non-zero values in `table`.
+    /// writing.
+    pub fn write_pointer_table(
+        &mut self,
+        offset: u64,
+        table: &[u64],
+        non_zero_flags: u64,
+    ) -> io::Result<()> {
+        self.file.seek(SeekFrom::Start(offset))?;
+        let mut buffer = BufWriter::with_capacity(table.len() * size_of::<u64>(), &self.file);
+        for addr in table {
+            let val = if *addr == 0 {
+                0
+            } else {
+                *addr | non_zero_flags
+            };
+            buffer.write_u64::<BigEndian>(val)?;
+        }
+        Ok(())
+    }
+
+    /// Read a refcount block from the file and returns a Vec containing the block.
+    /// Always returns a cluster's worth of data.
+    pub fn read_refcount_block(&mut self, offset: u64) -> io::Result<Vec<u16>> {
+        let count = self.cluster_size / size_of::<u16>() as u64;
+        let mut table = vec![0; count as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_u16_into::<BigEndian>(&mut table)?;
+        Ok(table)
+    }
+
+    /// Writes a refcount block to the file.
+    pub fn write_refcount_block(&mut self, offset: u64, table: &[u16]) -> io::Result<()> {
+        self.file.seek(SeekFrom::Start(offset))?;
+        let mut buffer = BufWriter::with_capacity(table.len() * size_of::<u16>(), &self.file);
+        for count in table {
+            buffer.write_u16::<BigEndian>(*count)?;
+        }
+        Ok(())
+    }
+
+    /// Allocates a new cluster at the end of the current file, return the address.
+    pub fn add_cluster_end(&mut self, max_valid_cluster_offset: u64) -> io::Result<Option<u64>> {
+        // Determine where the new end of the file should be and set_len, which
+        // translates to truncate(2).
+        let file_end: u64 = self.file.seek(SeekFrom::End(0))?;
+        let new_cluster_address: u64 = (file_end + self.cluster_size - 1) & !self.cluster_mask;
+
+        if new_cluster_address > max_valid_cluster_offset {
+            return Ok(None);
+        }
+
+        self.file.set_len(new_cluster_address + self.cluster_size)?;
+
+        Ok(Some(new_cluster_address))
+    }
+
+    /// Returns a reference to the underlying file.
+    pub fn file(&self) -> &File {
+        &self.file
+    }
+
+    /// Returns a mutable reference to the underlying file.
+    pub fn file_mut(&mut self) -> &mut File {
+        &mut self.file
+    }
+
+    /// Returns the size of the file's clusters.
+    pub fn cluster_size(&self) -> u64 {
+        self.cluster_size
+    }
+
+    /// Returns the offset of `address` within a cluster.
+    pub fn cluster_offset(&self, address: u64) -> u64 {
+        address & self.cluster_mask
+    }
+}

qcow/src/refcount.rs

@@ -0,0 +1,253 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std;
+use std::fmt::{self, Display};
+use std::io;
+
+use libc::EINVAL;
+
+use crate::qcow_raw_file::QcowRawFile;
+use crate::vec_cache::{CacheMap, Cacheable, VecCache};
+
+#[derive(Debug)]
+pub enum Error {
+    /// `EvictingCache` - Error writing a refblock from the cache to disk.
+    EvictingRefCounts(io::Error),
+    /// `InvalidIndex` - Address requested isn't within the range of the disk.
+    InvalidIndex,
+    /// `NeedCluster` - Handle this error by reading the cluster and calling the function again.
+    NeedCluster(u64),
+    /// `NeedNewCluster` - Handle this error by allocating a cluster and calling the function again.
+    NeedNewCluster,
+    /// `ReadingRefCounts` - Error reading the file in to the refcount cache.
+    ReadingRefCounts(io::Error),
+}
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+impl Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use self::Error::*;
+
+        match self {
+            EvictingRefCounts(e) => write!(
+                f,
+                "failed to write a refblock from the cache to disk: {}",
+                e
+            ),
+            InvalidIndex => write!(f, "address requested is not within the range of the disk"),
+            NeedCluster(addr) => write!(f, "cluster with addr={} needs to be read", addr),
+            NeedNewCluster => write!(f, "new cluster needs to be allocated for refcounts"),
+            ReadingRefCounts(e) => {
+                write!(f, "failed to read the file into the refcount cache: {}", e)
+            }
+        }
+    }
+}
+
+/// Represents the refcount entries for an open qcow file.
+#[derive(Debug)]
+pub struct RefCount {
+    ref_table: VecCache<u64>,
+    refcount_table_offset: u64,
+    refblock_cache: CacheMap<VecCache<u16>>,
+    refcount_block_entries: u64, // number of refcounts in a cluster.
+    cluster_size: u64,
+    max_valid_cluster_offset: u64,
+}
+
+impl RefCount {
+    /// Creates a `RefCount` from `file`, reading the refcount table from `refcount_table_offset`.
+    /// `refcount_table_entries` specifies the number of refcount blocks used by this image.
+    /// `refcount_block_entries` indicates the number of refcounts in each refcount block.
+    /// Each refcount table entry points to a refcount block.
+    pub fn new(
+        raw_file: &mut QcowRawFile,
+        refcount_table_offset: u64,
+        refcount_table_entries: u64,
+        refcount_block_entries: u64,
+        cluster_size: u64,
+    ) -> io::Result<RefCount> {
+        let ref_table = VecCache::from_vec(raw_file.read_pointer_table(
+            refcount_table_offset,
+            refcount_table_entries,
+            None,
+        )?);
+        let max_valid_cluster_index = (ref_table.len() as u64) * refcount_block_entries - 1;
+        let max_valid_cluster_offset = max_valid_cluster_index * cluster_size;
+        Ok(RefCount {
+            ref_table,
+            refcount_table_offset,
+            refblock_cache: CacheMap::new(50),
+            refcount_block_entries,
+            cluster_size,
+            max_valid_cluster_offset,
+        })
+    }
+
+    /// Returns the number of refcounts per block.
+    pub fn refcounts_per_block(&self) -> u64 {
+        self.refcount_block_entries
+    }
+
+    /// Returns the maximum valid cluster offset in the raw file for this refcount table.
+    pub fn max_valid_cluster_offset(&self) -> u64 {
+        self.max_valid_cluster_offset
+    }
+
+    /// Returns `NeedNewCluster` if a new cluster needs to be allocated for refcounts. If an
+    /// existing cluster needs to be read, `NeedCluster(addr)` is returned. The Caller should
+    /// allocate a cluster or read the required one and call this function again with the cluster.
+    /// On success, an optional address of a dropped cluster is returned. The dropped cluster can
+    /// be reused for other purposes.
+    pub fn set_cluster_refcount(
+        &mut self,
+        raw_file: &mut QcowRawFile,
+        cluster_address: u64,
+        refcount: u16,
+        mut new_cluster: Option<(u64, VecCache<u16>)>,
+    ) -> Result<Option<u64>> {
+        let (table_index, block_index) = self.get_refcount_index(cluster_address);
+
+        let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
+
+        // Fill the cache if this block isn't yet there.
+        if !self.refblock_cache.contains_key(table_index) {
+            // Need a new cluster
+            if let Some((addr, table)) = new_cluster.take() {
+                self.ref_table[table_index] = addr;
+                let ref_table = &self.ref_table;
+                self.refblock_cache
+                    .insert(table_index, table, |index, evicted| {
+                        raw_file.write_refcount_block(ref_table[index], evicted.get_values())
+                    })
+                    .map_err(Error::EvictingRefCounts)?;
+            } else {
+                if block_addr_disk == 0 {
+                    return Err(Error::NeedNewCluster);
+                }
+                return Err(Error::NeedCluster(block_addr_disk));
+            }
+        }
+
+        // Unwrap is safe here as the entry was filled directly above.
+        let dropped_cluster = if !self.refblock_cache.get(table_index).unwrap().dirty() {
+            // Free the previously used block and use a new one. Writing modified counts to new
+            // blocks keeps the on-disk state consistent even if it's out of date.
+            if let Some((addr, _)) = new_cluster.take() {
+                self.ref_table[table_index] = addr;
+                Some(block_addr_disk)
+            } else {
+                return Err(Error::NeedNewCluster);
+            }
+        } else {
+            None
+        };
+
+        self.refblock_cache.get_mut(table_index).unwrap()[block_index] = refcount;
+        Ok(dropped_cluster)
+    }
+
+    /// Flush the dirty refcount blocks. This must be done before flushing the table that points to
+    /// the blocks.
+    pub fn flush_blocks(&mut self, raw_file: &mut QcowRawFile) -> io::Result<()> {
+        // Write out all dirty L2 tables.
+        for (table_index, block) in self.refblock_cache.iter_mut().filter(|(_k, v)| v.dirty()) {
+            let addr = self.ref_table[*table_index];
+            if addr != 0 {
+                raw_file.write_refcount_block(addr, block.get_values())?;
+            } else {
+                return Err(std::io::Error::from_raw_os_error(EINVAL));
+            }
+            block.mark_clean();
+        }
+        Ok(())
+    }
+
+    /// Flush the refcount table that keeps the address of the refcounts blocks.
+    /// Returns true if the table changed since the previous `flush_table()` call.
+    pub fn flush_table(&mut self, raw_file: &mut QcowRawFile) -> io::Result<bool> {
+        if self.ref_table.dirty() {
+            raw_file.write_pointer_table(
+                self.refcount_table_offset,
+                &self.ref_table.get_values(),
+                0,
+            )?;
+            self.ref_table.mark_clean();
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    /// Gets the refcount for a cluster with the given address.
+    pub fn get_cluster_refcount(
+        &mut self,
+        raw_file: &mut QcowRawFile,
+        address: u64,
+    ) -> Result<u16> {
+        let (table_index, block_index) = self.get_refcount_index(address);
+        let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
+        if block_addr_disk == 0 {
+            return Ok(0);
+        }
+        if !self.refblock_cache.contains_key(table_index) {
+            let table = VecCache::from_vec(
+                raw_file
+                    .read_refcount_block(block_addr_disk)
+                    .map_err(Error::ReadingRefCounts)?,
+            );
+            let ref_table = &self.ref_table;
+            self.refblock_cache
+                .insert(table_index, table, |index, evicted| {
+                    raw_file.write_refcount_block(ref_table[index], evicted.get_values())
+                })
+                .map_err(Error::EvictingRefCounts)?;
+        }
+        Ok(self.refblock_cache.get(table_index).unwrap()[block_index])
+    }
+
+    /// Returns the refcount table for this file. This is only useful for debugging.
+    pub fn ref_table(&self) -> &[u64] {
+        &self.ref_table.get_values()
+    }
+
+    /// Returns the refcounts stored in the given block.
+    pub fn refcount_block(
+        &mut self,
+        raw_file: &mut QcowRawFile,
+        table_index: usize,
+    ) -> Result<Option<&[u16]>> {
+        let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
+        if block_addr_disk == 0 {
+            return Ok(None);
+        }
+        if !self.refblock_cache.contains_key(table_index) {
+            let table = VecCache::from_vec(
+                raw_file
+                    .read_refcount_block(block_addr_disk)
+                    .map_err(Error::ReadingRefCounts)?,
+            );
+            // TODO(dgreid) - closure needs to return an error.
+            let ref_table = &self.ref_table;
+            self.refblock_cache
+                .insert(table_index, table, |index, evicted| {
+                    raw_file.write_refcount_block(ref_table[index], evicted.get_values())
+                })
+                .map_err(Error::EvictingRefCounts)?;
+        }
+        // The index must exist as it was just inserted if it didn't already.
+        Ok(Some(
+            self.refblock_cache.get(table_index).unwrap().get_values(),
+        ))
+    }
+
+    // Gets the address of the refcount block and the index into the block for the given address.
+    fn get_refcount_index(&self, address: u64) -> (usize, usize) {
+        let block_index = (address / self.cluster_size) % self.refcount_block_entries;
+        let refcount_table_index = (address / self.cluster_size) / self.refcount_block_entries;
+        (refcount_table_index as usize, block_index as usize)
+    }
+}

qcow/src/vec_cache.rs

@@ -0,0 +1,185 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::collections::hash_map::IterMut;
+use std::collections::HashMap;
+use std::io;
+use std::ops::{Index, IndexMut};
+use std::slice::SliceIndex;
+
+/// Trait that allows for checking if an implementor is dirty. Useful for types that are cached so
+/// it can be checked if they need to be committed to disk.
+pub trait Cacheable {
+    /// Used to check if the item needs to be written out or if it can be discarded.
+    fn dirty(&self) -> bool;
+}
+
+#[derive(Debug)]
+/// Represents a vector that implements the `Cacheable` trait so it can be held in a cache.
+pub struct VecCache<T: 'static + Copy + Default> {
+    vec: Box<[T]>,
+    dirty: bool,
+}
+
+impl<T: 'static + Copy + Default> VecCache<T> {
+    /// Creates a `VecCache` that can hold `count` elements.
+    pub fn new(count: usize) -> VecCache<T> {
+        VecCache {
+            vec: vec![Default::default(); count].into_boxed_slice(),
+            dirty: true,
+        }
+    }
+
+    /// Creates a `VecCache` from the passed in `vec`.
+    pub fn from_vec(vec: Vec<T>) -> VecCache<T> {
+        VecCache {
+            vec: vec.into_boxed_slice(),
+            dirty: false,
+        }
+    }
+
+    pub fn get<I>(&self, index: I) -> Option<&<I as SliceIndex<[T]>>::Output>
+    where
+        I: SliceIndex<[T]>,
+    {
+        self.vec.get(index)
+    }
+
+    /// Gets a reference to the underlying vector.
+    pub fn get_values(&self) -> &[T] {
+        &self.vec
+    }
+
+    /// Mark this cache element as clean.
+    pub fn mark_clean(&mut self) {
+        self.dirty = false;
+    }
+
+    /// Returns the number of elements in the vector.
+    pub fn len(&self) -> usize {
+        self.vec.len()
+    }
+}
+
+impl<T: 'static + Copy + Default> Cacheable for VecCache<T> {
+    fn dirty(&self) -> bool {
+        self.dirty
+    }
+}
+
+impl<T: 'static + Copy + Default> Index<usize> for VecCache<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &T {
+        self.vec.index(index)
+    }
+}
+
+impl<T: 'static + Copy + Default> IndexMut<usize> for VecCache<T> {
+    fn index_mut(&mut self, index: usize) -> &mut T {
+        self.dirty = true;
+        self.vec.index_mut(index)
+    }
+}
+
+#[derive(Debug)]
+pub struct CacheMap<T: Cacheable> {
+    capacity: usize,
+    map: HashMap<usize, T>,
+}
+
+impl<T: Cacheable> CacheMap<T> {
+    pub fn new(capacity: usize) -> Self {
+        CacheMap {
+            capacity,
+            map: HashMap::with_capacity(capacity),
+        }
+    }
+
+    pub fn contains_key(&self, key: usize) -> bool {
+        self.map.contains_key(&key)
+    }
+
+    pub fn get(&self, index: usize) -> Option<&T> {
+        self.map.get(&index)
+    }
+
+    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
+        self.map.get_mut(&index)
+    }
+
+    pub fn iter_mut(&mut self) -> IterMut<usize, T> {
+        self.map.iter_mut()
+    }
+
+    // Check if the refblock cache is full and we need to evict.
+    pub fn insert<F>(&mut self, index: usize, block: T, write_callback: F) -> io::Result<()>
+    where
+        F: FnOnce(usize, T) -> io::Result<()>,
+    {
+        if self.map.len() == self.capacity {
+            // TODO(dgreid) - smarter eviction strategy.
+            let to_evict = *self.map.iter().nth(0).unwrap().0;
+            if let Some(evicted) = self.map.remove(&to_evict) {
+                if evicted.dirty() {
+                    write_callback(to_evict, evicted)?;
+                }
+            }
+        }
+        self.map.insert(index, block);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct NumCache(pub u64);
+    impl Cacheable for NumCache {
+        fn dirty(&self) -> bool {
+            true
+        }
+    }
+
+    #[test]
+    fn evicts_when_full() {
+        let mut cache = CacheMap::<NumCache>::new(3);
+        let mut evicted = None;
+        cache
+            .insert(0, NumCache(5), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(1, NumCache(6), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(2, NumCache(7), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(3, NumCache(8), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert!(evicted.is_some());
+
+        // Check that three of the four items inserted are still there and that the most recently
+        // inserted is one of them.
+        let num_items = (0..=3).filter(|k| cache.contains_key(&k)).count();
+        assert_eq!(num_items, 3);
+        assert!(cache.contains_key(&3));
+    }
+}