use std::sync::Arc;
use super::{
DESCRIPTOR,
cache::VhdCache,
constants::DEFAULT_SECTOR_SIZE,
dynamic_header::VhdDynamicHeader,
footer::{VhdDiskType, VhdFooter},
parser::{ParsedVhd, VhdBatLayout, parse},
};
use crate::{
ByteSource, ByteSourceCapabilities, ByteSourceHandle, ByteSourceSeekCost, Error, Result,
SourceHints, images::Image,
};
const MAX_BITMAP_CACHE_ENTRIES: usize = 64;
const MAX_BLOCK_CACHE_ENTRIES: usize = 64;
const BITMAP_CACHE_BUDGET_BYTES: usize = 4 * 1024 * 1024;
const BLOCK_CACHE_BUDGET_BYTES: usize = 64 * 1024 * 1024;
pub struct VhdImage {
source: ByteSourceHandle,
footer: VhdFooter,
dynamic_header: Option<VhdDynamicHeader>,
bat: VhdBatLayout,
parent_locator_paths: Arc<[String]>,
parent_image: Option<ByteSourceHandle>,
bitmap_cache: VhdCache<Vec<u8>>,
block_cache: VhdCache<Vec<u8>>,
}
impl VhdImage {
pub fn open(source: ByteSourceHandle) -> Result<Self> {
let parsed = parse(source.clone())?;
if parsed.footer.disk_type == VhdDiskType::Differential {
return Err(Error::invalid_source_reference(
"differential vhd images require source hints and a related-source resolver".to_string(),
));
}
Self::from_parsed(source, parsed, None)
}
pub fn open_with_hints(source: ByteSourceHandle, _hints: SourceHints<'_>) -> Result<Self> {
let parsed = parse(source.clone())?;
let parent_image = if parsed.footer.disk_type == VhdDiskType::Differential {
let resolver = _hints.resolver().ok_or_else(|| {
Error::invalid_source_reference(
"differential vhd images require a related-source resolver".to_string(),
)
})?;
let identity = _hints.source_identity().ok_or_else(|| {
Error::invalid_source_reference(
"differential vhd images require a source identity hint".to_string(),
)
})?;
let mut candidates = parsed.parent_locator_paths.clone();
if let Some(header) = &parsed.dynamic_header
&& !header.parent_name.is_empty()
{
candidates.push(header.parent_name.clone());
}
let mut parent_resolution = None;
for candidate in candidates {
if let Some(resolution) = resolve_parent_candidate(resolver, identity, &candidate)? {
parent_resolution = Some(resolution);
break;
}
}
let (parent_source, parent_path) = parent_resolution
.ok_or_else(|| Error::not_found("unable to resolve the parent vhd image"))?;
Some(Arc::new(Self::open_with_hints(
parent_source,
SourceHints::new()
.with_resolver(resolver)
.with_source_identity(&crate::SourceIdentity::new(parent_path)),
)?) as ByteSourceHandle)
} else {
None
};
Self::from_parsed(source, parsed, parent_image)
}
fn from_parsed(
source: ByteSourceHandle, parsed: ParsedVhd, parent_image: Option<ByteSourceHandle>,
) -> Result<Self> {
let (bitmap_cache_capacity, block_cache_capacity) = if let Some(header) = &parsed.dynamic_header
{
let bitmap_size = usize::try_from(header.sector_bitmap_size()?)
.map_err(|_| Error::invalid_range("vhd sector bitmap size is too large"))?;
let block_size = usize::try_from(header.block_size)
.map_err(|_| Error::invalid_range("vhd block size is too large"))?;
(
bounded_cache_capacity(
bitmap_size,
BITMAP_CACHE_BUDGET_BYTES,
MAX_BITMAP_CACHE_ENTRIES,
),
bounded_cache_capacity(
block_size,
BLOCK_CACHE_BUDGET_BYTES,
MAX_BLOCK_CACHE_ENTRIES,
),
)
} else {
(1, 1)
};
Ok(Self {
source,
footer: parsed.footer,
dynamic_header: parsed.dynamic_header,
bat: parsed.block_allocation_table,
parent_locator_paths: Arc::from(parsed.parent_locator_paths),
parent_image,
bitmap_cache: VhdCache::new(bitmap_cache_capacity),
block_cache: VhdCache::new(block_cache_capacity),
})
}
pub fn footer(&self) -> &VhdFooter {
&self.footer
}
pub fn dynamic_header(&self) -> Option<&VhdDynamicHeader> {
self.dynamic_header.as_ref()
}
pub fn parent_locator_paths(&self) -> &[String] {
&self.parent_locator_paths
}
fn bat_entry(&self, block_index: u64) -> Result<u32> {
if block_index >= u64::from(self.bat.entry_count) {
return Err(Error::invalid_range(format!(
"vhd block index {block_index} is out of bounds"
)));
}
let entry_offset = self
.bat
.file_offset
.checked_add(
block_index
.checked_mul(4)
.ok_or_else(|| Error::invalid_range("vhd BAT entry offset overflow"))?,
)
.ok_or_else(|| Error::invalid_range("vhd BAT entry offset overflow"))?;
let mut data = [0u8; 4];
self.source.read_exact_at(entry_offset, &mut data)?;
Ok(u32::from_be_bytes(data))
}
fn read_bitmap(&self, block_index: u64) -> Result<Arc<Vec<u8>>> {
let header = self
.dynamic_header
.as_ref()
.ok_or_else(|| Error::invalid_format("vhd dynamic header is missing"))?;
let sector_bitmap_size = usize::try_from(header.sector_bitmap_size()?)
.map_err(|_| Error::invalid_range("vhd sector bitmap size is too large"))?;
let bat_entry = self.bat_entry(block_index)?;
if bat_entry == u32::MAX {
return Err(Error::invalid_format(
"vhd sparse block bitmap was requested".to_string(),
));
}
let bitmap_offset = u64::from(bat_entry)
.checked_mul(u64::from(DEFAULT_SECTOR_SIZE))
.ok_or_else(|| Error::invalid_range("vhd bitmap offset overflow"))?;
self.bitmap_cache.get_or_load(block_index, || {
let data = self
.source
.read_bytes_at(bitmap_offset, sector_bitmap_size)?;
Ok(Arc::new(data))
})
}
fn read_block(&self, block_index: u64) -> Result<Option<Arc<Vec<u8>>>> {
let header = self
.dynamic_header
.as_ref()
.ok_or_else(|| Error::invalid_format("vhd dynamic header is missing"))?;
let bat_entry = self.bat_entry(block_index)?;
if bat_entry == u32::MAX {
return Ok(None);
}
let sector_bitmap_size = header.sector_bitmap_size()?;
let block_offset = u64::from(bat_entry)
.checked_mul(u64::from(DEFAULT_SECTOR_SIZE))
.and_then(|offset| offset.checked_add(sector_bitmap_size))
.ok_or_else(|| Error::invalid_range("vhd block offset overflow"))?;
let block_size = usize::try_from(header.block_size)
.map_err(|_| Error::invalid_range("vhd block size is too large"))?;
self
.block_cache
.get_or_load(block_index, || {
let data = self.source.read_bytes_at(block_offset, block_size)?;
Ok(Arc::new(data))
})
.map(Some)
}
}
fn bounded_cache_capacity(entry_size: usize, byte_budget: usize, max_entries: usize) -> usize {
if entry_size == 0 || max_entries == 0 {
return 1;
}
(byte_budget / entry_size).max(1).min(max_entries)
}
impl ByteSource for VhdImage {
fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
if offset >= self.footer.current_size || buf.is_empty() {
return Ok(0);
}
if self.footer.disk_type == VhdDiskType::Fixed {
let max = usize::try_from(self.footer.current_size - offset)
.map_err(|_| Error::invalid_range("vhd remaining size is too large"))?
.min(buf.len());
return self.source.read_at(offset, &mut buf[..max]);
}
let header = self
.dynamic_header
.as_ref()
.ok_or_else(|| Error::invalid_format("vhd dynamic header is missing"))?;
let block_size = u64::from(header.block_size);
let mut copied = 0usize;
while copied < buf.len() {
let absolute_offset = offset
.checked_add(copied as u64)
.ok_or_else(|| Error::invalid_range("vhd read offset overflow"))?;
if absolute_offset >= self.footer.current_size {
break;
}
let block_index = absolute_offset / block_size;
let within_block = absolute_offset % block_size;
let block_available = usize::try_from(
block_size
.checked_sub(within_block)
.ok_or_else(|| Error::invalid_range("vhd block range underflow"))?,
)
.map_err(|_| Error::invalid_range("vhd block size is too large"))?
.min(
usize::try_from(self.footer.current_size - absolute_offset)
.map_err(|_| Error::invalid_range("vhd remaining size is too large"))?,
);
let sector_available = usize::try_from(
u64::from(DEFAULT_SECTOR_SIZE) - (within_block % u64::from(DEFAULT_SECTOR_SIZE)),
)
.map_err(|_| Error::invalid_range("vhd sector size is too large"))?;
let available = block_available
.min(buf.len() - copied)
.min(sector_available);
let block = match self.read_block(block_index)? {
Some(block) => block,
None => {
if let Some(parent_image) = &self.parent_image {
parent_image.read_exact_at(absolute_offset, &mut buf[copied..copied + available])?;
} else {
buf[copied..copied + available].fill(0);
}
copied += available;
continue;
}
};
let bitmap = self.read_bitmap(block_index)?;
let sector_index = usize::try_from(within_block / u64::from(DEFAULT_SECTOR_SIZE))
.map_err(|_| Error::invalid_range("vhd sector index overflow"))?;
let sector_present = {
let byte = *bitmap.get(sector_index / 8).ok_or_else(|| {
Error::invalid_format("vhd sector bitmap does not cover the requested sector")
})?;
let bit = 7 - (sector_index % 8);
(byte & (1 << bit)) != 0
};
if !sector_present {
if let Some(parent_image) = &self.parent_image {
parent_image.read_exact_at(absolute_offset, &mut buf[copied..copied + available])?;
} else {
buf[copied..copied + available].fill(0);
}
copied += available;
continue;
}
let block_offset = usize::try_from(within_block)
.map_err(|_| Error::invalid_range("vhd block offset is too large"))?;
buf[copied..copied + available]
.copy_from_slice(&block[block_offset..block_offset + available]);
copied += available;
}
Ok(copied)
}
fn size(&self) -> Result<u64> {
Ok(self.footer.current_size)
}
fn capabilities(&self) -> ByteSourceCapabilities {
let mut capabilities = ByteSourceCapabilities::concurrent(ByteSourceSeekCost::Cheap);
if let Some(header) = &self.dynamic_header
&& let Ok(size) = usize::try_from(header.block_size)
{
capabilities = capabilities.with_preferred_chunk_size(size);
}
capabilities
}
fn telemetry_name(&self) -> &'static str {
"image.vhd"
}
}
impl Image for VhdImage {
fn descriptor(&self) -> crate::FormatDescriptor {
DESCRIPTOR
}
fn logical_sector_size(&self) -> Option<u32> {
Some(DEFAULT_SECTOR_SIZE)
}
fn is_sparse(&self) -> bool {
matches!(
self.footer.disk_type,
VhdDiskType::Dynamic | VhdDiskType::Differential
)
}
fn has_backing_chain(&self) -> bool {
self.parent_image.is_some()
}
}
fn resolve_parent_candidate(
resolver: &dyn crate::RelatedSourceResolver, identity: &crate::SourceIdentity, candidate: &str,
) -> Result<Option<(ByteSourceHandle, crate::RelatedPathBuf)>> {
if let Ok(relative) = crate::RelatedPathBuf::from_relative_path(candidate)
&& let Some(parent) = identity.logical_path().parent()
{
let joined = parent.join(&relative);
if let Some(source) = resolver.resolve(&crate::RelatedSourceRequest::new(
crate::RelatedSourcePurpose::BackingFile,
joined.clone(),
))? {
return Ok(Some((source, joined)));
}
}
let file_name = candidate.rsplit(['\\', '/']).next().unwrap_or(candidate);
let sibling = identity.sibling_path(file_name)?;
Ok(
resolver
.resolve(&crate::RelatedSourceRequest::new(
crate::RelatedSourcePurpose::BackingFile,
sibling.clone(),
))?
.map(|source| (source, sibling)),
)
}
#[cfg(test)]
mod tests {
use std::{collections::HashMap, path::Path, sync::Arc};
use super::*;
use crate::{RelatedSourceRequest, RelatedSourceResolver, SourceIdentity};
struct MemDataSource {
data: Vec<u8>,
}
impl ByteSource for MemDataSource {
fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
let offset = offset as usize;
if offset >= self.data.len() {
return Ok(0);
}
let read = buf.len().min(self.data.len() - offset);
buf[..read].copy_from_slice(&self.data[offset..offset + read]);
Ok(read)
}
fn size(&self) -> Result<u64> {
Ok(self.data.len() as u64)
}
}
struct Resolver {
files: HashMap<String, ByteSourceHandle>,
}
impl RelatedSourceResolver for Resolver {
fn resolve(&self, request: &RelatedSourceRequest) -> Result<Option<ByteSourceHandle>> {
Ok(self.files.get(&request.path.to_string()).cloned())
}
}
fn sample_source(relative_path: &str) -> ByteSourceHandle {
let path = Path::new(env!("CARGO_MANIFEST_DIR"))
.join("formats")
.join(relative_path);
Arc::new(MemDataSource {
data: std::fs::read(path).unwrap(),
})
}
#[test]
fn scales_cache_capacity_to_the_block_budget() {
assert_eq!(
bounded_cache_capacity(2 * 1024 * 1024, BLOCK_CACHE_BUDGET_BYTES, 64),
32
);
assert_eq!(
bounded_cache_capacity(128 * 1024 * 1024, BLOCK_CACHE_BUDGET_BYTES, 64),
1
);
}
#[test]
fn opens_dynamic_vhd_metadata() {
let image = VhdImage::open(sample_source("vhd/ext2.vhd")).unwrap();
assert_eq!(image.footer().disk_type, VhdDiskType::Dynamic);
assert_eq!(image.size().unwrap(), 4_212_736);
assert_eq!(image.dynamic_header().unwrap().block_count, 3);
}
#[test]
fn reads_full_ext2_dynamic_vhd_fixture() {
let image = VhdImage::open(sample_source("vhd/ext2.vhd")).unwrap();
let mut raw = std::fs::read(
Path::new(env!("CARGO_MANIFEST_DIR"))
.join("formats")
.join("ext/ext2.raw"),
)
.unwrap();
raw.resize(image.size().unwrap() as usize, 0);
assert_eq!(image.read_all().unwrap(), raw);
}
#[test]
fn reads_full_ntfs_dynamic_vhd_fixture() {
let image = VhdImage::open(sample_source("vhd/ntfs-dynamic.vhd")).unwrap();
let mut mbr_signature = [0u8; 2];
let mut ntfs_oem = [0u8; 8];
image.read_exact_at(510, &mut mbr_signature).unwrap();
image.read_exact_at(128 * 512 + 3, &mut ntfs_oem).unwrap();
assert_eq!(&mbr_signature, &[0x55, 0xAA]);
assert_eq!(&ntfs_oem, b"NTFS ");
}
#[test]
fn reads_differential_vhd_via_parent_resolution() {
let child = sample_source("vhd/ntfs-differential.vhd");
let parent = sample_source("vhd/ntfs-parent.vhd");
let resolver = Resolver {
files: HashMap::from([("vhd/ntfs-parent.vhd".to_string(), parent)]),
};
let identity = SourceIdentity::from_relative_path("vhd/ntfs-differential.vhd").unwrap();
let image = VhdImage::open_with_hints(
child,
SourceHints::new()
.with_resolver(&resolver)
.with_source_identity(&identity),
)
.unwrap();
let mut mbr_signature = [0u8; 2];
let mut ntfs_oem = [0u8; 8];
image.read_exact_at(510, &mut mbr_signature).unwrap();
image.read_exact_at(128 * 512 + 3, &mut ntfs_oem).unwrap();
assert!(image.has_backing_chain());
assert_eq!(&mbr_signature, &[0x55, 0xAA]);
assert_eq!(&ntfs_oem, b"NTFS ");
}
}
crate::images::driver::impl_image_data_source!(VhdImage);