Version 3.2.1

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

CHANGELOG.md

@@ -7,6 +7,16 @@
     to update the actual links at the bottom of the file.
 -->
 
+### Version 3.2.1
+
+* Increase hash tree and FEC size limits to accommodate partition images up to 8 GiB ([Issue #291], [PR #293])
+
+### Version 3.2.0
+
+* Fix potential infinite loop when interrupting avbroot at the right moment to a bug in the bzip2-rs library ([Issue #285], [PR #287])
+* Update all dependencies and fix new clippy lints ([PR #288])
+* Add support for adding the custom AVB public key to the list of trusted keys for DSU (booting signed GSIs) ([Discussion #286], [PR #289])
+
 ### Version 3.1.3
 
 * Build universal binary for macOS ([Issue #278], [PR #279])
@@ -170,6 +180,7 @@ Behind-the-scenes changes:
 <!-- Do not manually edit the lines below. Use `cargo xtask update-changelog` to regenerate. -->
 [Discussion #195]: https://github.com/chenxiaolong/avbroot/discussions/195
 [Discussion #235]: https://github.com/chenxiaolong/avbroot/discussions/235
+[Discussion #286]: https://github.com/chenxiaolong/avbroot/discussions/286
 [Issue #138]: https://github.com/chenxiaolong/avbroot/issues/138
 [Issue #144]: https://github.com/chenxiaolong/avbroot/issues/144
 [Issue #145]: https://github.com/chenxiaolong/avbroot/issues/145
@@ -186,6 +197,8 @@ Behind-the-scenes changes:
 [Issue #225]: https://github.com/chenxiaolong/avbroot/issues/225
 [Issue #265]: https://github.com/chenxiaolong/avbroot/issues/265
 [Issue #278]: https://github.com/chenxiaolong/avbroot/issues/278
+[Issue #285]: https://github.com/chenxiaolong/avbroot/issues/285
+[Issue #291]: https://github.com/chenxiaolong/avbroot/issues/291
 [PR #130]: https://github.com/chenxiaolong/avbroot/pull/130
 [PR #132]: https://github.com/chenxiaolong/avbroot/pull/132
 [PR #133]: https://github.com/chenxiaolong/avbroot/pull/133
@@ -267,3 +280,7 @@ Behind-the-scenes changes:
 [PR #276]: https://github.com/chenxiaolong/avbroot/pull/276
 [PR #277]: https://github.com/chenxiaolong/avbroot/pull/277
 [PR #279]: https://github.com/chenxiaolong/avbroot/pull/279
+[PR #287]: https://github.com/chenxiaolong/avbroot/pull/287
+[PR #288]: https://github.com/chenxiaolong/avbroot/pull/288
+[PR #289]: https://github.com/chenxiaolong/avbroot/pull/289
+[PR #293]: https://github.com/chenxiaolong/avbroot/pull/293

Cargo.toml

@@ -4,7 +4,7 @@ members = ["avbroot", "e2e", "fuzz", "xtask"]
 resolver = "2"
 
 [workspace.package]
-version = "3.1.3"
+version = "3.2.1"
 license = "GPL-3.0-only"
 edition = "2021"
 repository = "https://github.com/chenxiaolong/avbroot"

README.md

@@ -363,6 +363,12 @@ The only behavior this changes is where the partition is read from. When using `
 
 This has no impact on what patches are applied. For example, when using Magisk, the root patch is applied to the boot partition, no matter if the partition came from the original `payload.bin` or from `--replace`.
 
+### Booting signed GSIs
+
+Android's [Dynamic System Updates (DSU)](https://developer.android.com/topic/dsu) feature uses a different root of trust than the regular system. Instead of using the bootloader's `avb_custom_key`, it obtains the trusted keys from the `first_stage_ramdisk/avb/*.avbpubkey` files inside the `init_boot` or `vendor_boot` ramdisk. These files are encoded in the same binary format as `avb_pkmd.bin`.
+
+avbroot can add the custom AVB public key to this directory by passing in `--dsu` when patching an OTA. This allows booting [Generic System Images (GSI)](https://developer.android.com/topic/generic-system-image) signed by the custom AVB key.
+
 ### Clearing vbmeta flags
 
 Some Android builds may ship with a root `vbmeta` image with the flags set such that AVB is effectively disabled. When avbroot encounters these images, the patching process will fail with a message like:

avbroot/Cargo.toml

@@ -10,7 +10,7 @@ publish = false
 
 [dependencies]
 anyhow = "1.0.75"
-base64 = "0.21.3"
+base64 = "0.22.1"
 bitflags = "2.4.1"
 bstr = "1.6.2"
 byteorder = "1.4.3"
@@ -24,7 +24,7 @@ ctrlc = "3.4.0"
 flate2 = "1.0.27"
 gf256 = { version = "0.3.0", features = ["rs"] }
 hex = { version = "0.4.3", features = ["serde"] }
-liblzma = "0.2.1"
+liblzma = "0.3.0"
 lz4_flex = "0.11.1"
 memchr = "2.6.0"
 num-bigint-dig = "0.8.4"
@@ -47,20 +47,18 @@ sha1 = "0.10.5"
 sha2 = "0.10.7"
 tempfile = "3.8.0"
 thiserror = "1.0.47"
-toml_edit = { version = "0.21.0", features = ["serde"] }
+toml_edit = { version = "0.22.9", features = ["serde"] }
 topological-sort = "0.2.2"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
 x509-cert = { version = "0.2.4", features = ["builder"] }
 
-# There's an upstream bug that causes an infinite loop in the write::BzDecoder
-# destructor if the decoder is fed invalid data. While this never happens during
-# normal operation, it is possible to run into this by running `ota extract`
-# against a `--stripped` OTA file.
-# https://github.com/alexcrichton/bzip2-rs/pull/99
+# There are multiple upstream bugs that cause infinite loops in the Drop
+# implementation of write::BzDecoder. Unfortunately, the project is no longer
+# maintained, so we have to maintain our own fork with the necessary fixes.
 [dependencies.bzip2]
-git = "https://github.com/jongiddy/bzip2-rs"
-rev = "2aefcb4d3634de1df226c73d93f758d65228bb8c"
+git = "https://github.com/chenxiaolong/bzip2-rs"
+rev = "6e0f9836ec87b19261461b6cc1772e14aff8e851"
 
 # https://github.com/zip-rs/zip/pull/383
 [dependencies.zip]
@@ -74,7 +72,7 @@ rustix = { version = "0.38.9", default-features = false, features = ["process"]
 
 [build-dependencies]
 prost-build = "0.12.1"
-protox = "0.5.0"
+protox = "0.6.0"
 
 [dev-dependencies]
 assert_matches = "1.5.0"

avbroot/src/cli/avb.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -94,15 +94,13 @@ fn write_info(path: &Path, info: &AvbInfo) -> Result<()> {
 
 /// Packing with insecure algorithms is intentionally not supported, so promote
 /// to a secure algorithm if needed.
-fn promote_insecure_hash_algorithm(algorithm: &str) -> &str {
+fn promote_insecure_hash_algorithm(algorithm: &mut String) {
     const INSECURE_ALGORITHMS: &[&str] = &["sha1"];
     const NEW_ALGORITHM: &str = "sha256";
 
-    if INSECURE_ALGORITHMS.contains(&algorithm) {
+    if INSECURE_ALGORITHMS.contains(&algorithm.as_str()) {
         warn!("Changing insecure hash algorithm {algorithm} to {NEW_ALGORITHM}");
-        NEW_ALGORITHM
-    } else {
-        algorithm
+        NEW_ALGORITHM.clone_into(algorithm);
     }
 }
 
@@ -193,13 +191,13 @@ fn write_raw_and_update(
 
     match info.header.appended_descriptor_mut()? {
         AppendedDescriptorMut::HashTree(d) => {
-            d.hash_algorithm = promote_insecure_hash_algorithm(&d.hash_algorithm).to_owned();
+            promote_insecure_hash_algorithm(&mut d.hash_algorithm);
             d.image_size = image_size;
             d.update(&raw_file, &raw_file, None, cancel_signal)
                 .context("Failed to update hash tree descriptor")?;
         }
         AppendedDescriptorMut::Hash(d) => {
-            d.hash_algorithm = promote_insecure_hash_algorithm(&d.hash_algorithm).to_owned();
+            promote_insecure_hash_algorithm(&mut d.hash_algorithm);
             d.image_size = image_size;
             raw_file.rewind()?;
             d.update(&mut raw_file, cancel_signal)
@@ -646,7 +644,7 @@ fn repack_subcommand(cli: &RepackCli, cancel_signal: &AtomicBool) -> Result<()>
         // Write new hash tree and FEC data instead of copying the original.
         // There could have been errors in the original FEC data itself.
         if let AppendedDescriptorMut::HashTree(d) = info.header.appended_descriptor_mut()? {
-            d.hash_algorithm = promote_insecure_hash_algorithm(&d.hash_algorithm).to_owned();
+            promote_insecure_hash_algorithm(&mut d.hash_algorithm);
             d.update(&file, &file, None, cancel_signal)?;
         }
 

avbroot/src/cli/ota.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2022-2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -10,6 +10,7 @@ use std::{
     fmt::Display,
     fs::{self, File},
     io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write},
+    mem,
     ops::Range,
     path::{Path, PathBuf},
     sync::{atomic::AtomicBool, Mutex},
@@ -31,14 +32,16 @@ use crate::{
     cli,
     crypto::{self, PassphraseSource},
     format::{
-        avb::Header,
-        avb::{self, Descriptor},
+        avb::{self, Descriptor, Header},
         ota::{self, SigningWriter, ZipEntry},
         padding,
         payload::{self, PayloadHeader, PayloadWriter},
     },
     patch::{
-        boot::{self, BootImagePatch, MagiskRootPatcher, OtaCertPatcher, PrepatchedImagePatcher},
+        boot::{
+            self, BootImagePatch, DsuPubKeyPatcher, MagiskRootPatcher, OtaCertPatcher,
+            PrepatchedImagePatcher,
+        },
         system,
     },
     protobuf::{
@@ -187,26 +190,17 @@ fn open_input_files(
 }
 
 /// Patch the boot images listed in `required_images`. Not every image is
-/// necessarily patched. An [`OtaCertPatcher`] is always applied to the boot
-/// image that contains the trusted OTA certificate list. If `root_patcher` is
-/// specified, then it is used to patch the boot image for root access. If the
-/// original image is signed, then it will be re-signed with `key_avb`.
+/// necessarily patched. Each patcher will determine which image it should
+/// target. If the original image is signed, then it will be re-signed with
+/// `key_avb`.
 fn patch_boot_images<'a, 'b: 'a>(
     required_images: &'b RequiredImages,
     input_files: &mut HashMap<String, InputFile>,
-    root_patcher: Option<Box<dyn BootImagePatch + Sync>>,
+    boot_patchers: Vec<Box<dyn BootImagePatch + Sync>>,
     key_avb: &RsaPrivateKey,
-    cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let input_files = Mutex::new(input_files);
-    let mut boot_patchers = Vec::<Box<dyn BootImagePatch + Sync>>::new();
-    boot_patchers.push(Box::new(OtaCertPatcher::new(cert_ota.clone())));
-
-    if let Some(p) = root_patcher {
-        boot_patchers.push(p);
-    }
-
     let boot_partitions = required_images.iter_boot().collect::<Vec<_>>();
 
     info!(
@@ -479,7 +473,7 @@ fn update_security_descriptors(
         // vbmeta is signed; Use a chain descriptor.
         match parent_descriptor {
             Descriptor::ChainPartition(pd) => {
-                pd.public_key = child_header.public_key.clone();
+                child_header.public_key.clone_into(&mut pd.public_key);
             }
             _ => {
                 bail!("{child_name} descriptor ({parent_type}) in {parent_name} must be a chain descriptor");
@@ -522,7 +516,7 @@ fn update_metadata_descriptors(parent_header: &mut Header, child_header: &Header
                 });
 
                 if let Some(pd) = parent_property {
-                    pd.value = cd.value.clone();
+                    cd.value.clone_into(&mut pd.value);
                 } else {
                     parent_header
                         .descriptors
@@ -542,7 +536,7 @@ fn update_metadata_descriptors(parent_header: &mut Header, child_header: &Header
                 });
 
                 if let Some(pd) = parent_property {
-                    pd.cmdline = cd.cmdline.clone();
+                    cd.cmdline.clone_into(&mut pd.cmdline);
                 } else {
                     parent_header
                         .descriptors
@@ -698,7 +692,7 @@ fn patch_ota_payload(
     payload: &(dyn ReadSeekReopen + Sync),
     writer: impl Write,
     external_images: &HashMap<String, PathBuf>,
-    root_patcher: Option<Box<dyn BootImagePatch + Sync>>,
+    boot_patchers: Vec<Box<dyn BootImagePatch + Sync>>,
     clear_vbmeta_flags: bool,
     key_avb: &RsaPrivateKey,
     key_ota: &RsaPrivateKey,
@@ -750,9 +744,8 @@ fn patch_ota_payload(
     patch_boot_images(
         &required_images,
         &mut input_files,
-        root_patcher,
+        boot_patchers,
         key_avb,
-        cert_ota,
         cancel_signal,
     )?;
 
@@ -893,7 +886,7 @@ fn patch_ota_zip(
     zip_reader: &mut ZipArchive<impl Read + Seek>,
     mut zip_writer: &mut ZipWriter<impl Write>,
     external_images: &HashMap<String, PathBuf>,
-    mut root_patch: Option<Box<dyn BootImagePatch + Sync>>,
+    mut boot_patchers: Vec<Box<dyn BootImagePatch + Sync>>,
     clear_vbmeta_flags: bool,
     key_avb: &RsaPrivateKey,
     key_ota: &RsaPrivateKey,
@@ -1015,7 +1008,7 @@ fn patch_ota_zip(
                     &mut writer,
                     external_images,
                     // There's only one payload in the OTA.
-                    root_patch.take(),
+                    mem::take(&mut boot_patchers),
                     clear_vbmeta_flags,
                     key_avb,
                     key_ota,
@@ -1223,8 +1216,11 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()
         external_images.insert(name.to_owned(), path.to_owned());
     }
 
-    let root_patcher = if let Some(magisk) = &cli.root.magisk {
-        let patcher: Box<dyn BootImagePatch + Sync> = Box::new(
+    let mut boot_patchers = Vec::<Box<dyn BootImagePatch + Sync>>::new();
+    boot_patchers.push(Box::new(OtaCertPatcher::new(cert_ota.clone())));
+
+    if let Some(magisk) = &cli.root.magisk {
+        boot_patchers.push(Box::new(
             MagiskRootPatcher::new(
                 magisk,
                 cli.magisk_preinit_device.as_deref(),
@@ -1232,21 +1228,20 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()
                 cli.ignore_magisk_warnings,
             )
             .context("Failed to create Magisk boot image patcher")?,
-        );
-
-        Some(patcher)
+        ));
     } else if let Some(prepatched) = &cli.root.prepatched {
-        let patcher: Box<dyn BootImagePatch + Sync> = Box::new(PrepatchedImagePatcher::new(
+        boot_patchers.push(Box::new(PrepatchedImagePatcher::new(
             prepatched,
             cli.ignore_prepatched_compat + 1,
-        ));
-
-        Some(patcher)
+        )));
     } else {
         assert!(cli.root.rootless);
-        None
     };
 
+    if cli.dsu {
+        boot_patchers.push(Box::new(DsuPubKeyPatcher::new(key_avb.to_public_key())));
+    }
+
     let raw_reader = File::open(&cli.input)
         .map(PSeekFile::new)
         .with_context(|| format!("Failed to open for reading: {:?}", cli.input))?;
@@ -1272,7 +1267,7 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()
         &mut zip_reader,
         &mut zip_writer,
         &external_images,
-        root_patcher,
+        boot_patchers,
         cli.clear_vbmeta_flags,
         &key_avb,
         &key_ota,
@@ -1841,6 +1836,10 @@ pub struct PatchCli {
     )]
     pub ignore_prepatched_compat: u8,
 
+    /// Add AVB public key to trusted keys for DSU.
+    #[arg(long, help_heading = HEADING_OTHER)]
+    pub dsu: bool,
+
     /// Forcibly clear vbmeta flags if they disable AVB.
     #[arg(long, help_heading = HEADING_OTHER)]
     pub clear_vbmeta_flags: bool,

avbroot/src/format/avb.rs

@@ -50,13 +50,30 @@ pub const FOOTER_MAGIC: [u8; 4] = *b"AVBf";
 /// for early fail. No individual field can actually be this size.
 pub const HEADER_MAX_SIZE: u64 = 64 * 1024;
 
-/// Maximum hash tree size. The current limit equals the hash tree size for a
-/// 4GiB image using SHA512 digests and a block size of 4096.
-pub const HASH_TREE_MAX_SIZE: u64 = 68_177_920;
+/// Maximum hash tree size. The current limit equals the hash tree size for an
+/// 8GiB image using SHA256 digests and a block size of 4096. This is equal to:
+///
+/// ```rust
+/// use avbroot::format::hashtree::HashTree;
+/// let size = HashTree::new(4096, &ring::digest::SHA256, b"")
+///     .compute_level_offsets(8 * 1024 * 1024 * 1024)
+///     .unwrap()
+///    .first()
+///    .map(|r| r.end)
+///    .unwrap_or(0);
+/// ```
+pub const HASH_TREE_MAX_SIZE: u64 = 67_637_248;
 
-/// Maximum FEC data size. The current limit equals the FEC data size for a 4GiB
-/// image using 2 parity bytes per codeword.
-pub const FEC_DATA_MAX_SIZE: u64 = 33_959_936;
+/// Maximum FEC data size. The current limit equals the FEC data size for an
+/// 8GiB image using 2 parity bytes per codeword. This is equal to:
+///
+/// ```rust
+/// use avbroot::format::fec::Fec;
+/// let size = Fec::new(8 * 1024 * 1024 * 1024, 4096, 2)
+///     .unwrap()
+///     .fec_size();
+/// ```
+pub const FEC_DATA_MAX_SIZE: u64 = 67_911_680;
 
 #[derive(Debug, Error)]
 pub enum Error {

avbroot/src/format/fec.rs

@@ -166,8 +166,8 @@ impl Fec {
             return Err(Error::UnsupportedParity(parity));
         }
 
-        let blocks = util::div_ceil(file_size, u64::from(block_size));
-        let rounds = util::div_ceil(blocks, u64::from(rs_k));
+        let blocks = file_size.div_ceil(u64::from(block_size));
+        let rounds = blocks.div_ceil(u64::from(rs_k));
 
         // Check upfront so we don't need to do checked multiplication later.
         rounds
@@ -196,7 +196,7 @@ impl Fec {
 
     /// Get the size of the FEC data needed to cover the entire file.
     #[inline]
-    fn fec_size(&self) -> usize {
+    pub fn fec_size(&self) -> usize {
         usize::from(self.parity()) * self.rounds as usize * self.block_size as usize
     }
 
@@ -225,7 +225,7 @@ impl Fec {
             let end_block = if range.end % block_size == 0 {
                 range.end / block_size
             } else {
-                util::div_ceil(range.end, block_size)
+                range.end.div_ceil(block_size)
             };
 
             for block in start_block..end_block {

avbroot/src/format/hashtree.rs

@@ -72,14 +72,14 @@ impl HashTree {
     /// tree data. The items are returned with the bottom level's offsets first
     /// in the list. Note that the bottom level is stored at the end of the hash
     /// tree data.
-    fn compute_level_offsets(&self, image_size: u64) -> Result<Vec<Range<usize>>> {
+    pub fn compute_level_offsets(&self, image_size: u64) -> Result<Vec<Range<usize>>> {
         let algorithm = self.salted_context.algorithm();
         let digest_size = algorithm.output_len().next_power_of_two();
         let mut ranges = vec![];
         let mut level_size = image_size;
 
         while level_size > u64::from(self.block_size) {
-            let blocks = util::div_ceil(level_size, u64::from(self.block_size));
+            let blocks = level_size.div_ceil(u64::from(self.block_size));
             level_size = blocks
                 .checked_mul(digest_size as u64)
                 .and_then(|s| padding::round(s, u64::from(self.block_size)))
@@ -124,7 +124,7 @@ impl HashTree {
             let end_block = if range.end % block_size == 0 {
                 range.end / block_size
             } else {
-                util::div_ceil(range.end, block_size)
+                range.end.div_ceil(block_size)
             };
 
             result.push(start_block..end_block);

avbroot/src/format/ota.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2022-2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -152,15 +152,15 @@ pub fn parse_legacy_metadata(data: &str) -> Result<OtaMetadata> {
             }
             "post-build-incremental" => {
                 let p = metadata.postcondition.get_or_insert_with(Default::default);
-                p.build_incremental = value.to_owned();
+                value.clone_into(&mut p.build_incremental);
             }
             "post-sdk-level" => {
                 let p = metadata.postcondition.get_or_insert_with(Default::default);
-                p.sdk_level = value.to_owned();
+                value.clone_into(&mut p.sdk_level);
             }
             "post-security-patch-level" => {
                 let p = metadata.postcondition.get_or_insert_with(Default::default);
-                p.security_patch_level = value.to_owned();
+                value.clone_into(&mut p.security_patch_level);
             }
             "post-timestamp" => {
                 let p = metadata.postcondition.get_or_insert_with(Default::default);
@@ -176,7 +176,7 @@ pub fn parse_legacy_metadata(data: &str) -> Result<OtaMetadata> {
             }
             "pre-build-incremental" => {
                 let p = metadata.precondition.get_or_insert_with(Default::default);
-                p.build_incremental = value.to_owned();
+                value.clone_into(&mut p.build_incremental);
             }
             "spl-downgrade" => metadata.spl_downgrade = parse_yes()?,
             k if k.ends_with("-property-files") => {

avbroot/src/format/payload.rs

@@ -918,7 +918,7 @@ pub fn compress_image(
         });
     }
 
-    let chunks_total = util::div_ceil(file_size, CHUNK_SIZE);
+    let chunks_total = file_size.div_ceil(CHUNK_SIZE);
     let mut bytes_compressed = 0;
     let mut context_uncompressed = Context::new(&ring::digest::SHA256);
     let mut operations = vec![];
@@ -1059,7 +1059,7 @@ pub fn compress_modified_image(
         return Err(Error::ExtentsNotInOrder);
     }
 
-    let groups_total = util::div_ceil(operations.len(), OPERATION_GROUP);
+    let groups_total = operations.len().div_ceil(OPERATION_GROUP);
     let mut bytes_compressed = 0;
     let mut context_uncompressed = Context::new(&ring::digest::SHA256);
     let mut modified_operations = vec![];

avbroot/src/patch/boot.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2022-2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -23,7 +23,7 @@ use liblzma::{
 use rayon::iter::{IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator};
 use regex::bytes::Regex;
 use ring::digest::Context;
-use rsa::RsaPrivateKey;
+use rsa::{RsaPrivateKey, RsaPublicKey};
 use thiserror::Error;
 use tracing::{debug, debug_span, trace, warn, Span};
 use x509_cert::Certificate;
@@ -667,6 +667,137 @@ impl BootImagePatch for OtaCertPatcher {
     }
 }
 
+/// Add the AVB public key to DSU's list of trusted keys for verifying GSIs.
+pub struct DsuPubKeyPatcher {
+    key: RsaPublicKey,
+}
+
+impl DsuPubKeyPatcher {
+    const FIRST_STAGE_PATH: &'static [u8] = b"first_stage_ramdisk";
+    const DSU_KEYS_PATH: &'static [u8] = b"first_stage_ramdisk/avb";
+    const AVBROOT_KEY_PATH: &'static [u8] = b"first_stage_ramdisk/avb/avbroot.avbpubkey";
+
+    pub fn new(key: RsaPublicKey) -> Self {
+        Self { key }
+    }
+
+    fn patch_ramdisk(&self, ramdisk: &mut Vec<u8>, cancel_signal: &AtomicBool) -> Result<bool> {
+        let (mut entries, ramdisk_format) = load_ramdisk(ramdisk, cancel_signal)?;
+        if !entries.iter_mut().any(|e| e.path == Self::FIRST_STAGE_PATH) {
+            return Ok(false);
+        }
+
+        if !entries.iter().any(|e| e.path == Self::DSU_KEYS_PATH) {
+            entries.push(CpioEntry::new_directory(Self::DSU_KEYS_PATH, 0o755));
+        }
+
+        let data = CpioEntryData::Data(avb::encode_public_key(&self.key)?);
+
+        if let Some(e) = entries
+            .iter_mut()
+            .find(|e| e.path == Self::AVBROOT_KEY_PATH)
+        {
+            e.data = data;
+        } else {
+            entries.push(CpioEntry::new_file(Self::AVBROOT_KEY_PATH, 0o644, data));
+        };
+
+        *ramdisk = save_ramdisk(&entries, ramdisk_format, cancel_signal)?;
+
+        Ok(true)
+    }
+}
+
+impl BootImagePatch for DsuPubKeyPatcher {
+    fn patcher_name(&self) -> &'static str {
+        "DsuPubKeyPatcher"
+    }
+
+    fn find_targets<'a>(
+        &self,
+        boot_images: &HashMap<&'a str, BootImageInfo>,
+        cancel_signal: &AtomicBool,
+    ) -> Result<Vec<&'a str>> {
+        let mut dsu_keys_targets = vec![];
+        let mut first_stage_targets = vec![];
+
+        'outer: for (name, info) in boot_images {
+            let ramdisks = match &info.boot_image {
+                BootImage::V0Through2(b) => slice::from_ref(&b.ramdisk),
+                BootImage::V3Through4(b) => slice::from_ref(&b.ramdisk),
+                BootImage::VendorV3Through4(b) => &b.ramdisks,
+            };
+
+            for ramdisk in ramdisks {
+                if ramdisk.is_empty() {
+                    continue;
+                }
+
+                let (entries, _) = load_ramdisk(ramdisk, cancel_signal)?;
+                let mut found = false;
+
+                for entry in entries {
+                    if entry.path == Self::DSU_KEYS_PATH {
+                        dsu_keys_targets.push(*name);
+                        found = true;
+                    } else if entry.path == Self::FIRST_STAGE_PATH {
+                        first_stage_targets.push(*name);
+                        found = true;
+                    }
+                }
+
+                if found {
+                    continue 'outer;
+                }
+            }
+        }
+
+        if !dsu_keys_targets.is_empty() {
+            // Most builds trust as least one DSU key. For these builds, add the
+            // user's key to the same directory.
+            if dsu_keys_targets.len() > 1 {
+                return Err(Error::Validation(format!(
+                    "DSU keys found in more than one boot image: {dsu_keys_targets:?}",
+                )));
+            }
+
+            Ok(dsu_keys_targets)
+        } else {
+            // For builds that don't trust any DSU keys, pick the first boot
+            // image that contains a first stage ramdisk directory.
+            if !first_stage_targets.is_empty() {
+                first_stage_targets.sort();
+                first_stage_targets.resize(1, "");
+            }
+
+            Ok(first_stage_targets)
+        }
+    }
+
+    fn patch(&self, boot_image: &mut BootImage, cancel_signal: &AtomicBool) -> Result<()> {
+        let ramdisks = match boot_image {
+            BootImage::V0Through2(b) => slice::from_mut(&mut b.ramdisk),
+            BootImage::V3Through4(b) => slice::from_mut(&mut b.ramdisk),
+            BootImage::VendorV3Through4(b) => &mut b.ramdisks,
+        };
+
+        for ramdisk in ramdisks {
+            if ramdisk.is_empty() {
+                continue;
+            }
+
+            if self.patch_ramdisk(ramdisk, cancel_signal)? {
+                return Ok(());
+            }
+        }
+
+        Err(Error::Validation(format!(
+            "No ramdisk contains {:?}",
+            Self::FIRST_STAGE_PATH.as_bstr(),
+        )))
+    }
+}
+
 /// Replace the boot image with a prepatched boot image if it is compatible.
 ///
 /// An image is compatible if all the non-size-related header fields are
@@ -1044,7 +1175,7 @@ pub fn patch_boot_images<'a>(
             let (mut writer, context) = hashing_writer.finish();
 
             descriptor.image_size = writer.stream_position()?;
-            descriptor.hash_algorithm = "sha256".to_owned();
+            "sha256".clone_into(&mut descriptor.hash_algorithm);
             descriptor.root_digest = context.finish().as_ref().to_vec();
 
             if !info.header.public_key.is_empty() {

avbroot/src/patch/system.rs

@@ -129,7 +129,7 @@ pub fn patch_system_image(
         return Err(Error::NoHashTreeDescriptor);
     };
 
-    let num_chunks = util::div_ceil(footer.original_image_size, CHUNK_SIZE);
+    let num_chunks = footer.original_image_size.div_ceil(CHUNK_SIZE);
     trace!("Parallel heuristics search for otacerts.zip with {num_chunks} chunks");
 
     let modified_ranges = (0..num_chunks)

avbroot/src/util.rs

@@ -49,16 +49,6 @@ pub fn parent_path(path: &Path) -> &Path {
     Path::new(".")
 }
 
-/// Since Rust's built-in .div_ceil() is still nightly-only.
-pub fn div_ceil<T: PrimInt>(dividend: T, divisor: T) -> T {
-    dividend / divisor
-        + if dividend % divisor != T::zero() {
-            T::one()
-        } else {
-            T::zero()
-        }
-}
-
 /// Sort and merge overlapping intervals.
 pub fn merge_overlapping<T>(sections: &[Range<T>]) -> Vec<Range<T>>
 where

deny.toml

@@ -1,8 +1,6 @@
 [advisories]
-vulnerability = "deny"
-unmaintained = "deny"
+version = 2
 yanked = "deny"
-notice = "deny"
 ignore = [
     # https://rustsec.org/advisories/RUSTSEC-2023-0071
     #
@@ -29,19 +27,18 @@ ignore = [
 ]
 
 [licenses]
+version = 2
 include-dev = true
-unlicensed = "deny"
 allow = [
     "Apache-2.0",
     "Apache-2.0 WITH LLVM-exception",
     "BSD-3-Clause",
+    "GPL-3.0",
     "ISC",
     "MIT",
     "OpenSSL",
     "Unicode-DFS-2016",
 ]
-copyleft = "allow"
-default = "deny"
 
 [[licenses.clarify]]
 name = "ring"
@@ -73,6 +70,6 @@ bypass = [
 unknown-registry = "deny"
 unknown-git = "deny"
 allow-git = [
+    "https://github.com/chenxiaolong/bzip2-rs",
     "https://github.com/chenxiaolong/zip",
-    "https://github.com/jongiddy/bzip2-rs",
 ]

e2e/Cargo.toml

@@ -18,7 +18,7 @@ ring = "0.17.0"
 rsa = "0.9.6"
 serde = { version = "1.0.188", features = ["derive"] }
 tempfile = "3.8.0"
-toml_edit = { version = "0.21.0", features = ["serde"] }
+toml_edit = { version = "0.22.9", features = ["serde"] }
 topological-sort = "0.2.2"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"

e2e/e2e.toml

@@ -22,7 +22,7 @@ data.kernel = true
 avb.signed = true
 data.type = "boot"
 data.version = "v4"
-data.ramdisks = ["init"]
+data.ramdisks = [["init", "first_stage"]]
 
 [profile.pixel_v4_gki.partitions.system]
 avb.signed = false
@@ -43,11 +43,11 @@ data.deps = ["system"]
 avb.signed = false
 data.type = "boot"
 data.version = "vendor_v4"
-data.ramdisks = ["otacerts"]
+data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v4_gki.hashes]
-original = "f9477a35e3b60a495e49431c61e3897f11775f453a6a9897ead568357c963618"
-patched = "d3436650b4e0c60688dafb1472fa5fe95e67d19a8fad2da4850ffaa739d44574"
+original = "24a0a62cc08b96563f4872aee2fdd4a84d1a977b55c326dd9d4ddd92a1d326ea"
+patched = "29889670efea78bace221742b19e8ace88b67137fc2f46dcd9dbdf67e4e42267"
 
 # Google Pixel 6a
 # What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4, 2 ramdisks)
@@ -77,11 +77,11 @@ data.deps = ["system"]
 avb.signed = false
 data.type = "boot"
 data.version = "vendor_v4"
-data.ramdisks = ["init_and_otacerts", "dlkm"]
+data.ramdisks = [["init", "otacerts", "first_stage", "dsu_key_dir"], ["dlkm"]]
 
 [profile.pixel_v4_non_gki.hashes]
-original = "021b4510bc244f5f686fbff89eb2058ec9c96a2949c2fe8caa7750a78d593225"
-patched = "0f90ae4c26a54a735d48e13e98bea73b6d1c026b0e95fa5b4b26179a1d6bdc86"
+original = "0e7d0924a68d46e00abe96abfa0e5f3a98d5d15a32bef7401b91fca9a19748e8"
+patched = "1a2f53d9ac3a1da75e5e7502110bda437c5a725764f16656c564d42460136279"
 
 # Google Pixel 4a 5G
 # What's unique: boot (boot v3) + vendor_boot (vendor v3)
@@ -91,7 +91,7 @@ avb.signed = true
 data.type = "boot"
 data.version = "v3"
 data.kernel = true
-data.ramdisks = ["init"]
+data.ramdisks = [["init"]]
 
 [profile.pixel_v3.partitions.system]
 avb.signed = false
@@ -112,11 +112,11 @@ data.deps = ["system"]
 avb.signed = false
 data.type = "boot"
 data.version = "vendor_v3"
-data.ramdisks = ["otacerts"]
+data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v3.hashes]
-original = "12221a69ff32e137d5f19b61f576fc6b33f0973c4a81da7722c640554ff4bc4e"
-patched = "0a92969bbd7cb30071a0799eb20028546d1cec3e6ec1ca4d7e1fe7776f1399fc"
+original = "533e6f233cb98c98c945044c2ee81a6069e66baee6f7dbcfbf7523795a11215e"
+patched = "1eeae9dba0302c2d469bd03c8bccdc0469c171204dbd676a20cb6620d2d11c5c"
 
 # Google Pixel 4a
 # What's unique: boot (boot v2)
@@ -126,7 +126,7 @@ avb.signed = false
 data.type = "boot"
 data.version = "v2"
 data.kernel = true
-data.ramdisks = ["init_and_otacerts"]
+data.ramdisks = [["init", "otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v2.partitions.system]
 avb.signed = false
@@ -144,5 +144,5 @@ data.type = "vbmeta"
 data.deps = ["system"]
 
 [profile.pixel_v2.hashes]
-original = "8b38d2d999b5b6e240e894f669e9e2643b3764c108d53bb7b02447da725e7c18"
-patched = "85b411947145e89cdc7f71b7109a12586f4dbddce3021f0f96172fc465c189d6"
+original = "958dfa428abd2901178b90147903d7857ed78d6c016f6cb3af30d024a22a8f9a"
+patched = "b0d18ef350ca7b6499b7de4b0182f4ac3be7288ad238ce888708643e750f7631"

e2e/src/config.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -7,7 +7,7 @@ use std::{collections::BTreeMap, fs, path::Path};
 
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
-use toml_edit::Document;
+use toml_edit::DocumentMut;
 
 #[derive(Serialize, Deserialize)]
 pub struct Sha256Hash(
@@ -41,7 +41,8 @@ pub struct Avb {
 pub enum RamdiskContent {
     Init,
     Otacerts,
-    InitAndOtacerts,
+    FirstStage,
+    DsuKeyDir,
     Dlkm,
 }
 
@@ -62,7 +63,7 @@ pub struct BootData {
     #[serde(default)]
     pub kernel: bool,
     #[serde(default)]
-    pub ramdisks: Vec<RamdiskContent>,
+    pub ramdisks: Vec<Vec<RamdiskContent>>,
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -120,12 +121,12 @@ pub struct Config {
     pub profile: BTreeMap<String, Profile>,
 }
 
-pub fn load_config(path: &Path) -> Result<(Config, Document)> {
+pub fn load_config(path: &Path) -> Result<(Config, DocumentMut)> {
     let contents =
         fs::read_to_string(path).with_context(|| format!("Failed to read config: {path:?}"))?;
     let config: Config = toml_edit::de::from_str(&contents)
         .with_context(|| format!("Failed to parse config: {path:?}"))?;
-    let document: Document = contents.parse().unwrap();
+    let document: DocumentMut = contents.parse().unwrap();
 
     Ok((config, document))
 }

e2e/src/main.rs

@@ -227,6 +227,14 @@ fn ramdisk_add_otacerts(entries: &mut Vec<CpioEntry>, cert_ota: &Certificate) ->
     Ok(())
 }
 
+fn ramdisk_add_first_stage(entries: &mut Vec<CpioEntry>) {
+    entries.push(CpioEntry::new_directory(b"first_stage_ramdisk", 0o755));
+}
+
+fn ramdisk_add_dsu_key_dir(entries: &mut Vec<CpioEntry>) {
+    entries.push(CpioEntry::new_directory(b"first_stage_ramdisk/avb", 0o755));
+}
+
 fn ramdisk_add_dlkm(entries: &mut Vec<CpioEntry>) {
     for path in [b"lib".as_slice(), b"lib/modules".as_slice()] {
         entries.push(CpioEntry::new_directory(path, 0o755));
@@ -245,25 +253,29 @@ fn ramdisk_add_dlkm(entries: &mut Vec<CpioEntry>) {
 }
 
 fn create_ramdisk(
-    content: RamdiskContent,
+    content_list: &[RamdiskContent],
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<Vec<u8>> {
     let mut entries = vec![];
 
-    match content {
-        RamdiskContent::Init => {
-            ramdisk_add_init(&mut entries);
-        }
-        RamdiskContent::Otacerts => {
-            ramdisk_add_otacerts(&mut entries, cert_ota)?;
-        }
-        RamdiskContent::InitAndOtacerts => {
-            ramdisk_add_init(&mut entries);
-            ramdisk_add_otacerts(&mut entries, cert_ota)?;
-        }
-        RamdiskContent::Dlkm => {
-            ramdisk_add_dlkm(&mut entries);
+    for content in content_list {
+        match content {
+            RamdiskContent::Init => {
+                ramdisk_add_init(&mut entries);
+            }
+            RamdiskContent::Otacerts => {
+                ramdisk_add_otacerts(&mut entries, cert_ota)?;
+            }
+            RamdiskContent::FirstStage => {
+                ramdisk_add_first_stage(&mut entries);
+            }
+            RamdiskContent::DsuKeyDir => {
+                ramdisk_add_dsu_key_dir(&mut entries);
+            }
+            RamdiskContent::Dlkm => {
+                ramdisk_add_dlkm(&mut entries);
+            }
         }
     }
 
@@ -298,7 +310,7 @@ fn create_boot_image(
     let ramdisks = boot_data
         .ramdisks
         .iter()
-        .map(|c| create_ramdisk(*c, cert_ota, cancel_signal))
+        .map(|c| create_ramdisk(c, cert_ota, cancel_signal))
         .collect::<Result<Vec<_>>>()?;
 
     let boot_image = match boot_data.version {
@@ -361,17 +373,20 @@ fn create_boot_image(
                     ramdisk_metas: boot_data
                         .ramdisks
                         .iter()
-                        .map(|c| match c {
-                            RamdiskContent::Dlkm => RamdiskMeta {
-                                ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_DLKM,
-                                ramdisk_name: "dlkm".to_owned(),
-                                board_id: Default::default(),
-                            },
-                            _ => RamdiskMeta {
-                                ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_PLATFORM,
-                                ramdisk_name: String::new(),
-                                board_id: Default::default(),
-                            },
+                        .map(|c_list| {
+                            if c_list.iter().any(|c| *c == RamdiskContent::Dlkm) {
+                                RamdiskMeta {
+                                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_DLKM,
+                                    ramdisk_name: "dlkm".to_owned(),
+                                    board_id: Default::default(),
+                                }
+                            } else {
+                                RamdiskMeta {
+                                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_PLATFORM,
+                                    ramdisk_name: String::new(),
+                                    board_id: Default::default(),
+                                }
+                            }
                         })
                         .collect(),
                     bootconfig: String::new(),
@@ -968,6 +983,7 @@ fn patch_image(
         keys.ota_pass_file.path().as_os_str(),
         OsStr::new("--cert-ota"),
         keys.ota_cert_file.path().as_os_str(),
+        OsStr::new("--dsu"),
     ];
     args.extend_from_slice(extra_args);
 

xtask/Cargo.toml

@@ -12,4 +12,4 @@ publish = false
 anyhow = "1.0.75"
 clap = { version = "4.4.1", features = ["derive"] }
 regex = { version = "1.9.4", default-features = false, features = ["perf", "std"] }
-toml_edit = "0.21.0"
+toml_edit = "0.22.9"

xtask/src/version.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -11,7 +11,7 @@ use std::{
 
 use anyhow::{bail, Result};
 use clap::Parser;
-use toml_edit::{value, Document};
+use toml_edit::{value, DocumentMut};
 
 use crate::WORKSPACE_DIR;
 
@@ -19,7 +19,7 @@ fn update_cargo_version(version: &str) -> Result<()> {
     let path = Path::new(WORKSPACE_DIR).join("Cargo.toml");
     let data = fs::read_to_string(&path)?;
 
-    let mut document: Document = data.parse()?;
+    let mut document: DocumentMut = data.parse()?;
     document["workspace"]["package"]["version"] = value(version);
 
     fs::write(path, document.to_string())?;