From 35279b7c7ba08de924e898f59a64b459d448d411 Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Sun, 6 Apr 2025 10:18:19 -0400
Subject: [PATCH] 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.
---
 test-distro/src/depmod.rs   | 29 +++++++----------------------
 test-distro/src/lib.rs      | 27 +++++++++++++++++++++++++++
 test-distro/src/modprobe.rs | 26 ++++----------------------
 3 files changed, 38 insertions(+), 44 deletions(-)

diff --git a/test-distro/src/depmod.rs b/test-distro/src/depmod.rs
index bb057f52..1d5e64aa 100644
--- a/test-distro/src/depmod.rs
+++ b/test-distro/src/depmod.rs
@@ -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(())
diff --git a/test-distro/src/lib.rs b/test-distro/src/lib.rs
index 6f4ca9e6..c18b2beb 100644
--- a/test-distro/src/lib.rs
+++ b/test-distro/src/lib.rs
@@ -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)
+}
diff --git a/test-distro/src/modprobe.rs b/test-distro/src/modprobe.rs
index 42edd81f..07300503 100644
--- a/test-distro/src/modprobe.rs
+++ b/test-distro/src/modprobe.rs
@@ -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");