patch/boot: Replace function pointers with opener trait

Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>

avbroot/src/cli/ota.rs

@@ -43,8 +43,8 @@ use crate::{
     },
     patch::{
         boot::{
-            self, BootImagePatch, DsuPubKeyPatcher, MagiskRootPatcher, OtaCertPatcher,
-            PrepatchedImagePatcher,
+            self, BootImageOpener, BootImagePatch, DsuPubKeyPatcher, MagiskRootPatcher,
+            OtaCertPatcher, PrepatchedImagePatcher,
         },
         system,
     },
@@ -52,8 +52,8 @@ use crate::{
         build::tools::releasetools::OtaMetadata, chromeos_update_engine::DeltaArchiveManifest,
     },
     stream::{
-        self, FromReader, HashingWriter, MutexFile, ReadAt, SectionReader, SectionReaderAt,
-        ToWriter, UserPosFile, WriteAt,
+        self, FromReader, HashingWriter, MutexFile, ReadAt, ReadSeek, SectionReader,
+        SectionReaderAt, ToWriter, UserPosFile, WriteAt, WriteSeek,
     },
     util,
 };
@@ -200,7 +200,6 @@ fn patch_boot_images(
     key_avb: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
-    let input_files = Mutex::new(input_files);
     let boot_partitions = required_images
         .iter()
         .filter(|(_, flags)| flags.contains(PartitionFlags::BOOT))
@@ -212,19 +211,26 @@ fn patch_boot_images(
         util::join(util::sort(boot_partitions.iter()), ", "),
     );
 
-    boot::patch_boot_images(
-        &boot_partitions,
-        |name| {
-            let locked = input_files.lock().unwrap();
+    struct Opener<'a>(Mutex<&'a mut HashMap<String, InputFile>>);
+
+    impl<'a> BootImageOpener for Opener<'a> {
+        fn open_original(&self, name: &str) -> io::Result<Box<dyn ReadSeek + Sync>> {
+            let locked = self.0.lock().unwrap();
             Ok(Box::new(locked[name].file.clone()))
-        },
-        |name| {
-            let mut locked = input_files.lock().unwrap();
+        }
+
+        fn open_replacement(&self, name: &str) -> io::Result<Box<dyn WriteSeek + Sync>> {
+            let mut locked = self.0.lock().unwrap();
             let input_file = locked.get_mut(name).unwrap();
             input_file.file = tempfile::tempfile().map(Arc::new)?;
             input_file.state = InputFileState::Modified;
             Ok(Box::new(input_file.file.clone()))
-        },
+        }
+    }
+
+    boot::patch_boot_images(
+        &boot_partitions,
+        &Opener(Mutex::new(input_files)),
         key_avb,
         boot_patchers,
         cancel_signal,
@@ -2047,13 +2053,20 @@ pub fn verify_subcommand(cli: &VerifyCli, cancel_signal: &AtomicBool) -> Result<
         .filter(|(_, flags)| flags.contains(PartitionFlags::BOOT))
         .map(|(name, _)| name.as_str())
         .collect::<Vec<_>>();
-    let boot_images = boot::load_boot_images(&boot_image_names, |name| {
-        let path = util::path_join_single(temp_dir.path(), format!("{name}.img"))
-            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
 
-        Ok(Box::new(File::open(path)?))
-    })
-    .context("Failed to load all boot images")?;
+    struct Opener<'a>(&'a Path);
+
+    impl<'a> BootImageOpener for Opener<'a> {
+        fn open_original(&self, name: &str) -> io::Result<Box<dyn ReadSeek + Sync>> {
+            let path = util::path_join_single(self.0, format!("{name}.img"))
+                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
+
+            Ok(Box::new(File::open(path)?))
+        }
+    }
+
+    let boot_images = boot::load_boot_images(&boot_image_names, &Opener(temp_dir.path()))
+        .context("Failed to load all boot images")?;
     let targets = OtaCertPatcher::new(ota_cert.clone())
         .find_targets(&boot_images, cancel_signal)
         .context("Failed to find boot image containing otacerts.zip")?;

avbroot/src/patch/boot.rs

@@ -1264,9 +1264,20 @@ fn save_boot_image(
     Ok(())
 }
 
+pub trait BootImageOpener {
+    fn open_original(&self, name: &str) -> io::Result<Box<dyn ReadSeek + Sync>>;
+
+    fn open_replacement(&self, name: &str) -> io::Result<Box<dyn WriteSeek + Sync>> {
+        Err(io::Error::new(
+            io::ErrorKind::NotFound,
+            format!("{name} boot image not found"),
+        ))
+    }
+}
+
 pub fn load_boot_images<'a>(
     names: &[&'a str],
-    open_input: impl Fn(&str) -> io::Result<Box<dyn ReadSeek>> + Sync,
+    opener: &(dyn BootImageOpener + Sync),
 ) -> TargetsResult<HashMap<&'a str, BootImageInfo>> {
     let parent_span = Span::current();
 
@@ -1274,8 +1285,9 @@ pub fn load_boot_images<'a>(
         .par_iter()
         .map(|&name| {
             let _span = debug_span!(parent: &parent_span, "image", name).entered();
-            let mut reader =
-                open_input(name).map_err(|e| TargetsError::Open(name.to_owned(), e))?;
+            let mut reader = opener
+                .open_original(name)
+                .map_err(|e| TargetsError::Open(name.to_owned(), e))?;
 
             let info =
                 load_boot_image(&mut reader).map_err(|e| TargetsError::Load(name.to_owned(), e))?;
@@ -1292,8 +1304,7 @@ pub fn load_boot_images<'a>(
 /// be opened from multiple threads, but at most once each.
 pub fn patch_boot_images<'a>(
     names: &[&'a str],
-    open_input: impl Fn(&str) -> io::Result<Box<dyn ReadSeek>> + Sync,
-    open_output: impl Fn(&str) -> io::Result<Box<dyn WriteSeek>> + Sync,
+    opener: &(dyn BootImageOpener + Sync),
     key: &RsaSigningKey,
     patchers: &[Box<dyn BootImagePatch + Sync>],
     cancel_signal: &AtomicBool,
@@ -1306,7 +1317,7 @@ pub fn patch_boot_images<'a>(
     }
 
     // Preparse all images. Some patchers need to inspect every candidate.
-    let mut images = load_boot_images(names, open_input)?;
+    let mut images = load_boot_images(names, opener)?;
 
     // Find the targets that each patcher wants to patch.
     let all_targets = patchers
@@ -1357,7 +1368,9 @@ pub fn patch_boot_images<'a>(
     // Resign and write new images.
     groups.par_iter_mut().try_for_each(|(&name, (info, _))| {
         let _span = debug_span!(parent: &parent_span, "image", name).entered();
-        let mut writer = open_output(name).map_err(|e| TargetsError::Open(name.to_owned(), e))?;
+        let mut writer = opener
+            .open_replacement(name)
+            .map_err(|e| TargetsError::Open(name.to_owned(), e))?;
 
         save_boot_image(&mut writer, info, key).map_err(|e| TargetsError::Save(name.to_owned(), e))
     })?;