split up large segments in compressed images (#8)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "avml"
-version = "0.1.5"
+version = "0.2.0"
 license = "MIT"
 description = "A portable volatile memory acquisition tool"
 authors = ["[email protected]"]

README.md

@@ -88,7 +88,7 @@ avml-convert ./compressed.lime ./uncompressed.lime
 
 ## To compress an uncompressed LiME image
 ```
-avml-convert --compress ./uncompressed.lime ./compressed.lime
+avml-convert --format lime_compressed ./uncompressed.lime ./compressed.lime
 ```
 
 # Usage

src/bin/avml-convert.rs

@@ -8,8 +8,10 @@ extern crate byteorder;
 extern crate elf;
 extern crate snap;
 
+use avml::ONE_MB;
 use clap::{App, Arg};
 use snap::Reader;
+use std::convert::TryFrom;
 use std::error::Error;
 use std::fs::metadata;
 use std::io::prelude::*;
@@ -31,11 +33,11 @@ fn convert(src: String, dst: String, compress: bool) -> Result<(), Box<dyn Error
 
         match header.version {
             1 => {
-                avml::image::copy_block(&new_header, &mut image.src, &mut image.dst)?;
+                avml::image::copy_block(new_header, &mut image.src, &mut image.dst)?;
             }
             2 => {
                 let mut reader = Reader::new(&image.src);
-                avml::image::copy_block(&new_header, &mut reader, &mut image.dst)?;
+                avml::image::copy_block(new_header, &mut reader, &mut image.dst)?;
                 image.src.seek(SeekFrom::Current(8))?;
             }
             _ => unimplemented!(),
@@ -45,15 +47,70 @@ fn convert(src: String, dst: String, compress: bool) -> Result<(), Box<dyn Error
     Ok(())
 }
 
-fn run_app() -> Result<(), Box<dyn Error>> {
+fn convert_to_raw(src: String, dst: String) -> Result<(), Box<dyn Error>> {
+    let src_len = metadata(&src)?.len();
+    let mut image = avml::image::Image::new(1, &src, &dst)?;
+
+    loop {
+        let current = image.src.seek(SeekFrom::Current(0))?;
+        if current >= src_len {
+            break;
+        }
+        let current_dst = image.dst.seek(SeekFrom::Current(0))?;
+
+        let header = avml::image::Header::read(&image.src)?;
+        let mut zeros = vec![0; ONE_MB];
+
+        let mut unmapped = usize::try_from(header.range.start - current_dst)?;
+        while unmapped > ONE_MB {
+            image.dst.write_all(&zeros)?;
+            unmapped -= ONE_MB;
+        }
+        if unmapped > 0 {
+            zeros.resize(unmapped, 0);
+            image.dst.write_all(&zeros)?;
+        }
+
+        let size = usize::try_from(header.range.end - header.range.start)?;
+
+        match header.version {
+            1 => {
+                avml::image::copy(size, &mut image.src, &mut image.dst)?;
+            }
+            2 => {
+                let mut reader = Reader::new(&image.src);
+                avml::image::copy(size, &mut reader, &mut image.dst)?;
+                image.src.seek(SeekFrom::Current(8))?;
+            }
+            _ => unimplemented!(),
+        }
+    }
+
+    Ok(())
+}
+
+arg_enum! {
+    #[allow(non_camel_case_types)]
+    pub enum OutputFormat {
+        raw,
+        lime,
+        lime_compressed
+    }
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let default_format = format!("{}", OutputFormat::lime);
     let args = App::new("avml-convert")
         .author(crate_authors!())
         .about("AVML compress/decompress tool")
         .version(crate_version!())
         .args(&[
-            Arg::with_name("compress")
-                .long("compress")
-                .help("compress pages via snappy"),
+            Arg::with_name("format")
+                .long("format")
+                .help("output format")
+                .takes_value(true)
+                .default_value(&default_format)
+                .possible_values(&OutputFormat::variants()),
             Arg::with_name("source")
                 .help("name of the source file to read to on local system")
                 .required(true),
@@ -63,20 +120,14 @@ fn run_app() -> Result<(), Box<dyn Error>> {
         ])
         .get_matches();
 
-    let compress = args.is_present("compress");
     let src = value_t!(args.value_of("source"), String)?;
     let dst = value_t!(args.value_of("destination"), String)?;
 
-    convert(src, dst, compress)?;
-    Ok(())
-}
+    let format = value_t!(args.value_of("format"), OutputFormat)?;
 
-fn main() {
-    ::std::process::exit(match run_app() {
-        Ok(_) => 0,
-        Err(err) => {
-            eprintln!("error: {:?}", err);
-            1
-        }
-    });
+    match format {
+        OutputFormat::raw => convert_to_raw(src, dst),
+        OutputFormat::lime => convert(src, dst, false),
+        OutputFormat::lime_compressed => convert(src, dst, true),
+    }
 }

src/bin/avml.rs

@@ -5,6 +5,9 @@
 extern crate clap;
 extern crate avml;
 
+#[cfg(feature = "blobstore")]
+use avml::ONE_MB;
+
 use clap::{App, Arg};
 use std::error::Error;
 #[cfg(any(feature = "blobstore", feature = "put"))]
@@ -100,7 +103,7 @@ fn get_mem(src: Option<&str>, dst: &str, version: u32) -> Result<(), Box<dyn Err
     Err(From::from("unable to read physical memory"))
 }
 
-fn run_app() -> Result<(), Box<dyn Error>> {
+fn main() -> Result<(), Box<dyn Error>> {
     let sources = vec!["/proc/kcore", "/dev/crash", "/dev/mem"];
     let args = App::new(crate_name!())
         .author(crate_authors!())
@@ -169,8 +172,8 @@ fn run_app() -> Result<(), Box<dyn Error>> {
             value_t!(args.value_of("sas_block_size"), usize)?
         } else {
             100
-        } * 1024
-            * 1024;
+        } * ONE_MB;
+
         if let Some(sas_url) = sas_url {
             avml::blobstore::upload_sas(&dst, sas_url, sas_block_size)?;
             delete = true;
@@ -186,13 +189,3 @@ fn run_app() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
-
-fn main() {
-    ::std::process::exit(match run_app() {
-        Ok(_) => 0,
-        Err(err) => {
-            eprintln!("error: {:?}", err);
-            1
-        }
-    });
-}

src/blobstore.rs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+use crate::ONE_MB;
+
 use azure::prelude::*;
 use azure_sdk_core::errors::AzureError;
 use azure_sdk_core::prelude::*;
@@ -9,6 +11,7 @@ use azure_sdk_storage_core::prelude::*;
 use byteorder::{LittleEndian, WriteBytesExt};
 use retry::{delay::jitter, delay::Exponential, retry, OperationResult};
 use std::cmp;
+use std::convert::TryFrom;
 use std::error::Error;
 use std::fs::File;
 use std::io::prelude::*;
@@ -17,7 +20,7 @@ use url::Url;
 
 const BACKOFF: u64 = 100;
 const BACKOFF_COUNT: usize = 100;
-const MAX_BLOCK_SIZE: usize = 1024 * 1024 * 100;
+const MAX_BLOCK_SIZE: usize = ONE_MB * 100;
 
 /// Converts the block index into an block_id
 fn to_id(count: u64) -> Result<Vec<u8>, Box<dyn Error>> {
@@ -52,7 +55,7 @@ pub fn upload_sas(filename: &str, sas: &str, block_size: usize) -> Result<(), Bo
 
     let mut core = Core::new()?;
     let mut file = File::open(filename)?;
-    let size = file.metadata()?.len() as usize;
+    let size = usize::try_from(file.metadata()?.len())?;
     let mut sent = 0;
     let mut blocks = BlockList { blocks: Vec::new() };
     let mut data = vec![0; block_size];

src/image.rs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 use byteorder::{ByteOrder, LittleEndian, ReadBytesExt};
+use std::convert::TryFrom;
 use std::error::Error;
 use std::fs::OpenOptions;
 use std::io::prelude::*;
@@ -65,7 +66,7 @@ where
     R: Read,
     W: Write,
 {
-    let mut buf = vec![0; PAGE_SIZE as usize];
+    let mut buf = vec![0; PAGE_SIZE];
     while size >= PAGE_SIZE {
         src.read_exact(&mut buf)?;
         dst.write_all(&buf)?;
@@ -76,24 +77,23 @@ where
         src.read_exact(&mut buf)?;
         dst.write_all(&buf)?;
     }
-
     Ok(())
 }
 
-pub fn copy_block<R, W>(header: &Header, src: &mut R, mut dst: &mut W) -> Result<(), Box<dyn Error>>
+fn copy_block_impl<R, W>(header: Header, src: &mut R, mut dst: &mut W) -> Result<(), Box<dyn Error>>
 where
     R: Read,
     W: Write + std::io::Seek,
 {
     header.write(dst)?;
-    let size = header.range.end - header.range.start;
+    let size = usize::try_from(header.range.end - header.range.start).expect("invalid block size");
     if header.version == 1 {
-        copy(size as usize, src, dst)?;
+        copy(size, src, dst)?;
     } else {
         let begin = dst.seek(SeekFrom::Current(0))?;
         {
             let mut snap_fh = snap::Writer::new(&mut dst);
-            copy(size as usize, src, &mut snap_fh)?;
+            copy(size, src, &mut snap_fh)?;
         }
         let end = dst.seek(SeekFrom::Current(0))?;
         let mut size_bytes = [0; 8];
@@ -103,6 +103,35 @@ where
     Ok(())
 }
 
+pub fn copy_block<R, W>(mut header: Header, src: &mut R, dst: &mut W) -> Result<(), Box<dyn Error>>
+where
+    R: Read,
+    W: Write + std::io::Seek,
+{
+    if header.version == 2 {
+        let max_size = u64::try_from(100 * 256 * PAGE_SIZE).expect("invalid max page size");
+        while header.range.end - header.range.start > max_size {
+            copy_block_impl(
+                Header {
+                    range: Range {
+                        start: header.range.start,
+                        end: header.range.start + max_size,
+                    },
+                    version: header.version,
+                },
+                src,
+                dst,
+            )?;
+            header.range.start += max_size;
+        }
+    }
+    if header.range.end > header.range.start {
+        copy_block_impl(header, src, dst)?;
+    }
+
+    Ok(())
+}
+
 pub struct Image {
     pub version: u32,
     pub src: std::fs::File,
@@ -134,7 +163,7 @@ impl Image {
             self.src.seek(SeekFrom::Start(offset))?;
         }
 
-        copy_block(&header, &mut self.src, &mut self.dst)?;
+        copy_block(header, &mut self.src, &mut self.dst)?;
         Ok(())
     }
 }

src/lib.rs

@@ -26,3 +26,5 @@ pub mod blobstore;
 
 #[cfg(feature = "put")]
 pub mod upload;
+
+pub const ONE_MB: usize = 1024 * 1024;

test/test-conversion.sh

@@ -11,8 +11,8 @@ CONVERT=${2:-target/x86_64-unknown-linux-musl/release/avml-convert}
 
 for SKU in $(cat ${IMAGES_TXT}); do
     ${CONVERT} ${SKU}.lime ${SKU}.uncompressed.lime
-    ${CONVERT} --compress ${SKU}.uncompressed.lime ${SKU}.recompressed.lime
-    ${CONVERT} --compress ${SKU}.lime ${SKU}.compressed.lime
+    ${CONVERT} --format lime_compressed ${SKU}.uncompressed.lime ${SKU}.recompressed.lime
+    ${CONVERT} --format lime_compressed ${SKU}.lime ${SKU}.compressed.lime
     diff -q ${SKU}.lime ${SKU}.compressed.lime
     diff -q ${SKU}.lime ${SKU}.recompressed.lime
-done
+done