Merge pull request #147 from chenxiaolong/atomicbool

Stopping spewing Arc everywhere for the cancel signal

avbroot/src/boot.rs

@@ -11,7 +11,7 @@ use std::{
     num::ParseIntError,
     ops::Range,
     path::{Path, PathBuf},
-    sync::{atomic::AtomicBool, Arc},
+    sync::atomic::AtomicBool,
 };
 
 use regex::bytes::Regex;
@@ -89,7 +89,7 @@ fn save_ramdisk(entries: &[CpioEntryNew], format: CompressedFormat) -> Result<Ve
 }
 
 pub trait BootImagePatcher {
-    fn patch(&self, boot_image: &mut BootImage, cancel_signal: &Arc<AtomicBool>) -> Result<()>;
+    fn patch(&self, boot_image: &mut BootImage, cancel_signal: &AtomicBool) -> Result<()>;
 }
 
 /// Root a boot image with Magisk.
@@ -254,7 +254,7 @@ impl MagiskRootPatcher {
 }
 
 impl BootImagePatcher for MagiskRootPatcher {
-    fn patch(&self, boot_image: &mut BootImage, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+    fn patch(&self, boot_image: &mut BootImage, cancel_signal: &AtomicBool) -> Result<()> {
         let zip_reader = File::open(&self.apk_path)?;
         let mut zip = ZipArchive::new(BufReader::new(zip_reader))?;
 
@@ -466,7 +466,7 @@ impl OtaCertPatcher {
 }
 
 impl BootImagePatcher for OtaCertPatcher {
-    fn patch(&self, boot_image: &mut BootImage, _cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+    fn patch(&self, boot_image: &mut BootImage, _cancel_signal: &AtomicBool) -> Result<()> {
         let patched_any = match boot_image {
             BootImage::V0Through2(b) => self.patch_ramdisk(&mut b.ramdisk)?,
             BootImage::V3Through4(b) => self.patch_ramdisk(&mut b.ramdisk)?,
@@ -557,7 +557,7 @@ impl PrepatchedImagePatcher {
 }
 
 impl BootImagePatcher for PrepatchedImagePatcher {
-    fn patch(&self, boot_image: &mut BootImage, _cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+    fn patch(&self, boot_image: &mut BootImage, _cancel_signal: &AtomicBool) -> Result<()> {
         let prepatched_image = {
             let raw_reader = File::open(&self.prepatched)?;
             BootImage::from_reader(BufReader::new(raw_reader))?
@@ -735,7 +735,7 @@ pub fn patch_boot(
     writer: impl Write + Seek,
     key: &RsaPrivateKey,
     patchers: &[Box<dyn BootImagePatcher + Send>],
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let (mut header, footer, image_size) = avb::load_image(&mut reader)?;
     let Some(footer) = footer else {

avbroot/src/cli/args.rs

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
-use std::sync::{atomic::AtomicBool, Arc};
+use std::sync::atomic::AtomicBool;
 
 use anyhow::Result;
 use clap::{Parser, Subcommand};
@@ -34,7 +34,7 @@ pub struct Cli {
     pub command: Command,
 }
 
-pub fn main(cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn main(cancel_signal: &AtomicBool) -> Result<()> {
     let cli = Cli::parse();
 
     match cli.command {

avbroot/src/cli/avb.rs

@@ -10,7 +10,7 @@ use std::{
     io::{self, BufReader},
     path::{Path, PathBuf},
     str,
-    sync::{atomic::AtomicBool, Arc},
+    sync::atomic::AtomicBool,
 };
 
 use anyhow::{anyhow, bail, Context, Result};
@@ -107,7 +107,7 @@ pub fn verify_headers(
 pub fn verify_descriptors(
     directory: &Path,
     descriptors: &HashMap<String, Descriptor>,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     descriptors
         .par_iter()
@@ -149,7 +149,7 @@ pub fn verify_descriptors(
         .collect()
 }
 
-pub fn avb_main(cli: &AvbCli, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn avb_main(cli: &AvbCli, cancel_signal: &AtomicBool) -> Result<()> {
     match &cli.command {
         AvbCommand::Dump(c) => {
             let raw_reader = File::open(&c.input)

avbroot/src/cli/ota.rs

@@ -11,7 +11,7 @@ use std::{
     fs::{self, File},
     io::{self, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
     path::{Path, PathBuf},
-    sync::{atomic::AtomicBool, Arc, Mutex},
+    sync::{atomic::AtomicBool, Mutex},
     time::Instant,
 };
 
@@ -143,7 +143,7 @@ fn open_input_streams(
     required_images: &HashMap<String, String>,
     external_images: &HashMap<String, PathBuf>,
     header: &PayloadHeader,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<HashMap<String, Box<dyn ReadSeek + Send>>> {
     let mut input_streams = HashMap::<String, Box<dyn ReadSeek + Send>>::new();
 
@@ -185,7 +185,7 @@ fn patch_boot_images(
     root_patcher: Option<Box<dyn BootImagePatcher + Send>>,
     key_avb: &RsaPrivateKey,
     cert_ota: &Certificate,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let mut boot_patchers = HashMap::<&str, Vec<Box<dyn BootImagePatcher + Send>>>::new();
     boot_patchers
@@ -426,7 +426,7 @@ fn compress_image(
     mut stream: &mut Box<dyn ReadSeek + Send>,
     header: &Mutex<PayloadHeader>,
     block_size: u32,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     stream.rewind()?;
 
@@ -460,7 +460,7 @@ fn patch_ota_payload(
     key_avb: &RsaPrivateKey,
     key_ota: &RsaPrivateKey,
     cert_ota: &Certificate,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<(String, u64)> {
     let header =
         PayloadHeader::from_reader(open_payload()?).context("Failed to load OTA payload header")?;
@@ -630,7 +630,7 @@ fn patch_ota_zip(
     key_avb: &RsaPrivateKey,
     key_ota: &RsaPrivateKey,
     cert_ota: &Certificate,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<(OtaMetadata, u64)> {
     let mut missing = BTreeSet::from([
         ota::PATH_METADATA_PB,
@@ -797,7 +797,7 @@ fn extract_ota_zip(
     payload_size: u64,
     header: &PayloadHeader,
     images: &BTreeSet<String>,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     for name in images {
         if Path::new(name).file_name() != Some(OsStr::new(name)) {
@@ -843,7 +843,7 @@ fn extract_ota_zip(
     Ok(())
 }
 
-pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()> {
     let output = cli.output.as_ref().map_or_else(
         || {
             let mut s = cli.input.clone().into_os_string();
@@ -1008,7 +1008,7 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &Arc<AtomicBool>) -> Resu
     Ok(())
 }
 
-pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &AtomicBool) -> Result<()> {
     let raw_reader = File::open(&cli.input)
         .map(PSeekFile::new)
         .with_context(|| format!("Failed to open for reading: {:?}", cli.input))?;
@@ -1067,7 +1067,7 @@ pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &Arc<AtomicBool>) ->
     Ok(())
 }
 
-pub fn verify_subcommand(cli: &VerifyCli, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn verify_subcommand(cli: &VerifyCli, cancel_signal: &AtomicBool) -> Result<()> {
     let raw_reader = File::open(&cli.input)
         .map(PSeekFile::new)
         .with_context(|| format!("Failed to open for reading: {:?}", cli.input))?;
@@ -1184,7 +1184,7 @@ pub fn verify_subcommand(cli: &VerifyCli, cancel_signal: &Arc<AtomicBool>) -> Re
     Ok(())
 }
 
-pub fn ota_main(cli: &OtaCli, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+pub fn ota_main(cli: &OtaCli, cancel_signal: &AtomicBool) -> Result<()> {
     match &cli.command {
         OtaCommand::Patch(c) => patch_subcommand(c, cancel_signal),
         OtaCommand::Extract(c) => extract_subcommand(c, cancel_signal),

avbroot/src/format/avb.rs

@@ -7,10 +7,7 @@ use std::{
     cmp, fmt,
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     str,
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
+    sync::atomic::{AtomicBool, Ordering},
 };
 
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
@@ -337,7 +334,7 @@ impl HashtreeDescriptor {
         block_size: u32,
         algorithm: &'static Algorithm,
         salt: &[u8],
-        cancel_signal: &Arc<AtomicBool>,
+        cancel_signal: &AtomicBool,
     ) -> io::Result<Vec<u8>> {
         // Each digest must be a power of 2.
         let digest_padding = algorithm.output_len.next_power_of_two() - algorithm.output_len;
@@ -384,7 +381,7 @@ impl HashtreeDescriptor {
         block_size: u32,
         algorithm: &'static Algorithm,
         salt: &[u8],
-        cancel_signal: &Arc<AtomicBool>,
+        cancel_signal: &AtomicBool,
     ) -> io::Result<Vec<u8>> {
         assert!(
             image_size > block_size as u64,
@@ -418,7 +415,7 @@ impl HashtreeDescriptor {
         block_size: u32,
         algorithm: &'static Algorithm,
         salt: &[u8],
-        cancel_signal: &Arc<AtomicBool>,
+        cancel_signal: &AtomicBool,
     ) -> io::Result<(Vec<u8>, Vec<u8>)> {
         // Small files are hashed directly, exactly like a hash descriptor.
         if image_size <= u64::from(block_size) {
@@ -486,7 +483,7 @@ impl HashtreeDescriptor {
     pub fn verify(
         &self,
         open_input: impl Fn() -> io::Result<Box<dyn ReadSeek>> + Sync,
-        cancel_signal: &Arc<AtomicBool>,
+        cancel_signal: &AtomicBool,
     ) -> Result<()> {
         let algorithm = match self.hash_algorithm.as_str() {
             "sha256" => &ring::digest::SHA256,
@@ -683,7 +680,7 @@ impl fmt::Debug for HashDescriptor {
 
 impl HashDescriptor {
     /// Verify the root hash against the input reader.
-    pub fn verify(&self, reader: impl Read, cancel_signal: &Arc<AtomicBool>) -> Result<()> {
+    pub fn verify(&self, reader: impl Read, cancel_signal: &AtomicBool) -> Result<()> {
         let algorithm = match self.hash_algorithm.as_str() {
             "sha256" => &ring::digest::SHA256,
             "sha512" => &ring::digest::SHA512,

avbroot/src/format/ota.rs

@@ -7,7 +7,7 @@ use std::{
     collections::BTreeMap,
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     iter,
-    sync::{atomic::AtomicBool, Arc},
+    sync::atomic::AtomicBool,
 };
 
 use cms::signed_data::SignedData;
@@ -452,10 +452,7 @@ fn parse_ota_sig(mut reader: impl Read + Seek) -> Result<(SignedData, u64)> {
 /// CMS signed attributes are intentionally not supported because AOSP recovery
 /// does not support them either. It expects the CMS [`SignedData`] structure to
 /// be used for nothing more than a raw signature transport mechanism.
-pub fn verify_ota(
-    mut reader: impl Read + Seek,
-    cancel_signal: &Arc<AtomicBool>,
-) -> Result<Certificate> {
+pub fn verify_ota(mut reader: impl Read + Seek, cancel_signal: &AtomicBool) -> Result<Certificate> {
     let (sd, hashed_size) = parse_ota_sig(&mut reader)?;
 
     // Make sure the certificate in the CMS structure matches the otacert zip

avbroot/src/format/payload.rs

@@ -6,7 +6,7 @@
 use std::{
     collections::{HashMap, HashSet},
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
-    sync::{atomic::AtomicBool, Arc},
+    sync::atomic::AtomicBool,
 };
 
 use base64::engine::general_purpose::STANDARD;
@@ -619,7 +619,7 @@ pub fn verify_payload(
     mut reader: impl Read + Seek,
     cert: &Certificate,
     properties_raw: &str,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let header = PayloadHeader::from_reader(&mut reader)?;
     reader.rewind()?;
@@ -753,7 +753,7 @@ pub fn apply_operation(
     block_size: u32,
     blob_offset: u64,
     op: &InstallOperation,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     for extent in &op.dst_extents {
         let start_block = extent
@@ -859,7 +859,7 @@ pub fn extract_image_to_memory(
     open_payload: impl Fn() -> io::Result<Box<dyn ReadSeek>> + Sync,
     header: &PayloadHeader,
     partition_name: &str,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<SharedCursor> {
     let partition = header
         .manifest
@@ -900,7 +900,7 @@ pub fn extract_images<'a>(
     open_output: impl Fn(&str) -> io::Result<Box<dyn WriteSeek>> + Sync,
     header: &PayloadHeader,
     partition_names: impl IntoIterator<Item = &'a str>,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let mut remaining = partition_names.into_iter().collect::<HashSet<_>>();
     // We parallelize at the operation level or else one thread might get stuck

avbroot/src/stream.rs

@@ -569,7 +569,7 @@ pub fn copy_n_inspect(
     mut writer: impl Write,
     mut size: u64,
     mut inspect: impl FnMut(&[u8]),
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> io::Result<()> {
     let mut buf = [0u8; 16384];
 
@@ -599,7 +599,7 @@ pub fn copy_n(
     reader: impl Read,
     writer: impl Write,
     size: u64,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> io::Result<()> {
     copy_n_inspect(reader, writer, size, |_| {}, cancel_signal)
 }
@@ -610,7 +610,7 @@ pub fn copy_n(
 pub fn copy(
     mut reader: impl Read,
     mut writer: impl Write,
-    cancel_signal: &Arc<AtomicBool>,
+    cancel_signal: &AtomicBool,
 ) -> io::Result<u64> {
     let mut buf = [0u8; 16384];
     let mut copied = 0;
@@ -640,10 +640,7 @@ pub fn copy(
 mod tests {
     use std::{
         io::{self, Cursor, Read, Seek, SeekFrom, Write},
-        sync::{
-            atomic::{AtomicBool, Ordering},
-            Arc,
-        },
+        sync::atomic::{AtomicBool, Ordering},
     };
 
     use ring::digest::Context;
@@ -869,7 +866,7 @@ mod tests {
 
     #[test]
     fn copy() {
-        let cancel_signal = Arc::new(AtomicBool::new(false));
+        let cancel_signal = AtomicBool::new(false);
         let mut reader = Cursor::new(b"foobar");
         let mut writer = Cursor::new([0u8; 6]);