block: merge qcow, vhdx and block_util into block crate

This commit merges crates `qcow`, `vhdx` and `block_util` into the
crate `block`, which can allow `qcow` to use functions from `block_util`
without introducing a circular crate dependency.

This commit is based on crosvm implementation:
https://chromium.googlesource.com/crosvm/crosvm/+/f2eecc4152eca8d395566cffa2c102ec090a152d

Signed-off-by: Yu Li <liyu.yukiteru@bytedance.com>

block/src/qcow/qcow_raw_file.rs

@@ -0,0 +1,150 @@
+// 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-BSD-3-Clause file.
+
+use super::RawFile;
+use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
+use std::io::{self, BufWriter, Seek, SeekFrom};
+use std::mem::size_of;
+use vmm_sys_util::write_zeroes::WriteZeroes;
+
+/// A qcow file. Allows reading/writing clusters and appending clusters.
+#[derive(Debug)]
+pub struct QcowRawFile {
+    file: RawFile,
+    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: RawFile, 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(std::mem::size_of_val(table), &mut 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(std::mem::size_of_val(table), &mut 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 mutable reference to the underlying file.
+    pub fn file_mut(&mut self) -> &mut RawFile {
+        &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
+    }
+
+    /// Zeros out a cluster in the file.
+    pub fn zero_cluster(&mut self, address: u64) -> io::Result<()> {
+        let cluster_size = self.cluster_size as usize;
+        self.file.seek(SeekFrom::Start(address))?;
+        self.file.write_zeroes(cluster_size)?;
+        Ok(())
+    }
+}
+
+impl Clone for QcowRawFile {
+    fn clone(&self) -> Self {
+        QcowRawFile {
+            file: self.file.try_clone().expect("QcowRawFile cloning failed"),
+            cluster_size: self.cluster_size,
+            cluster_mask: self.cluster_mask,
+        }
+    }
+}

block/src/qcow/raw_file.rs

@@ -0,0 +1,354 @@
+// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Portions Copyright 2017 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-BSD-3-Clause file.
+//
+// Copyright © 2020 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
+
+use libc::c_void;
+use std::alloc::{alloc_zeroed, dealloc, Layout};
+use std::convert::TryInto;
+use std::fs::{File, Metadata};
+use std::io::{self, Read, Seek, SeekFrom, Write};
+use std::os::unix::io::{AsRawFd, RawFd};
+use std::slice;
+use vmm_sys_util::{seek_hole::SeekHole, write_zeroes::PunchHole, write_zeroes::WriteZeroesAt};
+
+#[derive(Debug)]
+pub struct RawFile {
+    file: File,
+    alignment: usize,
+    position: u64,
+}
+
+const BLK_ALIGNMENTS: [usize; 2] = [512, 4096];
+
+fn is_valid_alignment(fd: RawFd, alignment: usize) -> bool {
+    let layout = Layout::from_size_align(alignment, alignment).unwrap();
+    // SAFETY: layout has non-zero size
+    let ptr = unsafe { alloc_zeroed(layout) };
+    assert!(!ptr.is_null());
+
+    // SAFETY: FFI call
+    let ret = unsafe {
+        ::libc::pread(
+            fd,
+            ptr as *mut c_void,
+            alignment,
+            alignment.try_into().unwrap(),
+        )
+    };
+
+    // SAFETY: ptr was allocated by alloc_zeroed with layout
+    unsafe { dealloc(ptr, layout) };
+
+    ret >= 0
+}
+
+impl RawFile {
+    pub fn new(file: File, direct_io: bool) -> Self {
+        // Assume no alignment restrictions if we aren't using O_DIRECT.
+        let mut alignment = 0;
+        if direct_io {
+            for align in &BLK_ALIGNMENTS {
+                if is_valid_alignment(file.as_raw_fd(), *align) {
+                    alignment = *align;
+                    break;
+                }
+            }
+        }
+        RawFile {
+            file,
+            alignment,
+            position: 0,
+        }
+    }
+
+    fn round_up(&self, offset: u64) -> u64 {
+        let align: u64 = self.alignment.try_into().unwrap();
+        ((offset + align - 1) / align) * align
+    }
+
+    fn round_down(&self, offset: u64) -> u64 {
+        let align: u64 = self.alignment.try_into().unwrap();
+        (offset / align) * align
+    }
+
+    fn is_aligned(&self, buf: &[u8]) -> bool {
+        if self.alignment == 0 {
+            return true;
+        }
+
+        let align64: u64 = self.alignment.try_into().unwrap();
+
+        (self.position % align64 == 0)
+            && ((buf.as_ptr() as usize) % self.alignment == 0)
+            && (buf.len() % self.alignment == 0)
+    }
+
+    pub fn set_len(&self, size: u64) -> std::io::Result<()> {
+        self.file.set_len(size)
+    }
+
+    pub fn metadata(&self) -> std::io::Result<Metadata> {
+        self.file.metadata()
+    }
+
+    pub fn try_clone(&self) -> std::io::Result<RawFile> {
+        Ok(RawFile {
+            file: self.file.try_clone().expect("RawFile cloning failed"),
+            alignment: self.alignment,
+            position: self.position,
+        })
+    }
+
+    pub fn sync_all(&self) -> std::io::Result<()> {
+        self.file.sync_all()
+    }
+
+    pub fn sync_data(&self) -> std::io::Result<()> {
+        self.file.sync_data()
+    }
+}
+
+impl Read for RawFile {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        if self.is_aligned(buf) {
+            match self.file.read(buf) {
+                Ok(r) => {
+                    self.position = self.position.checked_add(r.try_into().unwrap()).unwrap();
+                    Ok(r)
+                }
+                Err(e) => Err(e),
+            }
+        } else {
+            let rounded_pos: u64 = self.round_down(self.position);
+            let file_offset: usize = self
+                .position
+                .checked_sub(rounded_pos)
+                .unwrap()
+                .try_into()
+                .unwrap();
+            let buf_len: usize = buf.len();
+            let rounded_len: usize = self
+                .round_up(
+                    file_offset
+                        .checked_add(buf_len)
+                        .unwrap()
+                        .try_into()
+                        .unwrap(),
+                )
+                .try_into()
+                .unwrap();
+
+            let layout = Layout::from_size_align(rounded_len, self.alignment).unwrap();
+            // SAFETY: layout has non-zero size
+            let tmp_ptr = unsafe { alloc_zeroed(layout) };
+            if tmp_ptr.is_null() {
+                return Err(io::Error::last_os_error());
+            }
+
+            // SAFETY: tmp_ptr is valid and at least rounded_len long
+            let tmp_buf = unsafe { slice::from_raw_parts_mut(tmp_ptr, rounded_len) };
+
+            // This can eventually replaced with read_at once its interface
+            // has been stabilized.
+            // SAFETY: FFI call. All parameters are valid.
+            let ret = unsafe {
+                ::libc::pread64(
+                    self.file.as_raw_fd(),
+                    tmp_buf.as_mut_ptr() as *mut c_void,
+                    tmp_buf.len(),
+                    rounded_pos.try_into().unwrap(),
+                )
+            };
+            if ret < 0 {
+                // SAFETY: tmp_ptr was allocated by alloc_zeroed with layout
+                unsafe { dealloc(tmp_ptr, layout) };
+                return Err(io::Error::last_os_error());
+            }
+
+            let read: usize = ret.try_into().unwrap();
+            if read < file_offset {
+                // SAFETY: tmp_ptr was allocated by alloc_zeroed with layout
+                unsafe { dealloc(tmp_ptr, layout) };
+                return Ok(0);
+            }
+
+            let mut to_copy = read - file_offset;
+            if to_copy > buf_len {
+                to_copy = buf_len;
+            }
+
+            buf.copy_from_slice(&tmp_buf[file_offset..(file_offset + buf_len)]);
+            // SAFETY: tmp_ptr was allocated by alloc_zeroed with layout
+            unsafe { dealloc(tmp_ptr, layout) };
+
+            self.seek(SeekFrom::Current(to_copy.try_into().unwrap()))
+                .unwrap();
+            Ok(to_copy)
+        }
+    }
+}
+
+impl Write for RawFile {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        if self.is_aligned(buf) {
+            match self.file.write(buf) {
+                Ok(r) => {
+                    self.position = self.position.checked_add(r.try_into().unwrap()).unwrap();
+                    Ok(r)
+                }
+                Err(e) => Err(e),
+            }
+        } else {
+            let rounded_pos: u64 = self.round_down(self.position);
+            let file_offset: usize = self
+                .position
+                .checked_sub(rounded_pos)
+                .unwrap()
+                .try_into()
+                .unwrap();
+            let buf_len: usize = buf.len();
+            let rounded_len: usize = self
+                .round_up(
+                    file_offset
+                        .checked_add(buf_len)
+                        .unwrap()
+                        .try_into()
+                        .unwrap(),
+                )
+                .try_into()
+                .unwrap();
+
+            let layout = Layout::from_size_align(rounded_len, self.alignment).unwrap();
+            // SAFETY: layout has non-zero size
+            let tmp_ptr = unsafe { alloc_zeroed(layout) };
+            if tmp_ptr.is_null() {
+                return Err(io::Error::last_os_error());
+            }
+
+            // SAFETY: tmp_ptr is at least rounded_len long
+            let tmp_buf = unsafe { slice::from_raw_parts_mut(tmp_ptr, rounded_len) };
+
+            // This can eventually replaced with read_at once its interface
+            // has been stabilized.
+            // SAFETY: FFI call
+            let ret = unsafe {
+                ::libc::pread64(
+                    self.file.as_raw_fd(),
+                    tmp_buf.as_mut_ptr() as *mut c_void,
+                    tmp_buf.len(),
+                    rounded_pos.try_into().unwrap(),
+                )
+            };
+            if ret < 0 {
+                // SAFETY: tmp_ptr was allocated by alloc_zeroed with layout
+                unsafe { dealloc(tmp_ptr, layout) };
+                return Err(io::Error::last_os_error());
+            };
+
+            tmp_buf[file_offset..(file_offset + buf_len)].copy_from_slice(buf);
+
+            // This can eventually replaced with write_at once its interface
+            // has been stabilized.
+            // SAFETY: FFI call
+            let ret = unsafe {
+                ::libc::pwrite64(
+                    self.file.as_raw_fd(),
+                    tmp_buf.as_ptr() as *const c_void,
+                    tmp_buf.len(),
+                    rounded_pos.try_into().unwrap(),
+                )
+            };
+
+            // SAFETY: tmp_ptr was allocated by alloc_zeroed with layout
+            unsafe { dealloc(tmp_ptr, layout) };
+
+            if ret < 0 {
+                return Err(io::Error::last_os_error());
+            }
+
+            let written: usize = ret.try_into().unwrap();
+            if written < file_offset {
+                Ok(0)
+            } else {
+                let mut to_seek = written - file_offset;
+                if to_seek > buf_len {
+                    to_seek = buf_len;
+                }
+
+                self.seek(SeekFrom::Current(to_seek.try_into().unwrap()))
+                    .unwrap();
+                Ok(to_seek)
+            }
+        }
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.file.sync_all()
+    }
+}
+
+impl Seek for RawFile {
+    fn seek(&mut self, newpos: SeekFrom) -> std::io::Result<u64> {
+        match self.file.seek(newpos) {
+            Ok(pos) => {
+                self.position = pos;
+                Ok(pos)
+            }
+            Err(e) => Err(e),
+        }
+    }
+}
+
+impl WriteZeroesAt for RawFile {
+    fn write_zeroes_at(&mut self, offset: u64, length: usize) -> std::io::Result<usize> {
+        self.file.write_zeroes_at(offset, length)
+    }
+}
+
+impl PunchHole for RawFile {
+    fn punch_hole(&mut self, offset: u64, length: u64) -> std::io::Result<()> {
+        self.file.punch_hole(offset, length)
+    }
+}
+
+impl SeekHole for RawFile {
+    fn seek_hole(&mut self, offset: u64) -> std::io::Result<Option<u64>> {
+        match self.file.seek_hole(offset) {
+            Ok(pos) => {
+                if let Some(p) = pos {
+                    self.position = p;
+                }
+                Ok(pos)
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    fn seek_data(&mut self, offset: u64) -> std::io::Result<Option<u64>> {
+        match self.file.seek_data(offset) {
+            Ok(pos) => {
+                if let Some(p) = pos {
+                    self.position = p;
+                }
+                Ok(pos)
+            }
+            Err(e) => Err(e),
+        }
+    }
+}
+
+impl Clone for RawFile {
+    fn clone(&self) -> Self {
+        RawFile {
+            file: self.file.try_clone().expect("RawFile cloning failed"),
+            alignment: self.alignment,
+            position: self.position,
+        }
+    }
+}

block/src/qcow/refcount.rs

@@ -0,0 +1,252 @@
+// 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-BSD-3-Clause file.
+
+use std::fmt::{self, Display};
+use std::io;
+
+use libc::EINVAL;
+
+use crate::qcow::{
+    qcow_raw_file::QcowRawFile,
+    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={addr} needs to be read"),
+            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(Clone, 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)
+    }
+}

block/src/qcow/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-BSD-3-Clause 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(Clone, 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(Clone, 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().next().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));
+    }
+}