diff --git a/avbroot/src/cli/ota.rs b/avbroot/src/cli/ota.rs index d9cdcbb..43d061b 100644 --- a/avbroot/src/cli/ota.rs +++ b/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>); + + impl<'a> BootImageOpener for Opener<'a> { + fn open_original(&self, name: &str) -> io::Result> { + 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> { + 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::>(); - 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> { + 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")?; diff --git a/avbroot/src/patch/boot.rs b/avbroot/src/patch/boot.rs index d991adf..b3f4d43 100644 --- a/avbroot/src/patch/boot.rs +++ b/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>; + + fn open_replacement(&self, name: &str) -> io::Result> { + 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> + Sync, + opener: &(dyn BootImageOpener + Sync), ) -> TargetsResult> { 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> + Sync, - open_output: impl Fn(&str) -> io::Result> + Sync, + opener: &(dyn BootImageOpener + Sync), key: &RsaSigningKey, patchers: &[Box], 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)) })?;