fstool 0.4.30

//! qcow2 refcount table + refcount blocks.
//!
//! The refcount table is an array of u64 BE pointers to refcount
//! *blocks*. Each refcount block carries an array of refcount entries
//! — one per physical cluster in the backing file. fstool only supports
//! `refcount_order = 4`, i.e. 16-bit refcounts (`cluster_size / 2`
//! entries per block).
//!
//! Allocation strategy (modify path):
//!
//! 1. Scan existing refcount blocks for a free entry (`refcount == 0`)
//!    starting at a hint cursor.
//! 2. If nothing free, grow the file: pick the next cluster past EOF,
//!    set its refcount to 1, and return it. If the refcount block that
//!    would track the new cluster doesn't exist yet, allocate THE
//!    refcount block at that same range (it tracks itself).
//!
//! The strategy assumes the refcount table itself never needs to grow.
//! That holds for any realistic image size: 64 KiB clusters give
//! 8192 table entries × 32768 cluster-refcounts per block = 256 M
//! clusters = 16 PiB of addressable storage.

use std::collections::HashMap;
use std::io::{Read, Seek, SeekFrom, Write};

use crate::Result;

use super::header::Header;

/// Bytes per refcount entry (we only support `refcount_order=4` → u16).
pub const REFCOUNT_ENTRY_BYTES: u64 = 2;

/// In-memory refcount management state.
pub struct Refcount {
    pub cluster_size: u64,
    pub table_offset: u64,
    pub table_clusters: u32,
    /// Pointers to refcount blocks (BE u64 on disk, native u64 here).
    /// `table.len() = table_clusters * cluster_size / 8`.
    pub table: Vec<u64>,
    /// Cached refcount blocks, keyed by physical block offset.
    pub block_cache: HashMap<u64, RefcountBlock>,
    pub block_cache_cap: usize,
    /// Hint: start the free-cluster scan from this cluster index.
    pub next_free_hint: u64,
    pub table_dirty: bool,
}

pub struct RefcountBlock {
    pub entries: Vec<u16>,
    pub dirty: bool,
}

impl Refcount {
    /// Entries per refcount block (`cluster_size / 2` for 16-bit refs).
    pub fn entries_per_block(&self) -> u64 {
        self.cluster_size / REFCOUNT_ENTRY_BYTES
    }

    /// Total table capacity (`table_clusters × cluster_size / 8`).
    pub fn table_capacity(&self) -> u64 {
        self.table_clusters as u64 * self.cluster_size / 8
    }

    /// Load the refcount table from disk. The refcount blocks themselves
    /// are loaded lazily on demand.
    pub fn load<F: Read + Seek>(file: &mut F, header: &Header) -> Result<Self> {
        let cluster_size = header.cluster_size();

        // `refcount_table_clusters` is attacker-controlled and the table is
        // allocated up front (`table_clusters * cluster_size` bytes) before any
        // bounds-checked read. Cap it before allocating: the table can never
        // legitimately be larger than the file that holds it.
        let table_bytes = (header.refcount_table_clusters as u64)
            .checked_mul(cluster_size)
            .ok_or_else(|| {
                crate::Error::InvalidImage(
                    "qcow2: refcount_table_clusters * cluster_size overflows".into(),
                )
            })?;
        let file_len = file.seek(SeekFrom::End(0))?;
        if table_bytes > file_len {
            return Err(crate::Error::InvalidImage(format!(
                "qcow2: refcount table ({table_bytes} bytes) exceeds file length {file_len}"
            )));
        }
        let table_entries = (table_bytes / 8) as usize;
        let mut raw = vec![0u8; table_entries * 8];
        file.seek(SeekFrom::Start(header.refcount_table_offset))?;
        file.read_exact(&mut raw)?;
        let table: Vec<u64> = raw
            .as_chunks::<8>()
            .0
            .iter()
            .map(|&c| u64::from_be_bytes(c))
            .collect();
        Ok(Self {
            cluster_size,
            table_offset: header.refcount_table_offset,
            table_clusters: header.refcount_table_clusters,
            table,
            block_cache: HashMap::new(),
            block_cache_cap: 32,
            next_free_hint: 0,
            table_dirty: false,
        })
    }

    /// Build a fresh in-memory refcount state for a brand-new image.
    /// The table and the (single, empty) refcount block exist only in
    /// memory; the caller writes them to disk at format time.
    ///
    /// `initial_clusters` is the set of physical cluster indices that
    /// should start with refcount=1 — i.e. the header / refcount table /
    /// refcount block / L1 clusters that the new image's layout uses.
    /// The freshly-created image seeds exactly one refcount block, so the
    /// whole initial layout has to fit inside it — `cluster_size / 2`
    /// clusters. Later growth allocates further blocks on demand
    /// ([`alloc_cluster`](Self::alloc_cluster)), so this bound binds only
    /// at creation. A layout that overflows it is an error rather than
    /// an out-of-bounds write.
    pub fn new_fresh(
        cluster_size: u64,
        table_offset: u64,
        first_refcount_block_offset: u64,
        initial_clusters: &[u64],
    ) -> Result<Self> {
        let table_capacity = cluster_size / 8;
        let mut table = vec![0u64; table_capacity as usize];
        table[0] = first_refcount_block_offset;
        let entries_per_block = cluster_size / REFCOUNT_ENTRY_BYTES;
        let mut entries = vec![0u16; entries_per_block as usize];
        for &c in initial_clusters {
            if c >= entries_per_block {
                return Err(crate::Error::Unsupported(format!(
                    "qcow2: the image's fixed layout reaches cluster {c}, past the \
                     {entries_per_block} one refcount block covers at cluster_size \
                     {cluster_size} — use a larger cluster size"
                )));
            }
            entries[c as usize] = 1;
        }
        let mut block_cache = HashMap::new();
        block_cache.insert(
            first_refcount_block_offset,
            RefcountBlock {
                entries,
                dirty: true,
            },
        );
        Ok(Self {
            cluster_size,
            table_offset,
            table_clusters: 1,
            table,
            block_cache,
            block_cache_cap: 32,
            next_free_hint: initial_clusters.iter().copied().max().unwrap_or(0) + 1,
            table_dirty: true,
        })
    }

    fn load_block<F: Read + Seek>(
        &mut self,
        file: &mut F,
        block_off: u64,
    ) -> Result<&mut RefcountBlock> {
        if !self.block_cache.contains_key(&block_off) {
            file.seek(SeekFrom::Start(block_off))?;
            let mut raw = vec![0u8; self.cluster_size as usize];
            file.read_exact(&mut raw)?;
            let entries: Vec<u16> = raw
                .as_chunks::<2>()
                .0
                .iter()
                .map(|&c| u16::from_be_bytes(c))
                .collect();
            // Drop a non-dirty cache entry if we're at the cap.
            if self.block_cache.len() >= self.block_cache_cap
                && let Some(k) = self
                    .block_cache
                    .iter()
                    .find(|(_, v)| !v.dirty)
                    .map(|(k, _)| *k)
            {
                self.block_cache.remove(&k);
            }
            self.block_cache.insert(
                block_off,
                RefcountBlock {
                    entries,
                    dirty: false,
                },
            );
        }
        Ok(self.block_cache.get_mut(&block_off).unwrap())
    }

    /// Allocate one physical cluster: find a free entry in an existing
    /// refcount block, bump it to 1, and return the cluster index. If no
    /// existing block has space, extend the file by allocating the next
    /// cluster past EOF (and a new refcount block for it if needed).
    /// Updates `file_len` to reflect any growth.
    pub fn alloc_cluster<F: Read + Write + Seek>(
        &mut self,
        file: &mut F,
        file_len: &mut u64,
    ) -> Result<u64> {
        // Pass 1: scan existing refcount blocks for refcount==0.
        let hint_block = self.next_free_hint / self.entries_per_block();
        let table_len = self.table.len();
        for sweep in 0..2 {
            // sweep 0: from hint to end; sweep 1: from 0 to hint.
            let (lo, hi) = if sweep == 0 {
                (hint_block as usize, table_len)
            } else {
                (0, hint_block as usize)
            };
            for block_idx in lo..hi {
                let block_off = self.table[block_idx];
                if block_off == 0 {
                    continue;
                }
                let entries_per_block = self.entries_per_block();
                let _block = self.load_block(file, block_off)?;
                // Scan starting from the hint within this block.
                let start_in_block = if sweep == 0 && block_idx == hint_block as usize {
                    (self.next_free_hint % entries_per_block) as usize
                } else {
                    0
                };
                let block = self.block_cache.get_mut(&block_off).unwrap();
                for (i, slot) in block.entries.iter_mut().enumerate().skip(start_in_block) {
                    if *slot == 0 {
                        *slot = 1;
                        block.dirty = true;
                        let cluster = block_idx as u64 * entries_per_block + i as u64;
                        self.next_free_hint = cluster + 1;
                        return Ok(cluster);
                    }
                }
            }
        }
        // Pass 2: grow the file. Start at the first cluster wholly past the
        // current EOF: when the file length is not cluster-aligned (a
        // compressed image whose last payload ends mid-cluster, say) the
        // cluster straddling EOF is in use and must not be handed out.
        let mut new_cluster = file_len.div_ceil(self.cluster_size);
        let entries_per_block = self.entries_per_block();
        loop {
            let block_idx = (new_cluster / entries_per_block) as usize;
            if block_idx >= table_len {
                return Err(crate::Error::Unsupported(
                    "qcow2: refcount table is full (image would exceed 16 PiB at default \
                     cluster size)"
                        .into(),
                ));
            }
            let block_off = self.table[block_idx];
            if block_off == 0 {
                // Need to allocate a new refcount block. It lives at the
                // next cluster past EOF (`new_cluster`), tracks itself, and
                // also tracks the data cluster we're handing out
                // (`new_cluster + 1`).
                let rcb_cluster = new_cluster;
                let data_cluster = new_cluster + 1;
                let next_block_idx = (data_cluster / entries_per_block) as usize;
                if next_block_idx != block_idx {
                    // The data cluster falls into the next refcount block —
                    // an unlikely edge at 2 GiB boundaries that we don't
                    // handle in v1.
                    return Err(crate::Error::Unsupported(
                        "qcow2: allocation across refcount-block boundary not implemented".into(),
                    ));
                }
                let rcb_off = rcb_cluster * self.cluster_size;
                let data_off = data_cluster * self.cluster_size;
                let mut entries = vec![0u16; entries_per_block as usize];
                entries[(rcb_cluster % entries_per_block) as usize] = 1;
                entries[(data_cluster % entries_per_block) as usize] = 1;
                self.block_cache.insert(
                    rcb_off,
                    RefcountBlock {
                        entries,
                        dirty: true,
                    },
                );
                self.table[block_idx] = rcb_off;
                self.table_dirty = true;
                *file_len = data_off + self.cluster_size;
                self.next_free_hint = data_cluster + 1;
                return Ok(data_cluster);
            }
            // Existing refcount block: the cluster past EOF should be free,
            // but an image may legitimately reference clusters beyond its
            // current length (a truncated preallocation, say). Never hand
            // out a cluster something already points at — move past it.
            let _block = self.load_block(file, block_off)?;
            let block = self.block_cache.get_mut(&block_off).unwrap();
            let idx = (new_cluster % entries_per_block) as usize;
            if block.entries[idx] != 0 {
                new_cluster += 1;
                continue;
            }
            block.entries[idx] = 1;
            block.dirty = true;
            *file_len = (new_cluster + 1) * self.cluster_size;
            self.next_free_hint = new_cluster + 1;
            return Ok(new_cluster);
        }
    }

    /// Decrement the refcount of every host cluster overlapping the byte
    /// range `[start_byte, start_byte + len)` by one, freeing (refcount → 0)
    /// any that reach zero. Used when copying a compressed cluster out: each
    /// compressed cluster contributes one reference to every host cluster its
    /// (byte-granular) payload touches, so several compressed clusters can
    /// share a host cluster with refcount > 1 — decrementing keeps that exact.
    pub fn release_range<F: Read + Write + Seek>(
        &mut self,
        file: &mut F,
        start_byte: u64,
        len: u64,
    ) -> Result<()> {
        if len == 0 {
            return Ok(());
        }
        let first = start_byte / self.cluster_size;
        let last = (start_byte + len - 1) / self.cluster_size;
        let epb = self.entries_per_block();
        for cluster in first..=last {
            let block_idx = (cluster / epb) as usize;
            if block_idx >= self.table.len() {
                continue;
            }
            let block_off = self.table[block_idx];
            if block_off == 0 {
                continue;
            }
            self.load_block(file, block_off)?;
            let block = self.block_cache.get_mut(&block_off).unwrap();
            let i = (cluster % epb) as usize;
            if block.entries[i] > 0 {
                block.entries[i] -= 1;
                block.dirty = true;
                if block.entries[i] == 0 && cluster < self.next_free_hint {
                    self.next_free_hint = cluster; // make it reusable
                }
            }
        }
        Ok(())
    }

    /// Write every dirty refcount block back, then the refcount table.
    pub fn flush<F: Write + Seek>(&mut self, file: &mut F) -> Result<()> {
        for (&off, block) in self.block_cache.iter_mut() {
            if !block.dirty {
                continue;
            }
            let mut raw = vec![0u8; self.cluster_size as usize];
            for (i, &e) in block.entries.iter().enumerate() {
                raw[i * 2..i * 2 + 2].copy_from_slice(&e.to_be_bytes());
            }
            file.seek(SeekFrom::Start(off))?;
            file.write_all(&raw)?;
            block.dirty = false;
        }
        if self.table_dirty {
            let mut raw = vec![0u8; self.table.len() * 8];
            for (i, &e) in self.table.iter().enumerate() {
                raw[i * 8..i * 8 + 8].copy_from_slice(&e.to_be_bytes());
            }
            file.seek(SeekFrom::Start(self.table_offset))?;
            file.write_all(&raw)?;
            self.table_dirty = false;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn fresh_state_marks_initial_clusters() {
        // Pretend cluster_size=512 to keep the test tiny; we just check
        // the bookkeeping. Initial clusters [0,1,2,3] are reserved.
        let r = Refcount::new_fresh(512, 512, 1024, &[0, 1, 2, 3]).unwrap();
        assert_eq!(r.table[0], 1024);
        let block = r.block_cache.get(&1024).unwrap();
        assert_eq!(block.entries[0], 1);
        assert_eq!(block.entries[1], 1);
        assert_eq!(block.entries[2], 1);
        assert_eq!(block.entries[3], 1);
        assert_eq!(block.entries[4], 0);
        assert_eq!(r.next_free_hint, 4);
    }

    #[test]
    fn alloc_picks_first_free_in_existing_block() {
        let mut r = Refcount::new_fresh(512, 512, 1024, &[0, 1, 2, 3]).unwrap();
        let mut file_len = 4 * 512;
        // No real I/O needed since the block is already in cache.
        let mut buf: Vec<u8> = Vec::new();
        let mut cur = Cursor::new(&mut buf);
        let c = r.alloc_cluster(&mut cur, &mut file_len).unwrap();
        assert_eq!(c, 4);
        let block = r.block_cache.get(&1024).unwrap();
        assert_eq!(block.entries[4], 1);
    }

    /// With every tracked cluster in use and a file length that is not
    /// cluster-aligned, the allocator must not hand out the cluster that
    /// straddles EOF — it is (partially) present and referenced — but the
    /// first one wholly past it, allocating a refcount block for it.
    #[test]
    fn grow_skips_the_cluster_straddling_eof() {
        let mut r = Refcount::new_fresh(512, 512, 1024, &[0, 1, 2, 3]).unwrap();
        // cluster_size 512 → one refcount block covers 256 clusters; mark
        // them all used so the scan in pass 1 finds nothing.
        r.block_cache.get_mut(&1024).unwrap().entries.fill(1);
        // The file ends 100 bytes into cluster 255.
        let mut file_len = 255 * 512 + 100;
        let mut buf: Vec<u8> = Vec::new();
        let mut cur = Cursor::new(&mut buf);
        let c = r.alloc_cluster(&mut cur, &mut file_len).unwrap();
        // Cluster 256 becomes the refcount block tracking the new range,
        // 257 is the data cluster; 255 stays untouched.
        assert_eq!(c, 257);
        assert_eq!(r.table[1], 256 * 512);
        assert_eq!(file_len, 258 * 512);
        let block = r.block_cache.get(&(256 * 512)).unwrap();
        assert_eq!(block.entries[0], 1);
        assert_eq!(block.entries[1], 1);
        assert_eq!(block.entries[2], 0);
    }

    /// A layout reaching past the single seeded refcount block must be a
    /// clean error, not an out-of-bounds write.
    #[test]
    fn rejects_a_layout_past_the_first_refcount_block() {
        // cluster_size 512 → one refcount block covers 256 clusters.
        assert!(Refcount::new_fresh(512, 512, 1024, &[0, 1, 255]).is_ok());
        let Err(err) = Refcount::new_fresh(512, 512, 1024, &[0, 1, 256]) else {
            panic!("a layout past the first refcount block should be refused");
        };
        assert!(matches!(err, crate::Error::Unsupported(_)), "{err}");
    }
}