From d3c72e8d19b3cdd3e3ccc858349180cb88e772b1 Mon Sep 17 00:00:00 2001 From: Brian Caswell Date: Tue, 25 Jan 2022 16:42:14 -0500 Subject: [PATCH] skip 1MB blocks of nulls when running with --compress --- src/bin/avml-convert.rs | 29 ++++++++++-- src/bin/avml.rs | 49 ++++++++++----------- src/image.rs | 84 ++++++++++++++++++++++++++++++----- src/iomem.rs | 98 +++++++++++++++++++++++++++++++++++------ 4 files changed, 205 insertions(+), 55 deletions(-) diff --git a/src/bin/avml-convert.rs b/src/bin/avml-convert.rs index 3128024..f6989ca 100644 --- a/src/bin/avml-convert.rs +++ b/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::>(); + + 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") - } } } diff --git a/src/bin/avml.rs b/src/bin/avml.rs index 826ac89..c013d55 100644 --- a/src/bin/avml.rs +++ b/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,29 +101,43 @@ fn kcore(ranges: &[Range], 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], 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::>(); + let mut image = avml::image::Image::new(version, mem, filename).with_context(|| { format!( "unable to create image. source:{} destination:{}", @@ -130,23 +145,8 @@ fn phys(ranges: &[Range], filename: &Path, mem: &Path, version: u32) -> Res 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], 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) diff --git a/src/image.rs b/src/image.rs index be6e089..a36d5e0 100644 --- a/src/image.rs +++ b/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, +} + impl Header { pub fn read(mut src: &File) -> Result { let magic = src @@ -95,7 +101,44 @@ where Ok(()) } -fn copy_block_impl(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(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(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(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(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) -> 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(()) } } diff --git a/src/iomem.rs b/src/iomem.rs index 32b1707..38590b8 100644 --- a/src/iomem.rs +++ b/src/iomem.rs @@ -2,10 +2,14 @@ // Licensed under the MIT License. use anyhow::{bail, Context, Result}; -use std::{fs::OpenOptions, io::prelude::*, path::Path}; +use std::{fs::OpenOptions, io::prelude::*, ops::Range, path::Path}; /// Parse /proc/iomem and return System RAM memory ranges -pub fn parse(path: &Path) -> Result>> { +pub fn parse() -> Result>> { + parse_file(Path::new("/proc/iomem")) +} + +fn parse_file(path: &Path) -> Result>> { let mut f = OpenOptions::new() .read(true) .open(path) @@ -40,13 +44,77 @@ pub fn parse(path: &Path) -> Result>> { Ok(ranges) } +#[must_use] +pub fn merge_ranges(mut ranges: Vec>) -> Vec> { + let mut result = vec![]; + ranges.sort_unstable_by_key(|r| r.start); + + while !ranges.is_empty() { + let mut range = ranges.remove(0); + while !ranges.is_empty() && range.end >= ranges[0].start { + let next = ranges.remove(0); + range = range.start..next.end; + } + result.push(range); + } + + result +} + +pub fn split_ranges(ranges: Vec>, max_size: u64) -> Result>> { + let mut result = vec![]; + + for mut range in ranges { + while range.end - range.start > max_size { + result.push(Range { + start: range.start, + end: range.start + max_size, + }); + range.start += max_size; + } + if !range.is_empty() { + result.push(range); + } + } + + Ok(result) +} + #[cfg(test)] mod tests { use super::*; #[test] - fn parse_iomem() { - let ranges = parse(Path::new("test/iomem.txt")).unwrap(); + fn test_merge_ranges() { + let result = merge_ranges(vec![0..3, 3..6, 7..10, 12..15]); + let expected = [0..6, 7..10, 12..15]; + assert_eq!(result, expected); + + let result = merge_ranges(vec![0..3, 3..6, 6..10]); + let expected = [0..10]; + assert_eq!(result, expected); + } + + #[test] + fn test_split_ranges() -> Result<()> { + let result = split_ranges(vec![0..30], 10)?; + let expected = [0..10, 10..20, 20..30]; + assert_eq!(result, expected); + + let result = split_ranges(vec![0..30], 7)?; + let expected = [0..7, 7..14, 14..21, 21..28, 28..30]; + assert_eq!(result, expected); + + let result = split_ranges(vec![0..10, 10..20, 20..30], 7)?; + let expected = [0..7, 7..10, 10..17, 17..20, 20..27, 27..30]; + assert_eq!(result, expected); + + Ok(()) + } + + #[test] + fn test_parse_iomem() -> Result<()> { + let ranges = parse_file(Path::new("test/iomem.txt"))?; let expected = [ 4096..654_335, 1_048_576..1_073_676_287, @@ -54,7 +122,7 @@ mod tests { ]; assert_eq!(ranges, expected); - let ranges = parse(Path::new("test/iomem-2.txt")).unwrap(); + let ranges = parse_file(Path::new("test/iomem-2.txt"))?; let expected = [ 4096..655_359, 1_048_576..1_055_838_207, @@ -64,7 +132,7 @@ mod tests { ]; assert_eq!(ranges, expected); - let ranges = parse(Path::new("test/iomem-3.txt")).unwrap(); + let ranges = parse_file(Path::new("test/iomem-3.txt"))?; let expected = [ 65_536..649_215, 1_048_576..2_146_303_999, @@ -72,16 +140,18 @@ mod tests { ]; assert_eq!(ranges, expected); - let ranges = parse(Path::new("test/iomem-4.txt")).unwrap(); + let ranges = parse_file(Path::new("test/iomem-4.txt"))?; let expected = [ - 4096..655359, - 1048576..1423523839, - 1423585280..1511186431, - 1780150272..1818623999, - 1818828800..1843613695, - 2071535616..2071986175, - 4294967296..414464344063, + 4_096..655_359, + 1_048_576..1_423_523_839, + 1_423_585_280..1_511_186_431, + 1_780_150_272..1_818_623_999, + 1_818_828_800..1_843_613_695, + 2_071_535_616..2_071_986_175, + 4_294_967_296..414_464_344_063, ]; assert_eq!(ranges, expected); + + Ok(()) } }