skip 1MB blocks of nulls when running with --compress (#45)

src/bin/avml-convert.rs

@@ -3,7 +3,11 @@
 
 use anyhow::{bail, Error, Result};
 use argh::FromArgs;
-use avml::ONE_MB;
+use avml::{
+    image::{Block, MAX_BLOCK_SIZE},
+    iomem::split_ranges,
+    ONE_MB,
+};
 use snap::read::FrameDecoder;
 use std::{
     convert::TryFrom,
@@ -86,6 +90,24 @@ fn convert_to_raw(src: &Path, dst: &Path) -> Result<()> {
     Ok(())
 }
 
+fn convert_from_raw(src: &Path, dst: &Path, compress: bool) -> Result<()> {
+    let src_len = metadata(&src)?.len();
+    let version = if compress { 2 } else { 1 };
+    let mut image = avml::image::Image::new(version, src, dst)?;
+
+    let ranges = split_ranges(vec![0..src_len], MAX_BLOCK_SIZE)?;
+
+    let blocks = ranges
+        .iter()
+        .map(|x| Block {
+            offset: x.start,
+            range: x.start..x.end,
+        })
+        .collect::<Vec<_>>();
+
+    image.write_blocks(&blocks)
+}
+
 #[derive(FromArgs)]
 /// AVML compress/decompress tool
 struct Config {
@@ -134,11 +156,10 @@ fn main() -> Result<()> {
         }
         (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) => bail!("no conversion required"),
-        (Format::Raw, Format::Lime | Format::LimeCompressed) => {
-            bail!("converting from raw not supported")
-        }
     }
 }

src/bin/avml.rs

@@ -3,6 +3,7 @@
 
 use anyhow::{anyhow, bail, Context, Result};
 use argh::FromArgs;
+use avml::image::Block;
 #[cfg(feature = "blobstore")]
 use avml::ONE_MB;
 use std::{
@@ -100,53 +101,52 @@ fn kcore(ranges: &[Range<u64>], filename: &Path, version: u32) -> Result<()> {
                 filename.display()
             )
         })?;
+
     let mut file = elf::File::open_stream(&mut image.src)
         .map_err(|e| anyhow!("unable to parse ELF structures from /proc/kcore: {:?}", e))?;
     file.phdrs.retain(|&x| x.progtype == elf::types::PT_LOAD);
     file.phdrs.sort_by(|a, b| a.vaddr.cmp(&b.vaddr));
     let start = file.phdrs[0].vaddr - ranges[0].start;
 
+    let mut blocks = vec![];
     for range in ranges {
         for phdr in &file.phdrs {
             if range.start == phdr.vaddr - start {
-                image.write_block(
-                    phdr.offset,
-                    Range {
-                        start: range.start,
-                        end: range.start + phdr.memsz,
-                    },
-                )?;
+                blocks.push(Block {
+                    offset: phdr.offset,
+                    range: range.start..range.start + phdr.memsz,
+                });
             }
         }
     }
+
+    image.write_blocks(&blocks)?;
     Ok(())
 }
 
 fn phys(ranges: &[Range<u64>], filename: &Path, mem: &Path, version: u32) -> Result<()> {
+    let is_crash = mem == Path::new("/dev/crash");
+    let blocks = ranges
+        .iter()
+        .map(|x| Block {
+            offset: x.start,
+            range: if is_crash {
+                x.start..((x.end >> 12) << 12)
+            } else {
+                x.start..x.end
+            },
+        })
+        .collect::<Vec<_>>();
+
     let mut image = avml::image::Image::new(version, mem, filename).with_context(|| {
         format!(
             "unable to create image. source:{} destination:{}",
             mem.display(),
             filename.display()
         )
     })?;
-    for range in ranges {
-        let end = if mem == Path::new("/dev/crash") {
-            (range.end >> 12) << 12
-        } else {
-            range.end
-        };
 
-        image
-            .write_block(
-                range.start,
-                Range {
-                    start: range.start,
-                    end,
-                },
-            )
-            .with_context(|| format!("unable to write block: {}:{}", range.start, end))?;
-    }
+    image.write_blocks(&blocks)?;
 
     Ok(())
 }
@@ -169,8 +169,7 @@ fn read_src(ranges: &[Range<u64>], src: &Source, dst: &Path, version: u32) -> Re
 }
 
 fn get_mem(src: Option<&Source>, dst: &Path, version: u32) -> Result<()> {
-    let ranges =
-        avml::iomem::parse(Path::new("/proc/iomem")).context("parsing /proc/iomem failed")?;
+    let ranges = avml::iomem::parse().context("unable to parse /proc/iomem")?;
 
     if let Some(src) = src {
         read_src(&ranges, src, dst, version)

src/image.rs

@@ -13,6 +13,7 @@ use std::{
     path::Path,
 };
 
+pub const MAX_BLOCK_SIZE: u64 = 0x1000 * 0x1000;
 const PAGE_SIZE: usize = 0x1000;
 const LIME_MAGIC: u32 = 0x4c69_4d45; // EMiL as u32le
 const AVML_MAGIC: u32 = 0x4c4d_5641; // AVML as u32le
@@ -23,6 +24,11 @@ pub struct Header {
     pub version: u32,
 }
 
+pub struct Block {
+    pub offset: u64,
+    pub range: Range<u64>,
+}
+
 impl Header {
     pub fn read(mut src: &File) -> Result<Self> {
         let magic = src
@@ -95,7 +101,44 @@ where
     Ok(())
 }
 
-fn copy_block_impl<R, W>(header: &Header, src: &mut R, mut dst: &mut W) -> Result<()>
+// read the entire block into memory, and only write it if it's not empty
+fn copy_if_nonzero<R, W>(header: &Header, src: &mut R, mut dst: &mut W) -> Result<()>
+where
+    R: Read,
+    W: Write + Seek,
+{
+    let size = usize::try_from(header.range.end - header.range.start)
+        .context("unable to create image range size")?;
+
+    let mut buf = vec![0; size];
+    src.read_exact(&mut buf)?;
+
+    // if the entire block is zero, we can skip it
+    if buf.iter().all(|x| x == &0) {
+        return Ok(());
+    }
+
+    header.write(dst)?;
+    if header.version == 1 {
+        dst.write_all(&buf)?;
+    } else {
+        let begin = dst
+            .seek(SeekFrom::Current(0))
+            .context("unable to seek to location")?;
+        {
+            let mut snap_fh = FrameEncoder::new(&mut dst);
+            snap_fh.write_all(&buf)?;
+        }
+        let end = dst.seek(SeekFrom::Current(0)).context("seek failed")?;
+        let mut size_bytes = [0; 8];
+        LittleEndian::write_u64_into(&[end - begin], &mut size_bytes);
+        dst.write_all(&size_bytes)
+            .context("write_all of size failed")?;
+    }
+    Ok(())
+}
+
+fn copy_large_block<R, W>(header: &Header, src: &mut R, mut dst: &mut W) -> Result<()>
 where
     R: Read,
     W: Write + Seek,
@@ -122,18 +165,28 @@ where
     Ok(())
 }
 
+fn copy_block_impl<R, W>(header: &Header, src: &mut R, dst: &mut W) -> Result<()>
+where
+    R: Read,
+    W: Write + Seek,
+{
+    if header.range.end - header.range.start > MAX_BLOCK_SIZE {
+        copy_large_block(header, src, dst)
+    } else {
+        copy_if_nonzero(header, src, dst)
+    }
+}
+
 pub fn copy_block<R, W>(mut header: Header, src: &mut R, dst: &mut W) -> Result<()>
 where
     R: Read,
     W: Write + Seek,
 {
     if header.version == 2 {
-        let max_size =
-            u64::try_from(100 * 256 * PAGE_SIZE).context("unable to create image range size")?;
-        while header.range.end - header.range.start > max_size {
+        while header.range.end - header.range.start > MAX_BLOCK_SIZE {
             let range = Range {
                 start: header.range.start,
-                end: header.range.start + max_size,
+                end: header.range.start + MAX_BLOCK_SIZE,
             };
             copy_block_impl(
                 &Header {
@@ -144,7 +197,7 @@ where
                 dst,
             )
             .with_context(|| format!("unable to copy block: {:?}", range))?;
-            header.range.start += max_size;
+            header.range.start += MAX_BLOCK_SIZE;
         }
     }
     if header.range.end > header.range.start {
@@ -177,20 +230,27 @@ impl Image {
         Ok(Self { version, src, dst })
     }
 
-    pub fn write_block(&mut self, offset: u64, range: Range<u64>) -> Result<()> {
+    pub fn write_blocks(&mut self, blocks: &[Block]) -> Result<()> {
+        for block in blocks {
+            self.write_block(block)?;
+        }
+        Ok(())
+    }
+
+    fn write_block(&mut self, block: &Block) -> Result<()> {
         let header = Header {
-            range: range.clone(),
+            range: block.range.clone(),
             version: self.version,
         };
 
-        if offset > 0 {
+        if block.offset > 0 {
             self.src
-                .seek(SeekFrom::Start(offset))
-                .with_context(|| format!("unable to seek to block: {}", offset))?;
+                .seek(SeekFrom::Start(block.offset))
+                .with_context(|| format!("unable to seek to block: {}", block.offset))?;
         }
 
         copy_block(header, &mut self.src, &mut self.dst)
-            .with_context(|| format!("unable to copy block: {:?}", range))?;
+            .with_context(|| format!("unable to copy block: {:?}", block.range))?;
         Ok(())
     }
 }