diff --git a/block/src/qcow_common.rs b/block/src/qcow_common.rs index eeea8bf51..195fb48aa 100644 --- a/block/src/qcow_common.rs +++ b/block/src/qcow_common.rs @@ -291,3 +291,114 @@ pub unsafe fn gather_from_iovecs(iovecs: &[libc::iovec], start: usize, len: usiz unsafe { gather_from_iovecs_into(iovecs, start, &mut result) }; result } + +#[cfg(test)] +pub(crate) mod unit_tests { + use std::fs::File; + use std::io::{Read, Seek, SeekFrom, Write}; + + use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; + use flate2::write::DeflateEncoder; + use flate2::Compression; + + const COMPRESSED_FLAG: u64 = 1 << 62; + const CLUSTER_USED_FLAG: u64 = 1 << 63; + const COMPRESSED_SECTOR_SIZE: u64 = 512; + + const HEADER_CLUSTER_BITS_OFFSET: u64 = 20; + const HEADER_L1_SIZE_OFFSET: u64 = 36; + const HEADER_L1_TABLE_OFFSET: u64 = 40; + + const L1_L2_ADDR_MASK: u64 = 0x00ff_ffff_ffff_fe00; + + fn make_compressed_l2_entry(host_offset: u64, compressed_len: usize, cluster_bits: u32) -> u64 { + let compressed_size_shift = 62 - (cluster_bits - 8); + let intra_sector_offset = host_offset & (COMPRESSED_SECTOR_SIZE - 1); + let total_bytes = compressed_len as u64 + intra_sector_offset; + let nsectors = total_bytes.div_ceil(COMPRESSED_SECTOR_SIZE); + let addr_part = host_offset & ((1 << compressed_size_shift) - 1); + let size_part = (nsectors - 1) << compressed_size_shift; + COMPRESSED_FLAG | size_part | addr_part + } + + /// Compress every allocated cluster in a QCOW2 image file in place. + /// + /// Walks L1 -> L2 tables, compresses each standard cluster with raw + /// deflate, appends the compressed payload at the end of the file, + /// and rewrites the L2 entry with the compressed layout. + pub fn compress_allocated_clusters(file: &mut File) { + file.seek(SeekFrom::Start(HEADER_CLUSTER_BITS_OFFSET)) + .unwrap(); + let cluster_bits = file.read_u32::().unwrap(); + let cluster_size = 1u64 << cluster_bits; + + file.seek(SeekFrom::Start(HEADER_L1_SIZE_OFFSET)).unwrap(); + let l1_size = file.read_u32::().unwrap(); + + file.seek(SeekFrom::Start(HEADER_L1_TABLE_OFFSET)).unwrap(); + let l1_table_offset = file.read_u64::().unwrap(); + + let entries_per_l2 = cluster_size / 8; + + let mut append_offset = file.seek(SeekFrom::End(0)).unwrap(); + append_offset = (append_offset + 511) & !511; + + for l1_idx in 0..l1_size as u64 { + let l1_entry_offset = l1_table_offset + l1_idx * 8; + file.seek(SeekFrom::Start(l1_entry_offset)).unwrap(); + let l1_entry = file.read_u64::().unwrap(); + + let l2_table_addr = l1_entry & L1_L2_ADDR_MASK; + if l2_table_addr == 0 { + continue; + } + + for l2_idx in 0..entries_per_l2 { + let l2_entry_offset = l2_table_addr + l2_idx * 8; + file.seek(SeekFrom::Start(l2_entry_offset)).unwrap(); + let l2_entry = file.read_u64::().unwrap(); + + if l2_entry & CLUSTER_USED_FLAG == 0 || l2_entry & COMPRESSED_FLAG != 0 { + continue; + } + + let host_cluster_addr = l2_entry & L1_L2_ADDR_MASK; + if host_cluster_addr == 0 { + continue; + } + + let mut cluster_data = vec![0u8; cluster_size as usize]; + file.seek(SeekFrom::Start(host_cluster_addr)).unwrap(); + file.read_exact(&mut cluster_data).unwrap(); + + let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(&cluster_data).unwrap(); + let compressed = encoder.finish().unwrap(); + + file.seek(SeekFrom::Start(append_offset)).unwrap(); + file.write_all(&compressed).unwrap(); + + // The L2 entry encodes the compressed size in units of + // 512 byte sectors. The reader decodes the sector count + // back and computes: nsectors * 512 - (addr & 511). + // Because addr is 512 aligned, this yields nsectors * 512 + // which rounds up to the next sector boundary. The file + // must contain enough bytes for that rounded up pread. + let padded_len = (compressed.len() + 511) & !511; + if padded_len > compressed.len() { + let padding = vec![0u8; padded_len - compressed.len()]; + file.write_all(&padding).unwrap(); + } + + let new_entry = + make_compressed_l2_entry(append_offset, compressed.len(), cluster_bits); + file.seek(SeekFrom::Start(l2_entry_offset)).unwrap(); + file.write_u64::(new_entry).unwrap(); + + append_offset += padded_len as u64; + } + } + + file.flush().unwrap(); + } +}