test-distro: extract common decompression code

Remove vector preallocation in the uncompressed case; the standard
library implementation of `io::Read` for `fs::File` already does this.
pull/530/merge
Tamir Duberstein 1 week ago
parent 680402be0c
commit 35279b7c7b

@ -6,16 +6,15 @@
use std::{
fs::File,
io::{BufWriter, Read, Write as _},
io::{BufWriter, Write as _},
path::PathBuf,
};
use anyhow::{Context as _, anyhow};
use clap::Parser;
use object::{Object, ObjectSection, ObjectSymbol, Section};
use test_distro::resolve_modules_dir;
use test_distro::{read_to_end, resolve_modules_dir};
use walkdir::WalkDir;
use xz2::read::XzDecoder;
#[derive(Parser)]
struct Args {
@ -65,26 +64,12 @@ fn main() -> anyhow::Result<()> {
continue;
};
let mut f =
File::open(path).with_context(|| format!("failed to open: {}", path.display()))?;
let stat = f
.metadata()
.with_context(|| format!("failed to get metadata for {}", path.display()))?;
let contents = read_to_end(path, compressed)
.with_context(|| format!("read_to_end({})", path.display()))?;
if compressed {
let mut decoder = XzDecoder::new(f);
// We don't know the size of the decompressed data, so we assume it's
// no more than twice the size of the compressed data.
let mut decompressed = Vec::with_capacity(stat.len() as usize * 2);
decoder.read_to_end(&mut decompressed)?;
read_aliases_from_module(&decompressed, module_name, &mut output)
} else {
let mut buf = Vec::with_capacity(stat.len() as usize);
f.read_to_end(&mut buf)
.with_context(|| format!("failed to read: {}", path.display()))?;
read_aliases_from_module(&buf, module_name, &mut output)
}
.with_context(|| format!("failed to read aliases from module {}", path.display()))?;
read_aliases_from_module(&contents, module_name, &mut output).with_context(|| {
format!("failed to read aliases from module {}", path.display())
})?;
}
}
Ok(())

@ -28,3 +28,30 @@ pub fn resolve_modules_dir() -> anyhow::Result<PathBuf> {
);
Ok(modules_dir)
}
pub fn read_to_end(path: &std::path::Path, compressed: bool) -> anyhow::Result<Vec<u8>> {
use std::io::Read as _;
let mut f = std::fs::File::open(path).context("open()")?;
let mut contents = Vec::new();
if compressed {
let stat = f.metadata().context("metadata()")?;
#[expect(clippy::manual_ok_err)]
let len = match usize::try_from(stat.len()) {
Ok(len) => Some(len),
Err(std::num::TryFromIntError { .. }) => None,
}
.and_then(|len| len.checked_mul(2))
.ok_or_else(|| anyhow::anyhow!("2 * {stat:?}.len() is too large to fit in a usize"))?;
contents.reserve(len);
xz2::read::XzDecoder::new(f).read_to_end(&mut contents)
} else {
f.read_to_end(&mut contents)
}
.context("read_to_end()")?;
Ok(contents)
}

@ -3,17 +3,13 @@
//! This implementation is incredibly naive and is only designed to work within
//! the constraints of the test environment. Not for production use.
use std::{
fs::File,
io::{BufRead as _, Read as _},
path::Path,
};
use std::{fs::File, io::BufRead as _, path::Path};
use anyhow::{Context as _, anyhow, bail};
use clap::Parser;
use glob::glob;
use nix::kmod::init_module;
use test_distro::resolve_modules_dir;
use test_distro::{read_to_end, resolve_modules_dir};
macro_rules! output {
($quiet:expr, $($arg:tt)*) => {
@ -61,28 +57,14 @@ fn try_main(quiet: bool, name: String) -> anyhow::Result<()> {
.context("glob error")?;
output!(quiet, "loading module: {}", module_path.display());
let mut f =
File::open(&module_path).with_context(|| format!("open(): {}", module_path.display()))?;
let stat = f
.metadata()
.with_context(|| format!("stat(): {}", module_path.display()))?;
let extension = module_path
.as_path()
.extension()
.ok_or_else(|| anyhow!("module has no extension: {}", module_path.display()))?;
let contents = if extension == "xz" {
output!(quiet, "decompressing module");
let mut decompressed = Vec::with_capacity(stat.len() as usize * 2);
xz2::read::XzDecoder::new(f).read_to_end(&mut decompressed)?;
decompressed
} else {
let mut contents: Vec<u8> = Vec::with_capacity(stat.len() as usize);
f.read_to_end(&mut contents)?;
contents
};
let contents = read_to_end(&module_path, extension == "xz")
.with_context(|| format!("read_to_end({})", module_path.display()))?;
if !contents.starts_with(&[0x7f, 0x45, 0x4c, 0x46]) {
bail!("module is not an valid ELF file");

Loading…
Cancel
Save