From cc6e35d4ee23a1dd61fd324a6f491917a5d4944c Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Wed, 20 May 2026 21:51:58 +0000 Subject: [PATCH] Replace numeric version field with Format enum Header and Image now carry a Format (Lime / AvmlCompressed) instead of a u32 version. The wire encoding (4-byte magic + 4-byte version) is unchanged, so on-disk .lime / .avml files read back identically. API changes: - image::Format enum, re-exported from the crate root - Header { format } and Snapshot::format(Format) replace the integer version field/builder - Image fields format and align_src are now pub(crate); src/dst remain pub for the convert binary's stream manipulation. Image::from_streams is the public constructor for non-file usage - Header::encode is infallible; Header::write no longer propagates a version error - Error::UnimplementedVersion is removed (UnsupportedFormat covers bad-magic / wrong-version pairs from disk) - avml-convert renames its CLI value-enum to CliFormat to avoid collision with image::Format Every format dispatch site (copy_if_nonzero, convert_block, convert_to_raw_image) is now exhaustive on the two variants. --- src/bin/avml-convert.rs | 118 +++++++++++++------------- src/bin/avml.rs | 6 +- src/image.rs | 180 ++++++++++++++++++++++++++-------------- src/lib.rs | 1 + src/snapshot.rs | 18 ++-- 5 files changed, 187 insertions(+), 136 deletions(-) diff --git a/src/bin/avml-convert.rs b/src/bin/avml-convert.rs index 39c97100..a2f32701 100644 --- a/src/bin/avml-convert.rs +++ b/src/bin/avml-convert.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -use avml::{Error, Result, image}; +use avml::{Error, Format, Result, image}; use clap::{Parser, ValueEnum}; use snap::read::FrameDecoder; use std::{ @@ -10,23 +10,22 @@ use std::{ path::{Path, PathBuf}, }; -fn convert(src: &Path, dst: &Path, compress: bool) -> Result<()> { +fn convert(src: &Path, dst: &Path, format: Format) -> Result<()> { let src_len = metadata(src) .map_err(|source| image::Error::Io { context: "unable to read source size", source, })? .len(); - let mut image = image::Image::::new(1, src, dst)?; - convert_image(&mut image, src_len, compress) + let mut image = image::Image::::new(format, src, dst)?; + convert_image(&mut image, src_len) } -fn convert_image(image: &mut image::Image, src_len: u64, compress: bool) -> Result<()> +fn convert_image(image: &mut image::Image, src_len: u64) -> Result<()> where R: Read + Seek, W: Write, { - image.version = if compress { 2 } else { 1 }; loop { let current = image .src @@ -50,7 +49,6 @@ where R: Read + Seek, W: Write + Seek, { - image.version = 1; loop { let current = image .src @@ -82,8 +80,8 @@ where let size = header.size()?; - match header.version { - 1 => { + match header.format { + Format::Lime => { let mut handle = (&mut image.src).take(size.try_into().map_err(image::Error::IntConversion)?); copy(&mut handle, &mut image.dst).map_err(|source| image::Error::Io { @@ -91,7 +89,7 @@ where source, })?; } - 2 => { + Format::AvmlCompressed => { let mut decoder = FrameDecoder::new(&mut image.src) .take(size.try_into().map_err(image::Error::IntConversion)?); copy(&mut decoder, &mut image.dst).map_err(|source| image::Error::Io { @@ -106,7 +104,6 @@ where source, })?; } - _ => return Err(image::Error::UnimplementedVersion.into()), } } @@ -120,21 +117,15 @@ fn convert_to_raw(src: &Path, dst: &Path) -> Result<()> { source, })? .len(); - let mut image = image::Image::::new(1, src, dst)?; + let mut image = image::Image::::new(Format::Lime, src, dst)?; convert_to_raw_image(&mut image, src_len) } -fn encode_raw_image( - image: &mut image::Image, - raw_len: u64, - compress: bool, -) -> Result<()> +fn encode_raw_image(image: &mut image::Image, raw_len: u64) -> Result<()> where R: Read + Seek, W: Write, { - image.version = if compress { 2 } else { 1 }; - let mut start = 0_u64; while start < raw_len { let end = raw_len.min(start.saturating_add(image::MAX_BLOCK_SIZE)); @@ -145,15 +136,15 @@ where Ok(()) } -fn convert_from_raw(src: &Path, dst: &Path, compress: bool) -> Result<()> { +fn convert_from_raw(src: &Path, dst: &Path, format: Format) -> Result<()> { let src_len = metadata(src) .map_err(|source| image::Error::Io { context: "unable to read source size", source, })? .len(); - let mut image = image::Image::::new(1, src, dst)?; - encode_raw_image(&mut image, src_len, compress) + let mut image = image::Image::::new(format, src, dst)?; + encode_raw_image(&mut image, src_len) } #[derive(Parser)] @@ -161,12 +152,12 @@ fn convert_from_raw(src: &Path, dst: &Path, compress: bool) -> Result<()> { #[command(version)] struct Config { /// specify output format - #[arg(long, value_enum, default_value_t = Format::LimeCompressed)] - source_format: Format, + #[arg(long, value_enum, default_value_t = CliFormat::LimeCompressed)] + source_format: CliFormat, /// specify output format - #[arg(long, value_enum, default_value_t = Format::Lime)] - format: Format, + #[arg(long, value_enum, default_value_t = CliFormat::Lime)] + format: CliFormat, /// name of the source file to read to on local system src: PathBuf, @@ -175,8 +166,8 @@ struct Config { dst: PathBuf, } -#[derive(ValueEnum, Clone)] -enum Format { +#[derive(ValueEnum, Clone, Copy, PartialEq, Eq)] +enum CliFormat { Raw, Lime, #[value(rename_all = "snake_case")] @@ -187,33 +178,36 @@ fn main() -> Result<()> { let config = Config::parse(); match (config.source_format, config.format) { - (Format::Lime | Format::LimeCompressed, Format::Raw) => { + (CliFormat::Lime | CliFormat::LimeCompressed, CliFormat::Raw) => { convert_to_raw(&config.src, &config.dst) } - (Format::Lime, Format::LimeCompressed) => convert(&config.src, &config.dst, true), - (Format::LimeCompressed, Format::Lime) => convert(&config.src, &config.dst, false), - (Format::Raw, Format::Lime) => convert_from_raw(&config.src, &config.dst, false), - (Format::Raw, Format::LimeCompressed) => convert_from_raw(&config.src, &config.dst, true), - (Format::Lime, Format::Lime) - | (Format::LimeCompressed, Format::LimeCompressed) - | (Format::Raw, Format::Raw) => Err(Error::NoConversionRequired), + (CliFormat::Lime, CliFormat::LimeCompressed) => { + convert(&config.src, &config.dst, Format::AvmlCompressed) + } + (CliFormat::LimeCompressed, CliFormat::Lime) => { + convert(&config.src, &config.dst, Format::Lime) + } + (CliFormat::Raw, CliFormat::Lime) => { + convert_from_raw(&config.src, &config.dst, Format::Lime) + } + (CliFormat::Raw, CliFormat::LimeCompressed) => { + convert_from_raw(&config.src, &config.dst, Format::AvmlCompressed) + } + (CliFormat::Lime, CliFormat::Lime) + | (CliFormat::LimeCompressed, CliFormat::LimeCompressed) + | (CliFormat::Raw, CliFormat::Raw) => Err(Error::NoConversionRequired), } } #[cfg(test)] mod tests { use crate::{convert_image, convert_to_raw_image, encode_raw_image}; - use avml::{Result, image}; + use avml::{Format, Result, image}; use rand::{Rng as _, SeedableRng as _, rngs::SmallRng}; use std::io::Cursor; - fn memory_image(src: &[u8]) -> image::Image, Cursor>> { - image::Image { - version: 1, - align_src: false, - src: Cursor::new(src), - dst: Cursor::new(Vec::new()), - } + fn memory_image(format: Format, src: &[u8]) -> image::Image, Cursor>> { + image::Image::from_streams(format, Cursor::new(src), Cursor::new(Vec::new())) } fn block_size() -> Result { @@ -249,43 +243,43 @@ mod tests { Ok(chunks.concat()) } - fn encode_raw(raw: &[u8], version: u32) -> Result> { - let mut image = memory_image(raw); + fn encode_raw(raw: &[u8], format: Format) -> Result> { + let mut image = memory_image(format, raw); let total = u64::try_from(raw.len()).map_err(image::Error::IntConversion)?; - encode_raw_image(&mut image, total, version == 2)?; + encode_raw_image(&mut image, total)?; Ok(image.dst.into_inner()) } - fn convert_encoded(encoded: &[u8], compress: bool) -> Result> { + fn convert_encoded(encoded: &[u8], format: Format) -> Result> { let encoded_len = u64::try_from(encoded.len()).map_err(image::Error::IntConversion)?; - let mut image = memory_image(encoded); - convert_image(&mut image, encoded_len, compress)?; + let mut image = memory_image(format, encoded); + convert_image(&mut image, encoded_len)?; Ok(image.dst.into_inner()) } fn decode_to_raw(encoded: &[u8]) -> Result> { let encoded_len = u64::try_from(encoded.len()).map_err(image::Error::IntConversion)?; - let mut image = memory_image(encoded); + let mut image = memory_image(Format::Lime, encoded); convert_to_raw_image(&mut image, encoded_len)?; Ok(image.dst.into_inner()) } - fn header_version(encoded: &[u8]) -> Result { - Ok(image::Header::read(Cursor::new(encoded))?.version) + fn header_format(encoded: &[u8]) -> Result { + Ok(image::Header::read(Cursor::new(encoded))?.format) } #[test] fn convert_sparse_raw_between_lime_and_compressed_formats() -> Result<()> { let raw = build_sparse_raw()?; - let lime = encode_raw(&raw, 1)?; - assert_eq!(header_version(&lime)?, 1); + let lime = encode_raw(&raw, Format::Lime)?; + assert_eq!(header_format(&lime)?, Format::Lime); assert_eq!(decode_to_raw(&lime)?, raw); - let compressed = convert_encoded(&lime, true)?; - assert_eq!(header_version(&compressed)?, 2); + let compressed = convert_encoded(&lime, Format::AvmlCompressed)?; + assert_eq!(header_format(&compressed)?, Format::AvmlCompressed); - let lime_roundtrip = convert_encoded(&compressed, false)?; - assert_eq!(header_version(&lime_roundtrip)?, 1); + let lime_roundtrip = convert_encoded(&compressed, Format::Lime)?; + assert_eq!(header_format(&lime_roundtrip)?, Format::Lime); assert_eq!(lime_roundtrip, lime); assert_eq!(decode_to_raw(&compressed)?, raw); @@ -301,11 +295,11 @@ mod tests { let block_size = block_size()?; raw.extend(vec![0; block_size]); - let lime = encode_raw(&raw, 1)?; + let lime = encode_raw(&raw, Format::Lime)?; assert_eq!(decode_to_raw(&lime)?, expected_raw); - let compressed = convert_encoded(&lime, true)?; - let lime_roundtrip = convert_encoded(&compressed, false)?; + let compressed = convert_encoded(&lime, Format::AvmlCompressed)?; + let lime_roundtrip = convert_encoded(&compressed, Format::Lime)?; assert_eq!(lime_roundtrip, lime); assert_eq!(decode_to_raw(&compressed)?, expected_raw); assert_eq!(decode_to_raw(&lime_roundtrip)?, expected_raw); diff --git a/src/bin/avml.rs b/src/bin/avml.rs index 0eb8e161..27430984 100644 --- a/src/bin/avml.rs +++ b/src/bin/avml.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -use avml::{Result, Snapshot, Source, iomem}; +use avml::{Format, Result, Snapshot, Source, iomem}; use clap::Parser; #[cfg(feature = "blobstore")] use core::num::NonZeroUsize; @@ -111,14 +111,14 @@ async fn upload(config: &Config) -> Result<()> { } fn acquire(config: &Config) -> Result<()> { - let version = if config.compress { 2 } else { 1 }; + let format = Format::from(config.compress); let ranges = iomem::parse()?; let snapshot = Snapshot::new(&config.filename, ranges) .source(config.source.clone()) .max_disk_usage_percentage(config.max_disk_usage_percentage) .max_disk_usage(config.max_disk_usage) - .version(version); + .format(format); snapshot.create()?; Ok(()) } diff --git a/src/image.rs b/src/image.rs index bc254039..da42a3f9 100644 --- a/src/image.rs +++ b/src/image.rs @@ -30,9 +30,6 @@ pub enum Error { #[error("file is too large")] TooLarge, - #[error("unimplemented version")] - UnimplementedVersion, - #[error("unsupported format")] UnsupportedFormat, @@ -49,6 +46,58 @@ pub enum Error { type Result = core::result::Result; +/// On-disk format for a memory snapshot. +/// +/// The wire encoding is unchanged across the two variants: a 32-bit +/// little-endian magic followed by a 32-bit little-endian version. This +/// enum exists so that internal code dispatches on a named format rather +/// than passing around bare integers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Format { + /// `LiME` v1: uncompressed memory blocks with `LiME` headers. + Lime, + /// AVML v2: Snappy-compressed memory blocks with AVML headers. + AvmlCompressed, +} + +impl Format { + const LIME_MAGIC: u32 = 0x4c69_4d45; // "LiME" + const AVML_MAGIC: u32 = 0x4c4d_5641; // "AVML" + + const fn magic(self) -> u32 { + match self { + Self::Lime => Self::LIME_MAGIC, + Self::AvmlCompressed => Self::AVML_MAGIC, + } + } + + const fn version(self) -> u32 { + match self { + Self::Lime => 1, + Self::AvmlCompressed => 2, + } + } + + fn from_wire(magic: u32, version: u32) -> Result { + match (magic, version) { + (Self::LIME_MAGIC, 1) => Ok(Self::Lime), + (Self::AVML_MAGIC, 2) => Ok(Self::AvmlCompressed), + _ => Err(Error::UnsupportedFormat), + } + } +} + +impl From for Format { + /// `true` selects the compressed AVML format; `false` selects `LiME`. + fn from(compress: bool) -> Self { + if compress { + Self::AvmlCompressed + } else { + Self::Lime + } + } +} + /// Largest block AVML emits in a single header. Ranges larger than this /// are split into `MAX_BLOCK_SIZE`-sized chunks before being written. /// @@ -60,13 +109,11 @@ type Result = core::result::Result; pub const MAX_BLOCK_SIZE: u64 = 0x1000 * 0x1000; const PAGE_SIZE: usize = 0x1000; const HEADER_LEN: usize = 32; -const LIME_MAGIC: u32 = 0x4c69_4d45; // EMiL as u32le -const AVML_MAGIC: u32 = 0x4c4d_5641; // AVML as u32le #[derive(Debug, Clone)] pub struct Header { pub range: Range, - pub version: u32, + pub format: Format, } #[derive(PartialEq, Eq, Debug, Clone)] @@ -111,42 +158,36 @@ impl Header { if padding != 0 { return Err(Error::InvalidPadding); } - if !(magic == LIME_MAGIC && version == 1 || magic == AVML_MAGIC && version == 2) { - return Err(Error::UnsupportedFormat); - } + let format = Format::from_wire(magic, version)?; Ok(Self { range: Range { start, end }, - version, + format, }) } - fn encode(&self) -> Result<[u8; 32]> { - let magic = match self.version { - 1 => LIME_MAGIC, - 2 => AVML_MAGIC, - _ => return Err(Error::UnimplementedVersion), - }; + fn encode(&self) -> [u8; HEADER_LEN] { let mut bytes = [0; HEADER_LEN]; - LittleEndian::write_u32_into(&[magic, self.version], &mut bytes[..8]); + LittleEndian::write_u32_into( + &[self.format.magic(), self.format.version()], + &mut bytes[..8], + ); LittleEndian::write_u64_into( &[self.range.start, self.range.end.saturating_sub(1), 0], &mut bytes[8..], ); - Ok(bytes) + bytes } /// Writes the header to the destination writer. /// /// # Errors - /// Returns an error if: - /// - The version is not supported - /// - The header cannot be written to the destination + /// Returns an error if the header cannot be written to the destination. pub fn write(&self, mut dst: W) -> Result<()> where W: Write, { - let bytes = self.encode()?; + let bytes = self.encode(); dst.write_all(&bytes).map_err(|source| Error::Io { context: "unable to write header", source, @@ -208,13 +249,29 @@ where } pub struct Image { - pub version: u32, - pub align_src: bool, + pub(crate) format: Format, + pub(crate) align_src: bool, pub src: R, pub dst: W, } impl Image { + /// Build an `Image` over arbitrary streams. + pub fn from_streams(format: Format, src: R, dst: W) -> Self { + Self { + format, + align_src: false, + src, + dst, + } + } + + /// The destination format this `Image` writes. + #[must_use] + pub fn format(&self) -> Format { + self.format + } + #[cfg(target_family = "windows")] fn open_dst(path: &Path) -> Result { OpenOptions::new() @@ -243,14 +300,15 @@ impl Image { }) } - /// Creates a new Image with the specified version, source filename, and destination filename. + /// Open `src_filename` for reading and `dst_filename` for writing, + /// producing an `Image` that emits the given destination `format`. /// /// # Errors /// Returns an error if: /// - The source file cannot be opened for reading /// - The destination file cannot be created or opened for writing pub fn new( - version: u32, + format: Format, src_filename: &Path, dst_filename: &Path, ) -> Result> { @@ -276,7 +334,7 @@ impl Image { let dst = Self::open_dst(dst_filename)?; Ok(Image:: { - version, + format, align_src, src, dst, @@ -318,7 +376,7 @@ impl Image { fn write_header(&mut self, range: Range) -> Result<()> { Header { range, - version: self.version, + format: self.format, } .write(&mut self.dst) } @@ -378,32 +436,35 @@ impl Image { } self.write_header(range.clone())?; - if self.version == 1 { - self.dst.write_all(&buf).map_err(|source| Error::Io { - context: "unable to write non-zero block", - source, - })?; - } else { - let mut encoder = SnapCountWriter::new(&mut self.dst); - encoder.write_all(&buf).map_err(|source| Error::Io { - context: "unable to write compressed block", - source, - })?; - encoder.finalize().map_err(|source| Error::Io { - context: "unable to finalize compressed block", - source, - })?; + match self.format { + Format::Lime => { + self.dst.write_all(&buf).map_err(|source| Error::Io { + context: "unable to write non-zero block", + source, + })?; + } + Format::AvmlCompressed => { + let mut encoder = SnapCountWriter::new(&mut self.dst); + encoder.write_all(&buf).map_err(|source| Error::Io { + context: "unable to write compressed block", + source, + })?; + encoder.finalize().map_err(|source| Error::Io { + context: "unable to finalize compressed block", + source, + })?; + } } Ok(()) } pub fn convert_block(&mut self) -> Result<()> { let header = self.read_header()?; - match header.version { - 1 => { + match header.format { + Format::Lime => { self.copy_block(header.range)?; } - 2 => { + Format::AvmlCompressed => { self.write_header(header.range.clone())?; { let size = range_len(header.range.clone()); @@ -420,7 +481,6 @@ impl Image { source, })?; } - _ => return Err(Error::UnimplementedVersion), } Ok(()) @@ -437,48 +497,44 @@ fn range_usize(value: Range) -> Result { #[cfg(test)] mod tests { + use super::{Format, Header, Image}; use core::ops::Range; use std::io::Cursor; #[test] - fn encode_header_v1() { + fn encode_header_lime() { let expected = b"\x45\x4d\x69\x4c\x01\x00\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\ \x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; - let header = super::Header { + let header = Header { range: Range { start: 0x1000, end: 0x20001, }, - version: 1, + format: Format::Lime, }; - assert!(matches!(header.encode(), Ok(x) if x == *expected)); + assert_eq!(header.encode(), *expected); } #[test] - fn encode_header_v2() { + fn encode_header_avml() { let expected = b"\x41\x56\x4d\x4c\x02\x00\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\ \x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; - let header = super::Header { + let header = Header { range: Range { start: 0x1000, end: 0x20001, }, - version: 2, + format: Format::AvmlCompressed, }; - assert!(matches!(header.encode(), Ok(x) if x == *expected)); + assert_eq!(header.encode(), *expected); } #[test] fn copy_block_skips_all_zero_ranges() -> super::Result<()> { - for version in [1, 2] { + for format in [Format::Lime, Format::AvmlCompressed] { let src = Cursor::new(vec![0; 0x4000]); let dst = Cursor::new(vec![]); - let mut image = super::Image { - version, - align_src: false, - src, - dst, - }; + let mut image = Image::from_streams(format, src, dst); image.copy_block(0..0x4000)?; assert!(image.dst.get_ref().is_empty()); } diff --git a/src/lib.rs b/src/lib.rs index 4a5dadb1..48ce70b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub use crate::upload::blobstore::{BlobUploader, DEFAULT_CONCURRENCY}; pub use crate::upload::http::put; pub use crate::{ errors::Error, + image::Format, snapshot::{Snapshot, Source}, }; diff --git a/src/snapshot.rs b/src/snapshot.rs index 9fbb5dd4..d5ff6f4b 100644 --- a/src/snapshot.rs +++ b/src/snapshot.rs @@ -5,7 +5,7 @@ use crate::disk_usage; use crate::{ errors::format_error, - image::{Block, Image}, + image::{Block, Format, Image}, }; use clap::ValueEnum; use core::{ @@ -197,7 +197,7 @@ pub struct Snapshot<'a> { source: Option, destination: &'a Path, memory_ranges: Vec>, - version: u32, + format: Format, max_disk_usage: Option, max_disk_usage_percentage: Option, } @@ -205,14 +205,14 @@ pub struct Snapshot<'a> { impl<'a> Snapshot<'a> { /// Create a new memory snapshot. /// - /// The default version implements the `LiME` format. + /// Defaults to the `LiME` format. #[must_use] pub fn new(destination: &'a Path, memory_ranges: Vec>) -> Self { Self { source: None, destination, memory_ranges, - version: 1, + format: Format::Lime, max_disk_usage: None, max_disk_usage_percentage: None, } @@ -246,10 +246,10 @@ impl<'a> Snapshot<'a> { Self { source, ..self } } - /// Specify the version of the snapshot format + /// Specify the snapshot format. #[must_use] - pub fn version(self, version: u32) -> Self { - Self { version, ..self } + pub fn format(self, format: Format) -> Self { + Self { format, ..self } } fn create_source(&self, src: &Source) -> Result<()> { @@ -394,7 +394,7 @@ impl<'a> Snapshot<'a> { } let mut image = - Image::::new(self.version, Path::new("/proc/kcore"), self.destination)?; + Image::::new(self.format, Path::new("/proc/kcore"), self.destination)?; self.check_disk_usage(&image)?; let file = elf::ElfStream::::open_stream(&mut image.src)?; @@ -470,7 +470,7 @@ impl<'a> Snapshot<'a> { }) .collect::>(); - let mut image = Image::::new(self.version, mem, self.destination)?; + let mut image = Image::::new(self.format, mem, self.destination)?; self.check_disk_usage(&image)?; image.write_blocks(&blocks)?;