block: qcow: Add support variable refcount widths

QCOW2 v3 specifies refcount_order 0-6 with
refcount_bits = 1 << refcount_order. Previously only 16-bit (order 4)
was supported.

Changes:
- RefcountBytes trait handles byte-aligned types (8/16/32/64-bit)
- Generic pack/unpack for sub-byte widths (1/2/4-bit)
- Function pointers for read/write selected at open time
- Internal refcount type widened from u16 to u64

Signed-off-by: Anatol Belski <anbelski@linux.microsoft.com>

block/src/qcow/mod.rs

@@ -186,7 +186,7 @@ const MAX_CLUSTER_BITS: u32 = 21;
 // This easily covers 1 TB files. When support for bigger files is needed the assumptions made to
 // keep these tables in RAM needs to be thrown out.
 const MAX_RAM_POINTER_TABLE_SIZE: u64 = 35_000_000;
-// Only support 2 byte refcounts, 2^refcount_order bits.
+// 16-bit refcounts.
 const DEFAULT_REFCOUNT_ORDER: u32 = 4;
 
 const V2_BARE_HEADER_SIZE: u32 = 72;
@@ -682,12 +682,10 @@ impl QcowHeader {
 
 fn max_refcount_clusters(refcount_order: u32, cluster_size: u32, num_clusters: u32) -> u64 {
     // Use u64 as the product of the u32 inputs can overflow.
-    let refcount_bytes = (0x01 << u64::from(refcount_order)) / 8;
+    let refcount_bits = 0x01u64 << u64::from(refcount_order);
-    let for_data = div_round_up_u64(
+    let cluster_bits = u64::from(cluster_size) * 8;
-        u64::from(num_clusters) * refcount_bytes,
+    let for_data = div_round_up_u64(u64::from(num_clusters) * refcount_bits, cluster_bits);
-        u64::from(cluster_size),
+    let for_refcounts = div_round_up_u64(for_data * refcount_bits, cluster_bits);
-    );
-    let for_refcounts = div_round_up_u64(for_data * refcount_bytes, u64::from(cluster_size));
     for_data + for_refcounts
 }
 
@@ -849,14 +847,13 @@ impl QcowFile {
         let backing_file =
             BackingFile::new(header.backing_file.as_ref(), direct_io, max_nesting_depth)?;
 
-        // Only support two byte refcounts.
+        // Validate refcount order to be 0..6
         let refcount_bits: u64 = 0x01u64
             .checked_shl(header.refcount_order)
             .ok_or(Error::UnsupportedRefcountOrder)?;
-        if refcount_bits != 16 {
+        if refcount_bits > 64 {
             return Err(Error::UnsupportedRefcountOrder);
         }
-        let refcount_bytes = refcount_bits.div_ceil(8);
 
         // Need at least one refcount cluster
         if header.refcount_table_clusters == 0 {
@@ -891,8 +888,8 @@ impl QcowFile {
             refcount_rebuild_required = true;
         }
 
-        let mut raw_file =
+        let mut raw_file = QcowRawFile::from(file, cluster_size, refcount_bits)
-            QcowRawFile::from(file, cluster_size).ok_or(Error::InvalidClusterSize)?;
+            .ok_or(Error::InvalidClusterSize)?;
         if refcount_rebuild_required {
             QcowFile::rebuild_refcounts(&mut raw_file, header.clone())?;
         }
@@ -928,7 +925,7 @@ impl QcowFile {
         if l1_clusters + refcount_clusters > MAX_RAM_POINTER_TABLE_SIZE {
             return Err(Error::TooManyRefcounts(refcount_clusters));
         }
-        let refcount_block_entries = cluster_size / refcount_bytes;
+        let refcount_block_entries = cluster_size * 8 / refcount_bits;
         let refcounts = RefCount::new(
             &mut raw_file,
             header.refcount_table_offset,
@@ -1067,7 +1064,7 @@ impl QcowFile {
     }
 
     /// Returns the `index`th refcount block from the file.
-    pub fn refcount_block(&mut self, index: usize) -> Result<Option<&[u16]>> {
+    pub fn refcount_block(&mut self, index: usize) -> Result<Option<&[u64]>> {
         self.refcounts
             .refcount_block(&mut self.raw_file, index)
             .map_err(Error::ReadingRefCountBlock)
@@ -1122,7 +1119,7 @@ impl QcowFile {
 
     /// Rebuild the reference count tables.
     fn rebuild_refcounts(raw_file: &mut QcowRawFile, header: QcowHeader) -> Result<()> {
-        fn add_ref(refcounts: &mut [u16], cluster_size: u64, cluster_address: u64) -> Result<()> {
+        fn add_ref(refcounts: &mut [u64], cluster_size: u64, cluster_address: u64) -> Result<()> {
             let idx = (cluster_address / cluster_size) as usize;
             if idx >= refcounts.len() {
                 return Err(Error::InvalidClusterIndex);
@@ -1132,13 +1129,13 @@ impl QcowFile {
         }
 
         // Add a reference to the first cluster (header plus extensions).
-        fn set_header_refcount(refcounts: &mut [u16], cluster_size: u64) -> Result<()> {
+        fn set_header_refcount(refcounts: &mut [u64], cluster_size: u64) -> Result<()> {
             add_ref(refcounts, cluster_size, 0)
         }
 
         // Add references to the L1 table clusters.
         fn set_l1_refcounts(
-            refcounts: &mut [u16],
+            refcounts: &mut [u64],
             header: &QcowHeader,
             cluster_size: u64,
         ) -> Result<()> {
@@ -1153,7 +1150,7 @@ impl QcowFile {
 
         // Traverse the L1 and L2 tables to find all reachable data clusters.
         fn set_data_refcounts(
-            refcounts: &mut [u16],
+            refcounts: &mut [u64],
             header: &QcowHeader,
             cluster_size: u64,
             raw_file: &mut QcowRawFile,
@@ -1192,7 +1189,7 @@ impl QcowFile {
 
         // Add references to the top-level refcount table clusters.
         fn set_refcount_table_refcounts(
-            refcounts: &mut [u16],
+            refcounts: &mut [u64],
             header: &QcowHeader,
             cluster_size: u64,
         ) -> Result<()> {
@@ -1211,7 +1208,7 @@ impl QcowFile {
         // This needs to be done last so that we have the correct refcounts for all other
         // clusters.
         fn alloc_refblocks(
-            refcounts: &mut [u16],
+            refcounts: &mut [u64],
             cluster_size: u64,
             refblock_clusters: u64,
         ) -> Result<Vec<u64>> {
@@ -1239,7 +1236,7 @@ impl QcowFile {
 
         // Write the updated reference count blocks and reftable.
         fn write_refblocks(
-            refcounts: &[u16],
+            refcounts: &[u64],
             mut header: QcowHeader,
             ref_table: &[u64],
             raw_file: &mut QcowRawFile,
@@ -1265,12 +1262,11 @@ impl QcowFile {
                 // If this is the last (partial) cluster, pad it out to a full refblock cluster.
                 if refblock.len() < refcount_block_entries as usize {
                     let refblock_padding =
-                        vec![0u16; refcount_block_entries as usize - refblock.len()];
+                        vec![0u64; refcount_block_entries as usize - refblock.len()];
+                    let byte_offset =
+                        refblock.len() as u64 * raw_file.cluster_size() / refcount_block_entries;
                     raw_file
-                        .write_refcount_block(
+                        .write_refcount_block(*refblock_addr + byte_offset, &refblock_padding)
-                            *refblock_addr + refblock.len() as u64 * 2,
-                            &refblock_padding,
-                        )
                         .map_err(Error::WritingHeader)?;
                 }
             }
@@ -1297,8 +1293,7 @@ impl QcowFile {
             .len();
 
         let refcount_bits = 1u64 << header.refcount_order;
-        let refcount_bytes = div_round_up_u64(refcount_bits, 8);
+        let refcount_block_entries = cluster_size * 8 / refcount_bits;
-        let refcount_block_entries = cluster_size / refcount_bytes;
         let pointers_per_cluster = cluster_size / size_of::<u64>() as u64;
         let data_clusters = div_round_up_u64(header.size, cluster_size);
         let l2_clusters = div_round_up_u64(data_clusters, pointers_per_cluster);
@@ -1554,7 +1549,7 @@ impl QcowFile {
         l1_index: usize,
         l2_index: usize,
         cluster_addr: u64,
-        set_refcounts: &mut Vec<(u64, u16)>,
+        set_refcounts: &mut Vec<(u64, u64)>,
     ) -> io::Result<()> {
         if !self.l2_cache.get(l1_index).unwrap().dirty() {
             // Free the previously used cluster if one exists. Modified tables are always
@@ -1804,7 +1799,7 @@ impl QcowFile {
     fn set_cluster_refcount_track_freed(
         &mut self,
         address: u64,
-        refcount: u16,
+        refcount: u64,
     ) -> std::io::Result<()> {
         let mut newly_unref = self.set_cluster_refcount(address, refcount)?;
         self.unref_clusters.append(&mut newly_unref);
@@ -1814,7 +1809,7 @@ impl QcowFile {
     // Set the refcount for a cluster with the given address.
     // Returns a list of any refblocks that can be reused, this happens when a refblock is moved,
     // the old location can be reused.
-    fn set_cluster_refcount(&mut self, address: u64, refcount: u16) -> std::io::Result<Vec<u64>> {
+    fn set_cluster_refcount(&mut self, address: u64, refcount: u64) -> std::io::Result<Vec<u64>> {
         let mut added_clusters = Vec::new();
         let mut unref_clusters = Vec::new();
         let mut refcount_set = false;
@@ -2655,7 +2650,7 @@ mod unit_tests {
     #[test]
     fn invalid_refcount_order() {
         let mut header = valid_header_v3();
-        header[99] = 2;
+        header[99] = 7;
         with_basic_file(&header, |disk_file: RawFile| {
             QcowFile::from(disk_file).expect_err("Invalid refcount order worked.");
         });
@@ -3449,8 +3444,9 @@ mod unit_tests {
         with_basic_file(&valid_header_v3(), |mut disk_file: RawFile| {
             let header = QcowHeader::new(&mut disk_file).expect("Failed to create Header.");
             let cluster_size = 65536;
-            let mut raw_file =
+            let refcount_bits = 1u64 << header.refcount_order;
-                QcowRawFile::from(disk_file, cluster_size).expect("Failed to create QcowRawFile.");
+            let mut raw_file = QcowRawFile::from(disk_file, cluster_size, refcount_bits)
+                .expect("Failed to create QcowRawFile.");
             QcowFile::rebuild_refcounts(&mut raw_file, header)
                 .expect("Failed to rebuild recounts.");
         });

block/src/qcow/qcow_raw_file.rs

@@ -4,7 +4,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
 
-use std::io::{self, BufWriter, Seek, SeekFrom, Write};
+use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
 use std::os::fd::{AsRawFd, RawFd};
 
@@ -13,25 +13,163 @@ use vmm_sys_util::write_zeroes::WriteZeroes;
 
 use super::RawFile;
 
+// Type aliases for the refcount read/write function pointers
+type RefcountReader = fn(&mut RawFile, usize) -> io::Result<Vec<u64>>;
+type RefcountWriter = fn(&mut RawFile, &[u64]) -> io::Result<()>;
+
+/// Big-endian file access trait.
+trait BeUint: Sized + Copy {
+    fn from_slice(bytes: &[u8]) -> u64;
+    fn write<W: Write>(w: &mut W, val: u64) -> io::Result<()>;
+}
+
+impl BeUint for u8 {
+    #[inline(always)]
+    fn from_slice(bytes: &[u8]) -> u64 {
+        bytes[0] as u64
+    }
+    #[inline(always)]
+    fn write<W: Write>(w: &mut W, val: u64) -> io::Result<()> {
+        w.write_u8(val as u8)
+    }
+}
+
+impl BeUint for u16 {
+    #[inline(always)]
+    fn from_slice(bytes: &[u8]) -> u64 {
+        u16::from_be_bytes([bytes[0], bytes[1]]) as u64
+    }
+    #[inline(always)]
+    fn write<W: Write>(w: &mut W, val: u64) -> io::Result<()> {
+        w.write_u16::<BigEndian>(val as u16)
+    }
+}
+
+impl BeUint for u32 {
+    #[inline(always)]
+    fn from_slice(bytes: &[u8]) -> u64 {
+        u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as u64
+    }
+    #[inline(always)]
+    fn write<W: Write>(w: &mut W, val: u64) -> io::Result<()> {
+        w.write_u32::<BigEndian>(val as u32)
+    }
+}
+
+impl BeUint for u64 {
+    #[inline(always)]
+    fn from_slice(bytes: &[u8]) -> u64 {
+        u64::from_be_bytes([
+            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
+        ])
+    }
+    #[inline(always)]
+    fn write<W: Write>(w: &mut W, val: u64) -> io::Result<()> {
+        w.write_u64::<BigEndian>(val)
+    }
+}
+
+/// Read byte-aligned refcounts.
+fn read_refcount<T: BeUint>(file: &mut RawFile, count: usize) -> io::Result<Vec<u64>> {
+    let bytes_per_entry = size_of::<T>();
+    let mut data = vec![0u8; count * bytes_per_entry];
+    file.read_exact(&mut data)?;
+    Ok(data
+        .chunks_exact(bytes_per_entry)
+        .map(T::from_slice)
+        .collect())
+}
+
+/// Write byte-aligned refcounts.
+fn write_refcount<T: BeUint>(file: &mut RawFile, table: &[u64]) -> io::Result<()> {
+    let bytes_per_entry = size_of::<T>();
+    let mut buffer = BufWriter::with_capacity(table.len() * bytes_per_entry, file);
+    for &val in table {
+        T::write(&mut buffer, val)?;
+    }
+    buffer.flush()
+}
+
+/// Read sub-byte refcounts. Bit 0 is the least significant bit.
+fn read_refcount_subbyte<const BITS: usize>(
+    file: &mut RawFile,
+    count: usize,
+) -> io::Result<Vec<u64>> {
+    const { assert!(BITS == 1 || BITS == 2 || BITS == 4) };
+    let entries_per_byte = 8 / BITS;
+    let mask = (1u64 << BITS) - 1;
+    let bytes_needed = count.div_ceil(entries_per_byte);
+    let mut bytes = vec![0u8; bytes_needed];
+    file.read_exact(&mut bytes)?;
+
+    let mut table = vec![0u64; count];
+    for (i, val) in table.iter_mut().enumerate() {
+        let byte_idx = i / entries_per_byte;
+        let bit_offset = (i % entries_per_byte) * BITS;
+        *val = (bytes[byte_idx] as u64 >> bit_offset) & mask;
+    }
+    Ok(table)
+}
+
+/// Write sub-byte refcounts. Bit 0 is the least significant bit.
+fn write_refcount_subbyte<const BITS: usize>(file: &mut RawFile, table: &[u64]) -> io::Result<()> {
+    const { assert!(BITS == 1 || BITS == 2 || BITS == 4) };
+    let entries_per_byte = 8 / BITS;
+    let mask = (1u64 << BITS) - 1;
+    let mut buffer = BufWriter::with_capacity(table.len().div_ceil(entries_per_byte), file);
+
+    for chunk in table.chunks(entries_per_byte) {
+        let mut byte = 0u8;
+        for (i, &val) in chunk.iter().enumerate() {
+            let bit_offset = i * BITS;
+            byte |= ((val & mask) << bit_offset) as u8;
+        }
+        buffer.write_u8(byte)?;
+    }
+    buffer.flush()
+}
+
 /// A qcow file. Allows reading/writing clusters and appending clusters.
 #[derive(Debug)]
 pub struct QcowRawFile {
     file: RawFile,
     cluster_size: u64,
     cluster_mask: u64,
+    refcount_block_entries: u64,
+    read_refcount_fn: RefcountReader,
+    write_refcount_fn: RefcountWriter,
 }
 
 impl QcowRawFile {
     /// Creates a `QcowRawFile` from the given `File`, `None` is returned if `cluster_size` is not
-    /// a power of two.
+    /// a power of two or refcount_bits is invalid.
-    pub fn from(file: RawFile, cluster_size: u64) -> Option<Self> {
+    pub fn from(file: RawFile, cluster_size: u64, refcount_bits: u64) -> Option<Self> {
         if !cluster_size.is_power_of_two() {
             return None;
         }
+
+        let (read_refcount_fn, write_refcount_fn): (RefcountReader, RefcountWriter) =
+            match refcount_bits {
+                1 => (read_refcount_subbyte::<1>, write_refcount_subbyte::<1>),
+                2 => (read_refcount_subbyte::<2>, write_refcount_subbyte::<2>),
+                4 => (read_refcount_subbyte::<4>, write_refcount_subbyte::<4>),
+                8 => (read_refcount::<u8>, write_refcount::<u8>),
+                16 => (read_refcount::<u16>, write_refcount::<u16>),
+                32 => (read_refcount::<u32>, write_refcount::<u32>),
+                64 => (read_refcount::<u64>, write_refcount::<u64>),
+                _ => return None,
+            };
+
+        // For sub-byte refcounts (1,2,4 bits), entries pack multiple per byte
+        let refcount_block_entries = cluster_size * 8 / refcount_bits;
+
         Some(QcowRawFile {
             file,
             cluster_size,
             cluster_mask: cluster_size - 1,
+            refcount_block_entries,
+            read_refcount_fn,
+            write_refcount_fn,
         })
     }
 
@@ -109,24 +247,17 @@ impl QcowRawFile {
 
     /// 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>> {
+    #[inline]
-        let count = self.cluster_size / size_of::<u16>() as u64;
+    pub fn read_refcount_block(&mut self, offset: u64) -> io::Result<Vec<u64>> {
-        let mut table = vec![0; count as usize];
         self.file.seek(SeekFrom::Start(offset))?;
-        self.file.read_u16_into::<BigEndian>(&mut table)?;
+        (self.read_refcount_fn)(&mut self.file, self.refcount_block_entries as usize)
-        Ok(table)
     }
 
     /// Writes a refcount block to the file.
-    pub fn write_refcount_block(&mut self, offset: u64, table: &[u16]) -> io::Result<()> {
+    #[inline]
+    pub fn write_refcount_block(&mut self, offset: u64, table: &[u64]) -> io::Result<()> {
         self.file.seek(SeekFrom::Start(offset))?;
-        let mut buffer = BufWriter::with_capacity(std::mem::size_of_val(table), &mut self.file);
+        (self.write_refcount_fn)(&mut self.file, table)
-
-        for count in table {
-            buffer.write_u16::<BigEndian>(*count)?;
-        }
-        buffer.flush()?;
-        Ok(())
     }
 
     /// Allocates a new cluster at the end of the current file, return the address.
@@ -191,6 +322,9 @@ impl Clone for QcowRawFile {
             file: self.file.try_clone().expect("QcowRawFile cloning failed"),
             cluster_size: self.cluster_size,
             cluster_mask: self.cluster_mask,
+            refcount_block_entries: self.refcount_block_entries,
+            read_refcount_fn: self.read_refcount_fn,
+            write_refcount_fn: self.write_refcount_fn,
         }
     }
 }

block/src/qcow/refcount.rs

@@ -38,7 +38,7 @@ pub type Result<T> = std::result::Result<T, Error>;
 pub struct RefCount {
     ref_table: VecCache<u64>,
     refcount_table_offset: u64,
-    refblock_cache: CacheMap<VecCache<u16>>,
+    refblock_cache: CacheMap<VecCache<u64>>,
     refcount_block_entries: u64, // number of refcounts in a cluster.
     cluster_size: u64,
     max_valid_cluster_offset: u64,
@@ -92,8 +92,8 @@ impl RefCount {
         &mut self,
         raw_file: &mut QcowRawFile,
         cluster_address: u64,
-        refcount: u16,
+        refcount: u64,
-        mut new_cluster: Option<(u64, VecCache<u16>)>,
+        mut new_cluster: Option<(u64, VecCache<u64>)>,
     ) -> Result<Option<u64>> {
         let (table_index, block_index) = self.get_refcount_index(cluster_address);
 
@@ -170,7 +170,7 @@ impl RefCount {
         &mut self,
         raw_file: &mut QcowRawFile,
         address: u64,
-    ) -> Result<u16> {
+    ) -> Result<u64> {
         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 {
@@ -202,7 +202,7 @@ impl RefCount {
         &mut self,
         raw_file: &mut QcowRawFile,
         table_index: usize,
-    ) -> Result<Option<&[u16]>> {
+    ) -> Result<Option<&[u64]>> {
         let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
         if block_addr_disk == 0 {
             return Ok(None);