Version 3.6.0

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

.github/workflows/ci.yml

@@ -13,24 +13,60 @@ concurrency:
 
 jobs:
   build:
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ matrix.artifact.os }}
     env:
       CARGO_TERM_COLOR: always
+      # https://github.com/rust-lang/rust/issues/78210
       RUSTFLAGS: -C strip=symbols -C target-feature=+crt-static
+      TARGETS: ${{ join(matrix.artifact.targets, ' ') || matrix.artifact.name }}
+      ANDROID_API: ${{ matrix.artifact.android_api }}
     strategy:
       fail-fast: false
       matrix:
-        os:
-          - ubuntu-latest
-          - windows-latest
-          - macos-latest
+        artifact:
+          - os: ubuntu-latest
+            name: x86_64-unknown-linux-gnu
+          - os: windows-latest
+            name: x86_64-pc-windows-msvc
+          - os: macos-latest
+            name: universal-apple-darwin
+            targets:
+              - aarch64-apple-darwin
+              - x86_64-apple-darwin
+            combine: lipo
+          # ubuntu-latest is not 24.04 yet and 22.04's qemu-user-static segfaults.
+          - os: ubuntu-24.04
+            name: aarch64-linux-android31
+            targets:
+              - aarch64-linux-android
+            android_api: '31'
     steps:
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           # For git describe
           fetch-depth: 0
 
+      - name: Install qemu-user-static
+        if: ${{ contains(matrix.artifact.name, 'android') }}
+        shell: bash
+        run: |
+          sudo apt-get -y update
+          sudo apt-get -y install qemu-user-static
+
+      - name: Set Android temporary directory
+        if: ${{ contains(matrix.artifact.name, 'android') }}
+        shell: bash
+        run: |
+          echo "TMPDIR=/tmp" >> "${GITHUB_ENV}"
+
+      - name: Install cargo-android
+        shell: bash
+        run: |
+          cargo install \
+              --git https://github.com/chenxiaolong/cargo-android \
+              --tag v0.1.1
+
       - name: Get version
         id: get_version
         shell: bash
@@ -40,14 +76,12 @@ jobs:
               | sed -E "s/^v//g;s/([^-]*-g)/r\1/;s/-/./g" \
               >> "${GITHUB_OUTPUT}"
 
-      - name: Get Rust target triple
-        id: get_target
+      - name: Install toolchains
         shell: bash
-        env:
-          RUSTC_BOOTSTRAP: '1'
         run: |
-          echo -n 'name=' >> "${GITHUB_OUTPUT}"
-          rustc -vV | sed -n 's|host: ||p' >> "${GITHUB_OUTPUT}"
+          for target in ${TARGETS}; do
+              rustup target add "${target}"
+          done
 
       - name: Cache Rust dependencies
         uses: Swatinem/rust-cache@v2
@@ -55,51 +89,75 @@ jobs:
       - name: Clippy
         shell: bash
         run: |
-          cargo clippy --release --workspace --features static \
-              --target ${{ steps.get_target.outputs.name }}
+          for target in ${TARGETS}; do
+              cargo android \
+                  clippy --release --workspace --features static \
+                  --target "${target}"
+          done
 
       - name: Build
         shell: bash
         run: |
-          cargo build --release --workspace --features static \
-              --target ${{ steps.get_target.outputs.name }}
+          for target in ${TARGETS}; do
+              cargo android \
+                  build --release --workspace --features static \
+                  --target "${target}"
+          done
 
       - name: Tests
         shell: bash
         run: |
-          cargo test --release --workspace --features static \
-              --target ${{ steps.get_target.outputs.name }}
+          for target in ${TARGETS}; do
+              cargo android \
+                  test --release --workspace --features static \
+                  --target "${target}"
+          done
 
       - name: End to end tests
         shell: bash
         run: |
-          cargo run --release -p e2e --features static \
-              --target ${{ steps.get_target.outputs.name }} \
-              -- test -a -c e2e/e2e.toml
+          for target in ${TARGETS}; do
+              cargo android \
+                  run --release -p e2e --features static \
+                  --target "${target}" \
+                  -- test -a -c e2e/e2e.toml
+          done
 
-      - name: Archive documentation
-        uses: actions/upload-artifact@v3
-        with:
-          name: avbroot-${{ steps.get_version.outputs.version }}-${{ steps.get_target.outputs.name }}
-          path: |
-            LICENSE
-            README.md
-
-      # Due to https://github.com/rust-lang/rust/issues/78210, we have to use
-      # the --target option, which puts all output files in a different path.
-      # Symlink that path to the normal output directory so that we don't need
-      # to specify the Rust triple everywhere.
-      - name: Symlink target directory
+      - name: Create output directory
         shell: bash
         run: |
           rm -rf target/output
-          ln -s ${{ steps.get_target.outputs.name }}/release target/output
 
-      # This is separate so we can have a flat directory structure.
+          case "${{ matrix.artifact.combine }}" in
+          lipo)
+              mkdir target/output
+              cmd=(lipo -output target/output/avbroot -create)
+              for target in ${TARGETS}; do
+                  cmd+=("target/${target}/release/avbroot")
+              done
+              "${cmd[@]}"
+              ;;
+          '')
+              ln -s "${TARGETS}/release" target/output
+              ;;
+          *)
+              echo >&2 "Unsupported combine argument"
+              exit 1
+              ;;
+          esac
+
+      # This is done to ensure a flat directory structure. The upload-artifact
+      # action no longer allows multiple uploads to the same destination.
+      - name: Copy documentation to target directory
+        shell: bash
+        run: cp LICENSE README.md target/output/
+
       - name: Archive executable
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: avbroot-${{ steps.get_version.outputs.version }}-${{ steps.get_target.outputs.name }}
+          name: avbroot-${{ steps.get_version.outputs.version }}-${{ matrix.artifact.name }}
           path: |
+            target/output/LICENSE
+            target/output/README.md
             target/output/avbroot
             target/output/avbroot.exe

.github/workflows/deny.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Run cargo-deny
         uses: EmbarkStudios/cargo-deny-action@v1

.github/workflows/release.yml

@@ -25,14 +25,12 @@ jobs:
           echo "version=${version}" >> "${GITHUB_OUTPUT}"
 
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Create release
-        uses: softprops/action-gh-release@v1
+        uses: softprops/action-gh-release@v2
         with:
-          token: ${{ secrets.GITHUB_TOKEN }}
           tag_name: v${{ steps.get_version.outputs.version }}
           name: Version ${{ steps.get_version.outputs.version }}
           body_path: RELEASE.md
           draft: true
-          prerelease: false

CHANGELOG.md

@@ -7,6 +7,69 @@
     to update the actual links at the bottom of the file.
 -->
 
+### Version 3.6.0
+
+* Add support for gzip compression when computing CoW size estimates ([Issue #332], [PR #333])
+  * This allows `--replace` to successfully replace dynamic partitions on legacy devices, like the Pixel 4a 5G
+* Minor code cleanup ([PR #334], [PR #335])
+
+### Version 3.5.0
+
+* Update all dependencies ([PR #329])
+* Add new unpack and pack commands for `payload.bin` files ([Issue #328], [PR #331])
+
+### Version 3.4.1
+
+* Update all dependencies ([PR #321])
+* Add support for Magisk 27006 ([PR #323])
+
+### Version 3.4.0
+
+* Fix (unreachable) minor error handling logic when attempting to use unsupported AVB signing algorithms ([PR #311])
+* Add support for performing signing operations with external programs ([Issue #310], [PR #312])
+  * See the linked issue for an example of how to sign with a Yubikey.
+
+### Version 3.3.0
+
+* Recompute CoW size estimate when replacing dynamic partitions ([Issue #306], [PR #307])
+  * Fixes out of space error when flashing a patched OTA that uses `--replace` to replace a dynamic partition (eg. `system`) with a larger or more incompressible image
+* Add `avbroot payload info` subcommand for inspecting `payload.bin` headers ([PR #309])
+
+### Version 3.2.3
+
+* Add prebuilt binary for Android (aarch64) ([PR #304])
+
+### Version 3.2.2
+
+* Add new `--recompute-size` option to `avbroot avb pack` to automatically recompute the image size for resizable images ([Discussion #294], [PR #296])
+* Add new `--output-info` option to `avbroot avb pack` to write a new `avb.toml` file containing computed values ([PR #297])
+* Add support for upcoming Magisk Canary 27003 ([Issue #301], [PR #268])
+
+### 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])
+
+### Version 3.1.2
+
+* Use `fastboot flashall` for initial setup to avoid needing to manually flash every partition ([PR #253])
+* Remove binary test files in the git repo and generate them at runtime ([Issue #265], [PR #276])
+* Fix portions of a couple error messages being incorrectly quoted ([PR #277])
+
+### Version 3.1.1
+
+* Cache salted SHA-256 contexts for a small performance improvement ([PR #257])
+* Fix loading certificates that have extra text outside of the marker lines ([PR #261])
+
 ### Version 3.1.0
 
 * The `OEMUnlockOnBoot` module has been split out to a separate repo ([Discussion #235], [PR #246])
@@ -155,6 +218,8 @@ 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
+[Discussion #294]: https://github.com/chenxiaolong/avbroot/discussions/294
 [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
@@ -169,6 +234,15 @@ Behind-the-scenes changes:
 [Issue #222]: https://github.com/chenxiaolong/avbroot/issues/222
 [Issue #223]: https://github.com/chenxiaolong/avbroot/issues/223
 [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
+[Issue #301]: https://github.com/chenxiaolong/avbroot/issues/301
+[Issue #306]: https://github.com/chenxiaolong/avbroot/issues/306
+[Issue #310]: https://github.com/chenxiaolong/avbroot/issues/310
+[Issue #328]: https://github.com/chenxiaolong/avbroot/issues/328
+[Issue #332]: https://github.com/chenxiaolong/avbroot/issues/332
 [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
@@ -242,5 +316,30 @@ Behind-the-scenes changes:
 [PR #246]: https://github.com/chenxiaolong/avbroot/pull/246
 [PR #247]: https://github.com/chenxiaolong/avbroot/pull/247
 [PR #251]: https://github.com/chenxiaolong/avbroot/pull/251
+[PR #253]: https://github.com/chenxiaolong/avbroot/pull/253
 [PR #255]: https://github.com/chenxiaolong/avbroot/pull/255
 [PR #256]: https://github.com/chenxiaolong/avbroot/pull/256
+[PR #257]: https://github.com/chenxiaolong/avbroot/pull/257
+[PR #261]: https://github.com/chenxiaolong/avbroot/pull/261
+[PR #268]: https://github.com/chenxiaolong/avbroot/pull/268
+[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
+[PR #296]: https://github.com/chenxiaolong/avbroot/pull/296
+[PR #297]: https://github.com/chenxiaolong/avbroot/pull/297
+[PR #304]: https://github.com/chenxiaolong/avbroot/pull/304
+[PR #307]: https://github.com/chenxiaolong/avbroot/pull/307
+[PR #309]: https://github.com/chenxiaolong/avbroot/pull/309
+[PR #311]: https://github.com/chenxiaolong/avbroot/pull/311
+[PR #312]: https://github.com/chenxiaolong/avbroot/pull/312
+[PR #321]: https://github.com/chenxiaolong/avbroot/pull/321
+[PR #323]: https://github.com/chenxiaolong/avbroot/pull/323
+[PR #329]: https://github.com/chenxiaolong/avbroot/pull/329
+[PR #331]: https://github.com/chenxiaolong/avbroot/pull/331
+[PR #333]: https://github.com/chenxiaolong/avbroot/pull/333
+[PR #334]: https://github.com/chenxiaolong/avbroot/pull/334
+[PR #335]: https://github.com/chenxiaolong/avbroot/pull/335

Cargo.toml

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

README.extra.md

@@ -32,6 +32,10 @@ This subcommand packs a new AVB image from the `avb.toml` file and, for appended
 
 Note that if the image is an appended image and its hash or hash tree descriptor uses an insecure algorithm, like `sha1`, then it will automatically be promoted to `sha256`.
 
+By default, for appended vbmeta images, the output image size will match the size of the original image that was unpacked. This size is specified by the `image_size` field in `avb.toml`. If the image is resizable (eg. `system`), then passing in `--recompute-size` will cause the `image_size` field to be ignored and the smallest possible output file that fits the raw image and AVB metadata will be built. This avoids wasting space if `raw.img` shrunk or allows the packing to work at all if `raw.img` grew. **Do not use this option for non-resizable images** (eg. `boot`) or else the device won't be able to boot.
+
+When packing an image, several of the fields in `avb.toml` may potentially be recomputed. To write a TOML file containing the new values, use `--output-info <output TOML>`. It is safe to overwrite the existing `avb.toml` if desired.
+
 ### Repacking an AVB image
 
 ```bash
@@ -248,3 +252,43 @@ avbroot hash-tree verify -i <input data file> -H <input hash tree file>
 ```
 
 This will check if the input file has any corrupted blocks. Currently, the command cannot report which specific blocks are corrupted, only whether the file is valid.
+
+## `avbroot payload`
+
+### Unpacking a payload binary
+
+```bash
+avbroot payload unpack -i <input payload>
+```
+
+This subcommand unpacks the payload header information to `payload.toml` and the partition images to the `payload_images` directory.
+
+Only full payload binaries can be unpacked. Delta payload binaries from incremental OTAs are not supported.
+
+### Packing a payload binary
+
+```bash
+avbroot payload pack -o <output payload> --key <OTA private key>
+```
+
+This subcommand packs a new payload binary from the `payload.toml` file and `payload_images` directory. Any `.img` files in the `payload_images` directory that don't have a corresponding entry in `payload.toml` are silently ignored.
+
+Packing a payload binary requires compressing all of the partition images, which is very CPU intensive. If re-signing an existing payload binary without making any other modifications is all that's needed, use the `repack` subcommand instead.
+
+### Repacking a payload binary
+
+```bash
+avbroot payload repack -i <input payload> -o <output payload> --key <OTA private key>
+```
+
+This subcommand is logically equivalent to `avbroot payload unpack` followed by `avbroot payload pack`, except significantly more efficient. Instead of decompressing and recompressing all partition images, the raw data is directly copied from the input payload binary.
+
+This is useful for re-signing a payload binary without making any other changes.
+
+### Showing payload header information
+
+```bash
+avbroot payload info -i <payload>
+```
+
+This subcommand shows all of the payload header fields (which will likely be extremely long).

README.md

@@ -1,5 +1,7 @@
 # avbroot
 
+(This page is also available in: [Russian (Русский)](./README.ru.md).)
+
 avbroot is a program for patching Android A/B-style OTA images for root access while preserving AVB (Android Verified Boot) using custom signing keys. It is compatible with both Magisk and KernelSU. If desired, it can also just re-sign an OTA without enabling root access.
 
 Having a good understanding of how AVB and A/B OTAs work is recommended prior to using avbroot. At the very least, please make sure the [warnings and caveats](#warnings-and-caveats) are well-understood to avoid the risk of hard bricking.
@@ -13,9 +15,9 @@ Having a good understanding of how AVB and A/B OTAs work is recommended prior to
   * `payload.bin` exists
   * `META-INF/com/android/metadata` (Android 10-11) or `META-INF/com/android/metadata.pb` (Android 12+) exists
 
-* The device must support using a custom public key for the bootloader's root of trust. This is normally done via the `fastboot flash avb_custom_key` command. All Pixel devices with unlockable bootloaders since the Pixel 2 support this. Other devices may support it as well, but there's no easy way to check without just trying it.
+* The device must support using a custom public key for the bootloader's root of trust. This is normally done via the `fastboot flash avb_custom_key` command.
 
-  * NOTE: Some OnePlus devices have a broken implementation where a custom public key can be set, but the device won't boot despite having proper signatures. Downgrading the bootloader to the version shipped with Android 11 might potentially help. This problem has been reported across multiple OnePlus models ([#186](https://github.com/chenxiaolong/avbroot/issues/186), [#195](https://github.com/chenxiaolong/avbroot/discussions/195), [#212](https://github.com/chenxiaolong/avbroot/issues/212)).
+  A list of devices known to work can be found in the issue tracker at [#299](https://github.com/chenxiaolong/avbroot/issues/299).
 
 ## Patches
 
@@ -25,7 +27,7 @@ avbroot applies the following patches to the partition images:
 
 * The `boot`, `recovery`, or `vendor_boot` image, depending on device, is patched to replace the OTA signature verification certificates with the custom OTA signing certificate. This allows future patched OTAs to be sideloaded from recovery mode after the bootloader has been locked. It also prevents accidental flashing of the original unpatched OTA.
 
-* The `system` or `my_engineering` image, depending on device, is also patched to replace the OTA signature verification certificates. This prevents the OS' system updater app from installing an unpatched OTA and also allows the use of custom OTA updater apps.
+* The `system` image is also patched to replace the OTA signature verification certificates. This prevents the OS' system updater app from installing an unpatched OTA and also allows the use of custom OTA updater apps.
 
 ## Warnings and Caveats
 
@@ -33,7 +35,11 @@ avbroot applies the following patches to the partition images:
 
     Repeat: **_ALWAYS leave `OEM unlocking` enabled if rooted._**
 
-* Any operation that causes an improperly-signed boot image to be flashed will result in the device being unbootable and unrecoverable without unlocking the bootloader again (and thus, triggering a data wipe). This includes the `Direct install` method for updating Magisk. Magisk updates **must** be done by repatching the OTA, not via the app.
+* Any operation that causes an improperly-signed boot image to be flashed will result in the device being unbootable and unrecoverable without unlocking the bootloader again (and thus, triggering a data wipe). A couple ways an improperly-signed boot image could be flashed include:
+
+    * The `Direct install` method for updating Magisk. Magisk updates **must** be done by repatching the OTA, not via the app.
+
+    * The `Uninstall Magisk` feature in Magisk. If root access is no longer needed, Magisk **must** be removed by repatching the OTA with the `--rootless` option, not via the app.
 
     If the boot image is ever modified, **do not reboot**. [Open an issue](https://github.com/chenxiaolong/avbroot/issues/new) for support and be very clear about what steps were done that lead to the situation. If Android is still running and root access works, it might be possible to recover without wiping and starting over.
 
@@ -124,47 +130,56 @@ When patching OTAs for multiple devices, generating unique keys for each device
     avbroot key generate-cert -k ota.key -o ota.crt
     ```
 
-The commands above are provided for convenience. avbroot is compatible with any standard PKCS8-encoded 4096-bit RSA private key and PEM-encoded X509 certificate, like those generated by openssl.
+The commands above are provided for convenience. avbroot is compatible with any standard PKCS#8-encoded 4096-bit RSA private key and PEM-encoded X509 certificate, like those generated by openssl.
 
-If you lose your AVB or OTA signing key, you will no longer be able to sign new OTA zips. You will have to generate new signing keys and unlock your bootloader again (including a data wipe). Follow the [Usage section](#usage) as if doing an initial setup.
+If you lose your AVB or OTA signing key, you will no longer be able to sign new OTA zips. You will have to generate new signing keys and unlock your bootloader again (triggering a data wipe). Follow the [Usage section](#usage) as if doing an initial setup.
 
 ## Initial setup
 
-1. Reboot into fastboot mode and unlock the bootloader if it isn't already unlocked. This will trigger a data wipe.
+1. Make sure that the version of fastboot is 34 or newer. Older versions have bugs that prevent the `fastboot flashall` command (required later) from working properly.
 
-2. When setting things up for the first time, the device must already be running the correct OS. Flash the original unpatched OTA if needed.
+    ```bash
+    fastboot --version
+    ```
 
-3. Extract the partition images from the patched OTA that are different from the original.
+2. Reboot into fastboot mode and unlock the bootloader if it isn't already unlocked. This will trigger a data wipe.
+
+    ```bash
+    fastboot flashing unlock
+    ```
+
+3. When setting things up for the first time, the device must already be running the correct OS. Flash the original unpatched OTA if needed.
+
+4. Extract the partition images from the patched OTA that are different from the original.
 
     ```bash
     avbroot ota extract \
         --input /path/to/ota.zip.patched \
-        --directory extracted
+        --directory extracted \
+        --fastboot
     ```
 
-4. Flash the partition images that were extracted.
+    If you prefer to extract and flash all OS partitions just to be safe, pass in `--all`.
 
-    For each partition inside `extracted/`, except for `system`, run:
+5. Flash the partition images that were extracted.
 
     ```bash
-    fastboot flash <partition> extracted/<partition>.img
+    ANDROID_PRODUCT_OUT=extracted fastboot flashall --skip-reboot
     ```
 
-    Then, reboot into recovery's fastbootd mode and flash `system`:
-
-    ```bash
-    fastboot reboot fastboot
-    fastboot flash system extracted/system.img
-    ```
-
-5. Set up the custom AVB public key in the bootloader.
+    Note that this only flashes the OS partitions. The bootloader and modem/radio partitions are left untouched due to fastboot limitations. If they are not already up to date or if unsure, after fastboot completes, follow the steps in the [updates section](#updates) to sideload the patched OTA once. Sideloading OTAs always ensures that all partitions are up to date.
+
+    Alternatively, for Pixel devices, running `flash-base.sh` from the factory image will also update the bootloader and modem.
+
+6. Set up the custom AVB public key in the bootloader after rebooting from fastbootd to bootloader.
 
     ```bash
+    fastboot reboot-bootloader
     fastboot erase avb_custom_key
     fastboot flash avb_custom_key /path/to/avb_pkmd.bin
     ```
 
-6. **[Optional]** Before locking the bootloader, reboot into Android once to confirm that everything is properly signed.
+7. **[Optional]** Before locking the bootloader, reboot into Android once to confirm that everything is properly signed.
 
     Install the Magisk or KernelSU app and run the following command:
 
@@ -178,13 +193,19 @@ If you lose your AVB or OTA signing key, you will no longer be able to sign new
     init: [libfs_avb]Returning avb_handle with status: Success
     ```
 
-7. Reboot back into fastboot and lock the bootloader. This will trigger a data wipe again.
+8. Reboot back into fastboot and lock the bootloader. This will trigger a data wipe again.
+
+    ```bash
+    fastboot flashing lock
+    ```
+
+    Confirm by pressing volume down and then power. Then reboot.
 
     Remember: **Do not uncheck `OEM unlocking`!**
 
     **WARNING**: If you are flashing CalyxOS, the setup wizard will [automatically turn off the `OEM unlocking` switch](https://github.com/CalyxOS/platform_packages_apps_SetupWizard/blob/7d2df25cedcbff83ddb608e628f9d97b38259c26/src/org/lineageos/setupwizard/SetupWizardApp.java#L135-L140). Make sure to manually reenable it again from Android's developer settings. Consider using the [`OEMUnlockOnBoot` module](https://github.com/chenxiaolong/OEMUnlockOnBoot) to automatically ensure OEM unlocking is enabled on every boot.
 
-8. That's it! To install future OS, Magisk, or KernelSU updates, see the [next section](#updates).
+9. That's it! To install future OS, Magisk, or KernelSU updates, see the [next section](#updates).
 
 ## Updates
 
@@ -352,6 +373,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:
@@ -409,6 +436,28 @@ avbroot ota extract \
     --all
 ```
 
+### Signing with an external program
+
+avbroot supports delegating all RSA signing operations to an external program with the `--signing-helper` option. When using this option, the `--key-avb` and `--key-ota` options must be given a public key instead of a private key.
+
+For each signing operation, avbroot will invoke the program with:
+
+```bash
+<helper> <algorithm> <public key>
+```
+
+The algorithm is one of `SHA{256,512}_RSA{2048,4096}` and the public key is what was passed to avbroot. The program can use the public key to find the corresponding private key (eg. on a hardware security module). avbroot will write a PKCS#1 v1.5 padded digest to `stdin` and the helper program is expected to perform a raw RSA signing operation and write the raw signature (octet string matching key size) to `stdout`.
+
+By default, this behavior is compatible with the `--signing_helper` option in AOSP's avbtool. However, avbroot additionally extends the arguments to support non-interactive use. If `--pass-{avb,ota}-file` or `--pass-{avb,ota}-env-var` are used, then the helper program will be invoked with two additional arguments that point to the password file or environment variable.
+
+```bash
+<helper> <algorithm> <public key> file <pass file>
+# or
+<helper> <algorithm> <public key> env <env file>
+```
+
+Note that avbroot will verify the signature returned by helper program against the public key. This ensures that the patching process will fail appropriately if the wrong private key was used.
+
 ## Building from source
 
 Make sure the [Rust toolchain](https://www.rust-lang.org/) is installed. Then run:
@@ -423,25 +472,19 @@ Debug builds work too, but they will run significantly slower (in the sha256 com
 
 By default, the executable links to the system's bzip2 and liblzma libraries, which are the only external libraries avbroot depends on. To compile and statically link these two libraries, pass in `--features static`.
 
+### Android cross-compilation
+
+To cross-compile for Android, install [cargo-android](https://github.com/chenxiaolong/cargo-android) and use the `cargo android` wrapper. To make a release build for aarch64, run:
+
+```bash
+cargo android build --release --target aarch64-linux-android
+```
+
+It is possible to run the tests if the host is running Linux, qemu-user-static is installed, and the executable is built with `RUSTFLAGS=-C target-feature=+crt-static` and `--features static`.
+
 ## Verifying digital signatures
 
-First, save the public key to a file listing the keys to be trusted. This is the same key listed in [the author's profile](https://github.com/chenxiaolong/).
-
-```bash
-echo 'avbroot ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDOe6/tBnO7xZhAWXRj3ApUYgn+XZ0wnQiXM8B7tPgv4' > avbroot_trusted_keys
-```
-
-Then, verify the signature of the zip file using the list of trusted keys.
-
-```bash
-ssh-keygen -Y verify -f avbroot_trusted_keys -I avbroot -n file -s <file>.zip.sig < <file>.zip
-```
-
-If the file is successfully verified, the output will be:
-
-```
-Good "file" signature for avbroot with ED25519 key SHA256:Ct0HoRyrFLrnF9W+A/BKEiJmwx7yWkgaW/JvghKrboA
-```
+To verify the digital signatures of the downloads, follow [the steps here](https://github.com/chenxiaolong/chenxiaolong/blob/master/VERIFY_SSH_SIGNATURES.md).
 
 ## Contributing
 

README.ru.md

@@ -0,0 +1,470 @@
+# avbroot
+
+avbroot – это программа для модификации OTA-образов Android A/B-формата с целью получения root-прав при сохранении прохождения AVB (Android Verified Boot) с использованием кастомных (пользовательских) ключей подписи. Она совместима как с Magisk, так и с KernelSU. При необходимости можно просто переподписать OTA, без получения root-доступа.
+
+Прежде чем использовать avbroot, рекомендуется иметь хорошее понимание того, как работают AVB и OTA в формате A/B. Как минимум, следует ознакомиться с [разделом предостережений,](#предостережения) чтобы избежать хардбрика устройства.
+
+**ПРИМЕЧАНИЕ:** avbroot 2.0 была переписана на Rust и больше не имеет в основе никакого кода AOSP, а CLI полностью обратно совместим. Тем не менее, старую реализацию на Python можно найти в одноименной ветке `python`.
+
+## Требования
+
+* Поддерживаются только устройства, использующие современную A/B-разметку. Это большинство девайсов, выпускаемых с Android 10 и новее (за исключением устройств от Samsung). Чтобы проверить, использует ли ваш телефон необходимую схему разметки, откройте zip-архив OTA и проверьте:
+
+  * наличие файла `payload.bin` (обычно находится в корне архива)
+  * наличие файла `META-INF/com/android/metadata` (Android 10-11) или `META-INF/com/android/metadata.pb` (Android 12+)
+
+* Устройство должно поддерживать установку пользовательского публичного ключа для подтверждения статуса доверия загрузчика. Обычно это производится с помощью команды `fastboot flash avb_custom_key`.
+
+  Список девайсов, на которых проверялась совместимость с указанным выше функционалом, находится здесь: [#299.](https://github.com/chenxiaolong/avbroot/issues/299)
+
+## Патчи
+
+avbroot модифицирует следующие образы:
+
+* `boot` или `init_boot`, в зависимости от устройства, модифицируется для получения root-доступа. В случае с Magisk, патч будет эквивалентен тому, что производится в самом приложении Magisk.
+
+* `boot`, `recovery` или `vendor_boot`, в зависимости от устройства, модифицируется для замены сертификата проверки подписи OTA на пользовательский. Это позволяет устанавливать будущие пропатченные OTA через режим Recovery уже после блокировки загрузчика, то есть в качестве обновления. Также это предотвращает случайную установку оригинального непропатченного OTA.
+
+* `system` тоже модифицируется для замены сертификата проверки подписи OTA. Это не позволит системному приложению обновлений ОС установить оригинальный непропатченный OTA и дает возможность использовать сторонние приложения для установки пропатченных OTA.
+
+## Предостережения
+
+* **Всегда оставляйте опцию `Заводской разблокировки`** (или OEM unlocking в англ.) **включенной при наличии root-прав с заблокированным загрузчиком.** Это очень важно. Доступ к root-правам потенциально позволяет перезаписать загрузочный раздел из-под системы, будь то сделано случайно или намеренно, файлом, который не был подписан должным образом. В таком случае, система и режим Recovery больше не смогут загрузиться, а команда `fastboot flashing unlock` будет недоступна, потому что параметр Заводской разблокировки отключен. То есть, это приведет к **_хардбрику устройства_**.
+
+    Повторюсь: **_ВСЕГДА оставляйте `Заводскую разблокировку` включенной при наличии root-прав._**
+
+* Любая операция, приводящая к прошивке некорректно подписанного загрузочного образа, приведет к тому, что устройство больше не сможет загрузиться в систему/режим Recovery, а для его восстановления потребуется повторная разблокировка загрузчика (и, следовательно, стирание всех пользовательских данных). К подобным операциям в том числе относятся:
+
+    * Метод `Прямой установки` для обновления Magisk. Magisk можно обновлять **только путем репатчинга OTA,** но не через его приложение.
+
+    * Функция `Удаление Magisk` в приложении Magisk. Если вам больше не нужен root-доступ, Magisk **должен быть удален путем репатчинга OTA** с использованием параметра `--rootless`, но не через его приложение.
+
+    Если в загрузочный раздел были внесены какие-либо изменения, **не перезагружайтесь**. Обратитесь за помощью, [открыв Issue,](https://github.com/chenxiaolong/avbroot/issues/new) и четко разъясните, какие конкретные действия привели к возникновению такой ситуации. Если Android всё еще работает и доступ к root-правам сохранился – вероятно, получится откатить изменения до исходного состояния, не стирая ваши данные.
+
+## Использование
+
+1. Убедитесь, что вы ознакомились и поняли указанные выше [предостережения.](#предостережения) 
+
+2. Скачайте последнюю версию со страницы [релизов.](https://github.com/chenxiaolong/avbroot/releases) Чтобы сверить цифровую подпись, см. раздел [проверки цифровых подписей.](#проверка-цифровых-подписей)
+
+    avbroot – это отдельный исполняемый файл. Он не требует установки и может быть запущен из любого места на диске.
+
+3. [Сгенерируйте ключи подписи.](#генерация-ключей)
+
+4. Пропатчите ОТА-архив с помощью команды:
+
+    ```bash
+    avbroot ota patch \
+        --input /путь/к/ota.zip \
+        --key-avb /путь/к/avb.key \
+        --key-ota /путь/к/ota.key \
+        --cert-ota /путь/к/ota.crt \
+    ```
+
+    Добавьте следующие аргументы в конец команды в зависимости от того, как вы хотите получить root-доступ.
+
+    * Для получения root-доступа с использованием Magisk:
+
+        ```bash
+        --magisk /путь/к/magisk.apk \
+        --magisk-preinit-device <имя>
+        ```
+
+        Если вы не знаете имени раздела предварительной инициализации Magisk, следуйте инструкции [в соответствующем разделе.](#предварительная-инициализация-устройства-для-magisk)
+
+        Если вы пропатчили загрузочный образ вручную через приложение Magisk (вместо автоматического идентичного патчинга через avbroot), используйте следующий аргумент:
+
+        ```bash
+        --prepatched /путь/к/magisk_patched-xxxxx_yyyyy.img
+        ```
+
+    * Для получения root-доступа с использованием KernelSU:
+
+        ```bash
+        --prepatched /путь/к/kernelsu_boot.img
+        ```
+
+    * Без root-доступа:
+
+        ```bash
+        --rootless
+        ```
+
+    Больше информации про существующие аргументы можно найти в разделе [расширенного использования.](#расширенное-использование)
+
+    Если название для `--output` не указывается, то готовый файл будет записан как `<название-ota-zip-в-input>.patched`.
+
+5. Готово! Для прошивки пропатченного OTA следуйте инструкции в разделе [первоначальной настройки.](#первоначальная-настройка) Для последующих обновлений тоже есть соответствующий [раздел обновлений.](#обновления)
+
+## Генерация ключей
+
+Во время патчинга OTA, avbroot подписывает несколько компонентов:
+
+* загрузочный образ (boot)
+* образ vbmeta
+* payload из OTA
+* сам архив OTA
+
+Первые два компонента подписываются ключом AVB, а последние два – ключом OTA. Можно использовать один и тот же ключ, однако в следующих шагах описано, как сгенерировать два отдельных.
+
+Если вы патчите OTA сразу для нескольких устройств, настоятельно рекомендуется генерировать уникальные ключи для каждого девайса – так вы защитите себя от случайной прошивки неподходящего OTA для другого телефона. 
+
+1. Сгенерируйте ключи подписи для AVB и OTA.
+
+    ```bash
+    avbroot key generate-key -o avb.key
+    avbroot key generate-key -o ota.key
+    ```
+
+2. Преобразуйте публичную часть ключа подписи AVB в формат метаданных публичного ключа AVB. Именно этот формат используется в загрузчике устройства для установки пользовательского ключа.
+
+    ```bash
+    avbroot key extract-avb -k avb.key -o avb_pkmd.bin
+    ```
+
+3. Сгенерируйте самоподписанный сертификат для ключа подписи OTA. Он используется режимом Recovery для проверки подписи OTA при сайдлоадинге обновления.
+
+    ```bash
+    avbroot key generate-cert -k ota.key -o ota.crt
+    ```
+
+avbroot совместим с любым стандартным 4096-битным приватным ключом RSA в кодировке PKCS8 и сертификатом X509 в кодировке PEM, например с теми, которые генерируются openssl.
+
+Если вы потеряете ключ(-и) подписи AVB или OTA, вы больше не сможете подписывать новые OTA-архивы. Придется генерировать новые ключи подписи и разблокировать загрузчик (что приведет к стиранию всех данных). В таком случае возвращайтесь к инструкции в разделе [использования.](#использование)
+
+## Первоначальная настройка
+
+1. Убедитесь, что вы используете утилиту fastboot версии 34 или новее. Предыдущие версии содержат баги, что не позволяют команде `fastboot flashall` (которая понадобится по ходу инструкции) работать правильно.
+
+    ```bash
+    fastboot --version
+    ```
+
+2. Перезагрузитесь в режим fastboot и разблокируйте загрузчик, если не сделали этого ранее. Это приведет к стиранию всех пользовательских данных.
+
+    ```bash
+    fastboot flashing unlock
+    ```
+
+3. Перед первой установкой, на устройстве уже должна быть установлена в оригинальном виде та прошивка, пропатченную версию которой вы собираетесь ставить. Если это не так, сначала установите оригинальную непропатченную OTA.
+
+4. Извлекаем из пропатченного OTA модифицированные образы:
+
+    ```bash
+    avbroot ota extract \
+        --input /путь/к/ota.zip.patched \
+        --directory extracted \
+        --fastboot
+    ```
+
+    Если вы на всякий случай хотите прошить вообще все разделы из ОТА, извлечь их можно, указав аргумент `--all`.
+
+5. Прошейте извлеченные образы разделов.
+
+    ```bash
+    ANDROID_PRODUCT_OUT=extracted fastboot flashall --skip-reboot
+    ```
+
+    Обратите внимание, что так прошиваются лишь те образы, что относятся к системе. Разделы загрузчика и модема же остаются нетронутыми из-за ограничений fastboot. Если они не обновлены до необходимой версии, или вы не уверены в этом, после прошивки перейдите к пункту [обновлений](#обновления) и установите пропатченный OTA-архив сайдлоадом в режиме Recovery. Прошивка полного OTA гарантирует, что абсолютно все разделы будут обновлены.
+
+    Для устройств Pixel есть ещё один вариант: запуск скрипта `flash-base.sh` из папки заводских образов (factory images) обновит загрузчик и модем.
+
+6. После перезагрузки из fastbootd в загрузчик (bootloader), установите пользовательский публичный ключ AVB в загрузчик:
+
+    ```bash
+    fastboot reboot-bootloader
+    fastboot erase avb_custom_key
+    fastboot flash avb_custom_key /путь/к/avb_pkmd.bin
+    ```
+
+7. **[Опционально]** Перед блокировкой загрузчика загрузитесь в систему, дабы убедиться, что все подписано правильно.
+
+    Установите приложение Magisk или KernelSU и выполните следующую команду:
+
+    ```bash
+    adb shell su -c 'dmesg | grep libfs_avb'
+    ```
+
+    Если AVB работает корректно, будет выведено следующее сообщение:
+
+    ```bash
+    init: [libfs_avb]Returning avb_handle with status: Success
+    ```
+
+8. Перезагрузитесь в fastboot и заблокируйте загрузчик. Это снова приведет к стиранию данных.
+
+    ```bash
+    fastboot flashing lock
+    ```
+
+    Подтвердите нажатием клавиш уменьшения громкости и включения, а после перезагрузитесь в систему.
+
+    Напоминаю: **не отключайте `Заводскую разблокировку`!**
+
+    **ПРЕДУПРЕЖДЕНИЕ**: Если вы прошили CalyxOS, мастер настройки [автоматически отключит опцию `Заводской разблокировки`.](https://github.com/CalyxOS/platform_packages_apps_SetupWizard/blob/7d2df25cedcbff83ddb608e628f9d97b38259c26/src/org/lineageos/setupwizard/SetupWizardApp.java#L135-L140) Не забудьте снова включить её вручную в настройках для разработчиков. Для перестраховки можете использовать [модуль `OEMUnlockOnBoot`,](https://github.com/chenxiaolong/OEMUnlockOnBoot) который автоматически включает пункт Заводской разблокировки при каждом запуске системы.
+
+9. Готово! Установка последующих обновлений системы, Magisk или KernelSU, описывается в [следующем разделе.](#обновления)
+
+## Обновления
+
+Обновления Android, Magisk и KernelSU выполняются одинаково – исключительно путем обновления или репатчинга того же самого OTA.
+
+1. Если Magisk или KernelSU обновились, сначала установите их новый `.apk`. Если вы случайно открыли приложение после обновления, убедитесь, что оно не начало прошивать загрузочный образ. Если появится предложение обновить сам загрузочный образ – отклоните его.
+
+2. Следуйте инструкции в разделе [использования,](#использование) чтобы пропатчить OTA уже с новым .apk Magisk'а/предварительно пропатченным образом с Magisk или KernelSU.
+
+3. Перезагрузитесь в режим Recovery. Если устройство повисло на сплеше с сообщением "No command", удерживайте кнопку питания, а затем нажмите кнопку увеличения громкости один раз.
+
+4. Обновитесь (Apply update from adb → `adb sideload <ota.zip.patched>`).
+
+5. Готово!
+
+## Возврат на заводскую прошивку
+
+Если вы хотите отказаться от использования avbroot и вернуться на стоковую прошивку:
+
+1. Перезагрузитесь в режим fastboot и разблокируйте загрузчик. Это приведет к стиранию всех пользовательских данных.
+
+2. Удалите пользовательский публичный ключ AVB.
+
+    ```bash
+    fastboot erase avb_custom_key
+    ```
+
+3. Прошейте стоковую прошивку. Готово.
+
+## OTA-обновления
+
+avbroot заменяет `/system/etc/security/otacerts.zip` в разделах системы и Recovery на новый архив, содержащий пользовательский сертификат подписи OTA. Это предотвращает случайную установку непропатченных OTA как при загрузке в Android, так и при сайдлоадинге через Recovery.
+
+Рекомендуется отключить приложение обновлений системы, чтобы оно не пыталось установить непропатченные OTA:
+
+* Стоковая прошивка: Отключите `Автоматические обновления системы` (Automatic system updates в англ.) в настройках для разработчиков.
+* Кастомная прошивка: Отключите приложение обновлений системы (или запретите ему доступ к Интернету) через Настройки -> Приложения -> Все приложения -> (меню/три точки) -> Показать системные -> (найдите приложение обновлений, например Обновления системы/Updater).
+
+Это особенно важно для некоторых кастомных прошивок, поскольку их фирменное приложение для обновления системы может уйти в бесконечный цикл, загружая OTA-обновление, а затем повторяя попытку загрузки и установки при неудачной проверке подписи.
+
+Если вы хотите поднять собственный сервер для ОТА-обновлений, вам может быть интересно приложение [Custota.](https://github.com/chenxiaolong/Custota)
+
+## Режим обслуживания
+
+Некоторые устройства поставляются с режимом обслуживания, который загружает систему с чистым образом `userdata`, благодаря чему специалист по ремонту может проводить диагностику устройства, не имея доступа к пользовательским данным владельца.
+
+Если на устройстве есть root-права, использовать этот режим небезопасно. Если у вас обычная сборка Magisk/KernelSU, не подписанная вашим собственным ключом, кто угодно может установить официальное приложение Magisk/KernelSU в режиме обслуживания и запросто получить root-права без какой-либо аутентификации.
+
+Потому, чтобы безопасно использовать режим обслуживания:
+
+1. Отключите root-доступ на устройстве, пропатчив OTA с аргументом `--rootless` (вместо `--magisk` или `--prepatched`) и прошив его.
+
+2. Включите режим обслуживания.
+
+3. Получив отремонтированное устройство обратно, выйдите из режима обслуживания.
+
+4. Прошейте рутированный OTA в обычном режиме.
+
+Поскольку удаление root-прав и повторное их получение выполняются путем перепрошивки OTA, данные устройства стёрты не будут.
+
+## Предварительная инициализация устройства для Magisk
+
+Magisk версии 25211 и новее требует наличие раздела, доступного для записи пользовательских правил SELinux, к которым необходимо обращаться на ранних этапах загрузки. Его можно определить только на реальном устройстве, поэтому avbroot требует указания точного названия с помощью аргумента `--magisk-preinit-device <имя>`. Чтобы получить имя раздела:
+
+1. Извлеките загрузочный образ из оригинального, непропатченного OTA:
+
+    ```bash
+    avbroot ota extract \
+        --input /path/to/ota.zip \
+        --directory . \
+        --boot-only
+    ```
+
+2. Теперь нужно пропатчить загрузочный образ с помощью приложения Magisk. Это **ДОЛЖНО** быть сделано именно на целевом устройстве или устройстве той же модели! Имя раздела будет неверным и не подойдет, если пропатчить образ на устройстве иной модели.
+
+    Приложение Magisk выведет в лог строку, подобную следующей:
+
+    ```
+    - Pre-init storage partition device ID: <имя>
+    ```
+
+    Также и avbroot может вывести информацию о разделе, обнаруженном Magisk, для этого выполните команду:
+
+    ```bash
+    avbroot boot magisk-info \
+        --image magisk_patched-*.img
+    ```
+
+    Имя раздела будет выведено как: `PREINITDEVICE=<имя>`.
+
+    Теперь, когда имя раздела известно, его нужно указать avbroot с помощью команды `--magisk-preinit-device <имя>`. Имя раздела стоит запомнить или сохранить где-нибудь на будущее, оно вряд ли изменится при обновлении Magisk.
+
+Если запустить приложение Magisk на целевом устройстве невозможно (например, телефон не загружается), пропатчите OTA с аргументом `--ignore-magisk-warnings` и прошейте его. Затем выполните указанные выше шаги и повторно пропатчите OTA, но уже с указанием аргумента `--magisk-preinit-device <имя>`.
+
+## Проверка OTA
+
+Чтобы проверить все подписи и хэши, связанные с установкой OTA и процессом загрузки AVB, выполните команду:
+
+```bash
+avbroot ota verify \
+    --input /путь/к/ota.zip \
+    --cert-ota /путь/к/ota.crt \
+    --public-key-avb /путь/к/avb_pkmd.bin
+```
+
+Эта команда работает для любого OTA, независимо от того, пропатчено оно или нет.
+
+Если опции `--cert-ota` и `--public-key-avb` не указаны, то подписи проверяются только на корректность, не проверяя, совпадают ли они внутри всех файлов.
+
+## Подсказки через Tab
+
+Поскольку avbroot имеет множество опций, будет удобно настроить подсказки с автозаполнением для используемой оболочки. Конфигурации генерируются в самом avbroot.
+
+#### bash
+
+Добавьте в `~/.bashrc`:
+
+```bash
+eval "$(avbroot completion -s bash)"
+```
+
+#### zsh
+
+Добавьте в `~/.zshrc`:
+
+```bash
+eval "$(avbroot completion -s zsh)"
+```
+
+#### fish
+
+Добавьте в `~/.config/fish/config.fish`:
+
+```bash
+avbroot completion -s fish | source
+```
+
+#### PowerShell
+
+Добавьте в загрузочный скрипт PowerShell (`profile.ps1`):
+
+```powershell
+Invoke-Expression (& avbroot completion -s powershell)
+```
+
+## Расширенное использование
+
+### Использование заранее пропатченного boot.img
+
+avbroot может подменить используемый загрузочный образ на заранее пропатченный (вместо того, чтобы самостоятельно применять патч). Это пригодится в случае, если у вас уже имеется пропатченный через приложение Magisk образ ядра или образ с поддержкой KernelSU. Для этого используйте аргумент `--prepatched <загрузочный образ>` вместо `--magisk <apk>`. То есть, указав `--prepatched`, avbroot пропустит применение патчинга Magisk'ом, но по-прежнему применит патч OTA-сертификата.
+
+Обратите внимание, что avbroot проверяет совместимость предварительно пропатченного образа с оригинальным. Например, если поля заголовка образа не совпадают, или вовсе указан иной, незагрузочный образ, то процесс патча будет прерван. Эти проверки, конечно, ничего не гарантируют, но должны предостеречь от случайного использования некорректного образа. Чтобы обойти базовые проверки безопасности, укажите аргумент `--ignore-prepatched-compat`. Если вы хотите убрать вообще все проверки (чего делать крайне не рекомендуется), укажите его дважды.
+
+### Пропуск патчей для root-доступа
+
+avbroot можно использовать для простого переподписания OTA, указав аргумент `--rootless` вместо `--magisk`/`--prepatched`. В таком случае пропатченный OTA не будет рутирован. Единственная модификация, которая будет применена – это замена сертификата проверки OTA, чтобы систему можно было обновлять с помощью будущих пропатченных OTA.
+
+### Подмена разделов
+
+avbroot поддерживает подмену целых образов в OTA, даже тех, что не являются загрузочными (например, `vendor_dlkm`). Образ можно заменить, указав аргумент `--replace <имя раздела> /путь/к/образу.img`.
+
+Единственное, что меняется – это то, откуда считывается раздел. При использовании `--replace` вместо использования образа раздела из оригинального `payload.bin` в OTA, он берется напрямую по указанному вами пути. Таким образом, заменяющие образы разделов должны иметь правильные колонтитулы vbmeta, соответствующие оригинальным.
+
+Это не влияет на ход применения пачтей. Например, при использовании Magisk, патч получения root-прав применяется к загрузочному образу одинаково, независимо от того, был ли он получен из оригинального `payload.bin` или это файл, указанный через `--replace`.
+
+### Очистка флагов vbmeta
+
+Некоторые сборки Android-прошивок могут поставляться с образом `vbmeta`, в котором флаги установлены таким образом, что AVB фактически отключен. Если avbroot сталкивается с такими образом, процесс патчинга завершается ошибкой с сообщением следующего типа:
+
+```
+Verified boot is disabled by vbmeta's header flags: 0x3
+```
+
+Чтобы принудительно включить AVB (очистив флаги), укажите аргумент `--clear-vbmeta-flags`.
+
+### Использование в неинтерактивном режиме
+
+По умолчанию avbroot интерактивно запрашивает пароли к приватным ключам. Чтобы запустить avbroot в неинтерактивном режиме, можно:
+
+* Предоставить пароли через файлы:
+
+    ```bash
+    avbroot ota patch \
+        --pass-avb-file /путь/к/avb.passphrase \
+        --pass-ota-file /путь/к/ota.passphrase \
+        <...>
+    ```
+
+    На Unix-подобных системах "файлы" могут быть каналами ("pipes"). В оболочках, поддерживающих подстановку процесса (bash, zsh и т. д.), пароль можно запросить с помощью команды (например, запрашивая у менеджера паролей).
+
+    ```bash
+    avbroot ota patch \
+        --pass-avb-file <(команда для запроса пароля AVB) \
+        --pass-ota-file <(команда для запроса пароля OTA) \
+        <...>
+    ```
+
+* Предоставить пароли через переменные среды. Это менее безопасно, поскольку любой процесс, запущенный от имени того же пользователя, может видеть значения переменных среды.
+
+    ```bash
+    export PASSPHRASE_AVB="пароль AVB"
+    export PASSPHRASE_OTA="пароль OTA"
+
+    avbroot ota patch \
+        --pass-avb-env-var PASSPHRASE_AVB \
+        --pass-ota-env-var PASSPHRASE_OTA \
+        <...>
+    ```
+
+* Использовать незашифрованные приватные ключи. Крайне не рекомендуется.
+
+### Извлечение всей OTA
+
+Чтобы извлечь все образы, содержащиеся в `payload.bin`, используйте команду:
+
+```bash
+avbroot ota extract \
+    --input /путь/к/ota.zip \
+    --directory extracted \
+    --all
+```
+
+## Сборка из исходного кода
+
+Убедитесь, что у вас установлен [набор инструментов Rust.](https://www.rust-lang.org/ru/) Затем выполните:
+
+```bash
+cargo build --release
+```
+
+Исполняемый файл будет записан в `target/release/avbroot`.
+
+Дебаг-сборки тоже работают, но они будут работать значительно медленнее (в вычислениях sha256), потому что оптимизации компилятора отключены.
+
+По умолчанию исполняемый файл ссылается на системные библиотеки bzip2 и liblzma, от которых зависит avbroot. Чтобы скомпилировать и статически связать эти две библиотеки, укажите аргумент `--features static`.
+
+## Проверка цифровых подписей
+
+Сначала сохраните открытый ключ в файл, перечисляющий ключи, которым можно доверять. Это тот же ключ, который указан [в профиле автора.](https://github.com/chenxiaolong/)
+
+```bash
+echo 'avbroot ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDOe6/tBnO7xZhAWXRj3ApUYgn+XZ0wnQiXM8B7tPgv4' > avbroot_trusted_keys
+```
+
+Затем проверьте подпись zip-файла, используя список доверенных ключей.
+
+```bash
+ssh-keygen -Y verify -f avbroot_trusted_keys -I avbroot -n file -s <file>.zip.sig < <file>.zip
+```
+
+Если файл успешно проверен, вывод будет следующим:
+
+```
+Good "file" signature for avbroot with ED25519 key SHA256:Ct0HoRyrFLrnF9W+A/BKEiJmwx7yWkgaW/JvghKrboA
+```
+
+## Вклад
+
+Буду рад вашему вкладу в разработку! Однако я вряд ли приму изменения для поддержки устройств, которые ведут себя значительно иначе, чем устройства Pixel.
+
+## Лицензия
+
+avbroot распространяется по лицензии GPLv3. Полный текст лицензии см. в [`LICENSE`.](./LICENSE)
+

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,14 +24,15 @@ 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"
+miniz_oxide = "0.8.0"
 num-bigint-dig = "0.8.4"
 num-traits = "0.2.16"
 phf = { version = "0.11.2", features = ["macros"] }
 pkcs8 = { version = "0.10.2", features = ["encryption", "pem"] }
-prost = "0.12.1"
+prost = "0.13.1"
 rand = "0.8.5"
 rayon = "1.7.0"
 regex = { version = "1.9.4", default-features = false, features = ["perf", "std"] }
@@ -47,20 +48,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]
@@ -73,8 +72,9 @@ features = ["deflate"]
 rustix = { version = "0.38.9", default-features = false, features = ["process"] }
 
 [build-dependencies]
-prost-build = "0.12.1"
-protox = "0.5.0"
+constcat = "0.5.0"
+prost-build = "0.13.1"
+protox = "0.7.0"
 
 [dev-dependencies]
 assert_matches = "1.5.0"

avbroot/build.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -30,8 +30,62 @@ fn main() {
 
     let file_descriptors = protox::compile(&protos, [&in_dir]).unwrap();
 
+    const CUE_AI: &str = ".chromeos_update_engine.ApexInfo";
+    const CUE_DAM: &str = ".chromeos_update_engine.DeltaArchiveManifest";
+    const CUE_DPG: &str = ".chromeos_update_engine.DynamicPartitionGroup";
+    const CUE_DPM: &str = ".chromeos_update_engine.DynamicPartitionMetadata";
+    const CUE_PU: &str = ".chromeos_update_engine.PartitionUpdate";
+    const CUE_VABCFS: &str = ".chromeos_update_engine.VABCFeatureSet";
+
+    const DERIVE_SERDE: &str = "#[derive(serde::Deserialize, serde::Serialize)]";
+    const SERDE_DEFAULT: &str = "#[serde(default)]";
+    const SERDE_SKIP: &str = "#[serde(skip)]";
+    const SERDE_SKIP_IF_VEC_EMPTY: &str = "#[serde(skip_serializing_if = \"Vec::is_empty\")]";
+
+    use constcat::concat as c;
+
     prost_build::Config::new()
         .btree_map(["."])
+        // Allow deserializing and serializing the types we care about.
+        .type_attribute(CUE_AI, DERIVE_SERDE)
+        .type_attribute(CUE_DAM, DERIVE_SERDE)
+        .type_attribute(CUE_DPG, DERIVE_SERDE)
+        .type_attribute(CUE_DPM, DERIVE_SERDE)
+        .type_attribute(CUE_PU, DERIVE_SERDE)
+        .type_attribute(CUE_VABCFS, DERIVE_SERDE)
+        // Allow default-initializing all fields.
+        .type_attribute(CUE_AI, SERDE_DEFAULT)
+        .type_attribute(CUE_DAM, SERDE_DEFAULT)
+        .type_attribute(CUE_DPG, SERDE_DEFAULT)
+        .type_attribute(CUE_DPM, SERDE_DEFAULT)
+        .type_attribute(CUE_PU, SERDE_DEFAULT)
+        .type_attribute(CUE_VABCFS, SERDE_DEFAULT)
+        // Don't serialize fields that define the structure of the payload
+        // binary and that we recompute during packing.
+        .field_attribute(c!(CUE_DAM, ".signatures_offset"), SERDE_SKIP)
+        .field_attribute(c!(CUE_DAM, ".signatures_size"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".operations"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".estimate_cow_size"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".old_partition_info"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".new_partition_info"), SERDE_SKIP)
+        // Don't serialize AVB 1.0 fields.
+        .field_attribute(c!(CUE_PU, ".hash_tree_data_extent"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".hash_tree_extent"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".hash_tree_algorithm"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".hash_tree_salt"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".fec_data_extent"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".fec_extent"), SERDE_SKIP)
+        .field_attribute(c!(CUE_PU, ".fec_roots"), SERDE_SKIP)
+        // Don't serialize fields for incremental OTAs.
+        .field_attribute(c!(CUE_PU, ".merge_operations"), SERDE_SKIP)
+        // Don't serialize fields for vendor-signed images, which update_engine
+        // doesn't support anyway.
+        .field_attribute(c!(CUE_PU, ".new_partition_signature"), SERDE_SKIP)
+        // Don't serialize empty lists.
+        .field_attribute(c!(CUE_DAM, ".apex_info"), SERDE_SKIP_IF_VEC_EMPTY)
+        .field_attribute(c!(CUE_DAM, ".partitions"), SERDE_SKIP_IF_VEC_EMPTY)
+        .field_attribute(c!(CUE_DPG, ".partition_names"), SERDE_SKIP_IF_VEC_EMPTY)
+        .field_attribute(c!(CUE_DPM, ".groups"), SERDE_SKIP_IF_VEC_EMPTY)
         .compile_fds(file_descriptors)
         .unwrap();
 }

avbroot/src/cli/args.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -15,7 +15,7 @@ use clap::{Parser, Subcommand, ValueEnum};
 use tracing::{debug, Level};
 use tracing_subscriber::fmt::{format::Writer, time::FormatTime};
 
-use crate::cli::{avb, boot, completion, cpio, fec, hashtree, key, ota};
+use crate::cli::{avb, boot, completion, cpio, fec, hashtree, key, ota, payload};
 
 #[allow(clippy::large_enum_variant)]
 #[derive(Debug, Subcommand)]
@@ -28,6 +28,7 @@ pub enum Command {
     HashTree(hashtree::HashTreeCli),
     Key(key::KeyCli),
     Ota(ota::OtaCli),
+    Payload(payload::PayloadCli),
     /// (Deprecated: Use `avbroot ota patch` instead.)
     Patch(ota::PatchCli),
     /// (Deprecated: Use `avbroot ota extract` instead.)
@@ -36,39 +37,6 @@ pub enum Command {
     MagiskInfo(boot::MagiskInfoCli),
 }
 
-#[derive(Debug, Clone, Copy, ValueEnum)]
-pub enum LogLevel {
-    Trace,
-    Debug,
-    Info,
-    Warn,
-    Error,
-}
-
-impl LogLevel {
-    fn as_level(self) -> Level {
-        match self {
-            Self::Trace => Level::TRACE,
-            Self::Debug => Level::DEBUG,
-            Self::Info => Level::INFO,
-            Self::Warn => Level::WARN,
-            Self::Error => Level::ERROR,
-        }
-    }
-}
-
-impl Default for LogLevel {
-    fn default() -> Self {
-        Self::Info
-    }
-}
-
-impl fmt::Display for LogLevel {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str(self.to_possible_value().ok_or(fmt::Error)?.get_name())
-    }
-}
-
 #[derive(Debug, Clone, Copy, ValueEnum)]
 pub enum LogFormat {
     Short,
@@ -95,8 +63,8 @@ pub struct Cli {
     pub command: Command,
 
     /// Lowest log message severity to output.
-    #[arg(long, global = true, value_name = "LEVEL", default_value_t)]
-    pub log_level: LogLevel,
+    #[arg(long, global = true, value_name = "LEVEL", default_value_t = Level::INFO)]
+    pub log_level: Level,
 
     /// Output format for log messages.
     #[arg(long, global = true, value_name = "FORMAT", default_value_t)]
@@ -123,11 +91,11 @@ impl FormatTime for ShortUptime {
     }
 }
 
-pub fn init_logging(log_level: LogLevel, log_format: LogFormat) {
+pub fn init_logging(log_level: Level, log_format: LogFormat) {
     let builder = tracing_subscriber::fmt()
         .with_writer(io::stderr)
         .with_ansi(io::stderr().is_terminal())
-        .with_max_level(log_level.as_level());
+        .with_max_level(log_level);
 
     match log_format {
         LogFormat::Short => {
@@ -163,6 +131,7 @@ pub fn main(logging_initialized: &AtomicBool, cancel_signal: &AtomicBool) -> Res
         Command::HashTree(c) => hashtree::hash_tree_main(&c, cancel_signal),
         Command::Key(c) => key::key_main(&c),
         Command::Ota(c) => ota::ota_main(&c, cancel_signal),
+        Command::Payload(c) => payload::payload_main(&c, cancel_signal),
         // Deprecated aliases.
         Command::Patch(c) => ota::patch_subcommand(&c, cancel_signal),
         Command::Extract(c) => ota::extract_subcommand(&c, cancel_signal),

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
  */
 
@@ -25,7 +25,7 @@ use serde::{Deserialize, Serialize};
 use tracing::{debug_span, info, warn, Span};
 
 use crate::{
-    crypto::{self, PassphraseSource},
+    crypto::{self, PassphraseSource, RsaSigningKey},
     format::avb::{
         self, AlgorithmType, AppendedDescriptorMut, AppendedDescriptorRef, Descriptor, Footer,
         HashTreeDescriptor, Header, KernelCmdlineDescriptor,
@@ -57,16 +57,22 @@ fn read_avb_image(path: &Path) -> Result<(AvbInfo, BufReader<File>)> {
     Ok((info, reader))
 }
 
-fn write_avb_image(file: PSeekFile, info: &mut AvbInfo) -> Result<()> {
+fn write_avb_image(file: PSeekFile, info: &mut AvbInfo, recompute_size: bool) -> Result<()> {
     let mut writer = BufWriter::new(file);
 
-    if let Some(f) = &mut info.footer {
-        avb::write_appended_image(&mut writer, &info.header, f, info.image_size)
-            .context("Failed to write appended AVB image")?;
+    info.image_size = if let Some(f) = &mut info.footer {
+        let image_size = if recompute_size {
+            None
+        } else {
+            Some(info.image_size)
+        };
+
+        avb::write_appended_image(&mut writer, &info.header, f, image_size)
+            .context("Failed to write appended AVB image")?
     } else {
         avb::write_root_image(&mut writer, &info.header, 4096)
-            .context("Failed to write root AVB image")?;
-    }
+            .context("Failed to write root AVB image")?
+    };
 
     writer.flush().context("Failed to flush writes")?;
 
@@ -94,15 +100,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 +197,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)
@@ -379,12 +383,26 @@ fn sign_or_clear(info: &mut AvbInfo, orig_header: &Header, key_group: &KeyGroup)
                 key_group.pass_file.as_deref(),
                 key_group.pass_env_var.as_deref(),
             );
-            let private_key = crypto::read_pem_key_file(key_path, &source)
-                .with_context(|| format!("Failed to load key: {key_path:?}"))?;
+            let signing_key = if let Some(helper) = &key_group.signing_helper {
+                let public_key = crypto::read_pem_public_key_file(key_path)
+                    .with_context(|| format!("Failed to load key: {key_path:?}"))?;
 
-            info.header.set_algo_for_key(&private_key)?;
+                RsaSigningKey::External {
+                    program: helper.clone(),
+                    public_key_file: key_path.clone(),
+                    public_key,
+                    passphrase_source: source,
+                }
+            } else {
+                let private_key = crypto::read_pem_key_file(key_path, &source)
+                    .with_context(|| format!("Failed to load key: {key_path:?}"))?;
+
+                RsaSigningKey::Internal(private_key)
+            };
+
+            info.header.set_algo_for_key(&signing_key)?;
             info.header
-                .sign(&private_key)
+                .sign(&signing_key)
                 .context("Failed to sign new AVB header")?;
         }
         SignAction::Clear => {
@@ -627,12 +645,16 @@ fn pack_subcommand(cli: &PackCli, cancel_signal: &AtomicBool) -> Result<()> {
 
     sign_or_clear(&mut info, &orig_header, &cli.key)?;
 
-    write_avb_image(file, &mut info)?;
+    write_avb_image(file, &mut info, cli.recompute_size)?;
 
     // We display the info at the very end after both the header and footer are
     // updated so that incorrect/incomplete information isn't shown.
     display_info(&cli.display, &info);
 
+    if let Some(path) = &cli.output_info {
+        write_info(path, &info)?;
+    }
+
     Ok(())
 }
 
@@ -646,7 +668,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)?;
         }
 
@@ -661,7 +683,7 @@ fn repack_subcommand(cli: &RepackCli, cancel_signal: &AtomicBool) -> Result<()>
 
     sign_or_clear(&mut info, &orig_header, &cli.key)?;
 
-    write_avb_image(file, &mut info)?;
+    write_avb_image(file, &mut info, false)?;
 
     // We display the info at the very end after both the header and footer are
     // updated so that incorrect/incomplete information isn't shown.
@@ -735,12 +757,15 @@ struct DisplayGroup {
 
 #[derive(Debug, Args)]
 struct KeyGroup {
-    /// Path to private key for signing.
+    /// Path to signing key.
     ///
-    /// A private key is needed if packing an image where the original header
+    /// A signing key is needed if packing an image where the original header
     /// was signed and the header needs to be modified (eg. for a new checksum).
-    /// If the header was originally not signed, then the private key is not
+    /// If the header was originally not signed, then the signing key is not
     /// used, unless --force is specified.
+    ///
+    /// This should normally be a private key. However, if --signing-helper is
+    /// used, then it should be a public key instead.
     #[arg(short, long, value_name = "FILE", value_parser)]
     key: Option<PathBuf>,
 
@@ -760,6 +785,15 @@ struct KeyGroup {
     /// File containing private key passphrase.
     #[arg(long, value_name = "FILE", value_parser, group = "pass")]
     pass_file: Option<PathBuf>,
+
+    /// External program for signing.
+    ///
+    /// If this option is specified, then --key must refer to a public key. The
+    /// program will be invoked as:
+    ///
+    /// <program> <algo> <public key> [file <pass file>|env <pass env>]
+    #[arg(long, value_name = "PROGRAM", value_parser)]
+    signing_helper: Option<PathBuf>,
 }
 
 /// Unpack an AVB image.
@@ -829,12 +863,28 @@ struct PackCli {
     #[arg(long, value_name = "FILE", value_parser, default_value = "avb.toml")]
     input_info: PathBuf,
 
+    /// Path to output AVB info TOML.
+    ///
+    /// If specified, the AVB info containing all recomputed fields will be
+    /// written to this file. This can point to the same file as --input-info.
+    #[arg(long, value_name = "FILE", value_parser)]
+    output_info: Option<PathBuf>,
+
     /// Path to input raw image.
     ///
     /// Appended AVB images require a raw image.
     #[arg(long, value_name = "FILE", value_parser, default_value = "raw.img")]
     input_raw: PathBuf,
 
+    /// Recompute image size.
+    ///
+    /// By default, it is assumed that the image has a fixed size specified by
+    /// the image_size top-level field in the AVB info TOML. If flag is passed,
+    /// then that field is ignored and the smallest possible output image will
+    /// be created.
+    #[arg(long)]
+    recompute_size: bool,
+
     #[command(flatten)]
     key: KeyGroup,
 

avbroot/src/cli/cpio.rs

@@ -303,7 +303,7 @@ struct UnpackCli {
     #[arg(short, long, value_name = "FILE", value_parser)]
     input: PathBuf,
 
-    /// Path to output cpio info TOML.
+    /// Path to output info TOML.
     #[arg(long, value_name = "FILE", value_parser, default_value = "cpio.toml")]
     output_info: PathBuf,
 

avbroot/src/cli/key.rs

@@ -55,6 +55,9 @@ pub fn key_main(cli: &KeyCli) -> Result<()> {
                     .with_context(|| format!("Failed to load key: {p:?}"))?;
 
                 private_key.to_public_key()
+            } else if let Some(p) = &c.input.public_key {
+                crypto::read_pem_public_key_file(p)
+                    .with_context(|| format!("Failed to load public key: {p:?}"))?
             } else if let Some(p) = &c.input.cert {
                 let certificate = crypto::read_pem_cert_file(p)
                     .with_context(|| format!("Failed to load certificate: {p:?}"))?;
@@ -93,6 +96,10 @@ struct PublicKeyInputGroup {
     #[arg(short, long, value_name = "FILE", value_parser)]
     key: Option<PathBuf>,
 
+    /// Path to public key.
+    #[arg(short, long, value_name = "FILE", value_parser, conflicts_with_all = ["pass_env_var", "pass_file"])]
+    public_key: Option<PathBuf>,
+
     /// Path to certificate.
     #[arg(short, long, value_name = "FILE", value_parser, conflicts_with_all = ["pass_env_var", "pass_file"])]
     cert: Option<PathBuf>,

avbroot/src/cli/mod.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -12,3 +12,4 @@ pub mod fec;
 pub mod hashtree;
 pub mod key;
 pub mod ota;
+pub mod payload;

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},
@@ -20,7 +21,6 @@ use cap_std::{ambient_authority, fs::Dir};
 use cap_tempfile::TempDir;
 use clap::{value_parser, ArgAction, Args, Parser, Subcommand};
 use rayon::{iter::IntoParallelRefIterator, prelude::ParallelIterator};
-use rsa::RsaPrivateKey;
 use tempfile::NamedTempFile;
 use topological_sort::TopologicalSort;
 use tracing::{debug_span, info, warn};
@@ -29,16 +29,18 @@ use zip::{write::FileOptions, CompressionMethod, ZipArchive, ZipWriter};
 
 use crate::{
     cli,
-    crypto::{self, PassphraseSource},
+    crypto::{self, PassphraseSource, RsaSigningKey},
     format::{
-        avb::Header,
-        avb::{self, Descriptor},
+        avb::{self, Descriptor, Header},
         ota::{self, SigningWriter, ZipEntry},
         padding,
-        payload::{self, PayloadHeader, PayloadWriter},
+        payload::{self, PayloadHeader, PayloadWriter, VabcAlgo},
     },
     patch::{
-        boot::{self, BootImagePatch, MagiskRootPatcher, OtaCertPatcher, PrepatchedImagePatcher},
+        boot::{
+            self, BootImagePatch, DsuPubKeyPatcher, MagiskRootPatcher, OtaCertPatcher,
+            PrepatchedImagePatcher,
+        },
         system,
     },
     protobuf::{
@@ -187,26 +189,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>>,
-    key_avb: &RsaPrivateKey,
-    cert_ota: &Certificate,
+    boot_patchers: Vec<Box<dyn BootImagePatch + Sync>>,
+    key_avb: &RsaSigningKey,
     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!(
@@ -247,7 +240,7 @@ fn patch_system_image<'a, 'b: 'a>(
     required_images: &'b RequiredImages,
     input_files: &mut HashMap<String, InputFile>,
     cert_ota: &Certificate,
-    key_avb: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<(&'b str, Vec<Range<u64>>)> {
     let Some(target) = required_images.iter_system().next() else {
@@ -388,7 +381,7 @@ fn get_vbmeta_patch_order(
     }
 
     if !missing.is_empty() {
-        warn!("Partitions aren't protected by AVB: {:?}", joined(missing));
+        warn!("Partitions aren't protected by AVB: {}", joined(missing));
     }
 
     // Ensure that there's only a single root of trust. Otherwise, there could
@@ -479,7 +472,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");
@@ -493,9 +486,7 @@ fn update_security_descriptors(
 /// Get the text before the first equal sign in the kernel command line if it is
 /// not empty.
 fn cmdline_prefix(cmdline: &str) -> Option<&str> {
-    let Some((prefix, _)) = cmdline.split_once('=') else {
-        return None;
-    };
+    let (prefix, _) = cmdline.split_once('=')?;
     if prefix.is_empty() {
         return None;
     }
@@ -524,7 +515,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
@@ -544,7 +535,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
@@ -573,7 +564,7 @@ fn update_vbmeta_headers(
     headers: &mut HashMap<String, Header>,
     order: &mut [(String, HashSet<String>)],
     clear_vbmeta_flags: bool,
-    key: &RsaPrivateKey,
+    key: &RsaSigningKey,
     block_size: u64,
 ) -> Result<()> {
     for (name, deps) in order {
@@ -632,7 +623,7 @@ fn update_vbmeta_headers(
 /// If `ranges` is [`None`], then the entire file is compressed. Otherwise, only
 /// the chunks containing the specified ranges are compressed. In the latter
 /// scenario, unmodified chunks must be copied from the original payload.
-fn compress_image(
+pub fn compress_image(
     name: &str,
     file: &mut PSeekFile,
     header: &mut PayloadHeader,
@@ -682,12 +673,43 @@ fn compress_image(
 
     info!("Compressing full image: {name}");
 
-    // Otherwise, compress the entire image.
-    let (partition_info, operations) =
-        payload::compress_image(&*file, &writer, name, block_size, cancel_signal)?;
+    // Otherwise, compress the entire image. If VABC is enabled, we need to
+    // update the CoW size estimate or else the CoW block device may run out of
+    // space during flashing.
+    let vabc_algo = if partition.estimate_cow_size.is_some() {
+        info!("Needs updated CoW size estimate: {name}");
+
+        // Only CoW v2 seems to exist in the wild currently, so that is all we
+        // support.
+        let Some(dpm) = &header.manifest.dynamic_partition_metadata else {
+            bail!("Dynamic partition metadata is missing");
+        };
+
+        if !dpm.vabc_enabled() {
+            bail!("Partition has CoW estimate, but VABC is disabled: {name}");
+        }
+
+        let cow_version = dpm.cow_version();
+        if dpm.cow_version() != 2 {
+            bail!("Unsupported CoW version: {cow_version}");
+        }
+
+        let compression = dpm.vabc_compression_param();
+        let Some(vabc_algo) = VabcAlgo::new(compression) else {
+            bail!("Unsupported VABC compression: {compression}");
+        };
+
+        Some(vabc_algo)
+    } else {
+        None
+    };
+
+    let (partition_info, operations, cow_estimate) =
+        payload::compress_image(&*file, &writer, name, block_size, vabc_algo, cancel_signal)?;
 
     partition.new_partition_info = Some(partition_info);
     partition.operations = operations;
+    partition.estimate_cow_size = cow_estimate;
 
     *file = writer;
 
@@ -700,22 +722,20 @@ 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,
+    key_avb: &RsaSigningKey,
+    key_ota: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<(String, u64)> {
-    let header = PayloadHeader::from_reader(payload.reopen_boxed()?)
+    let mut header = PayloadHeader::from_reader(payload.reopen_boxed()?)
         .context("Failed to load OTA payload header")?;
     if !header.is_full_ota() {
         bail!("Payload is a delta OTA, not a full OTA");
     }
 
-    let header = Mutex::new(header);
-    let mut header_locked = header.lock().unwrap();
-    let all_partitions = header_locked
+    let all_partitions = header
         .manifest
         .partitions
         .iter()
@@ -733,7 +753,7 @@ fn patch_ota_payload(
     // Determine what images need to be patched. For simplicity, we pre-read all
     // vbmeta images since they're tiny. They're discarded later if the they
     // don't need to be modified.
-    let required_images = RequiredImages::new(&header_locked.manifest);
+    let required_images = RequiredImages::new(&header.manifest);
     let vbmeta_images = required_images.iter_vbmeta().collect::<HashSet<_>>();
 
     // The set of source images to be inserted into the new payload, replacing
@@ -745,16 +765,15 @@ fn patch_ota_payload(
         payload,
         &required_images,
         external_images,
-        &header_locked,
+        &header,
         cancel_signal,
     )?;
 
     patch_boot_images(
         &required_images,
         &mut input_files,
-        root_patcher,
+        boot_patchers,
         key_avb,
-        cert_ota,
         cancel_signal,
     )?;
 
@@ -788,7 +807,7 @@ fn patch_ota_payload(
         &mut vbmeta_order,
         clear_vbmeta_flags,
         key_avb,
-        header_locked.manifest.block_size().into(),
+        header.manifest.block_size().into(),
     )?;
 
     // Unmodified vbmeta images no longer need to be kept around either.
@@ -800,7 +819,7 @@ fn patch_ota_payload(
             let modified_operations = compress_image(
                 &name,
                 &mut input_file.file,
-                &mut header_locked,
+                &mut header,
                 // We can only perform the optimization of avoiding
                 // recompression if the image came from the original payload.
                 if name == system_target && !external_images.contains_key(&name) {
@@ -818,7 +837,7 @@ fn patch_ota_payload(
 
     info!("Generating new OTA payload");
 
-    let mut payload_writer = PayloadWriter::new(writer, header_locked.clone(), key_ota.clone())
+    let mut payload_writer = PayloadWriter::new(writer, header.clone(), key_ota.clone())
         .context("Failed to write payload header")?;
     let mut orig_payload_reader = payload.reopen_boxed().context("Failed to open payload")?;
 
@@ -836,7 +855,7 @@ fn patch_ota_payload(
 
         let pi = payload_writer.partition_index().unwrap();
         let oi = payload_writer.operation_index().unwrap();
-        let orig_partition = &header_locked.manifest.partitions[pi];
+        let orig_partition = &header.manifest.partitions[pi];
         let orig_operation = &orig_partition.operations[oi];
         let data_offset = orig_operation
             .data_offset
@@ -866,7 +885,7 @@ fn patch_ota_payload(
 
         // Otherwise, copy from the original payload.
         let data_offset = data_offset
-            .checked_add(header_locked.blob_offset)
+            .checked_add(header.blob_offset)
             .ok_or_else(|| anyhow!("data_offset overflow in partition #{pi} operation #{oi}"))?;
 
         orig_payload_reader
@@ -882,7 +901,7 @@ fn patch_ota_payload(
         .with_context(|| format!("Failed to copy from original payload: {name}"))?;
     }
 
-    let (_, properties, metadata_size) = payload_writer
+    let (_, _, properties, metadata_size) = payload_writer
         .finish()
         .context("Failed to finalize payload")?;
 
@@ -895,10 +914,10 @@ 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,
+    key_avb: &RsaSigningKey,
+    key_ota: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<(OtaMetadata, u64)> {
@@ -916,7 +935,7 @@ fn patch_ota_zip(
     }
 
     if !missing.is_empty() {
-        bail!("Missing entries in OTA zip: {:?}", joined(missing));
+        bail!("Missing entries in OTA zip: {}", joined(missing));
     } else if !paths.contains(ota::PATH_METADATA) && !paths.contains(ota::PATH_METADATA_PB) {
         bail!(
             "Neither legacy nor protobuf OTA metadata files exist: {:?}, {:?}",
@@ -1017,7 +1036,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,
@@ -1073,7 +1092,7 @@ fn patch_ota_zip(
     Ok((metadata, payload_metadata_size.unwrap()))
 }
 
-fn extract_ota_zip(
+pub fn extract_payload(
     raw_reader: &PSeekFile,
     directory: &Dir,
     payload_offset: u64,
@@ -1199,10 +1218,38 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()
         cli.pass_ota_env_var.as_deref(),
     );
 
-    let key_avb = crypto::read_pem_key_file(&cli.key_avb, &source_avb)
-        .with_context(|| format!("Failed to load key: {:?}", cli.key_avb))?;
-    let key_ota = crypto::read_pem_key_file(&cli.key_ota, &source_ota)
-        .with_context(|| format!("Failed to load key: {:?}", cli.key_ota))?;
+    let (key_avb, key_ota) = if let Some(helper) = &cli.signing_helper {
+        let public_key_avb = crypto::read_pem_public_key_file(&cli.key_avb)
+            .with_context(|| format!("Failed to load key: {:?}", cli.key_avb))?;
+        let public_key_ota = crypto::read_pem_public_key_file(&cli.key_ota)
+            .with_context(|| format!("Failed to load key: {:?}", cli.key_ota))?;
+
+        let key_avb = RsaSigningKey::External {
+            program: helper.clone(),
+            public_key_file: cli.key_avb.clone(),
+            public_key: public_key_avb,
+            passphrase_source: source_avb,
+        };
+        let key_ota = RsaSigningKey::External {
+            program: helper.clone(),
+            public_key_file: cli.key_ota.clone(),
+            public_key: public_key_ota,
+            passphrase_source: source_ota,
+        };
+
+        (key_avb, key_ota)
+    } else {
+        let private_key_avb = crypto::read_pem_key_file(&cli.key_avb, &source_avb)
+            .with_context(|| format!("Failed to load key: {:?}", cli.key_avb))?;
+        let private_key_ota = crypto::read_pem_key_file(&cli.key_ota, &source_ota)
+            .with_context(|| format!("Failed to load key: {:?}", cli.key_ota))?;
+
+        let key_avb = RsaSigningKey::Internal(private_key_avb);
+        let key_ota = RsaSigningKey::Internal(private_key_ota);
+
+        (key_avb, key_ota)
+    };
+
     let cert_ota = crypto::read_pem_cert_file(&cli.cert_ota)
         .with_context(|| format!("Failed to load certificate: {:?}", cli.cert_ota))?;
 
@@ -1225,8 +1272,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(),
@@ -1234,21 +1284,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))?;
@@ -1274,7 +1323,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,
@@ -1346,11 +1395,12 @@ pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &AtomicBool) -> Resul
         .with_context(|| format!("Failed to open for reading: {:?}", cli.input))?;
     let mut zip = ZipArchive::new(BufReader::new(raw_reader.reopen()?))
         .with_context(|| format!("Failed to read zip: {:?}", cli.input))?;
-    let payload_entry = zip
-        .by_name(ota::PATH_PAYLOAD)
-        .with_context(|| format!("Failed to open zip entry: {:?}", ota::PATH_PAYLOAD))?;
-    let payload_offset = payload_entry.data_start();
-    let payload_size = payload_entry.size();
+    let (payload_offset, payload_size) = {
+        let entry = zip
+            .by_name(ota::PATH_PAYLOAD)
+            .with_context(|| format!("Failed to open zip entry: {}", ota::PATH_PAYLOAD))?;
+        (entry.data_start(), entry.size())
+    };
 
     // Open the payload data directly.
     let mut payload_reader = SectionReader::new(
@@ -1393,7 +1443,7 @@ pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &AtomicBool) -> Resul
     let directory = Dir::open_ambient_dir(&cli.directory, authority)
         .with_context(|| format!("Failed to open directory: {:?}", cli.directory))?;
 
-    extract_ota_zip(
+    extract_payload(
         &raw_reader,
         &directory,
         payload_offset,
@@ -1403,6 +1453,144 @@ pub fn extract_subcommand(cli: &ExtractCli, cancel_signal: &AtomicBool) -> Resul
         cancel_signal,
     )?;
 
+    if cli.fastboot {
+        const ANDROID_INFO: &str = "android-info.txt";
+        const FASTBOOT_INFO: &str = "fastboot-info.txt";
+
+        // Generate android-info.txt, which is always required for fastboot's
+        // flashall subcommand. We only add a basic device check to avoid
+        // accidental flashes on the wrong device.
+
+        let mut metadata_entry = zip
+            .by_name(ota::PATH_METADATA_PB)
+            .with_context(|| format!("Failed to open zip entry: {:?}", ota::PATH_METADATA_PB))?;
+        let mut metadata_raw = vec![];
+        metadata_entry
+            .read_to_end(&mut metadata_raw)
+            .with_context(|| format!("Failed to read OTA metadata: {}", ota::PATH_METADATA_PB))?;
+        let metadata = ota::parse_protobuf_metadata(&metadata_raw)
+            .with_context(|| format!("Failed to parse OTA metadata: {}", ota::PATH_METADATA_PB))?;
+
+        let device = metadata
+            .precondition
+            .as_ref()
+            .and_then(|p| p.device.first())
+            .ok_or_else(|| anyhow!("Device codename not found in OTA metadata"))?;
+
+        directory
+            .write(ANDROID_INFO, format!("require board={device}\n"))
+            .with_context(|| format!("Failed to write file: {ANDROID_INFO}"))?;
+
+        // Find out which images can be flashed with fastboot. The bootloader
+        // (and potentially modem) partitions need to be flashed as a whole and
+        // an OTA doesn't contain sufficient information to generate the
+        // required combined file.
+        let mut flashable_images = BTreeSet::new();
+
+        for name in &unique_images {
+            let file = directory
+                .open(format!("{name}.img"))
+                .with_context(|| format!("Failed to open image for reading: {name}"))?;
+
+            match avb::load_image(file) {
+                Ok(_) => {
+                    flashable_images.insert(name);
+                }
+                // Treat images without AVB metadata as bootloader partitions.
+                Err(avb::Error::InvalidHeaderMagic(_)) => continue,
+                Err(e) => return Err(e).with_context(|| format!("Failed to load image: {name}")),
+            }
+        }
+
+        // Generate fastboot-info.txt to be able to control how exactly the
+        // images are flashed. This solves two problems:
+        //
+        // 1. fastboot flashall, by default, expects super_empty.img to exist
+        //    for A/B devices. If it doesn't exist, it assumes that all images
+        //    are meant for non-dynamic partitions and skips the fastbootd
+        //    reboot. We cannot generate a proper super_empty.img because it
+        //    requires knowing the number of super partitions along with their
+        //    names and sizes. This information is not available in an OTA. The
+        //    fastboot info file allows us to control when to reboot to
+        //    fastbootd and flash each partition.
+        //
+        //    This approach is a bit slower because we have to reboot into
+        //    fastbootd. When the device only has a single super partition,
+        //    fastboot will normally locally combine super_empty.img with the
+        //    individual partitions to form a full super image and then flash it
+        //    via the bootloader's fastboot mode.
+        //
+        //    Also, because we can't generate super_empty.img, if the super
+        //    partition metadata on the device somehow becomes corrupted,
+        //    running fastboot flash all using avbroot's extracted images isn't
+        //    sufficient for recovering the device. The user will need to copy
+        //    that file from the factory image into the output directory and add
+        //    an `update-super` line after `reboot fastboot` the fastboot info
+        //    file.
+        //
+        // 2. fastboot flashall does not look at all .img files in a directory.
+        //    Instead, it has a hardcoded list of partitions known to AOSP. This
+        //    means obscure OEM-specific partitions are just silently ignored.
+        //    Using a fastboot info file with explicit flash instructions avoids
+        //    this problem entirely.
+
+        let mut dynamic = BTreeSet::new();
+
+        if let Some(dpm) = &header.manifest.dynamic_partition_metadata {
+            for group in &dpm.groups {
+                for name in &group.partition_names {
+                    if unique_images.contains(name) {
+                        dynamic.insert(name);
+                    }
+                }
+            }
+        }
+
+        let flash_command = |name| {
+            let (prefix, suffix) = if flashable_images.contains(name) {
+                ("", "")
+            } else {
+                ("#", " # Bootloader/modem images cannot be flashed")
+            };
+
+            let extra_arg = if *name == "vbmeta" {
+                "--apply-vbmeta "
+            } else {
+                ""
+            };
+
+            format!("{prefix}flash {extra_arg}{name}{suffix}\n")
+        };
+
+        let mut fastboot_info = String::new();
+
+        fastboot_info.push_str("# Generated by avbroot\n");
+        fastboot_info.push_str("version 1\n");
+
+        fastboot_info.push_str("# Flash non-dynamic partitions\n");
+        for name in &unique_images {
+            if !dynamic.contains(name) {
+                fastboot_info.push_str(&flash_command(name));
+            }
+        }
+
+        fastboot_info.push_str("# Reboot to fastbootd\n");
+        fastboot_info.push_str("reboot fastboot\n");
+
+        fastboot_info.push_str("# Flash dynamic partitions\n");
+        for name in &dynamic {
+            fastboot_info.push_str(&flash_command(name));
+        }
+
+        fastboot_info.push_str("# Wipe data when -w flag is used\n");
+        fastboot_info.push_str("if-wipe erase userdata\n");
+        fastboot_info.push_str("if-wipe erase metadata\n");
+
+        directory
+            .write(FASTBOOT_INFO, fastboot_info)
+            .with_context(|| format!("Failed to write file: {FASTBOOT_INFO}"))?;
+    }
+
     Ok(())
 }
 
@@ -1467,7 +1655,7 @@ pub fn verify_subcommand(cli: &VerifyCli, cancel_signal: &AtomicBool) -> Result<
         .cloned()
         .collect::<BTreeSet<_>>();
 
-    extract_ota_zip(
+    extract_payload(
         &raw_reader,
         &temp_dir,
         pf_payload.offset,
@@ -1588,7 +1776,10 @@ pub struct PatchCli {
     #[arg(short, long, value_name = "FILE", value_parser, help_heading = HEADING_PATH)]
     pub output: Option<PathBuf>,
 
-    /// Private key for signing vbmeta images.
+    /// Signing key for vbmeta headers.
+    ///
+    /// This should normally be a private key. However, if --signing-helper is
+    /// used, then it should be a public key instead.
     #[arg(
         long,
         alias = "privkey-avb",
@@ -1598,7 +1789,10 @@ pub struct PatchCli {
     )]
     pub key_avb: PathBuf,
 
-    /// Private key for signing the OTA.
+    /// Signing key for the OTA.
+    ///
+    /// This should normally be a private key. However, if --signing-helper is
+    /// used, then it should be a public key instead.
     #[arg(
         long,
         alias = "privkey-ota",
@@ -1656,13 +1850,22 @@ pub struct PatchCli {
     )]
     pub pass_ota_file: Option<PathBuf>,
 
+    /// External program for signing.
+    ///
+    /// If this option is specified, then --key-avb and --key-ota must refer to
+    /// public keys. The program will be invoked as:
+    ///
+    /// <program> <algo> <public key> [file <pass file>|env <pass env>]
+    #[arg(long, value_name = "PROGRAM", value_parser, help_heading = HEADING_KEY)]
+    pub signing_helper: Option<PathBuf>,
+
     /// Use partition image from a file instead of the original payload.
     #[arg(
         long,
         value_names = ["PARTITION", "FILE"],
         value_parser = value_parser!(OsString),
         num_args = 2,
-        help_heading = HEADING_PATH,
+        help_heading = HEADING_PATH
     )]
     pub replace: Vec<OsString>,
 
@@ -1704,6 +1907,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,
@@ -1712,6 +1919,7 @@ pub struct PatchCli {
     #[arg(
         long,
         value_name = "PARTITION",
+        hide = true,
         help_heading = HEADING_OTHER
     )]
     pub boot_partition: Option<String>,
@@ -1737,8 +1945,12 @@ pub struct ExtractCli {
     pub boot_only: bool,
 
     /// (Deprecated: no longer needed)
-    #[arg(long, value_name = "PARTITION")]
+    #[arg(long, value_name = "PARTITION", hide = true)]
     pub boot_partition: Option<String>,
+
+    /// Generate fastboot info files.
+    #[arg(long)]
+    pub fastboot: bool,
 }
 
 /// Verify signatures of an OTA.

avbroot/src/cli/payload.rs

@@ -0,0 +1,473 @@
+/*
+ * SPDX-FileCopyrightText: 2024 Andrew Gunnerson
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+use std::{
+    collections::HashMap,
+    ffi::{OsStr, OsString},
+    fs::{self, File},
+    io::{BufReader, BufWriter, Seek, SeekFrom},
+    path::{Path, PathBuf},
+    sync::atomic::AtomicBool,
+};
+
+use anyhow::{anyhow, bail, Context, Result};
+use cap_std::{ambient_authority, fs::Dir};
+use clap::{Args, Parser, Subcommand};
+use tracing::info;
+
+use crate::{
+    cli::ota,
+    crypto::{self, PassphraseSource, RsaSigningKey},
+    format::payload::{PayloadHeader, PayloadWriter},
+    stream::{self, FromReader, PSeekFile},
+};
+
+fn open_reader(path: &Path) -> Result<(BufReader<File>, PayloadHeader)> {
+    let mut reader = File::open(path)
+        .map(BufReader::new)
+        .with_context(|| format!("Failed to open payload for reading: {path:?}"))?;
+    let header = PayloadHeader::from_reader(&mut reader)
+        .with_context(|| format!("Failed to read payload header: {path:?}"))?;
+    if !header.is_full_ota() {
+        bail!("Payload is a delta OTA, not a full OTA");
+    }
+
+    Ok((reader, header))
+}
+
+fn open_writer(
+    path: &Path,
+    header: PayloadHeader,
+    key: RsaSigningKey,
+) -> Result<PayloadWriter<BufWriter<File>>> {
+    let writer = File::create(path)
+        .map(BufWriter::new)
+        .with_context(|| format!("Failed to open payload for writing: {path:?}"))?;
+    let payload_writer = PayloadWriter::new(writer, header, key)
+        .with_context(|| format!("Failed to write payload header: {path:?}"))?;
+
+    Ok(payload_writer)
+}
+
+fn read_info(path: &Path) -> Result<PayloadHeader> {
+    let data = fs::read_to_string(path)
+        .with_context(|| format!("Failed to read payload info TOML: {path:?}"))?;
+    let info = toml_edit::de::from_str(&data)
+        .with_context(|| format!("Failed to parse payload info TOML: {path:?}"))?;
+
+    Ok(info)
+}
+
+fn write_info(path: &Path, manifest: &PayloadHeader) -> Result<()> {
+    let data = toml_edit::ser::to_string_pretty(manifest)
+        .with_context(|| format!("Failed to serialize payload info TOML: {path:?}"))?;
+    fs::write(path, data)
+        .with_context(|| format!("Failed to write payload info TOML: {path:?}"))?;
+
+    Ok(())
+}
+
+fn display_header(cli: &PayloadCli, header: &PayloadHeader) {
+    if !cli.quiet {
+        println!("{header:#?}");
+    }
+}
+
+fn load_key(group: &KeyGroup) -> Result<RsaSigningKey> {
+    let source = PassphraseSource::new(
+        &group.key,
+        group.pass_file.as_deref(),
+        group.pass_env_var.as_deref(),
+    );
+    let signing_key = if let Some(helper) = &group.signing_helper {
+        let public_key = crypto::read_pem_public_key_file(&group.key)
+            .with_context(|| format!("Failed to load key: {:?}", group.key))?;
+
+        RsaSigningKey::External {
+            program: helper.clone(),
+            public_key_file: group.key.clone(),
+            public_key,
+            passphrase_source: source,
+        }
+    } else {
+        let private_key = crypto::read_pem_key_file(&group.key, &source)
+            .with_context(|| format!("Failed to load key: {:?}", group.key))?;
+
+        RsaSigningKey::Internal(private_key)
+    };
+
+    Ok(signing_key)
+}
+
+fn unpack_subcommand(
+    payload_cli: &PayloadCli,
+    cli: &UnpackCli,
+    cancel_signal: &AtomicBool,
+) -> Result<()> {
+    let (mut reader, header) = open_reader(&cli.input)?;
+    let payload_size = reader
+        .seek(SeekFrom::End(0))
+        .with_context(|| format!("Failed to get file size: {:?}", cli.input))?;
+
+    display_header(payload_cli, &header);
+
+    write_info(&cli.output_info, &header)?;
+
+    let authority = ambient_authority();
+    Dir::create_ambient_dir_all(&cli.output_images, authority)
+        .with_context(|| format!("Failed to create directory: {:?}", cli.output_images))?;
+    let directory = Dir::open_ambient_dir(&cli.output_images, authority)
+        .with_context(|| format!("Failed to open directory: {:?}", cli.output_images))?;
+
+    ota::extract_payload(
+        &PSeekFile::new(reader.into_inner()),
+        &directory,
+        0,
+        payload_size,
+        &header,
+        &header
+            .manifest
+            .partitions
+            .iter()
+            .map(|p| &p.partition_name)
+            .cloned()
+            .collect(),
+        cancel_signal,
+    )?;
+
+    Ok(())
+}
+
+fn pack_subcommand(
+    payload_cli: &PayloadCli,
+    cli: &PackCli,
+    cancel_signal: &AtomicBool,
+) -> Result<()> {
+    let signing_key = load_key(&cli.key)?;
+
+    let mut header = read_info(&cli.input_info)?;
+
+    let authority = ambient_authority();
+    let directory = Dir::open_ambient_dir(&cli.input_images, authority)
+        .with_context(|| format!("Failed to open directory: {:?}", cli.input_images))?;
+
+    for p in &header.manifest.partitions {
+        let name = &p.partition_name;
+
+        if Path::new(name).file_name() != Some(OsStr::new(name)) {
+            bail!("Unsafe partition name: {name}");
+        }
+    }
+
+    // Pre-open all of the image files.
+    let input_files = header
+        .manifest
+        .partitions
+        .iter()
+        .map(|p| {
+            let path = format!("{}.img", p.partition_name);
+            let file = directory
+                .open(&path)
+                .map(|f| PSeekFile::new(f.into_std()))
+                .with_context(|| format!("Failed to open file: {path:?}"))?;
+
+            Ok((p.partition_name.clone(), file))
+        })
+        .collect::<Result<HashMap<_, _>>>()?;
+
+    // Compress the images and compute the list of install operations for
+    // insertion into the payload header. The compressed data is stored in new
+    // temp files and the original input files are dropped.
+    let mut compressed_files = input_files
+        .into_iter()
+        .map(|(name, mut input_file)| {
+            ota::compress_image(&name, &mut input_file, &mut header, None, cancel_signal)
+                .with_context(|| format!("Failed to compress image: {name}"))?;
+
+            Ok((name, input_file))
+        })
+        .collect::<Result<HashMap<_, _>>>()?;
+
+    info!("Generating new OTA payload");
+
+    // Now we can write the actual payload. With everything precomputed, this is
+    // mostly just a simple copy.
+    let mut payload_writer = open_writer(&cli.output, header.clone(), signing_key)?;
+
+    while payload_writer
+        .begin_next_operation()
+        .context("Failed to begin next payload blob entry")?
+    {
+        let name = payload_writer.partition().unwrap().partition_name.clone();
+        let operation = payload_writer.operation().unwrap();
+
+        let Some(data_length) = operation.data_length else {
+            // Otherwise, this is a ZERO/DISCARD operation.
+            continue;
+        };
+
+        let pi = payload_writer.partition_index().unwrap();
+        let oi = payload_writer.operation_index().unwrap();
+        let orig_partition = &header.manifest.partitions[pi];
+        let orig_operation = &orig_partition.operations[oi];
+        let data_offset = orig_operation
+            .data_offset
+            .ok_or_else(|| anyhow!("Missing data_offset in partition #{pi} operation #{oi}"))?;
+
+        // The compressed chunks are laid out sequentially and data_offset is
+        // set to the offset within that file.
+        let Some(input_file) = compressed_files.get_mut(&name) else {
+            unreachable!("Compressed data not found for image: {name}");
+        };
+
+        input_file
+            .seek(SeekFrom::Start(data_offset))
+            .with_context(|| format!("Failed to seek image: {name}"))?;
+
+        stream::copy_n(input_file, &mut payload_writer, data_length, cancel_signal)
+            .with_context(|| format!("Failed to copy from replacement image: {name}"))?;
+    }
+
+    let (_, header, properties, _) = payload_writer
+        .finish()
+        .context("Failed to finalize payload")?;
+
+    // Display the header information now that it has been finalized.
+    display_header(payload_cli, &header);
+
+    // Optionally, write payload_properties.txt.
+    if let Some(path) = &cli.output_properties {
+        fs::write(path, properties)
+            .with_context(|| format!("Failed to write payload properties: {path:?}"))?;
+    }
+
+    Ok(())
+}
+
+fn repack_subcommand(
+    payload_cli: &PayloadCli,
+    cli: &RepackCli,
+    cancel_signal: &AtomicBool,
+) -> Result<()> {
+    let signing_key = load_key(&cli.key)?;
+
+    let (mut reader, header) = open_reader(&cli.input)?;
+
+    info!("Generating new OTA payload");
+
+    let mut payload_writer = open_writer(&cli.output, header.clone(), signing_key)?;
+
+    while payload_writer
+        .begin_next_operation()
+        .context("Failed to begin next payload blob entry")?
+    {
+        let name = payload_writer.partition().unwrap().partition_name.clone();
+        let operation = payload_writer.operation().unwrap();
+
+        let Some(data_length) = operation.data_length else {
+            // Otherwise, this is a ZERO/DISCARD operation.
+            continue;
+        };
+
+        let pi = payload_writer.partition_index().unwrap();
+        let oi = payload_writer.operation_index().unwrap();
+        let orig_partition = &header.manifest.partitions[pi];
+        let orig_operation = &orig_partition.operations[oi];
+        let data_offset = orig_operation
+            .data_offset
+            .ok_or_else(|| anyhow!("Missing data_offset in partition #{pi} operation #{oi}"))?;
+
+        // Directly copy blobs from the original payload.
+        let data_offset = data_offset
+            .checked_add(header.blob_offset)
+            .ok_or_else(|| anyhow!("data_offset overflow in partition #{pi} operation #{oi}"))?;
+
+        reader
+            .seek(SeekFrom::Start(data_offset))
+            .with_context(|| format!("Failed to seek original payload to {data_offset}"))?;
+
+        stream::copy_n(&mut reader, &mut payload_writer, data_length, cancel_signal)
+            .with_context(|| format!("Failed to copy from original payload: {name}"))?;
+    }
+
+    let (_, header, properties, _) = payload_writer
+        .finish()
+        .context("Failed to finalize payload")?;
+
+    // Display the header information now that it has been finalized.
+    display_header(payload_cli, &header);
+
+    // Optionally, write payload_properties.txt.
+    if let Some(path) = &cli.output_properties {
+        fs::write(path, properties)
+            .with_context(|| format!("Failed to write payload properties: {path:?}"))?;
+    }
+
+    Ok(())
+}
+
+fn info_subcommand(payload_cli: &PayloadCli, cli: &InfoCli) -> Result<()> {
+    let (_, header) = open_reader(&cli.input)?;
+
+    display_header(payload_cli, &header);
+
+    Ok(())
+}
+
+pub fn payload_main(cli: &PayloadCli, cancel_signal: &AtomicBool) -> Result<()> {
+    match &cli.command {
+        PayloadCommand::Unpack(c) => unpack_subcommand(cli, c, cancel_signal),
+        PayloadCommand::Pack(c) => pack_subcommand(cli, c, cancel_signal),
+        PayloadCommand::Repack(c) => repack_subcommand(cli, c, cancel_signal),
+        PayloadCommand::Info(c) => info_subcommand(cli, c),
+    }
+}
+
+#[derive(Debug, Args)]
+struct KeyGroup {
+    /// Path to signing key.
+    ///
+    /// This should normally be a private key. However, if --signing-helper is
+    /// used, then it should be a public key instead.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    key: PathBuf,
+
+    /// Environment variable containing private key passphrase.
+    #[arg(long, value_name = "ENV_VAR", value_parser, group = "pass")]
+    pass_env_var: Option<OsString>,
+
+    /// File containing private key passphrase.
+    #[arg(long, value_name = "FILE", value_parser, group = "pass")]
+    pass_file: Option<PathBuf>,
+
+    /// External program for signing.
+    ///
+    /// If this option is specified, then --key must refer to a public key. The
+    /// program will be invoked as:
+    ///
+    /// <program> <algo> <public key> [file <pass file>|env <pass env>]
+    #[arg(long, value_name = "PROGRAM", value_parser)]
+    signing_helper: Option<PathBuf>,
+}
+
+/// Unpack a payload binary.
+///
+/// Each partition is extracted to `<partition name>.img` in the output images
+/// directory. The payload header metadata is written to the info TOML file.
+///
+/// If any partition names are unsafe to use in a path, the extraction process
+/// will fail and exit. Extracted files are never written outside of the tree
+/// directory, even if an external process tries to interfere.
+#[derive(Debug, Parser)]
+struct UnpackCli {
+    /// Path to input payload binary.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    input: PathBuf,
+
+    /// Path to output info TOML.
+    #[arg(
+        long,
+        value_name = "FILE",
+        value_parser,
+        default_value = "payload.toml"
+    )]
+    output_info: PathBuf,
+
+    /// Path to output images directory.
+    #[arg(
+        long,
+        value_name = "DIR",
+        value_parser,
+        default_value = "payload_images"
+    )]
+    output_images: PathBuf,
+}
+
+/// Pack a payload binary.
+///
+/// The new payload binary will *only* contain images listed in the info TOML
+/// file. Extra images in the input images directory that aren't listed will be
+/// silently ignored. Images are added to the payload in the order that they are
+/// listed in the info TOML file.
+#[derive(Debug, Parser)]
+struct PackCli {
+    /// Path to output payload binary.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    output: PathBuf,
+
+    /// Path to output payload properties file.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    output_properties: Option<PathBuf>,
+
+    /// Path to input info TOML.
+    #[arg(
+        long,
+        value_name = "FILE",
+        value_parser,
+        default_value = "payload.toml"
+    )]
+    input_info: PathBuf,
+
+    /// Path to input images directory.
+    #[arg(
+        long,
+        value_name = "DIR",
+        value_parser,
+        default_value = "payload_images"
+    )]
+    input_images: PathBuf,
+
+    #[command(flatten)]
+    key: KeyGroup,
+}
+
+/// Repack a payload binary.
+///
+/// This command is equivalent to running `unpack` and `pack`, except without
+/// storing the unpacked data to disk nor recompressing the partition images.
+#[derive(Debug, Parser)]
+struct RepackCli {
+    /// Path to input payload binary.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    input: PathBuf,
+
+    /// Path to output payload binary.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    output: PathBuf,
+
+    /// Path to output payload properties file.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    output_properties: Option<PathBuf>,
+
+    #[command(flatten)]
+    key: KeyGroup,
+}
+
+/// Display payload information.
+#[derive(Debug, Parser)]
+struct InfoCli {
+    /// Path to input payload file.
+    #[arg(short, long, value_name = "FILE", value_parser)]
+    input: PathBuf,
+}
+
+#[derive(Debug, Subcommand)]
+enum PayloadCommand {
+    Unpack(UnpackCli),
+    Pack(PackCli),
+    Repack(RepackCli),
+    Info(InfoCli),
+}
+
+/// Inspect OTA payloads.
+#[derive(Debug, Parser)]
+pub struct PayloadCli {
+    #[command(subcommand)]
+    command: PayloadCommand,
+
+    /// Don't print payload header information.
+    #[arg(short, long, global = true)]
+    quiet: bool,
+}

avbroot/src/crypto.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -9,6 +9,7 @@ use std::{
     fs::{self, File, OpenOptions},
     io::{self, BufReader, BufWriter, Read, Write},
     path::{Path, PathBuf},
+    process::{Command, ExitStatus, Stdio},
     time::Duration,
 };
 
@@ -22,12 +23,16 @@ use cms::{
 };
 use pkcs8::{
     pkcs5::{pbes2, scrypt},
-    DecodePrivateKey, EncodePrivateKey, EncodePublicKey, EncryptedPrivateKeyInfo, LineEnding,
-    PrivateKeyInfo,
+    DecodePrivateKey, DecodePublicKey, EncodePrivateKey, EncodePublicKey, EncryptedPrivateKeyInfo,
+    LineEnding, PrivateKeyInfo,
 };
 use rand::RngCore;
-use rsa::{pkcs1v15::SigningKey, Pkcs1v15Sign, RsaPrivateKey, RsaPublicKey};
-use sha2::Sha256;
+use rsa::{
+    pkcs1v15::SigningKey, traits::PublicKeyParts, Pkcs1v15Sign, RsaPrivateKey, RsaPublicKey,
+};
+use serde::{Deserialize, Serialize};
+use sha1::Sha1;
+use sha2::{Digest, Sha256, Sha512};
 use thiserror::Error;
 use x509_cert::{
     builder::{Builder, CertificateBuilder, Profile},
@@ -40,6 +45,20 @@ use x509_cert::{
 
 #[derive(Debug, Error)]
 pub enum Error {
+    #[error("Signature algorithm not supported: {0:?}")]
+    UnsupportedAlgorithm(SignatureAlgorithm),
+    #[error("RSA key size ({}) not supported", .0 * 8)]
+    UnsupportedKey(usize),
+    #[error("Invalid digest length ({0} bytes) for {1:?}")]
+    InvalidDigestLength(usize, SignatureAlgorithm),
+    #[error("Invalid signature length ({0} bytes) for {1:?}")]
+    InvalidSignatureLength(usize, SignatureAlgorithm),
+    #[error("Failed to run command: {0}")]
+    CommandSpawnFailed(String, #[source] io::Error),
+    #[error("Command failed with status: {1}: {0}")]
+    CommandExecutionFailed(String, ExitStatus),
+    #[error("Signature from signing helper does not match public key: {0:?}")]
+    SigningHelperBadSignature(PathBuf),
     #[error("Passphrases do not match")]
     ConfirmPassphrase,
     #[error("Failed to read environment variable: {0:?}")]
@@ -54,6 +73,10 @@ pub enum Error {
     SaveKeyEncrypted(#[source] pkcs8::Error),
     #[error("Failed to save unencrypted private key")]
     SaveKeyUnencrypted(#[source] pkcs8::Error),
+    #[error("Failed to RSA sign digest")]
+    RsaSign(#[source] rsa::Error),
+    #[error("Failed to RSA verify signature")]
+    RsaVerify(#[source] rsa::Error),
     #[error("X509 error")]
     X509(#[from] x509_cert::builder::Error),
     #[error("SPKI error")]
@@ -68,6 +91,34 @@ pub enum Error {
 
 type Result<T> = std::result::Result<T, Error>;
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
+pub enum SignatureAlgorithm {
+    Sha1WithRsa,
+    Sha256WithRsa,
+    Sha512WithRsa,
+}
+
+impl SignatureAlgorithm {
+    /// Length of digest required by the signing algorithm.
+    pub fn digest_len(self) -> usize {
+        match self {
+            Self::Sha1WithRsa => Sha1::output_size(),
+            Self::Sha256WithRsa => Sha256::output_size(),
+            Self::Sha512WithRsa => Sha512::output_size(),
+        }
+    }
+
+    /// Compute the digest of the specified data.
+    pub fn hash(self, data: &[u8]) -> Vec<u8> {
+        match self {
+            Self::Sha1WithRsa => Sha1::digest(data).to_vec(),
+            Self::Sha256WithRsa => Sha256::digest(data).to_vec(),
+            Self::Sha512WithRsa => Sha512::digest(data).to_vec(),
+        }
+    }
+}
+
+#[derive(Clone)]
 pub enum PassphraseSource {
     Prompt(String),
     EnvVar(OsString),
@@ -110,6 +161,181 @@ impl PassphraseSource {
     }
 }
 
+fn check_key_size(size: usize) -> Result<()> {
+    // RustCrypto does not support 8192-bit keys.
+    if size > 4096 / 8 {
+        return Err(Error::UnsupportedKey(size));
+    }
+
+    Ok(())
+}
+
+/// Copied from rsa-0.9.6 since the function is not exported.
+fn pkcs1v15_sign_pad(prefix: &[u8], hashed: &[u8], k: usize) -> rsa::Result<Vec<u8>> {
+    let hash_len = hashed.len();
+    let t_len = prefix.len() + hashed.len();
+    if k < t_len + 11 {
+        return Err(rsa::Error::MessageTooLong);
+    }
+
+    // EM = 0x00 || 0x01 || PS || 0x00 || T
+    let mut em = vec![0xff; k];
+    em[0] = 0;
+    em[1] = 1;
+    em[k - t_len - 1] = 0;
+    em[k - t_len..k - hash_len].copy_from_slice(prefix);
+    em[k - hash_len..k].copy_from_slice(hashed);
+
+    Ok(em)
+}
+
+#[derive(Clone)]
+pub enum RsaSigningKey {
+    Internal(RsaPrivateKey),
+    External {
+        program: PathBuf,
+        public_key_file: PathBuf,
+        public_key: RsaPublicKey,
+        passphrase_source: PassphraseSource,
+    },
+}
+
+impl RsaSigningKey {
+    /// Size of key in bytes.
+    pub fn size(&self) -> usize {
+        match self {
+            Self::Internal(key) => key.size(),
+            Self::External { public_key, .. } => public_key.size(),
+        }
+    }
+
+    /// Get the public key portion of the signing key.
+    pub fn to_public_key(&self) -> RsaPublicKey {
+        match self {
+            RsaSigningKey::Internal(key) => key.to_public_key(),
+            RsaSigningKey::External { public_key, .. } => public_key.clone(),
+        }
+    }
+
+    /// Sign the digest with the specified signature algorithm.
+    pub fn sign(&self, algo: SignatureAlgorithm, digest: &[u8]) -> Result<Vec<u8>> {
+        if digest.len() != algo.digest_len() {
+            return Err(Error::InvalidDigestLength(digest.len(), algo));
+        }
+
+        check_key_size(self.size())?;
+
+        let scheme = match algo {
+            // We don't support signing with insecure algorithms.
+            SignatureAlgorithm::Sha1WithRsa => return Err(Error::UnsupportedAlgorithm(algo)),
+            SignatureAlgorithm::Sha256WithRsa => Pkcs1v15Sign::new::<Sha256>(),
+            SignatureAlgorithm::Sha512WithRsa => Pkcs1v15Sign::new::<Sha512>(),
+        };
+
+        match self {
+            Self::Internal(key) => key.sign(scheme, digest).map_err(Error::RsaSign),
+            Self::External {
+                program,
+                public_key,
+                public_key_file,
+                passphrase_source,
+            } => {
+                let key_bits = public_key.size() * 8;
+                let algo_str = match algo {
+                    SignatureAlgorithm::Sha1WithRsa => unreachable!(),
+                    SignatureAlgorithm::Sha256WithRsa => format!("SHA256_RSA{key_bits}"),
+                    SignatureAlgorithm::Sha512WithRsa => format!("SHA512_RSA{key_bits}"),
+                };
+
+                let mut command = Command::new(program);
+                command.arg(algo_str);
+                command.arg(public_key_file);
+
+                match passphrase_source {
+                    PassphraseSource::Prompt(_) => {}
+                    PassphraseSource::EnvVar(v) => {
+                        command.arg("env");
+                        command.arg(v);
+                    }
+                    PassphraseSource::File(p) => {
+                        command.arg("file");
+                        command.arg(p);
+                    }
+                }
+
+                command.stdin(Stdio::piped());
+                command.stdout(Stdio::piped());
+                command.stderr(Stdio::inherit());
+
+                let mut child = command
+                    .spawn()
+                    .map_err(|e| Error::CommandSpawnFailed(format!("{command:?}"), e))?;
+
+                // We don't bother with spawning a thread. The pipe capacity on
+                // all major OSs is significantly larger than the digest, so we
+                // don't risk deadlocking even if the process doesn't read from
+                // stdin.
+                //
+                // Pipe capacities:
+                // * Linux: 64 KiB
+                // * macOS: 4 KiB, 16 KiB (usually), or 64 KiB
+                // * Windows: 4 KiB
+
+                let padded_digest = pkcs1v15_sign_pad(&scheme.prefix, digest, public_key.size())?;
+                child.stdin.as_mut().unwrap().write_all(&padded_digest)?;
+
+                let child = child.wait_with_output()?;
+
+                if !child.status.success() {
+                    return Err(Error::CommandExecutionFailed(
+                        format!("{command:?}"),
+                        child.status,
+                    ));
+                } else if child.stdout.len() != self.size() {
+                    return Err(Error::InvalidSignatureLength(child.stdout.len(), algo));
+                }
+
+                // Check that the helper signed with the proper key.
+                if let Err(e) = self.to_public_key().verify_sig(algo, digest, &child.stdout) {
+                    return match e {
+                        Error::RsaVerify(_) => {
+                            Err(Error::SigningHelperBadSignature(public_key_file.clone()))
+                        }
+                        e => Err(e),
+                    };
+                }
+
+                Ok(child.stdout)
+            }
+        }
+    }
+}
+
+pub trait RsaPublicKeyExt {
+    fn verify_sig(&self, algo: SignatureAlgorithm, digest: &[u8], signature: &[u8]) -> Result<()>;
+}
+
+impl RsaPublicKeyExt for RsaPublicKey {
+    /// Verify the signature against the specified key.
+    fn verify_sig(&self, algo: SignatureAlgorithm, digest: &[u8], signature: &[u8]) -> Result<()> {
+        // Check this explicitly so we can provide a better error message.
+        if digest.len() != algo.digest_len() {
+            return Err(Error::InvalidDigestLength(digest.len(), algo));
+        }
+
+        check_key_size(self.size())?;
+
+        let scheme = match algo {
+            SignatureAlgorithm::Sha1WithRsa => Pkcs1v15Sign::new::<Sha1>(),
+            SignatureAlgorithm::Sha256WithRsa => Pkcs1v15Sign::new::<Sha256>(),
+            SignatureAlgorithm::Sha512WithRsa => Pkcs1v15Sign::new::<Sha512>(),
+        };
+
+        self.verify(scheme, digest, signature)
+            .map_err(Error::RsaVerify)
+    }
+}
+
 /// Generate an 4096-bit RSA key pair.
 pub fn generate_rsa_key_pair() -> Result<RsaPrivateKey> {
     let mut rng = rand::thread_rng();
@@ -164,6 +390,9 @@ fn reformat_pem(data: &[u8]) -> Result<Vec<u8>> {
             continue;
         } else if line.starts_with(b"-----BEGIN CERTIFICATE-----") {
             inside_base64 = true;
+
+            result.extend_from_slice(line);
+            result.push(b'\n');
         } else if line.starts_with(b"-----END CERTIFICATE-----") {
             inside_base64 = false;
 
@@ -173,13 +402,13 @@ fn reformat_pem(data: &[u8]) -> Result<Vec<u8>> {
             }
 
             base64.clear();
+
+            result.extend_from_slice(line);
+            result.push(b'\n');
         } else if inside_base64 {
             base64.extend_from_slice(line);
             continue;
         }
-
-        result.extend_from_slice(line);
-        result.push(b'\n');
     }
 
     if inside_base64 {
@@ -225,6 +454,16 @@ pub fn write_pem_cert_file(path: &Path, cert: &Certificate) -> Result<()> {
     write_pem_cert(writer, cert)
 }
 
+/// Read PEM-encoded PKCS8 public key from a reader.
+pub fn read_pem_public_key(mut reader: impl Read) -> Result<RsaPublicKey> {
+    let mut data = String::new();
+    reader.read_to_string(&mut data)?;
+
+    let key = RsaPublicKey::from_public_key_pem(&data)?;
+
+    Ok(key)
+}
+
 /// Write PEM-encoded PKCS8 public key to a writer.
 pub fn write_pem_public_key(mut writer: impl Write, key: &RsaPublicKey) -> Result<()> {
     let data = key.to_public_key_pem(LineEnding::LF)?;
@@ -234,6 +473,14 @@ pub fn write_pem_public_key(mut writer: impl Write, key: &RsaPublicKey) -> Resul
     Ok(())
 }
 
+/// Read PEM-encoded PKCS8 public key from a file.
+pub fn read_pem_public_key_file(path: &Path) -> Result<RsaPublicKey> {
+    let file = File::open(path)?;
+    let reader = BufReader::new(file);
+
+    read_pem_public_key(reader)
+}
+
 /// Write PEM-encoded PKCS8 public key to a file.
 pub fn write_pem_public_key_file(path: &Path, key: &RsaPublicKey) -> Result<()> {
     let file = File::create(path)?;
@@ -348,7 +595,7 @@ pub fn get_public_key(cert: &Certificate) -> Result<RsaPublicKey> {
 }
 
 /// Check if a certificate matches a private key.
-pub fn cert_matches_key(cert: &Certificate, key: &RsaPrivateKey) -> Result<bool> {
+pub fn cert_matches_key(cert: &Certificate, key: &RsaSigningKey) -> Result<bool> {
     let public_key = get_public_key(cert)?;
 
     Ok(key.to_public_key() == public_key)
@@ -386,12 +633,11 @@ pub fn get_cms_certs(sd: &SignedData) -> Vec<Certificate> {
 /// a transport mechanism for a raw signature. Thus, we need to ensure that the
 /// signature covers nothing but the raw data.
 pub fn cms_sign_external(
-    key: &RsaPrivateKey,
+    key: &RsaSigningKey,
     cert: &Certificate,
     digest: &[u8],
 ) -> Result<ContentInfo> {
-    let scheme = Pkcs1v15Sign::new::<Sha256>();
-    let signature = key.sign(scheme, digest)?;
+    let signature = key.sign(SignatureAlgorithm::Sha256WithRsa, digest)?;
 
     let digest_algorithm = AlgorithmIdentifierOwned {
         oid: const_oid::db::rfc5912::ID_SHA_256,

avbroot/src/format/avb.rs

@@ -1,10 +1,10 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
 use std::{
-    cmp, fmt,
+    fmt,
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     ops::Range,
     str,
@@ -16,12 +16,12 @@ use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use num_bigint_dig::{ModInverse, ToBigInt};
 use num_traits::{Pow, ToPrimitive};
 use ring::digest::{Algorithm, Context};
-use rsa::{traits::PublicKeyParts, BigUint, Pkcs1v15Sign, RsaPrivateKey, RsaPublicKey};
+use rsa::{traits::PublicKeyParts, BigUint, RsaPublicKey};
 use serde::{Deserialize, Serialize};
-use sha2::{Digest, Sha256, Sha512};
 use thiserror::Error;
 
 use crate::{
+    crypto::{self, RsaPublicKeyExt, RsaSigningKey, SignatureAlgorithm},
     escape,
     format::{
         fec::{self, Fec},
@@ -29,8 +29,8 @@ use crate::{
         padding,
     },
     stream::{
-        self, CountingReader, FromReader, ReadDiscardExt, ReadSeekReopen, ReadStringExt, ToWriter,
-        WriteSeekReopen, WriteStringExt, WriteZerosExt,
+        self, CountingReader, CountingWriter, FromReader, ReadDiscardExt, ReadSeekReopen,
+        ReadStringExt, ToWriter, WriteSeekReopen, WriteStringExt, WriteZerosExt,
     },
     util,
 };
@@ -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 {
@@ -86,12 +103,9 @@ pub enum Error {
     UnsupportedAlgorithm(AlgorithmType),
     #[error("Hashing algorithm not supported: {0:?}")]
     UnsupportedHashAlgorithm(String),
-    #[error("Incorrect key size ({key_size} bytes) for algorithm {algo:?} ({} bytes)", algo.public_key_len())]
-    IncorrectKeySize {
-        key_size: usize,
-        algo: AlgorithmType,
-    },
-    #[error("RSA key size (0) is not compatible with any AVB signing algorithm")]
+    #[error("Incorrect key size ({}) for algorithm {1:?}", .0 * 8)]
+    IncorrectKeySize(usize, AlgorithmType),
+    #[error("RSA key size ({}) is not compatible with any AVB signing algorithm", .0 * 8)]
     UnsupportedKey(usize),
     #[error("Hash tree does not immediately follow image data")]
     HashTreeGap,
@@ -103,16 +117,18 @@ pub enum Error {
     MismatchedFecBlockSizes { data: u32, hash: u32 },
     #[error("Must have exactly one hash or hash tree descriptor")]
     NoAppendedDescriptor,
-    #[error("Failed to RSA sign digest")]
-    RsaSign(#[source] rsa::Error),
-    #[error("Failed to RSA verify signature")]
-    RsaVerify(#[source] rsa::Error),
-    #[error("{0} byte image size is too small to fit header or footer")]
-    ImageSizeTooSmall(u64),
+    #[error("{0} byte image size is too small to fit header")]
+    TooSmallForHeader(u64),
+    #[error("{0} byte image size is too small to fit footer")]
+    TooSmallForFooter(u64),
+    #[error("Crypto error")]
+    Crypto(#[from] crypto::Error),
     #[error("Hash tree error")]
     HashTree(#[from] hashtree::Error),
     #[error("FEC error")]
     Fec(#[from] fec::Error),
+    #[error("RSA error")]
+    Rsa(#[from] rsa::Error),
     #[error("I/O error")]
     Io(#[from] io::Error),
 }
@@ -137,6 +153,7 @@ pub enum AlgorithmType {
     Sha512Rsa2048,
     Sha512Rsa4096,
     Sha512Rsa8192,
+    #[serde(untagged)]
     Unknown(u32),
 }
 
@@ -167,18 +184,24 @@ impl AlgorithmType {
         }
     }
 
-    pub fn hash_len(self) -> usize {
+    pub fn to_digest_algorithm(self) -> Option<SignatureAlgorithm> {
         match self {
-            Self::None | Self::Unknown(_) => 0,
             Self::Sha256Rsa2048 | Self::Sha256Rsa4096 | Self::Sha256Rsa8192 => {
-                Sha256::output_size()
+                Some(SignatureAlgorithm::Sha256WithRsa)
             }
             Self::Sha512Rsa2048 | Self::Sha512Rsa4096 | Self::Sha512Rsa8192 => {
-                Sha512::output_size()
+                Some(SignatureAlgorithm::Sha512WithRsa)
             }
+            _ => None,
         }
     }
 
+    pub fn digest_len(self) -> usize {
+        self.to_digest_algorithm()
+            .map(|a| a.digest_len())
+            .unwrap_or_default()
+    }
+
     pub fn signature_len(self) -> usize {
         match self {
             Self::None | Self::Unknown(_) => 0,
@@ -198,49 +221,36 @@ impl AlgorithmType {
     }
 
     pub fn hash(self, data: &[u8]) -> Vec<u8> {
-        match self {
-            Self::None | Self::Unknown(_) => vec![],
-            Self::Sha256Rsa2048 | Self::Sha256Rsa4096 | Self::Sha256Rsa8192 => {
-                Sha256::digest(data).to_vec()
-            }
-            Self::Sha512Rsa2048 | Self::Sha512Rsa4096 | Self::Sha512Rsa8192 => {
-                Sha512::digest(data).to_vec()
-            }
-        }
-    }
-
-    pub fn sign(self, key: &RsaPrivateKey, digest: &[u8]) -> Result<Vec<u8>> {
-        let signature = match self {
-            Self::None | Self::Unknown(_) => vec![],
-            Self::Sha256Rsa2048 | Self::Sha256Rsa4096 | Self::Sha256Rsa8192 => {
-                let scheme = Pkcs1v15Sign::new::<Sha256>();
-                key.sign(scheme, digest).map_err(Error::RsaSign)?
-            }
-            Self::Sha512Rsa2048 | Self::Sha512Rsa4096 | Self::Sha512Rsa8192 => {
-                let scheme = Pkcs1v15Sign::new::<Sha512>();
-                key.sign(scheme, digest).map_err(Error::RsaSign)?
-            }
+        let Some(algo) = self.to_digest_algorithm() else {
+            return vec![];
         };
 
-        Ok(signature)
+        algo.hash(data)
+    }
+
+    pub fn sign(self, key: &RsaSigningKey, digest: &[u8]) -> Result<Vec<u8>> {
+        let Some(algo) = self.to_digest_algorithm() else {
+            return if self == Self::None {
+                Ok(vec![])
+            } else {
+                Err(Error::UnsupportedAlgorithm(self))
+            };
+        };
+
+        key.sign(algo, digest).map_err(|e| e.into())
     }
 
     pub fn verify(self, key: &RsaPublicKey, digest: &[u8], signature: &[u8]) -> Result<()> {
-        match self {
-            Self::None | Self::Unknown(_) => {}
-            Self::Sha256Rsa2048 | Self::Sha256Rsa4096 | Self::Sha256Rsa8192 => {
-                let scheme = Pkcs1v15Sign::new::<Sha256>();
-                key.verify(scheme, digest, signature)
-                    .map_err(Error::RsaVerify)?;
-            }
-            Self::Sha512Rsa2048 | Self::Sha512Rsa4096 | Self::Sha512Rsa8192 => {
-                let scheme = Pkcs1v15Sign::new::<Sha512>();
-                key.verify(scheme, digest, signature)
-                    .map_err(Error::RsaVerify)?;
-            }
-        }
+        let Some(algo) = self.to_digest_algorithm() else {
+            return if self == Self::None {
+                Ok(())
+            } else {
+                Err(Error::UnsupportedAlgorithm(self))
+            };
+        };
 
-        Ok(())
+        key.verify_sig(algo, digest, signature)
+            .map_err(|e| e.into())
     }
 }
 
@@ -1442,7 +1452,7 @@ impl Header {
         result.ok_or(Error::NoAppendedDescriptor)
     }
 
-    pub fn set_algo_for_key(&mut self, key: &RsaPrivateKey) -> Result<()> {
+    pub fn set_algo_for_key(&mut self, key: &RsaSigningKey) -> Result<()> {
         let key_raw = encode_public_key(&key.to_public_key())?;
 
         for algo in [AlgorithmType::Sha256Rsa2048, AlgorithmType::Sha256Rsa4096] {
@@ -1462,30 +1472,17 @@ impl Header {
         self.public_key_metadata.clear();
     }
 
-    pub fn sign(&mut self, key: &RsaPrivateKey) -> Result<()> {
+    pub fn sign(&mut self, key: &RsaSigningKey) -> Result<()> {
         let key_raw = encode_public_key(&key.to_public_key())?;
 
-        // RustCrypto does not support 8192-bit keys.
-        match self.algorithm_type {
-            AlgorithmType::Sha256Rsa8192
-            | AlgorithmType::Sha512Rsa8192
-            | AlgorithmType::Unknown(_) => {
-                return Err(Error::UnsupportedAlgorithm(self.algorithm_type));
-            }
-            _ => {}
-        }
-
         if key_raw.len() != self.algorithm_type.public_key_len() {
-            return Err(Error::IncorrectKeySize {
-                key_size: key_raw.len(),
-                algo: self.algorithm_type,
-            });
+            return Err(Error::IncorrectKeySize(key.size(), self.algorithm_type));
         }
 
         // The public key and the sizes of the hash and signature are included
         // in the data that's about to be signed.
         self.public_key = key_raw;
-        self.hash.resize(self.algorithm_type.hash_len(), 0);
+        self.hash.resize(self.algorithm_type.digest_len(), 0);
         self.signature
             .resize(self.algorithm_type.signature_len(), 0);
 
@@ -1506,18 +1503,16 @@ impl Header {
     /// and return the public key. If the header is not signed, then `None` is
     /// returned.
     pub fn verify(&self) -> Result<Option<RsaPublicKey>> {
-        // RustCrypto does not support 8192-bit keys.
-        match self.algorithm_type {
-            AlgorithmType::None => return Ok(None),
-            a @ AlgorithmType::Sha256Rsa8192
-            | a @ AlgorithmType::Sha512Rsa8192
-            | a @ AlgorithmType::Unknown(_) => return Err(Error::UnsupportedAlgorithm(a)),
-            _ => {}
-        }
-
         // Reconstruct the public key.
         let public_key = decode_public_key(&self.public_key)?;
 
+        if self.public_key.len() != self.algorithm_type.public_key_len() {
+            return Err(Error::IncorrectKeySize(
+                public_key.size(),
+                self.algorithm_type,
+            ));
+        }
+
         let mut without_auth_writer = Cursor::new(Vec::new());
         self.to_writer_internal(&mut without_auth_writer, true)?;
         let without_auth = without_auth_writer.into_inner();
@@ -1806,8 +1801,7 @@ pub fn decode_public_key(data: &[u8]) -> Result<RsaPublicKey> {
     reader.read_exact(&mut modulus_raw)?;
 
     let modulus = BigUint::from_bytes_be(&modulus_raw);
-    let public_key =
-        RsaPublicKey::new(modulus, BigUint::from(65537u32)).map_err(Error::RsaVerify)?;
+    let public_key = RsaPublicKey::new(modulus, BigUint::from(65537u32))?;
 
     Ok(public_key)
 }
@@ -1834,97 +1828,95 @@ pub fn load_image(mut reader: impl Read + Seek) -> Result<(Header, Option<Footer
     Ok((header, footer, image_size))
 }
 
-/// Write a vbmeta header to the specified writer. If a footer is specified, it
-/// will be used as the basis of the newly written footer, with the original
-/// image size, vbmeta header offset, and vbmeta header size fields updated
-/// appropriately.
-///
-/// The writer must not have an existing vbmeta header or footer.
-fn write_image_internal(
-    mut writer: impl Write + Seek,
-    header: &Header,
-    footer: Option<&mut Footer>,
-    image_size: Option<u64>,
-    block_size: u64,
-) -> Result<()> {
-    let eof_image_size = if footer.is_some() {
-        match header.appended_descriptor()? {
-            AppendedDescriptorRef::HashTree(d) => d
-                .image_size
-                .checked_add(d.tree_size)
-                .and_then(|s| s.checked_add(d.fec_size))
-                .ok_or_else(|| Error::FieldOutOfBounds("eof_image_size"))?,
-            AppendedDescriptorRef::Hash(d) => d.image_size,
-        }
-    } else {
-        0
-    };
-
-    writer.seek(SeekFrom::Start(eof_image_size))?;
-
-    // The header must be block-aligned.
-    let vbmeta_offset = if block_size > 0 {
-        let padding_size = padding::write_zeros(&mut writer, block_size)?;
-        eof_image_size
-            .checked_add(padding_size)
-            .ok_or_else(|| Error::FieldOutOfBounds("vbmeta_offset"))?
-    } else {
-        eof_image_size
-    };
-
-    header.to_writer(&mut writer)?;
-    let vbmeta_end = writer.stream_position()?;
-
-    if let Some(s) = image_size {
-        let footer_space = if footer.is_some() {
-            cmp::max(block_size, Footer::SIZE as u64)
-        } else {
-            0
-        };
-
-        if s < footer_space || vbmeta_end > s - footer_space {
-            return Err(Error::ImageSizeTooSmall(s));
-        }
-    }
-
-    if block_size > 0 {
-        padding::write_zeros(&mut writer, block_size)?;
-    }
-
-    if let Some(f) = footer {
-        let footer_offset = image_size.unwrap() - Footer::SIZE as u64;
-        writer.seek(SeekFrom::Start(footer_offset))?;
-
-        let original_image_size = match header.appended_descriptor()? {
-            AppendedDescriptorRef::HashTree(d) => d.image_size,
-            AppendedDescriptorRef::Hash(d) => d.image_size,
-        };
-
-        f.original_image_size = original_image_size;
-        f.vbmeta_offset = vbmeta_offset;
-        f.vbmeta_size = vbmeta_end - vbmeta_offset;
-
-        f.to_writer(&mut writer)?;
-    }
-
-    Ok(())
-}
-
 /// Write a vbmeta header to the specified writer. This is meant for writing
 /// vbmeta partition images, not appended vbmeta images. The writer must refer
-/// to an empty file.
-pub fn write_root_image(writer: impl Write + Seek, header: &Header, block_size: u64) -> Result<()> {
-    write_image_internal(writer, header, None, None, block_size)
+/// to an empty file. Returns the size of the new file.
+pub fn write_root_image(writer: impl Write, header: &Header, block_size: u64) -> Result<u64> {
+    let mut counting_writer = CountingWriter::new(writer);
+
+    header.to_writer(&mut counting_writer)?;
+    padding::write_zeros(&mut counting_writer, block_size)?;
+
+    Ok(counting_writer.stream_position()?)
 }
 
 /// Write a vbmeta header and footer to the specified writer. This is meant for
 /// appending vbmeta data to existing partition data, not writing vbmeta images.
+/// If `image_size` is specified, then the writer is guaranteed to not grow
+/// past that size and an error is returned if the header and footer won't fit.
+/// Otherwise, the writer will grow to the necessary size. Returns the size of
+/// the new file.
 pub fn write_appended_image(
-    writer: impl Write + Seek,
+    mut writer: impl Write + Seek,
     header: &Header,
     footer: &mut Footer,
-    image_size: u64,
-) -> Result<()> {
+    image_size: Option<u64>,
+) -> Result<u64> {
     // avbtool hardcodes a 4096 block size for appended non-sparse images.
-    write_image_internal(writer, header, Some(footer), Some(image_size), 4096)
+    const BLOCK_SIZE: u64 = 4096;
+
+    // Logical image size, excluding the AVB header and footer.
+    let logical_image_size = match header.appended_descriptor()? {
+        AppendedDescriptorRef::HashTree(d) => d
+            .image_size
+            .checked_add(d.tree_size)
+            .and_then(|s| s.checked_add(d.fec_size))
+            .ok_or_else(|| Error::FieldOutOfBounds("logical_image_size"))?,
+        AppendedDescriptorRef::Hash(d) => d.image_size,
+    };
+
+    writer.seek(SeekFrom::Start(logical_image_size))?;
+
+    // The header start offset must be block aligned.
+    let header_offset = {
+        let padding_size = padding::write_zeros(&mut writer, BLOCK_SIZE)?;
+        logical_image_size
+            .checked_add(padding_size)
+            .ok_or_else(|| Error::FieldOutOfBounds("header_offset"))?
+    };
+
+    // The header lives at the beginning of the empty space.
+    let mut header_buf = Cursor::new(Vec::new());
+    header.to_writer(&mut header_buf)?;
+    let header_size = header_buf.stream_position()?;
+    let header_padding = padding::write_zeros(&mut header_buf, BLOCK_SIZE)?;
+    let header_end_padded = header_offset
+        .checked_add(header_size)
+        .and_then(|s| s.checked_add(header_padding))
+        .ok_or_else(|| Error::FieldOutOfBounds("header_end_padded"))?;
+
+    if let Some(s) = image_size {
+        if header_end_padded > s {
+            return Err(Error::TooSmallForHeader(s));
+        }
+    }
+
+    writer.write_all(&header_buf.into_inner())?;
+
+    // The footer lives in its own separate block at the end of the empty space.
+    let footer_end = if let Some(s) = image_size {
+        if s - header_end_padded < BLOCK_SIZE {
+            return Err(Error::TooSmallForFooter(s));
+        }
+
+        s
+    } else {
+        header_end_padded
+            .checked_add(BLOCK_SIZE)
+            .ok_or_else(|| Error::FieldOutOfBounds("footer_end"))?
+    };
+
+    let footer_offset = footer_end - Footer::SIZE as u64;
+    writer.seek(SeekFrom::Start(footer_offset))?;
+
+    footer.original_image_size = match header.appended_descriptor()? {
+        AppendedDescriptorRef::HashTree(d) => d.image_size,
+        AppendedDescriptorRef::Hash(d) => d.image_size,
+    };
+    footer.vbmeta_offset = header_offset;
+    footer.vbmeta_size = header_size;
+
+    footer.to_writer(&mut writer)?;
+
+    Ok(footer_end)
 }

avbroot/src/format/bootimage.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -12,11 +12,11 @@ use std::{
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use num_traits::ToPrimitive;
 use ring::digest::Context;
-use rsa::RsaPrivateKey;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::{
+    crypto::RsaSigningKey,
     format::{
         avb::{self, Descriptor, Header},
         padding,
@@ -765,7 +765,7 @@ impl BootImageV3Through4 {
     /// Sign the boot image with a legacy VTS signature. Returns true if the
     /// image was successfully signed. Returns false if there's no vbmeta
     /// structure to sign in [`V4Extra::signature`].
-    pub fn sign(&mut self, key: &RsaPrivateKey) -> Result<bool> {
+    pub fn sign(&mut self, key: &RsaSigningKey) -> Result<bool> {
         let mut context = Context::new(&ring::digest::SHA256);
         let image_size;
 

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

@@ -52,18 +52,19 @@ pub enum Error {
 
 type Result<T> = std::result::Result<T, Error>;
 
-pub struct HashTree<'a> {
+pub struct HashTree {
     block_size: u32,
-    algorithm: &'static Algorithm,
-    salt: &'a [u8],
+    salted_context: Context,
 }
 
-impl<'a> HashTree<'a> {
-    pub fn new(block_size: u32, algorithm: &'static Algorithm, salt: &'a [u8]) -> Self {
+impl HashTree {
+    pub fn new(block_size: u32, algorithm: &'static Algorithm, salt: &[u8]) -> Self {
+        let mut salted_context = Context::new(algorithm);
+        salted_context.update(salt);
+
         Self {
             block_size,
-            algorithm,
-            salt,
+            salted_context,
         }
     }
 
@@ -71,13 +72,14 @@ impl<'a> HashTree<'a> {
     /// 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>>> {
-        let digest_size = self.algorithm.output_len().next_power_of_two();
+    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)))
@@ -122,7 +124,7 @@ impl<'a> HashTree<'a> {
             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);
@@ -143,8 +145,8 @@ impl<'a> HashTree<'a> {
         cancel_signal: &AtomicBool,
     ) -> io::Result<()> {
         // Each digest must be a power of 2.
-        let digest_padding =
-            self.algorithm.output_len().next_power_of_two() - self.algorithm.output_len();
+        let algorithm = self.salted_context.algorithm();
+        let digest_padding = algorithm.output_len().next_power_of_two() - algorithm.output_len();
         let mut buf = vec![0u8; self.block_size as usize];
 
         while size > 0 {
@@ -157,8 +159,7 @@ impl<'a> HashTree<'a> {
             // with padding.
             buf[n..].fill(0);
 
-            let mut context = Context::new(self.algorithm);
-            context.update(self.salt);
+            let mut context = self.salted_context.clone();
             context.update(&buf);
 
             // Add the digest to the tree level. Each tree node must be a power
@@ -191,7 +192,8 @@ impl<'a> HashTree<'a> {
         );
 
         // Parallelize in larger chunks to avoid too much seek thrashing.
-        let digest_size = self.algorithm.output_len().next_power_of_two();
+        let algorithm = self.salted_context.algorithm();
+        let digest_size = algorithm.output_len().next_power_of_two();
         let multiplier = 1024u64;
 
         level_data
@@ -222,7 +224,8 @@ impl<'a> HashTree<'a> {
         level_data: &mut [u8],
         cancel_signal: &AtomicBool,
     ) -> io::Result<()> {
-        let digest_size = self.algorithm.output_len().next_power_of_two();
+        let algorithm = self.salted_context.algorithm();
+        let digest_size = algorithm.output_len().next_power_of_two();
 
         level_data
             .par_chunks_exact_mut(digest_size)
@@ -263,8 +266,7 @@ impl<'a> HashTree<'a> {
             let mut buf = vec![0u8; image_size as usize];
             reader.read_exact(&mut buf)?;
 
-            let mut context = Context::new(self.algorithm);
-            context.update(self.salt);
+            let mut context = self.salted_context.clone();
             context.update(&buf);
             let digest = context.finish();
 
@@ -309,8 +311,7 @@ impl<'a> HashTree<'a> {
         }
 
         // Calculate the root hash.
-        let mut context = Context::new(self.algorithm);
-        context.update(self.salt);
+        let mut context = self.salted_context.clone();
         context.update(&hash_tree_data[level_offsets.last().unwrap().clone()]);
         let root_hash = context.finish().as_ref().to_vec();
 
@@ -402,8 +403,9 @@ impl<'a> HashTree<'a> {
 
         if hash_tree_data != actual_hash_tree_data {
             // These are multiple megabytes, so only report the hashes.
-            let expected = ring::digest::digest(self.algorithm, hash_tree_data);
-            let actual = ring::digest::digest(self.algorithm, &actual_hash_tree_data);
+            let algorithm = self.salted_context.algorithm();
+            let expected = ring::digest::digest(algorithm, hash_tree_data);
+            let actual = ring::digest::digest(algorithm, &actual_hash_tree_data);
 
             return Err(Error::InvalidHashTree {
                 expected: hex::encode(expected),

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
  */
 
@@ -15,15 +15,12 @@ use const_oid::{db::rfc5912, ObjectIdentifier};
 use memchr::memmem;
 use prost::Message;
 use ring::digest::Context;
-use rsa::{Pkcs1v15Sign, RsaPrivateKey};
-use sha1::Sha1;
-use sha2::Sha256;
 use thiserror::Error;
 use x509_cert::{der::Encode, Certificate};
 use zip::{result::ZipError, write::FileOptions, CompressionMethod, ZipArchive, ZipWriter};
 
 use crate::{
-    crypto,
+    crypto::{self, RsaPublicKeyExt, RsaSigningKey, SignatureAlgorithm},
     format::payload::{self, PayloadHeader},
     protobuf::build::tools::releasetools::{ota_metadata::OtaType, OtaMetadata},
     stream::{self, FromReader, HashingReader, HashingWriter},
@@ -88,8 +85,6 @@ pub enum Error {
     Spki(#[from] pkcs8::spki::Error),
     #[error("x509 DER error")]
     Der(#[from] x509_cert::der::Error),
-    #[error("RSA error")]
-    Rsa(#[from] rsa::Error),
     #[error("Zip error")]
     Zip(#[from] ZipError),
     #[error("I/O error")]
@@ -152,15 +147,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 +171,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") => {
@@ -586,12 +581,12 @@ pub fn verify_ota(mut reader: impl Read + Seek, cancel_signal: &AtomicBool) -> R
     reader.seek(SeekFrom::Start(0))?;
 
     // We support SHA1 for verification only.
-    let (algorithm, scheme) = if signer.digest_alg.oid == rfc5912::ID_SHA_256 {
-        (&ring::digest::SHA256, Pkcs1v15Sign::new::<Sha256>())
+    let (algorithm, algo) = if signer.digest_alg.oid == rfc5912::ID_SHA_256 {
+        (&ring::digest::SHA256, SignatureAlgorithm::Sha256WithRsa)
     } else {
         (
             &ring::digest::SHA1_FOR_LEGACY_USE_ONLY,
-            Pkcs1v15Sign::new::<Sha1>(),
+            SignatureAlgorithm::Sha1WithRsa,
         )
     };
 
@@ -603,7 +598,7 @@ pub fn verify_ota(mut reader: impl Read + Seek, cancel_signal: &AtomicBool) -> R
     let digest = context.finish();
 
     // Verify the signature against the public key.
-    public_key.verify(scheme, digest.as_ref(), signer.signature.as_bytes())?;
+    public_key.verify_sig(algo, digest.as_ref(), signer.signature.as_bytes())?;
 
     Ok(cert.clone())
 }
@@ -661,7 +656,7 @@ impl<W: Write> SigningWriter<W> {
         }
     }
 
-    pub fn finish(mut self, key: &RsaPrivateKey, cert: &Certificate) -> Result<W> {
+    pub fn finish(mut self, key: &RsaSigningKey, cert: &Certificate) -> Result<W> {
         if self.used < self.queue.len() {
             return Err(
                 io::Error::new(io::ErrorKind::InvalidData, "Too small to contain EOCD").into(),

avbroot/src/format/payload.rs

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: 2022-2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
 use std::{
     collections::{HashMap, HashSet},
+    fmt,
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     ops::Range,
     sync::atomic::AtomicBool,
@@ -26,13 +27,12 @@ use rayon::{
     prelude::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator},
 };
 use ring::digest::{Context, Digest};
-use rsa::{traits::PublicKeyParts, Pkcs1v15Sign, RsaPrivateKey};
-use sha2::Sha256;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use x509_cert::Certificate;
 
 use crate::{
-    crypto,
+    crypto::{self, RsaPublicKeyExt, RsaSigningKey, SignatureAlgorithm},
     protobuf::chromeos_update_engine::{
         install_operation::Type, signatures::Signature, DeltaArchiveManifest, Extent,
         InstallOperation, PartitionInfo, PartitionUpdate, Signatures,
@@ -100,19 +100,19 @@ pub enum Error {
     ProtobufDecode(#[from] prost::DecodeError),
     #[error("XZ stream error")]
     XzStream(#[from] liblzma::stream::Error),
-    #[error("RSA error")]
-    Rsa(#[from] rsa::Error),
     #[error("I/O error")]
     Io(#[from] io::Error),
 }
 
 type Result<T> = std::result::Result<T, Error>;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct PayloadHeader {
     pub version: u64,
     pub manifest: DeltaArchiveManifest,
+    #[serde(skip)]
     pub metadata_signature_size: u32,
+    #[serde(skip)]
     pub blob_offset: u64,
 }
 
@@ -175,9 +175,8 @@ impl<R: Read> FromReader<R> for PayloadHeader {
 
 /// Sign `digest` with `key` and return a [`Signatures`] protobuf struct with
 /// the signature padded to the maximum size.
-fn sign_digest(digest: &[u8], key: &RsaPrivateKey) -> Result<Signatures> {
-    let scheme = Pkcs1v15Sign::new::<Sha256>();
-    let mut digest_signed = key.sign(scheme, digest)?;
+fn sign_digest(digest: &[u8], key: &RsaSigningKey) -> Result<Signatures> {
+    let mut digest_signed = key.sign(SignatureAlgorithm::Sha256WithRsa, digest)?;
     assert!(
         digest_signed.len() <= key.size(),
         "Signature exceeds maximum size",
@@ -214,8 +213,7 @@ fn verify_digest(digest: &[u8], signatures: &Signatures, cert: &Certificate) ->
         };
         let without_padding = &data[..size as usize];
 
-        let scheme = Pkcs1v15Sign::new::<Sha256>();
-        match public_key.verify(scheme, digest, without_padding) {
+        match public_key.verify_sig(SignatureAlgorithm::Sha256WithRsa, digest, without_padding) {
             Ok(_) => return Ok(()),
             Err(e) => last_error = Some(e),
         }
@@ -292,7 +290,7 @@ pub struct PayloadWriter<W: Write> {
     h_partial: Context,
     /// Includes signatures (hashes are for properties file).
     h_full: Context,
-    key: RsaPrivateKey,
+    key: RsaSigningKey,
 }
 
 /// Write data to a writer and one or more hashers.
@@ -315,7 +313,7 @@ impl<W: Write> PayloadWriter<W> {
     /// fields are ignored and internally recomputed to guarantee that there are
     /// no gaps. All partitions' install operation data is written to the blob
     /// section in order.
-    pub fn new(mut inner: W, mut header: PayloadHeader, key: RsaPrivateKey) -> Result<Self> {
+    pub fn new(mut inner: W, mut header: PayloadHeader, key: RsaSigningKey) -> Result<Self> {
         let mut blob_size = 0;
 
         // The blob must contain all data in sequential order with no gaps.
@@ -390,12 +388,12 @@ impl<W: Write> PayloadWriter<W> {
     }
 
     /// Finalize the payload. If this function is not called, the payload will
-    /// be left in an incomplete state. Returns the original writer, the
-    /// contents that should be written for `payload_properties.txt` and the
-    /// length of the header + manifest + manifest signature sections (for
-    /// constructing the `payload_metadata.bin` OTA metadata property files
+    /// be left in an incomplete state. Returns the original writer, the final
+    /// header, the contents that should be written for `payload_properties.txt`
+    /// and the length of the header + manifest + manifest signature sections
+    /// (for constructing the `payload_metadata.bin` OTA metadata property files
     /// entry).
-    pub fn finish(mut self) -> Result<(W, String, u64)> {
+    pub fn finish(mut self) -> Result<(W, PayloadHeader, String, u64)> {
         // Append payload signature.
         let payload_partial_hash = self.h_partial.clone().finish();
         let payload_sig = sign_digest(payload_partial_hash.as_ref(), &self.key)?;
@@ -419,7 +417,7 @@ impl<W: Write> PayloadWriter<W> {
             self.metadata_size as u64,
         );
 
-        Ok((self.inner, properties, metadata_with_sig_size))
+        Ok((self.inner, self.header, properties, metadata_with_sig_size))
     }
 
     /// Prepare for writing the next source data blob corresponding to an
@@ -890,6 +888,52 @@ fn compress_chunk(raw_data: &[u8], cancel_signal: &AtomicBool) -> Result<(Vec<u8
     Ok((data, digest_compressed))
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
+pub enum VabcAlgo {
+    Lz4,
+    Gzip,
+}
+
+impl VabcAlgo {
+    pub fn new(name: &str) -> Option<Self> {
+        match name {
+            "lz4" => Some(Self::Lz4),
+            "gz" => Some(Self::Gzip),
+            _ => None,
+        }
+    }
+
+    fn compressed_size(&self, mut raw_data: &[u8], block_size: u32) -> u64 {
+        let mut total = 0;
+
+        while !raw_data.is_empty() {
+            let n = raw_data.len().min(block_size as usize);
+            let compressed = match self {
+                Self::Lz4 => lz4_flex::block::compress(&raw_data[..n]),
+                // We use the miniz_oxide backend for flate2, but flate2 doesn't
+                // expose a nice function for compressing to a vec, so just use
+                // miniz_oxide directly.
+                Self::Gzip => miniz_oxide::deflate::compress_to_vec_zlib(&raw_data[..n], 9),
+            };
+
+            total += compressed.len().min(n) as u64;
+
+            raw_data = &raw_data[n..];
+        }
+
+        total
+    }
+}
+
+impl fmt::Display for VabcAlgo {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Lz4 => f.write_str("lz4"),
+            Self::Gzip => f.write_str("gz"),
+        }
+    }
+}
+
 /// Compress the image and return the corresponding information to insert into
 /// the payload manifest's [`PartitionUpdate`] instance. The uncompressed data
 /// is split into 2 MiB chunks, which are read and compressed in parallel, and
@@ -897,13 +941,21 @@ fn compress_chunk(raw_data: &[u8], cancel_signal: &AtomicBool) -> Result<(Vec<u8
 /// a corresponding [`InstallOperation`] in the return value. The caller must
 /// update [`InstallOperation::data_offset`] in each operation manually because
 /// the initial values are relative to 0.
+///
+/// If `vabc_algo` is set, the VABC CoW v2 size estimate will be computed. The
+/// caller must update [`PartitionUpdate::estimate_cow_size`] with this value or
+/// else update_engine may fail to flash the partition due to running out of
+/// space on the CoW block device. CoW v2 + other algorithms and also CoW v3 are
+/// currently unsupported because there currently are no known OTAs that use
+/// those configurations.
 pub fn compress_image(
     input: &(dyn ReadSeekReopen + Sync),
     output: &(dyn WriteSeekReopen + Sync),
     partition_name: &str,
     block_size: u32,
+    vabc_algo: Option<VabcAlgo>,
     cancel_signal: &AtomicBool,
-) -> Result<(PartitionInfo, Vec<InstallOperation>)> {
+) -> Result<(PartitionInfo, Vec<InstallOperation>, Option<u64>)> {
     const CHUNK_SIZE: u64 = 2 * 1024 * 1024;
     const CHUNK_GROUP: u64 = 32;
 
@@ -918,9 +970,10 @@ 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 cow_estimate = 0;
     let mut operations = vec![];
 
     // Read the file one group at a time. This allows for some parallelization
@@ -957,8 +1010,13 @@ pub fn compress_image(
         let mut compressed_data_group = uncompressed_data_group
             .into_par_iter()
             .map(
-                |(raw_offset, raw_data)| -> Result<(Vec<u8>, InstallOperation)> {
+                |(raw_offset, raw_data)| -> Result<(Vec<u8>, InstallOperation, u64)> {
                     let (data, digest_compressed) = compress_chunk(&raw_data, cancel_signal)?;
+                    let cow_size = if let Some(algo) = vabc_algo {
+                        algo.compressed_size(&raw_data, block_size)
+                    } else {
+                        0
+                    };
 
                     let extent = Extent {
                         start_block: Some(raw_offset / u64::from(block_size)),
@@ -971,19 +1029,20 @@ pub fn compress_image(
                     operation.dst_extents.push(extent);
                     operation.data_sha256_hash = Some(digest_compressed.as_ref().to_vec());
 
-                    Ok((data, operation))
+                    Ok((data, operation, cow_size))
                 },
             )
             .collect::<Result<Vec<_>>>()?;
 
-        for (data, operation) in &mut compressed_data_group {
+        for (data, operation, cow_size) in &mut compressed_data_group {
             operation.data_offset = Some(bytes_compressed);
             bytes_compressed += data.len() as u64;
+            cow_estimate += *cow_size;
         }
 
         let group_operations = compressed_data_group
             .into_par_iter()
-            .map(|(data, operation)| -> Result<InstallOperation> {
+            .map(|(data, operation, _)| -> Result<InstallOperation> {
                 let mut writer = output.reopen_boxed()?;
                 writer.seek(SeekFrom::Start(operation.data_offset.unwrap()))?;
                 writer.write_all(&data)?;
@@ -1001,7 +1060,30 @@ pub fn compress_image(
         hash: Some(digest_uncompressed.as_ref().to_vec()),
     };
 
-    Ok((partition_info, operations))
+    let cow_estimate = if vabc_algo.is_some() {
+        // lz4_flex and miniz_oxide usually compress better than the lz4 and
+        // zlib implementations used by libsnapshot_cow. Make up for this by
+        // adding percentage-based overhead.
+        cow_estimate += cow_estimate / 100;
+
+        // We also need to account for constant overhead, especially with
+        // smaller partitions. We can match what delta_generator normally adds
+        // in CowWriterV2::InitPos() exactly. Since we only ever create full
+        // OTAs, we can assume that all CoW operations are kCowReplaceOp.
+
+        // sizeof(CowHeader).
+        cow_estimate += 38;
+        // header_.buffer_size (equal to BUFFER_REGION_DEFAULT_SIZE).
+        cow_estimate += 2 * 1024 * 1024;
+        // CowOptions::cluster_ops * sizeof(CowOperationV2).
+        cow_estimate += 200 * 20;
+
+        Some(cow_estimate)
+    } else {
+        None
+    };
+
+    Ok((partition_info, operations, cow_estimate))
 }
 
 fn extents_sorted(operations: &[InstallOperation]) -> bool {
@@ -1059,7 +1141,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,14 +23,14 @@ use liblzma::{
 use rayon::iter::{IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator};
 use regex::bytes::Regex;
 use ring::digest::Context;
-use rsa::RsaPrivateKey;
+use rsa::RsaPublicKey;
 use thiserror::Error;
 use tracing::{debug, debug_span, trace, warn, Span};
 use x509_cert::Certificate;
 use zip::{result::ZipError, ZipArchive};
 
 use crate::{
-    crypto,
+    crypto::{self, RsaSigningKey},
     format::{
         avb::{self, AppendedDescriptorMut, Footer, Header},
         bootimage::{self, BootImage, BootImageExt, RamdiskMeta},
@@ -70,7 +70,11 @@ pub enum Error {
     #[error("XZ stream error")]
     XzStream(#[from] liblzma::stream::Error),
     #[error("Zip error")]
-    Zip(#[from] ZipError),
+    Zip(#[source] ZipError),
+    #[error("Zip error for entry name: {0:?}")]
+    ZipEntryName(String, #[source] ZipError),
+    #[error("Zip error for entry index #{0}")]
+    ZipEntryIndex(usize, #[source] ZipError),
     #[error("I/O error")]
     Io(#[from] io::Error),
     #[error("File I/O error")]
@@ -157,6 +161,14 @@ impl MagiskRootPatcher {
     const VER_XZ_BACKUP: Range<u32> =
         26403..Self::VERS_SUPPORTED[Self::VERS_SUPPORTED.len() - 1].end;
 
+    const ZIP_INIT_LD: &'static str = "lib/arm64-v8a/libinit-ld.so";
+    const ZIP_LIBMAGISK: &'static str = "lib/arm64-v8a/libmagisk.so";
+    const ZIP_LIBMAGISK32: &'static str = "lib/armeabi-v7a/libmagisk32.so";
+    const ZIP_LIBMAGISK64: &'static str = "lib/arm64-v8a/libmagisk64.so";
+    const ZIP_MAGISKINIT: &'static str = "lib/arm64-v8a/libmagiskinit.so";
+    const ZIP_STUB: &'static str = "assets/stub.apk";
+    const ZIP_UTIL_FUNCTIONS: &'static str = "assets/util_functions.sh";
+
     pub fn new(
         path: &Path,
         preinit_device: Option<&str>,
@@ -208,8 +220,10 @@ impl MagiskRootPatcher {
     fn get_version(path: &Path) -> Result<u32> {
         let reader = File::open(path).map_err(|e| Error::File(path.to_owned(), e))?;
         let reader = BufReader::new(reader);
-        let mut zip = ZipArchive::new(reader)?;
-        let entry = zip.by_name("assets/util_functions.sh")?;
+        let mut zip = ZipArchive::new(reader).map_err(Error::Zip)?;
+        let entry = zip
+            .by_name(Self::ZIP_UTIL_FUNCTIONS)
+            .map_err(|e| Error::ZipEntryName(Self::ZIP_UTIL_FUNCTIONS.to_owned(), e))?;
         let mut entry = BufReader::new(entry);
         let mut line = String::new();
 
@@ -384,7 +398,7 @@ impl BootImagePatch for MagiskRootPatcher {
     fn patch(&self, boot_image: &mut BootImage, cancel_signal: &AtomicBool) -> Result<()> {
         let zip_reader =
             File::open(&self.apk_path).map_err(|e| Error::File(self.apk_path.clone(), e))?;
-        let mut zip = ZipArchive::new(BufReader::new(zip_reader))?;
+        let mut zip = ZipArchive::new(BufReader::new(zip_reader)).map_err(Error::Zip)?;
 
         // Load the first ramdisk. If it doesn't exist, we have to generate one
         // from scratch.
@@ -413,7 +427,9 @@ impl BootImagePatch for MagiskRootPatcher {
 
         // Add magiskinit.
         {
-            let mut zip_entry = zip.by_name("lib/arm64-v8a/libmagiskinit.so")?;
+            let mut zip_entry = zip
+                .by_name(Self::ZIP_MAGISKINIT)
+                .map_err(|e| Error::ZipEntryName(Self::ZIP_MAGISKINIT.to_owned(), e))?;
             let mut data = vec![];
             zip_entry.read_to_end(&mut data)?;
 
@@ -424,28 +440,40 @@ impl BootImagePatch for MagiskRootPatcher {
             ));
         }
 
-        // Add xz-compressed magisk32 and magisk64. We currently unconditionally
-        // include magisk32 because the boot image itself doesn't contain
-        // sufficient information to determine if a device is 64-bit only.
         let mut xz_files = HashMap::<&str, &[u8]>::new();
-        xz_files.insert(
-            "lib/armeabi-v7a/libmagisk32.so",
-            b"overlay.d/sbin/magisk32.xz",
-        );
-        xz_files.insert(
-            "lib/arm64-v8a/libmagisk64.so",
-            b"overlay.d/sbin/magisk64.xz",
-        );
+        if zip.file_names().any(|n| n == Self::ZIP_LIBMAGISK) {
+            // Newer Magisk versions only include a single binary for the target
+            // ABI in the ramdisk. fb5ee86615ed3df830e8538f8b39b1b133caea34.
+            debug!("Single libmagisk");
+            xz_files.insert(Self::ZIP_LIBMAGISK, b"overlay.d/sbin/magisk.xz");
+        } else {
+            // Older Magisk versions include the 64-bit binary and, optionally,
+            // the 32-bit binary if the device supports it. We unconditionally
+            // include the magisk32 because the boot image itself doesn't have
+            // sufficient information to determine if a device is 64-bit only.
+            debug!("Split libmagisk32/libmagisk64");
+            xz_files.insert(Self::ZIP_LIBMAGISK32, b"overlay.d/sbin/magisk32.xz");
+            xz_files.insert(Self::ZIP_LIBMAGISK64, b"overlay.d/sbin/magisk64.xz");
+        }
 
         // Add stub apk, which only exists after Magisk commit
         // ad0e6511e11ebec65aa9b5b916e1397342850319.
-        if zip.file_names().any(|n| n == "assets/stub.apk") {
+        if zip.file_names().any(|n| n == Self::ZIP_STUB) {
             debug!("Magisk stub found");
-            xz_files.insert("assets/stub.apk", b"overlay.d/sbin/stub.xz");
+            xz_files.insert(Self::ZIP_STUB, b"overlay.d/sbin/stub.xz");
+        }
+
+        // Add init-ld, which only exists after Magisk commit
+        // 33aebb59763b6ec27209563035303700e998633d
+        if zip.file_names().any(|n| n == Self::ZIP_INIT_LD) {
+            debug!("Magisk init-ld found");
+            xz_files.insert(Self::ZIP_INIT_LD, b"overlay.d/sbin/init-ld.xz");
         }
 
         for (source, target) in xz_files {
-            let reader = zip.by_name(source)?;
+            let reader = zip
+                .by_name(source)
+                .map_err(|e| Error::ZipEntryName(source.to_owned(), e))?;
             let buf = Self::xz_compress(reader, cancel_signal)?;
 
             entries.push(CpioEntry::new_file(target, 0o644, CpioEntryData::Data(buf)));
@@ -564,10 +592,12 @@ impl OtaCertPatcher {
                 continue;
             };
 
-            let mut zip = ZipArchive::new(Cursor::new(&data))?;
+            let mut zip = ZipArchive::new(Cursor::new(&data)).map_err(Error::Zip)?;
 
             for index in 0..zip.len() {
-                let zip_entry = zip.by_index(index)?;
+                let zip_entry = zip
+                    .by_index(index)
+                    .map_err(|e| Error::ZipEntryIndex(index, e))?;
                 if !zip_entry.name().ends_with(".x509.pem") {
                     debug!("Skipping invalid entry path: {}", zip_entry.name());
                     continue;
@@ -667,6 +697,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
@@ -969,7 +1130,7 @@ 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,
-    key: &RsaPrivateKey,
+    key: &RsaSigningKey,
     patchers: &[Box<dyn BootImagePatch + Sync>],
     cancel_signal: &AtomicBool,
 ) -> Result<HashSet<&'a str>> {
@@ -1044,7 +1205,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() {
@@ -1053,7 +1214,12 @@ pub fn patch_boot_images<'a>(
                 info.header.sign(key)?;
             }
 
-            avb::write_appended_image(writer, &info.header, &mut info.footer, info.image_size)?;
+            avb::write_appended_image(
+                writer,
+                &info.header,
+                &mut info.footer,
+                Some(info.image_size),
+            )?;
 
             Ok(())
         })

avbroot/src/patch/system.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -11,13 +11,13 @@ use std::{
 
 use memchr::memmem;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
-use rsa::RsaPrivateKey;
 use thiserror::Error;
 use tracing::{debug, debug_span, trace, Span};
 use x509_cert::Certificate;
 use zip::ZipArchive;
 
 use crate::{
+    crypto::RsaSigningKey,
     format::{
         avb::{self, AppendedDescriptorMut, Footer},
         ota,
@@ -111,7 +111,7 @@ pub fn patch_system_image(
     input: &(dyn ReadSeekReopen + Sync),
     output: &(dyn WriteSeekReopen + Sync),
     certificate: &Certificate,
-    key: &RsaPrivateKey,
+    key: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<(Vec<Range<u64>>, Vec<Range<u64>>)> {
     // This must be a multiple of normal filesystem block sizes (eg. 4 KiB).
@@ -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)
@@ -207,7 +207,7 @@ pub fn patch_system_image(
     }
 
     let writer = output.reopen_boxed()?;
-    avb::write_appended_image(writer, &header, &mut footer, image_size)?;
+    avb::write_appended_image(writer, &header, &mut footer, Some(image_size))?;
 
     let AppendedDescriptorMut::HashTree(descriptor) = header.appended_descriptor_mut()? else {
         return Err(Error::NoHashTreeDescriptor);

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

avbroot/tests/avb.rs

@@ -1,10 +1,10 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
 use std::{
-    io::{Cursor, Read, Seek, SeekFrom, Write},
+    io::{Cursor, Read, Seek, Write},
     sync::atomic::AtomicBool,
 };
 
@@ -14,11 +14,16 @@ use rsa::RsaPrivateKey;
 
 use avbroot::{
     self,
-    format::avb::{self, AppendedDescriptorMut, AppendedDescriptorRef},
-    stream::{self, SharedCursor},
+    crypto::RsaSigningKey,
+    format::avb::{
+        self, AlgorithmType, AppendedDescriptorMut, AppendedDescriptorRef,
+        ChainPartitionDescriptor, Descriptor, Footer, HashDescriptor, HashTreeDescriptor, Header,
+        KernelCmdlineDescriptor, PropertyDescriptor,
+    },
+    stream::SharedCursor,
 };
 
-fn get_test_key() -> RsaPrivateKey {
+fn get_test_key() -> RsaSigningKey {
     let data = include_str!(concat!(
         env!("CARGO_WORKSPACE_DIR"),
         "/e2e/keys/TEST_KEY_DO_NOT_USE_avb.key",
@@ -28,157 +33,439 @@ fn get_test_key() -> RsaPrivateKey {
         "/e2e/keys/TEST_KEY_DO_NOT_USE_avb.passphrase",
     ));
 
-    RsaPrivateKey::from_pkcs8_encrypted_pem(data, passphrase.trim_end()).unwrap()
+    let key = RsaPrivateKey::from_pkcs8_encrypted_pem(data, passphrase.trim_end()).unwrap();
+    RsaSigningKey::Internal(key)
+}
+
+fn repeat_str(s: &str, max_len: usize) -> String {
+    assert!(!s.is_empty());
+
+    let mut result = s.repeat(max_len / s.len());
+    result.push_str(&s[..max_len % s.len()]);
+
+    result
+}
+
+fn repeat_array<const N: usize>(data: &[u8]) -> [u8; N] {
+    assert!(!data.is_empty());
+
+    let mut result = [0u8; N];
+
+    for i in 0..N / data.len() {
+        result[i * data.len()..][..data.len()].copy_from_slice(data);
+    }
+
+    let remain = N % data.len();
+    result[N - remain..].copy_from_slice(&data[..remain]);
+
+    result
 }
 
 #[test]
 fn round_trip_root_image() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/vbmeta_root.img",
-    ));
-    let reader = Cursor::new(data);
-
-    let (mut header, footer, _) = avb::load_image(reader).unwrap();
-    assert_matches!(footer, None);
+    let mut header = Header {
+        required_libavb_version_major: 1,
+        required_libavb_version_minor: 0,
+        algorithm_type: AlgorithmType::Sha256Rsa4096,
+        hash: vec![],       // autogenerated
+        signature: vec![],  // autogenerated
+        public_key: vec![], // autogenerated
+        public_key_metadata: vec![
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
+            0x0e, 0x0f,
+        ],
+        descriptors: vec![
+            Descriptor::Property(PropertyDescriptor {
+                key: "foobar".to_owned(),
+                value: b"Invalid UTF-8: \xFF".to_vec(),
+            }),
+            Descriptor::HashTree(HashTreeDescriptor {
+                dm_verity_version: 1,
+                image_size: 4096,
+                tree_offset: 0,
+                tree_size: 2048,
+                data_block_size: 4096,
+                hash_block_size: 4096,
+                fec_num_roots: 1,
+                fec_offset: 2048,
+                fec_size: 2048,
+                hash_algorithm: "sha512".to_owned(),
+                partition_name: "hashtreed_partition".to_owned(),
+                salt: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(8),
+                root_digest: [0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10].repeat(8),
+                flags: 0,
+                reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+            }),
+            Descriptor::Hash(HashDescriptor {
+                image_size: 6,
+                hash_algorithm: "sha256".to_owned(),
+                partition_name: "hashed_partition".to_owned(),
+                salt: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(4),
+                root_digest: [0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10].repeat(4),
+                flags: 0,
+                reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+            }),
+            Descriptor::KernelCmdline(KernelCmdlineDescriptor {
+                flags: 1,
+                cmdline: "foobar".to_owned(),
+            }),
+            Descriptor::ChainPartition(ChainPartitionDescriptor {
+                rollback_index_location: 1,
+                partition_name: "chained_partition".to_owned(),
+                public_key: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(129),
+                flags: 0xfedcba98,
+                reserved: repeat_array(&[0x76, 0x54, 0x32, 0x10, 0xfe, 0xdc, 0xba, 0x98]),
+            }),
+        ],
+        rollback_index: 1677974400,
+        flags: 0,
+        rollback_index_location: 0,
+        release_string: repeat_str("MaxLength", 48),
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
 
+    // Sign the header.
     let key = get_test_key();
-
+    header.sign(&key).unwrap();
     assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
 
-    // Clear out the signature-related fields and re-sign.
-    header.hash.clear();
-    header.signature.clear();
-    header.public_key.clear();
-    header.sign(&key).unwrap();
-
+    // Write vbmeta structures.
     let mut writer = Cursor::new(Vec::new());
     avb::write_root_image(&mut writer, &header, 64).unwrap();
-    let new_data = writer.into_inner();
+    let data = writer.into_inner();
 
-    assert_eq!(data, new_data.as_slice());
+    // Verify checksum of the output.
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        [
+            0xc4, 0xa5, 0xda, 0x3e, 0x09, 0xa2, 0xc8, 0x70, 0xcb, 0xf0, 0x96, 0x79, 0x0e, 0x1e,
+            0x80, 0xae, 0x5e, 0x37, 0x81, 0x27, 0x24, 0xc3, 0x6c, 0xa9, 0x42, 0x9e, 0x2c, 0xb1,
+            0x81, 0xad, 0xce, 0xee, 0x8d, 0x4f, 0x76, 0x45, 0x54, 0xc1, 0x31, 0x6a, 0xa7, 0x81,
+            0x5c, 0x59, 0xa8, 0xe8, 0x76, 0xab, 0xed, 0x5b, 0x07, 0x07, 0x38, 0xdd, 0x09, 0x86,
+            0x05, 0x39, 0x23, 0x2d, 0x7b, 0xcc, 0x57, 0x06,
+        ],
+    );
+
+    // Parse the generated image.
+    let mut reader = Cursor::new(&data);
+    let (new_header, new_footer, new_image_size) = avb::load_image(&mut reader).unwrap();
+    assert_matches!(new_footer, None);
+
+    assert_eq!(new_header, header);
+    assert_eq!(new_image_size, data.len() as u64);
 }
 
 #[test]
 fn round_trip_appended_hash_image() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/vbmeta_appended_hash.img",
-    ));
-    let mut reader = Cursor::new(data);
-    let cancel_signal = AtomicBool::new(false);
-
-    let (mut header, footer, image_size) = avb::load_image(&mut reader).unwrap();
-    let mut footer = footer.unwrap();
-
-    let key = get_test_key();
-
-    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
-
-    // Verify the digest.
-    match header.appended_descriptor().unwrap() {
-        AppendedDescriptorRef::HashTree(_) => panic!("Expected hash descriptor"),
-        AppendedDescriptorRef::Hash(d) => {
-            reader.rewind().unwrap();
-            d.verify(&mut reader, &cancel_signal).unwrap();
-        }
-    }
+    let image_size = 12288;
+    let raw_data = b"foobar";
+    let mut header = Header {
+        required_libavb_version_major: 1,
+        required_libavb_version_minor: 0,
+        algorithm_type: AlgorithmType::Sha256Rsa4096,
+        hash: vec![],       // autogenerated
+        signature: vec![],  // autogenerated
+        public_key: vec![], // autogenerated
+        public_key_metadata: vec![
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
+            0x0e, 0x0f,
+        ],
+        descriptors: vec![
+            Descriptor::Property(PropertyDescriptor {
+                key: "foobar".to_owned(),
+                value: b"Invalid UTF-8: \xFF".to_vec(),
+            }),
+            Descriptor::Hash(HashDescriptor {
+                image_size: raw_data.len() as u64,
+                hash_algorithm: "sha256".to_owned(),
+                partition_name: "vbmeta_appended_hash".to_owned(),
+                salt: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(4),
+                root_digest: vec![], // autogenerated
+                flags: 0,
+                reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+            }),
+        ],
+        rollback_index: 1677974400,
+        flags: 0,
+        rollback_index_location: 0,
+        release_string: repeat_str("MaxLength", 48),
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
+    let mut footer = Footer {
+        version_major: 1,
+        version_minor: 0,
+        original_image_size: 0, // autogenerated
+        vbmeta_offset: 0,       // autogenerated
+        vbmeta_size: 0,         // autogenerated
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
 
     let mut writer = Cursor::new(Vec::new());
+    let cancel_signal = AtomicBool::new(false);
 
-    // Copy the partition data.
-    reader.seek(SeekFrom::Start(0)).unwrap();
-    stream::copy_n(
-        &mut reader,
-        &mut writer,
-        footer.original_image_size,
-        &cancel_signal,
-    )
-    .unwrap();
+    // Write the raw partition data.
+    writer.write_all(raw_data).unwrap();
 
-    // Regenerate the digest.
+    // Regenerate the raw image digest.
     match header.appended_descriptor_mut().unwrap() {
         AppendedDescriptorMut::HashTree(_) => panic!("Expected hash descriptor"),
         AppendedDescriptorMut::Hash(d) => {
-            d.root_digest.clear();
             writer.rewind().unwrap();
             d.update(&mut writer, &cancel_signal).unwrap();
         }
     }
 
-    // Clear out the signature-related fields and re-sign.
-    header.hash.clear();
-    header.signature.clear();
-    header.public_key.clear();
+    // Verify the raw image digest.
+    match header.appended_descriptor().unwrap() {
+        AppendedDescriptorRef::HashTree(_) => panic!("Expected hash descriptor"),
+        AppendedDescriptorRef::Hash(d) => {
+            writer.rewind().unwrap();
+            d.verify(&mut writer, &cancel_signal).unwrap();
+        }
+    }
+
+    // Sign the header.
+    let key = get_test_key();
     header.sign(&key).unwrap();
+    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
 
-    // Write new vbmeta structures.
-    avb::write_appended_image(&mut writer, &header, &mut footer, image_size).unwrap();
-    let new_data = writer.into_inner();
+    // Write vbmeta structures.
+    avb::write_appended_image(&mut writer, &header, &mut footer, Some(image_size)).unwrap();
+    let data = writer.into_inner();
 
-    assert_eq!(data, new_data.as_slice());
+    // Verify checksum of the output.
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        [
+            0x09, 0x98, 0x0c, 0x9d, 0x11, 0x50, 0xde, 0xb1, 0x55, 0x3b, 0x00, 0x76, 0xbe, 0x25,
+            0xfd, 0xe6, 0x46, 0x22, 0xbd, 0x9a, 0x05, 0x86, 0xea, 0x07, 0x4d, 0x8f, 0x7b, 0x15,
+            0x36, 0x20, 0x0d, 0xf0, 0x7e, 0x96, 0xd2, 0x58, 0xde, 0xf2, 0xa6, 0x91, 0x6d, 0x01,
+            0x7b, 0x03, 0x96, 0x70, 0xf8, 0x3b, 0x76, 0x74, 0xf0, 0xbf, 0x47, 0xe0, 0xd2, 0xd4,
+            0x5d, 0xbf, 0xb7, 0x9c, 0xf5, 0xf8, 0xaf, 0x3c,
+        ],
+    );
+
+    // Parse the generated image.
+    let mut reader = Cursor::new(&data);
+    let (new_header, new_footer, new_image_size) = avb::load_image(&mut reader).unwrap();
+    let new_footer = new_footer.unwrap();
+
+    assert_eq!(new_header, header);
+    assert_eq!(new_footer, footer);
+    assert_eq!(new_image_size, image_size);
 }
 
 #[test]
-fn round_trip_appended_hash_tree_image() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/vbmeta_appended_hash_tree.img",
-    ));
-    let mut reader = SharedCursor::default();
-    reader.write_all(data).unwrap();
-    let cancel_signal = AtomicBool::new(false);
-
-    let (mut header, footer, image_size) = avb::load_image(&mut reader).unwrap();
-    let mut footer = footer.unwrap();
-
-    let key = get_test_key();
-
-    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
-
-    // Verify the hash tree and FEC data.
-    match header.appended_descriptor().unwrap() {
-        AppendedDescriptorRef::HashTree(d) => d.verify(&reader, &cancel_signal).unwrap(),
-        AppendedDescriptorRef::Hash(_) => panic!("Expected hash tree descriptor"),
-    }
+fn round_trip_appended_hash_tree_image_fixed_size() {
+    let image_size = 32768;
+    let raw_data: [u8; 8192] = repeat_array(b"foobar");
+    let mut header = Header {
+        required_libavb_version_major: 1,
+        required_libavb_version_minor: 0,
+        algorithm_type: AlgorithmType::Sha256Rsa4096,
+        hash: vec![],       // autogenerated
+        signature: vec![],  // autogenerated
+        public_key: vec![], // autogenerated
+        public_key_metadata: vec![
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
+            0x0e, 0x0f,
+        ],
+        descriptors: vec![
+            Descriptor::Property(PropertyDescriptor {
+                key: "foobar".to_owned(),
+                value: b"Invalid UTF-8: \xFF".to_vec(),
+            }),
+            Descriptor::HashTree(HashTreeDescriptor {
+                dm_verity_version: 1,
+                image_size: raw_data.len() as u64,
+                tree_offset: 0, // autogenerated
+                tree_size: 0,   // autogenerated
+                data_block_size: 4096,
+                hash_block_size: 4096,
+                fec_num_roots: 2,
+                fec_offset: 0, // autogenerated
+                fec_size: 0,   // autogenerated
+                hash_algorithm: "sha256".to_owned(),
+                partition_name: "vbmeta_appended_hash_tree".to_owned(),
+                salt: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(4),
+                root_digest: vec![], // autogenerated
+                flags: 0,
+                reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+            }),
+        ],
+        rollback_index: 1677974400,
+        flags: 0,
+        rollback_index_location: 0,
+        release_string: repeat_str("MaxLength", 48),
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
+    let mut footer = Footer {
+        version_major: 1,
+        version_minor: 0,
+        original_image_size: 0, // autogenerated
+        vbmeta_offset: 0,       // autogenerated
+        vbmeta_size: 0,         // autogenerated
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
 
     let mut writer = SharedCursor::default();
+    let cancel_signal = AtomicBool::new(false);
 
-    // Copy the partition data, excluding the hash tree and FEC data.
-    reader.seek(SeekFrom::Start(0)).unwrap();
-    stream::copy_n(
-        &mut reader,
-        &mut writer,
-        footer.original_image_size,
-        &cancel_signal,
-    )
-    .unwrap();
+    // Write the raw partition data.
+    writer.write_all(&raw_data).unwrap();
 
-    // Regenerate the hash tree and FEC data.
+    // Generate and write the hash tree and FEC data.
     match header.appended_descriptor_mut().unwrap() {
         AppendedDescriptorMut::HashTree(d) => {
-            d.root_digest.clear();
-            d.tree_offset = 0;
-            d.tree_size = 0;
-            d.fec_offset = 0;
-            d.fec_size = 0;
-
             d.update(&writer, &writer, None, &cancel_signal).unwrap();
         }
         AppendedDescriptorMut::Hash(_) => panic!("Expected hash tree descriptor"),
     }
 
-    // Clear out the signature-related fields and re-sign.
-    header.hash.clear();
-    header.signature.clear();
-    header.public_key.clear();
+    // Verify the hash tree and FEC data.
+    match header.appended_descriptor_mut().unwrap() {
+        AppendedDescriptorMut::HashTree(d) => {
+            d.verify(&writer, &cancel_signal).unwrap();
+        }
+        AppendedDescriptorMut::Hash(_) => panic!("Expected hash tree descriptor"),
+    }
+
+    // Sign the header.
+    let key = get_test_key();
     header.sign(&key).unwrap();
+    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
 
-    // Write new vbmeta structures.
-    avb::write_appended_image(&mut writer, &header, &mut footer, image_size).unwrap();
-    let mut new_data = Vec::new();
+    // Write vbmeta structures.
+    avb::write_appended_image(&mut writer, &header, &mut footer, Some(image_size)).unwrap();
+    let mut data = Vec::new();
     writer.rewind().unwrap();
-    writer.read_to_end(&mut new_data).unwrap();
+    writer.read_to_end(&mut data).unwrap();
 
-    assert_eq!(data, new_data.as_slice());
+    // Verify checksum of the output.
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        [
+            0xb5, 0x56, 0x65, 0x81, 0x5a, 0x16, 0x65, 0xa9, 0xa6, 0xc6, 0x9e, 0x41, 0x89, 0x9f,
+            0xe9, 0xbc, 0xea, 0x59, 0x4d, 0x14, 0x8a, 0x9e, 0x2b, 0x13, 0xa0, 0x3a, 0x8e, 0xd4,
+            0x59, 0xcd, 0x74, 0xe7, 0x99, 0xbd, 0xa3, 0x58, 0x4b, 0x84, 0xf2, 0x04, 0xe2, 0x12,
+            0x48, 0xfe, 0x4f, 0x67, 0x1f, 0x2a, 0xaa, 0x22, 0x51, 0x19, 0x83, 0x95, 0xa8, 0x03,
+            0xf5, 0x87, 0x12, 0x05, 0x8e, 0x14, 0xd9, 0xbd
+        ],
+    );
+
+    // Parse the generated image.
+    let mut reader = Cursor::new(&data);
+    let (new_header, new_footer, new_image_size) = avb::load_image(&mut reader).unwrap();
+    let new_footer = new_footer.unwrap();
+
+    assert_eq!(new_header, header);
+    assert_eq!(new_footer, footer);
+    assert_eq!(new_image_size, image_size);
+}
+
+#[test]
+fn round_trip_appended_hash_tree_image_minimum_size() {
+    let raw_data: [u8; 8192] = repeat_array(b"foobar");
+    let mut header = Header {
+        required_libavb_version_major: 1,
+        required_libavb_version_minor: 0,
+        algorithm_type: AlgorithmType::Sha256Rsa4096,
+        hash: vec![],       // autogenerated
+        signature: vec![],  // autogenerated
+        public_key: vec![], // autogenerated
+        public_key_metadata: vec![
+            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
+            0x0e, 0x0f,
+        ],
+        descriptors: vec![
+            Descriptor::Property(PropertyDescriptor {
+                key: "foobar".to_owned(),
+                value: b"Invalid UTF-8: \xFF".to_vec(),
+            }),
+            Descriptor::HashTree(HashTreeDescriptor {
+                dm_verity_version: 1,
+                image_size: raw_data.len() as u64,
+                tree_offset: 0, // autogenerated
+                tree_size: 0,   // autogenerated
+                data_block_size: 4096,
+                hash_block_size: 4096,
+                fec_num_roots: 2,
+                fec_offset: 0, // autogenerated
+                fec_size: 0,   // autogenerated
+                hash_algorithm: "sha256".to_owned(),
+                partition_name: "vbmeta_appended_hash_tree".to_owned(),
+                salt: [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef].repeat(4),
+                root_digest: vec![], // autogenerated
+                flags: 0,
+                reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+            }),
+        ],
+        rollback_index: 1677974400,
+        flags: 0,
+        rollback_index_location: 0,
+        release_string: repeat_str("MaxLength", 48),
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
+    let mut footer = Footer {
+        version_major: 1,
+        version_minor: 0,
+        original_image_size: 0, // autogenerated
+        vbmeta_offset: 0,       // autogenerated
+        vbmeta_size: 0,         // autogenerated
+        reserved: repeat_array(&[0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]),
+    };
+
+    let mut writer = SharedCursor::default();
+    let cancel_signal = AtomicBool::new(false);
+
+    // Write the raw partition data.
+    writer.write_all(&raw_data).unwrap();
+
+    // Generate and write the hash tree and FEC data.
+    match header.appended_descriptor_mut().unwrap() {
+        AppendedDescriptorMut::HashTree(d) => {
+            d.update(&writer, &writer, None, &cancel_signal).unwrap();
+        }
+        AppendedDescriptorMut::Hash(_) => panic!("Expected hash tree descriptor"),
+    }
+
+    // Verify the hash tree and FEC data.
+    match header.appended_descriptor_mut().unwrap() {
+        AppendedDescriptorMut::HashTree(d) => {
+            d.verify(&writer, &cancel_signal).unwrap();
+        }
+        AppendedDescriptorMut::Hash(_) => panic!("Expected hash tree descriptor"),
+    }
+
+    // Sign the header.
+    let key = get_test_key();
+    header.sign(&key).unwrap();
+    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
+
+    // Write vbmeta structures.
+    avb::write_appended_image(&mut writer, &header, &mut footer, None).unwrap();
+    let mut data = Vec::new();
+    writer.rewind().unwrap();
+    writer.read_to_end(&mut data).unwrap();
+
+    // Verify checksum of the output.
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        [
+            0xd6, 0x69, 0x19, 0x6a, 0x36, 0xc8, 0x1c, 0xe9, 0xc4, 0x85, 0xbe, 0xff, 0x43, 0xb1,
+            0x9f, 0xd4, 0x1d, 0x6c, 0xf9, 0xd2, 0xf3, 0xa6, 0x5f, 0x66, 0x41, 0xd5, 0xf3, 0xfd,
+            0x28, 0xdb, 0x14, 0x67, 0xc6, 0xa8, 0xef, 0xc4, 0xd4, 0x67, 0x6c, 0xb8, 0x66, 0xbb,
+            0x56, 0x5a, 0x4a, 0xf5, 0xd8, 0x92, 0x7c, 0x42, 0xbc, 0x47, 0xdb, 0x94, 0x38, 0x15,
+            0x4b, 0x2d, 0xd0, 0x28, 0x1f, 0xd1, 0x45, 0xa9,
+        ],
+    );
+
+    // Parse the generated image.
+    let mut reader = Cursor::new(&data);
+    let (new_header, new_footer, new_image_size) = avb::load_image(&mut reader).unwrap();
+    let new_footer = new_footer.unwrap();
+
+    assert_eq!(new_header, header);
+    assert_eq!(new_footer, footer);
+    assert_eq!(new_image_size, 28672);
 }

avbroot/tests/bootimage.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -7,13 +7,20 @@ use std::io::Cursor;
 
 use avbroot::{
     self,
-    format::bootimage::{BootImage, BootImageExt},
+    crypto::RsaSigningKey,
+    format::{
+        avb::{AlgorithmType, Descriptor, HashDescriptor, Header},
+        bootimage::{
+            self, BootImage, BootImageExt, BootImageV0Through2, BootImageV3Through4, RamdiskMeta,
+            V1Extra, V2Extra, V4Extra, VendorBootImageV3Through4, VendorV4Extra,
+        },
+    },
     stream::{FromReader, ToWriter},
 };
 use pkcs8::DecodePrivateKey;
 use rsa::RsaPrivateKey;
 
-fn get_test_key() -> RsaPrivateKey {
+fn get_test_key() -> RsaSigningKey {
     let data = include_str!(concat!(
         env!("CARGO_WORKSPACE_DIR"),
         "/e2e/keys/TEST_KEY_DO_NOT_USE_avb.key",
@@ -23,103 +30,322 @@ fn get_test_key() -> RsaPrivateKey {
         "/e2e/keys/TEST_KEY_DO_NOT_USE_avb.passphrase",
     ));
 
-    RsaPrivateKey::from_pkcs8_encrypted_pem(data, passphrase.trim_end()).unwrap()
+    let key = RsaPrivateKey::from_pkcs8_encrypted_pem(data, passphrase.trim_end()).unwrap();
+    RsaSigningKey::Internal(key)
 }
 
-fn round_trip(data: &[u8], expected_version: u32) {
-    let reader = Cursor::new(data);
-    let mut image = BootImage::from_reader(reader).unwrap();
+fn repeat(s: &str, max_len: usize) -> String {
+    assert!(!s.is_empty());
 
+    let mut result = s.repeat(max_len / s.len());
+    result.push_str(&s[..max_len % s.len()]);
+
+    result
+}
+
+fn round_trip(image: &BootImage, sha512: &[u8; 64], expected_version: u32) {
     assert_eq!(image.header_version(), expected_version);
 
-    match &mut image {
-        BootImage::V3Through4(b) => {
-            let should_sign = b
-                .v4_extra
-                .as_ref()
-                .map_or(false, |v4| v4.signature.is_some());
-            let key = get_test_key();
-            let signed = b.sign(&key).unwrap();
-            assert_eq!(signed, should_sign);
-        }
-        _ => {}
-    }
-
     let mut writer = Cursor::new(Vec::new());
     image.to_writer(&mut writer).unwrap();
-    let new_data = writer.into_inner();
+    let data = writer.into_inner();
 
-    assert_eq!(data, new_data);
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        sha512,
+    );
+
+    let reader = Cursor::new(data);
+    let new_image = BootImage::from_reader(reader).unwrap();
+
+    assert_eq!(&new_image, image);
 }
 
 #[test]
 fn round_trip_v0() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v0.img",
-    ));
-    round_trip(data, 0);
+    let image = BootImage::V0Through2(BootImageV0Through2 {
+        kernel_addr: 0x01234567,
+        ramdisk_addr: 0x89abcdef,
+        second_addr: 0x02468ace,
+        tags_addr: 0x13579bdf,
+        page_size: 4096,
+        os_version: 0x76543210,
+        name: repeat("Name", 16),
+        cmdline: repeat("Cmdline", 512),
+        id: [
+            0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff, 0xffeeddcc, 0xbbaa9988, 0x77665544,
+            0x33221100,
+        ],
+        extra_cmdline: repeat("ExtraCmdline", 1024),
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+        second: b"second data".to_vec(),
+        v1_extra: None,
+        v2_extra: None,
+    });
+    let sha512 = [
+        0x23, 0x65, 0x0b, 0xfa, 0x7a, 0x09, 0x0a, 0xdf, 0xdd, 0x9a, 0x6c, 0x03, 0xfa, 0xc5, 0xe1,
+        0xfa, 0x27, 0x65, 0xa0, 0x94, 0xef, 0xa2, 0x0c, 0xc5, 0x3e, 0xd9, 0x67, 0x7d, 0x88, 0x7b,
+        0xb3, 0x48, 0x39, 0xab, 0x28, 0x77, 0x7b, 0x18, 0xec, 0x60, 0xe0, 0xb7, 0x0d, 0x15, 0x26,
+        0xb2, 0xd4, 0x27, 0x25, 0x92, 0x5c, 0x7b, 0x0b, 0x5c, 0xf7, 0xed, 0x27, 0x6c, 0x39, 0xb5,
+        0xb7, 0x44, 0xbb, 0xec,
+    ];
+
+    round_trip(&image, &sha512, 0);
 }
 
 #[test]
 fn round_trip_v1() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v1.img",
-    ));
-    round_trip(data, 1);
+    let image = BootImage::V0Through2(BootImageV0Through2 {
+        kernel_addr: 0x01234567,
+        ramdisk_addr: 0x89abcdef,
+        second_addr: 0x02468ace,
+        tags_addr: 0x13579bdf,
+        page_size: 4096,
+        os_version: 0x76543210,
+        name: repeat("Name", 16),
+        cmdline: repeat("Cmdline", 512),
+        id: [
+            0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff, 0xffeeddcc, 0xbbaa9988, 0x77665544,
+            0x33221100,
+        ],
+        extra_cmdline: repeat("ExtraCmdline", 1024),
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+        second: b"second data".to_vec(),
+        v1_extra: Some(V1Extra {
+            recovery_dtbo_offset: 0x0123456789abcdef,
+            recovery_dtbo: b"recovery_dtbo data".to_vec(),
+        }),
+        v2_extra: None,
+    });
+    let sha512 = [
+        0x37, 0x8e, 0xf1, 0xf0, 0xb8, 0x44, 0x0f, 0x9e, 0x16, 0xc0, 0x15, 0x98, 0xa2, 0xb5, 0x06,
+        0x63, 0x59, 0xf4, 0x91, 0xb6, 0x28, 0x03, 0xe6, 0xdc, 0xd2, 0x0d, 0xd7, 0x49, 0x33, 0x63,
+        0x91, 0xd4, 0xa8, 0x24, 0xff, 0xb0, 0x5f, 0x99, 0x2a, 0x9a, 0xb3, 0x66, 0x81, 0x41, 0x69,
+        0xb0, 0xbc, 0xe2, 0x5b, 0x33, 0x3f, 0x39, 0x6a, 0xa8, 0xbd, 0xe1, 0x15, 0x3e, 0x51, 0x5a,
+        0x2a, 0x9d, 0x23, 0x90,
+    ];
+
+    round_trip(&image, &sha512, 1);
 }
 
 #[test]
 fn round_trip_v2() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v2.img",
-    ));
-    round_trip(data, 2);
+    let image = BootImage::V0Through2(BootImageV0Through2 {
+        kernel_addr: 0x01234567,
+        ramdisk_addr: 0x89abcdef,
+        second_addr: 0x02468ace,
+        tags_addr: 0x13579bdf,
+        page_size: 4096,
+        os_version: 0x76543210,
+        name: repeat("Name", 16),
+        cmdline: repeat("Cmdline", 512),
+        id: [
+            0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff, 0xffeeddcc, 0xbbaa9988, 0x77665544,
+            0x33221100,
+        ],
+        extra_cmdline: repeat("ExtraCmdline", 1024),
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+        second: b"second data".to_vec(),
+        v1_extra: Some(V1Extra {
+            recovery_dtbo_offset: 0x0123456789abcdef,
+            recovery_dtbo: b"recovery_dtbo data".to_vec(),
+        }),
+        v2_extra: Some(V2Extra {
+            dtb_addr: 0xfedcba9876543210,
+            dtb: b"dtb data".to_vec(),
+        }),
+    });
+    let sha512 = [
+        0x04, 0x24, 0x5b, 0xb7, 0x07, 0x82, 0xa9, 0x08, 0x68, 0xb9, 0xc9, 0x65, 0x1f, 0x53, 0xd7,
+        0x6c, 0xcf, 0xf3, 0x48, 0x58, 0x9a, 0xd4, 0xb1, 0xf3, 0xd8, 0x6f, 0x95, 0x10, 0x70, 0x2f,
+        0x53, 0x30, 0x60, 0x60, 0xe4, 0x68, 0xd9, 0x84, 0xe8, 0x0a, 0xf3, 0x12, 0xb3, 0xa3, 0x1b,
+        0x06, 0x88, 0x2f, 0x5d, 0x34, 0x5a, 0xea, 0x8f, 0xbb, 0x54, 0x49, 0x3c, 0xc1, 0x9c, 0xc6,
+        0x24, 0x80, 0x03, 0xde,
+    ];
+
+    round_trip(&image, &sha512, 2);
 }
 
 #[test]
 fn round_trip_v3() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v3.img",
-    ));
-    round_trip(data, 3);
+    let image = BootImage::V3Through4(BootImageV3Through4 {
+        os_version: 0x01234567,
+        reserved: [0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff],
+        cmdline: repeat("Cmdline", 1536),
+        v4_extra: None,
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+    });
+    let sha512 = [
+        0x30, 0xea, 0x77, 0x0a, 0xd3, 0x24, 0x6a, 0x3f, 0xf8, 0xdf, 0xe6, 0xd9, 0x5a, 0xa1, 0xd3,
+        0xa4, 0x3b, 0x8a, 0x13, 0x39, 0x5e, 0x58, 0x24, 0x3e, 0x71, 0x31, 0x78, 0xa1, 0x2c, 0xad,
+        0x1d, 0xca, 0x24, 0x12, 0xf5, 0xfb, 0x2c, 0x48, 0xa5, 0x3d, 0xc0, 0x38, 0x55, 0xb6, 0xfd,
+        0xd3, 0x30, 0xe0, 0x69, 0x11, 0x28, 0xd7, 0x29, 0xda, 0x2e, 0x5a, 0x49, 0x5c, 0x39, 0x1d,
+        0xb9, 0xdb, 0x53, 0xe1,
+    ];
+
+    round_trip(&image, &sha512, 3);
 }
 
 #[test]
 fn round_trip_v4() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v4.img",
-    ));
-    round_trip(data, 4);
+    let image = BootImage::V3Through4(BootImageV3Through4 {
+        os_version: 0x01234567,
+        reserved: [0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff],
+        cmdline: repeat("Cmdline", 1536),
+        v4_extra: Some(V4Extra { signature: None }),
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+    });
+    let sha512 = [
+        0xa8, 0x1d, 0x2b, 0x78, 0x22, 0x45, 0x0b, 0xe7, 0xc2, 0x3a, 0xd8, 0xda, 0x95, 0x49, 0x77,
+        0x18, 0xd0, 0x7b, 0x9b, 0x7f, 0xc7, 0xf6, 0x48, 0xb4, 0x2d, 0x85, 0x6d, 0xe3, 0x5a, 0xa3,
+        0x24, 0xb6, 0x94, 0x56, 0xb9, 0x07, 0x84, 0xdb, 0x50, 0x01, 0xca, 0x6c, 0x86, 0x26, 0x32,
+        0x79, 0x0c, 0xc5, 0x70, 0xcf, 0xcc, 0x7f, 0xc3, 0x5b, 0x96, 0x56, 0x23, 0x5c, 0xd0, 0x50,
+        0xb0, 0x98, 0xdb, 0x4a,
+    ];
+
+    round_trip(&image, &sha512, 4);
 }
 
 #[test]
 fn round_trip_v4_vts() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/boot_v4_vts.img",
-    ));
-    round_trip(data, 4);
+    let mut header = Header {
+        required_libavb_version_major: 1,
+        required_libavb_version_minor: 0,
+        algorithm_type: AlgorithmType::Sha256Rsa4096,
+        hash: vec![],       // autogenerated
+        signature: vec![],  // autogenerated
+        public_key: vec![], // autogenerated
+        public_key_metadata: vec![],
+        descriptors: vec![Descriptor::Hash(HashDescriptor {
+            image_size: 12288,
+            hash_algorithm: "sha256".to_owned(),
+            partition_name: "boot".to_owned(),
+            salt: vec![0x64, 0x30, 0x30, 0x64, 0x66, 0x30, 0x30, 0x64],
+            root_digest: vec![
+                0xab, 0xd5, 0x48, 0x3e, 0x11, 0xe7, 0x94, 0x0c, 0xb9, 0xbf, 0x38, 0x75, 0x87, 0xa4,
+                0xa1, 0x65, 0x99, 0x81, 0xa1, 0xb8, 0x39, 0x62, 0xb7, 0xc1, 0xfa, 0xf1, 0xb0, 0xcd,
+                0x63, 0x07, 0xd4, 0x49,
+            ],
+            flags: 0,
+            reserved: [0; 60],
+        })],
+        rollback_index: 0,
+        flags: 0,
+        rollback_index_location: 0,
+        release_string: "avbtool 1.2.0".to_owned(),
+        reserved: [0; 80],
+    };
+
+    let key = get_test_key();
+    header.sign(&key).unwrap();
+    assert_eq!(header.verify().unwrap().unwrap(), key.to_public_key());
+
+    let image = BootImage::V3Through4(BootImageV3Through4 {
+        os_version: 0x01234567,
+        reserved: [0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff],
+        cmdline: repeat("Cmdline", 1536),
+        v4_extra: Some(V4Extra {
+            signature: Some(header),
+        }),
+        kernel: b"kernel data".to_vec(),
+        ramdisk: b"ramdisk data".to_vec(),
+    });
+    let sha512 = [
+        0x19, 0x47, 0x15, 0x3c, 0x1f, 0x62, 0x84, 0xee, 0xbc, 0x16, 0x9e, 0x5a, 0xf2, 0x45, 0x2a,
+        0xf7, 0x40, 0xc2, 0x18, 0x7f, 0x23, 0xb2, 0xa4, 0x20, 0x10, 0xdf, 0xb1, 0x5c, 0xf2, 0x7f,
+        0x6f, 0x79, 0x22, 0x1d, 0x29, 0x27, 0x78, 0xea, 0xb3, 0x9e, 0x1f, 0xfe, 0xeb, 0xc8, 0x9f,
+        0xe6, 0xef, 0xce, 0xa2, 0x28, 0x0b, 0x05, 0x1d, 0x52, 0xff, 0xab, 0xd4, 0x6f, 0x87, 0x11,
+        0xd9, 0xb6, 0x5f, 0x2e,
+    ];
+
+    round_trip(&image, &sha512, 4);
 }
 
 #[test]
 fn round_trip_vendor_v3() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/vendor_v3.img",
-    ));
-    round_trip(data, 3);
+    let image = BootImage::VendorV3Through4(VendorBootImageV3Through4 {
+        page_size: 4096,
+        kernel_addr: 0x01234567,
+        ramdisk_addr: 0x89abcdef,
+        cmdline: repeat("Cmdline", 2048),
+        tags_addr: 0xfedcba98,
+        name: repeat("Name", 16),
+        dtb: b"dtb data".to_vec(),
+        dtb_addr: 0x76543210,
+        ramdisks: vec![b"ramdisk data".to_vec()],
+        v4_extra: None,
+    });
+    let sha512 = [
+        0x17, 0x18, 0xb9, 0x67, 0x4c, 0x82, 0x71, 0x98, 0x6a, 0x8a, 0xb8, 0x85, 0x3c, 0x77, 0x9e,
+        0x27, 0xeb, 0xce, 0x2a, 0x23, 0x04, 0x63, 0x7c, 0x94, 0xd4, 0xad, 0x1f, 0x3c, 0xee, 0x7e,
+        0x41, 0x8b, 0xa8, 0xd9, 0x35, 0xec, 0xf2, 0xc1, 0x52, 0x3a, 0xd9, 0x5b, 0xbe, 0x63, 0xe8,
+        0x00, 0xd2, 0x23, 0x4e, 0x37, 0x76, 0x31, 0x5a, 0xfc, 0x63, 0x43, 0x32, 0x34, 0x30, 0xf6,
+        0x3e, 0x2e, 0x3e, 0x66,
+    ];
+
+    round_trip(&image, &sha512, 3);
 }
 
 #[test]
 fn round_trip_vendor_v4() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/vendor_v4.img",
-    ));
-    round_trip(data, 4);
+    let board_id = [
+        0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff, 0xffeeddcc, 0xbbaa9988, 0x77665544,
+        0x33221100, 0x004488cc, 0x115599dd, 0x2266aaee, 0x3377bbff, 0xffbb7733, 0xeeaa6622,
+        0xdd995511, 0xcc884400,
+    ];
+    let image = BootImage::VendorV3Through4(VendorBootImageV3Through4 {
+        page_size: 2048,
+        kernel_addr: 0x01234567,
+        ramdisk_addr: 0x89abcdef,
+        cmdline: repeat("Cmdline", 2048),
+        tags_addr: 0xfedcba98,
+        name: repeat("Name", 16),
+        dtb: b"dtb data".to_vec(),
+        dtb_addr: 0x76543210,
+        ramdisks: vec![
+            b"ramdisk 0 data".to_vec(),
+            b"ramdisk 1 data".to_vec(),
+            b"ramdisk 2 data".to_vec(),
+            b"ramdisk 3 data".to_vec(),
+        ],
+        v4_extra: Some(VendorV4Extra {
+            ramdisk_metas: vec![
+                RamdiskMeta {
+                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_NONE,
+                    ramdisk_name: repeat("None", 32),
+                    board_id,
+                },
+                RamdiskMeta {
+                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_PLATFORM,
+                    ramdisk_name: repeat("Platform", 32),
+                    board_id,
+                },
+                RamdiskMeta {
+                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_RECOVERY,
+                    ramdisk_name: repeat("Recovery", 32),
+                    board_id,
+                },
+                RamdiskMeta {
+                    ramdisk_type: bootimage::VENDOR_RAMDISK_TYPE_DLKM,
+                    ramdisk_name: repeat("Dlkm", 32),
+                    board_id,
+                },
+            ],
+            bootconfig: "bootconfig data".to_owned(),
+        }),
+    });
+    let sha512 = [
+        0x0e, 0x3f, 0x86, 0x9e, 0xad, 0x98, 0xbb, 0x53, 0xc7, 0xc4, 0x3f, 0xb8, 0xc6, 0x06, 0xdc,
+        0xb2, 0xe5, 0x47, 0x66, 0xe3, 0xaf, 0x2c, 0xa4, 0x91, 0x8d, 0x4b, 0xc5, 0x70, 0x1e, 0x51,
+        0x19, 0x23, 0x7c, 0xab, 0x40, 0x24, 0x95, 0xef, 0xc8, 0x65, 0xdb, 0x5f, 0x0a, 0x41, 0x93,
+        0xff, 0x6c, 0x22, 0xb4, 0x9a, 0xe2, 0x20, 0xc1, 0x95, 0xa0, 0x3c, 0xc2, 0x13, 0xdb, 0xc8,
+        0x24, 0x33, 0x77, 0x75,
+    ];
+
+    round_trip(&image, &sha512, 4);
 }

avbroot/tests/cpio.rs

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
@@ -7,22 +7,64 @@ use std::io::{self, Cursor};
 
 use avbroot::{
     self,
-    format::cpio::{CpioEntryType, CpioReader, CpioWriter},
+    format::cpio::{CpioEntry, CpioEntryData, CpioEntryType, CpioReader, CpioWriter},
     util,
 };
 
+fn generate_archive() -> Vec<u8> {
+    let writer = Cursor::new(Vec::new());
+    let mut cpio_writer = CpioWriter::new(writer, false);
+
+    for entry in [
+        CpioEntry::new_symlink(b"symlink", b"target"),
+        CpioEntry::new_directory(b"directory", 0o755),
+        CpioEntry::new_file(b"file", 0o644, CpioEntryData::Data(b"foobar".to_vec())),
+        CpioEntry {
+            path: b"reserved".to_vec(),
+            data: CpioEntryData::Size(0),
+            inode: 12345,
+            file_type: CpioEntryType::Reserved,
+            file_mode: 0o4777,
+            uid: 12345678,
+            gid: 87654321,
+            nlink: 2,
+            mtime: 1700000000,
+            dev_maj: 2222,
+            dev_min: 3333,
+            rdev_maj: 4444,
+            rdev_min: 5555,
+            crc32: 0xfedcba09,
+        },
+    ] {
+        cpio_writer.start_entry(&entry).unwrap();
+    }
+
+    let writer = cpio_writer.finish().unwrap();
+    let data = writer.into_inner();
+
+    assert_eq!(
+        ring::digest::digest(&ring::digest::SHA512, &data).as_ref(),
+        [
+            0xb0, 0x51, 0xac, 0x28, 0x6f, 0x78, 0xe2, 0xe7, 0x45, 0xa0, 0x52, 0x7c, 0xff, 0x42,
+            0x30, 0x55, 0xbd, 0x64, 0x7d, 0x4e, 0xb8, 0xe6, 0x95, 0xe5, 0x9b, 0xd1, 0x13, 0xd6,
+            0x43, 0x0e, 0x32, 0xb2, 0x4e, 0x62, 0xa4, 0x55, 0x64, 0x48, 0xb7, 0x32, 0x26, 0x57,
+            0x75, 0x07, 0xf5, 0xa6, 0x0f, 0x18, 0xc3, 0x9e, 0x9f, 0x06, 0xdb, 0xa4, 0xf7, 0xeb,
+            0x5e, 0x8f, 0xce, 0xd0, 0x2b, 0x54, 0x39, 0x57
+        ],
+    );
+
+    data
+}
+
 #[test]
 fn round_trip_archive() {
-    let data = include_bytes!(concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/data/archive.cpio",
-    ));
+    let data = generate_archive();
     assert_ne!(data.len() % 512, 0);
 
     for pad_to_block_size in [false, true] {
         println!("Pad to block size: {pad_to_block_size}");
 
-        let reader = Cursor::new(data);
+        let reader = Cursor::new(&data);
         let mut cpio_reader = CpioReader::new(reader, false);
 
         let writer = Cursor::new(Vec::new());
@@ -40,7 +82,7 @@ fn round_trip_archive() {
         let new_data = writer.get_ref().as_slice();
 
         if pad_to_block_size {
-            assert!(new_data.starts_with(data));
+            assert!(new_data.starts_with(&data));
             assert!(util::is_zero(&new_data[data.len()..]));
             assert_eq!(new_data.len() % 512, 0);
         } else {

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"
@@ -65,14 +62,12 @@ include-workspace = true
 bypass = [
     # Copies of unmodified crashwrangler objects for old macOS versions.
     { name = "honggfuzz", allow-globs = ["honggfuzz/third_party/mac/CrashReport_*.o"] },
-    # Test files for liblzma's test suite
-    { name = "liblzma-sys", allow-globs = ["xz/tests/compress_prepared_bcj_*"] },
 ]
 
 [sources]
 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

@@ -15,10 +15,10 @@ clap = { version = "4.4.1", features = ["derive"] }
 ctrlc = "3.4.0"
 hex = { version = "0.4.3", features = ["serde"] }
 ring = "0.17.0"
-rsa = "0.9.6"
+rsa = { version = "0.9.6", features = ["hazmat"] }
 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

@@ -12,6 +12,9 @@ security_patch_level = "2024-01-01"
 # Google Pixel 7 Pro
 # What's unique: init_boot (boot v4) + vendor_boot (vendor v4)
 
+[profile.pixel_v4_gki]
+vabc_algo = "Lz4"
+
 [profile.pixel_v4_gki.partitions.boot]
 avb.signed = true
 data.type = "boot"
@@ -22,7 +25,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,15 +46,18 @@ 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 = "c00f891f941f3dddb28966f7b07f3acea773bee104dace82b37c2d1341f09422"
+patched = "ce9d8ee97828d233809742a5d3f23aa27b042675b1935ca9e3df0592c55788fd"
 
 # Google Pixel 6a
 # What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4, 2 ramdisks)
 
+[profile.pixel_v4_non_gki]
+vabc_algo = "Lz4"
+
 [profile.pixel_v4_non_gki.partitions.boot]
 avb.signed = true
 data.type = "boot"
@@ -77,21 +83,24 @@ 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 = "4d692bc777b568b0626d3c08d2e6f83f1b472db5ad903486daaec6a78d0cc26e"
+patched = "e27673e4f30933710c11d51f0e73849068cbe9bc9f54e6076bdd93f9a5c8ea0a"
 
 # Google Pixel 4a 5G
 # What's unique: boot (boot v3) + vendor_boot (vendor v3)
 
+[profile.pixel_v3]
+vabc_algo = "Lz4"
+
 [profile.pixel_v3.partitions.boot]
 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,21 +121,24 @@ 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 = "f432dc7931520feb238474aa707dd5299747562ffe6129f3f763b5f11ac473ab"
+patched = "3850a2e73bd783a1ec4a70c59f37d2374e017c20df7ab4b591182b14d187c18e"
 
 # Google Pixel 4a
 # What's unique: boot (boot v2)
 
+[profile.pixel_v2]
+vabc_algo = "Gzip"
+
 [profile.pixel_v2.partitions.boot]
 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 +156,5 @@ data.type = "vbmeta"
 data.deps = ["system"]
 
 [profile.pixel_v2.hashes]
-original = "8b38d2d999b5b6e240e894f669e9e2643b3764c108d53bb7b02447da725e7c18"
-patched = "85b411947145e89cdc7f71b7109a12586f4dbddce3021f0f96172fc465c189d6"
+original = "1b45235b58054009cc496f6c3ee11d3dc16ed5c388c861761e26a6fce83103a0"
+patched = "193b2dc70dd465d686f35c7b7f74d2cc1b06a55e48cf5c2e4df0f667e03032fc"

e2e/src/cli.rs

@@ -3,10 +3,11 @@
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
-use std::path::PathBuf;
+use std::{ffi::OsString, path::PathBuf};
 
-use avbroot::cli::args::{LogFormat, LogLevel};
-use clap::{Args, Parser, Subcommand};
+use avbroot::cli::args::LogFormat;
+use clap::{Args, Parser, Subcommand, ValueEnum};
+use tracing::Level;
 
 #[derive(Debug, Args)]
 pub struct ProfileGroup {
@@ -69,10 +70,34 @@ pub struct Cli {
     pub command: Command,
 
     /// Lowest log message severity to output.
-    #[arg(long, global = true, value_name = "LEVEL", default_value_t)]
-    pub log_level: LogLevel,
+    #[arg(long, global = true, value_name = "LEVEL", default_value_t = Level::INFO)]
+    pub log_level: Level,
 
     /// Output format for log messages.
     #[arg(long, global = true, value_name = "FORMAT", default_value_t = LogFormat::Medium)]
     pub log_format: LogFormat,
 }
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
+pub enum PassSource {
+    Env,
+    File,
+}
+
+#[derive(Debug, Parser)]
+pub struct HelperCli {
+    /// Signature algorithm.
+    pub algorithm: String,
+
+    /// Public key.
+    #[arg(value_name = "FILE", value_parser)]
+    pub public_key: PathBuf,
+
+    /// Non-interactive password source.
+    #[arg(value_name = "SOURCE")]
+    pub pass_source: PassSource,
+
+    /// Non-interactive password source value.
+    #[arg(value_name = "VALUE", value_parser)]
+    pub pass_source_value: OsString,
+}

e2e/src/config.rs

@@ -1,13 +1,14 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-License-Identifier: GPL-3.0-only
  */
 
 use std::{collections::BTreeMap, fs, path::Path};
 
 use anyhow::{Context, Result};
+use avbroot::format::payload::VabcAlgo;
 use serde::{Deserialize, Serialize};
-use toml_edit::Document;
+use toml_edit::DocumentMut;
 
 #[derive(Serialize, Deserialize)]
 pub struct Sha256Hash(
@@ -41,7 +42,8 @@ pub struct Avb {
 pub enum RamdiskContent {
     Init,
     Otacerts,
-    InitAndOtacerts,
+    FirstStage,
+    DsuKeyDir,
     Dlkm,
 }
 
@@ -62,7 +64,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)]
@@ -108,6 +110,7 @@ pub struct Partition {
 #[derive(Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct Profile {
+    pub vabc_algo: Option<VabcAlgo>,
     pub partitions: BTreeMap<String, Partition>,
     pub hashes: Hashes,
 }
@@ -120,12 +123,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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
  * SPDX-FileCopyrightText: 2023 Pascal Roeleven
  * SPDX-License-Identifier: GPL-3.0-only
  */
@@ -9,10 +9,11 @@ mod config;
 
 use std::{
     collections::{BTreeMap, BTreeSet},
-    ffi::OsStr,
+    env,
+    ffi::{OsStr, OsString},
     fs::{self, File},
-    io::{self, BufReader, BufWriter, Cursor, Seek, SeekFrom, Write},
-    path::Path,
+    io::{self, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
+    path::{Path, PathBuf},
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc,
@@ -22,7 +23,7 @@ use std::{
 use anyhow::{anyhow, bail, Context, Result};
 use avbroot::{
     cli::ota::{ExtractCli, PatchCli, VerifyCli},
-    crypto::{self, PassphraseSource},
+    crypto::{self, PassphraseSource, RsaSigningKey},
     format::{
         avb::{
             self, AlgorithmType, ChainPartitionDescriptor, Descriptor, Footer, HashDescriptor,
@@ -48,15 +49,15 @@ use avbroot::{
     stream::{self, CountingWriter, HashingReader, PSeekFile, Reopen, ToWriter},
 };
 use clap::Parser;
-use rsa::RsaPrivateKey;
-use tempfile::{NamedTempFile, TempDir};
+use rsa::{rand_core::OsRng, traits::PublicKeyParts, BigUint};
+use tempfile::TempDir;
 use topological_sort::TopologicalSort;
 use tracing::{info, info_span};
 use x509_cert::Certificate;
 use zip::{write::FileOptions, CompressionMethod, ZipWriter};
 
 use crate::{
-    cli::{Cli, Command, ListCli, ProfileGroup, TestCli},
+    cli::{Cli, Command, HelperCli, ListCli, PassSource, ProfileGroup, TestCli},
     config::{
         Avb, BootData, BootVersion, Config, Data, DmVerityContent, DmVerityData, OtaInfo,
         Partition, Profile, RamdiskContent, VbmetaData,
@@ -100,7 +101,7 @@ fn append_avb(
     avb: &Avb,
     hash_tree: bool,
     ota_info: &OtaInfo,
-    key_avb: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
     let image_size = file.seek(SeekFrom::End(0))?;
@@ -196,7 +197,7 @@ fn append_avb(
         // Give enough free space for changes from patching.
         .max(1024 * 1024);
 
-    avb::write_appended_image(file, &header, &mut footer, full_image_size)?;
+    avb::write_appended_image(file, &header, &mut footer, Some(full_image_size))?;
 
     Ok(())
 }
@@ -227,6 +228,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 +254,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);
+            }
         }
     }
 
@@ -286,7 +299,7 @@ fn create_boot_image(
     avb: &Avb,
     boot_data: &BootData,
     ota_info: &OtaInfo,
-    key_avb: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
@@ -298,7 +311,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 +374,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(),
@@ -410,7 +426,7 @@ fn create_dm_verity_image(
     avb: &Avb,
     dm_verity_data: &DmVerityData,
     ota_info: &OtaInfo,
-    key_avb: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
@@ -439,7 +455,7 @@ fn create_vbmeta_image(
     avb: &Avb,
     vbmeta_data: &VbmetaData,
     inputs: &BTreeMap<String, PSeekFile>,
-    key: &RsaPrivateKey,
+    key: &RsaSigningKey,
 ) -> Result<()> {
     let mut descriptors = Vec::new();
 
@@ -491,7 +507,7 @@ fn create_vbmeta_image(
 fn create_partition_images(
     partitions: &BTreeMap<String, Partition>,
     ota_info: &OtaInfo,
-    key_avb: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<BTreeMap<String, PSeekFile>> {
@@ -561,7 +577,8 @@ fn create_payload(
     partitions: &BTreeMap<String, Partition>,
     inputs: &BTreeMap<String, PSeekFile>,
     ota_info: &OtaInfo,
-    key_ota: &RsaPrivateKey,
+    profile: &Profile,
+    key_ota: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<(String, u64)> {
     let dynamic_partitions_names = partitions
@@ -578,8 +595,14 @@ fn create_payload(
             .map(PSeekFile::new)
             .with_context(|| format!("Failed to create temp file for: {name}"))?;
 
-        let (partition_info, operations) =
-            payload::compress_image(file, &writer, name, 4096, cancel_signal)?;
+        let vabc_algo = if dynamic_partitions_names.contains(name) {
+            profile.vabc_algo
+        } else {
+            None
+        };
+
+        let (partition_info, operations, cow_estimate) =
+            payload::compress_image(file, &writer, name, 4096, vabc_algo, cancel_signal)?;
 
         compressed.insert(name, writer);
 
@@ -602,7 +625,7 @@ fn create_payload(
             fec_roots: None,
             version: None,
             merge_operations: vec![],
-            estimate_cow_size: None,
+            estimate_cow_size: cow_estimate,
         });
     }
 
@@ -623,7 +646,7 @@ fn create_payload(
                 }],
                 snapshot_enabled: Some(true),
                 vabc_enabled: Some(true),
-                vabc_compression_param: Some("gz".to_owned()),
+                vabc_compression_param: profile.vabc_algo.map(|a| a.to_string()),
                 cow_version: Some(2),
                 vabc_feature_set: None,
             }),
@@ -666,7 +689,7 @@ fn create_payload(
             .with_context(|| format!("Failed to copy from image: {name}"))?;
     }
 
-    let (_, properties, metadata_size) = payload_writer
+    let (_, _, properties, metadata_size) = payload_writer
         .finish()
         .context("Failed to finalize payload")?;
 
@@ -677,8 +700,8 @@ fn create_ota(
     output: &Path,
     ota_info: &OtaInfo,
     profile: &Profile,
-    key_avb: &RsaPrivateKey,
-    key_ota: &RsaPrivateKey,
+    key_avb: &RsaSigningKey,
+    key_ota: &RsaSigningKey,
     cert_ota: &Certificate,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
@@ -724,6 +747,7 @@ fn create_ota(
                     &profile.partitions,
                     &inputs,
                     ota_info,
+                    profile,
                     key_ota,
                     cancel_signal,
                 )
@@ -808,12 +832,16 @@ fn create_fake_magisk(output: &Path) -> Result<()> {
 
     for path in [
         "assets/stub.apk",
+        "lib/arm64-v8a/libinit-ld.so",
         "lib/arm64-v8a/libmagisk64.so",
         "lib/arm64-v8a/libmagiskinit.so",
+        "lib/armeabi-v7a/libinit-ld.so",
         "lib/armeabi-v7a/libmagisk32.so",
         "lib/armeabi-v7a/libmagiskinit.so",
+        "lib/x86/libinit-ld.so",
         "lib/x86/libmagisk32.so",
         "lib/x86/libmagiskinit.so",
+        "lib/x86_64/libinit-ld.so",
         "lib/x86_64/libmagisk64.so",
         "lib/x86_64/libmagiskinit.so",
     ] {
@@ -829,15 +857,18 @@ fn create_fake_magisk(output: &Path) -> Result<()> {
 }
 
 struct KeySet {
-    avb_key: RsaPrivateKey,
-    ota_key: RsaPrivateKey,
+    avb_key: RsaSigningKey,
+    ota_key: RsaSigningKey,
     ota_cert: Certificate,
-    avb_key_file: NamedTempFile,
-    avb_pass_file: NamedTempFile,
-    avb_pkmd_file: NamedTempFile,
-    ota_key_file: NamedTempFile,
-    ota_pass_file: NamedTempFile,
-    ota_cert_file: NamedTempFile,
+    _key_dir: TempDir,
+    avb_key_file: PathBuf,
+    avb_public_key_file: PathBuf,
+    avb_pass_file: PathBuf,
+    avb_pkmd_file: PathBuf,
+    ota_key_file: PathBuf,
+    ota_public_key_file: PathBuf,
+    ota_pass_file: PathBuf,
+    ota_cert_file: PathBuf,
 }
 
 macro_rules! new_keys_with_prefix {
@@ -886,13 +917,8 @@ macro_rules! new_keys_with_prefix {
 }
 
 impl KeySet {
-    fn write_temp(data: &[u8]) -> Result<NamedTempFile> {
-        let mut temp_file =
-            NamedTempFile::new().context("Failed to create temp file for test keys")?;
-        temp_file
-            .write_all(data)
-            .with_context(|| format!("Failed to write test key data: {:?}", temp_file.path()))?;
-        Ok(temp_file)
+    fn write(path: &Path, data: &[u8]) -> Result<()> {
+        fs::write(path, data).with_context(|| format!("Failed to write test key data: {path:?}"))
     }
 
     fn new_with_data(
@@ -903,36 +929,54 @@ impl KeySet {
         ota_pass: &[u8],
         ota_cert: &[u8],
     ) -> Result<Self> {
-        let avb_key_file = Self::write_temp(avb_key)?;
-        let avb_pass_file = Self::write_temp(avb_pass)?;
-        let avb_pkmd_file = Self::write_temp(avb_pkmd)?;
-        let ota_key_file = Self::write_temp(ota_key)?;
-        let ota_pass_file = Self::write_temp(ota_pass)?;
-        let ota_cert_file = Self::write_temp(ota_cert)?;
+        let key_dir = TempDir::new().context("Failed to create temp directory")?;
+        let avb_key_file = key_dir.path().join("avb.key");
+        let avb_public_key_file = key_dir.path().join("avb.public.key");
+        let avb_pass_file = key_dir.path().join("avb.passphrase");
+        let avb_pkmd_file = key_dir.path().join("avb_pkmd.bin");
+        let ota_key_file = key_dir.path().join("ota.key");
+        let ota_public_key_file = key_dir.path().join("ota.public.key");
+        let ota_pass_file = key_dir.path().join("ota.passphrase");
+        let ota_cert_file = key_dir.path().join("ota.crt");
+
+        Self::write(&avb_key_file, avb_key)?;
+        Self::write(&avb_pass_file, avb_pass)?;
+        Self::write(&avb_pkmd_file, avb_pkmd)?;
+        Self::write(&ota_key_file, ota_key)?;
+        Self::write(&ota_pass_file, ota_pass)?;
+        Self::write(&ota_cert_file, ota_cert)?;
 
         let avb_key = crypto::read_pem_key_file(
-            avb_key_file.path(),
-            &PassphraseSource::File(avb_pass_file.path().to_owned()),
+            &avb_key_file,
+            &PassphraseSource::File(avb_pass_file.clone()),
         )
+        .map(RsaSigningKey::Internal)
         .context("Failed to load AVB test key")?;
 
         let ota_key = crypto::read_pem_key_file(
-            ota_key_file.path(),
-            &PassphraseSource::File(ota_pass_file.path().to_owned()),
+            &ota_key_file,
+            &PassphraseSource::File(ota_pass_file.clone()),
         )
+        .map(RsaSigningKey::Internal)
         .context("Failed to load OTA test key")?;
 
-        let ota_cert = crypto::read_pem_cert_file(ota_cert_file.path())
-            .context("Failed to load OTA test cert")?;
+        crypto::write_pem_public_key_file(&avb_public_key_file, &avb_key.to_public_key())?;
+        crypto::write_pem_public_key_file(&ota_public_key_file, &ota_key.to_public_key())?;
+
+        let ota_cert =
+            crypto::read_pem_cert_file(&ota_cert_file).context("Failed to load OTA test cert")?;
 
         Ok(Self {
             avb_key,
             ota_key,
             ota_cert,
+            _key_dir: key_dir,
             avb_key_file,
+            avb_public_key_file,
             avb_pass_file,
             avb_pkmd_file,
             ota_key_file,
+            ota_public_key_file,
             ota_pass_file,
             ota_cert_file,
         })
@@ -945,12 +989,20 @@ impl KeySet {
 fn patch_image(
     input_file: &Path,
     output_file: &Path,
+    system_image_file: &Path,
     extra_args: &[&OsStr],
     keys: &KeySet,
+    signing_helper: bool,
     cancel_signal: &AtomicBool,
 ) -> Result<()> {
     info!("Patching OTA: {input_file:?} -> {output_file:?}");
 
+    let (avb_key_file, ota_key_file) = if signing_helper {
+        (&keys.avb_public_key_file, &keys.ota_public_key_file)
+    } else {
+        (&keys.avb_key_file, &keys.ota_key_file)
+    };
+
     // We're intentionally using the CLI interface.
     let mut args: Vec<&OsStr> = vec![
         OsStr::new("patch"),
@@ -958,17 +1010,29 @@ fn patch_image(
         input_file.as_os_str(),
         OsStr::new("--output"),
         output_file.as_os_str(),
+        OsStr::new("--replace"),
+        OsStr::new("system"),
+        system_image_file.as_os_str(),
         OsStr::new("--key-avb"),
-        keys.avb_key_file.path().as_os_str(),
+        avb_key_file.as_os_str(),
         OsStr::new("--pass-avb-file"),
-        keys.avb_pass_file.path().as_os_str(),
+        keys.avb_pass_file.as_os_str(),
         OsStr::new("--key-ota"),
-        keys.ota_key_file.path().as_os_str(),
+        ota_key_file.as_os_str(),
         OsStr::new("--pass-ota-file"),
-        keys.ota_pass_file.path().as_os_str(),
+        keys.ota_pass_file.as_os_str(),
         OsStr::new("--cert-ota"),
-        keys.ota_cert_file.path().as_os_str(),
+        keys.ota_cert_file.as_os_str(),
+        OsStr::new("--dsu"),
     ];
+
+    let argv0: OsString;
+    if signing_helper {
+        argv0 = env::args_os().next().unwrap();
+        args.push(OsStr::new("--signing-helper"));
+        args.push(&argv0);
+    }
+
     args.extend_from_slice(extra_args);
 
     let cli = PatchCli::try_parse_from(args)?;
@@ -1000,9 +1064,9 @@ fn verify_image(input_file: &Path, keys: &KeySet, cancel_signal: &AtomicBool) ->
         OsStr::new("--input"),
         input_file.as_os_str(),
         OsStr::new("--public-key-avb"),
-        keys.avb_pkmd_file.path().as_os_str(),
+        keys.avb_pkmd_file.as_os_str(),
         OsStr::new("--cert-ota"),
-        keys.ota_cert_file.path().as_os_str(),
+        keys.ota_cert_file.as_os_str(),
     ])?;
     avbroot::cli::ota::verify_subcommand(&cli, cancel_signal)?;
 
@@ -1103,11 +1167,18 @@ fn test_subcommand(cli: &TestCli, cancel_signal: &AtomicBool) -> Result<()> {
             .with_context(|| format!("[{name}] Failed to verify original OTA hash"))?;
 
         // Patch once using Magisk.
+        extract_image(&out_original, &profile_dir, cancel_signal)
+            .with_context(|| format!("[{name}] Failed to extract OTA"))?;
+
+        let system_image = profile_dir.join("system.img");
+
         patch_image(
             &out_original,
             &out_magisk,
+            &system_image,
             &args_magisk,
             &test_keys,
+            false,
             cancel_signal,
         )
         .with_context(|| format!("[{name}] Failed to patch OTA"))?;
@@ -1133,8 +1204,10 @@ fn test_subcommand(cli: &TestCli, cancel_signal: &AtomicBool) -> Result<()> {
         patch_image(
             &out_original,
             &out_prepatched,
+            &system_image,
             &args_prepatched,
             &test_keys,
+            true,
             cancel_signal,
         )
         .with_context(|| format!("[{name}] Failed to patch OTA"))?;
@@ -1159,7 +1232,87 @@ fn list_subcommand(cli: &ListCli) -> Result<()> {
     Ok(())
 }
 
+/// A basic --signing-helper implementation that just signs via RustCrypto.
+fn helper_mode() -> Result<()> {
+    let cli = HelperCli::parse();
+
+    let private_key_path = {
+        let parent = cli.public_key.parent().unwrap_or(Path::new("."));
+        let name = cli
+            .public_key
+            .file_name()
+            .and_then(|n| n.to_str())
+            .ok_or_else(|| anyhow!("Bad filename: {:?}", cli.public_key))?;
+
+        parent.join(name.replace(".public", ""))
+    };
+    let source = match cli.pass_source {
+        PassSource::Env => PassphraseSource::EnvVar(cli.pass_source_value),
+        PassSource::File => PassphraseSource::File(cli.pass_source_value.into()),
+    };
+    let private_key = crypto::read_pem_key_file(&private_key_path, &source)
+        .with_context(|| format!("Failed to load private key: {private_key_path:?}"))?;
+    let public_key = crypto::read_pem_public_key_file(&cli.public_key)
+        .with_context(|| format!("Failed to load public key: {:?}", cli.public_key))?;
+
+    if private_key.to_public_key() != public_key {
+        bail!("Private key does not match public key");
+    }
+
+    let (hash_algo, key_algo) = cli
+        .algorithm
+        .split_once('_')
+        .ok_or_else(|| anyhow!("Unknown algorithm: {:?}", cli.algorithm))?;
+
+    if key_algo != format!("RSA{}", private_key.size() * 8) {
+        bail!(
+            "{key_algo} does not match key size ({})",
+            private_key.size() * 8
+        );
+    } else if hash_algo != "SHA256" && hash_algo != "SHA512" {
+        bail!("Unknown hash algorithm: {hash_algo}");
+    }
+
+    let mut padded_digest = vec![];
+    io::stdin()
+        .read_to_end(&mut padded_digest)
+        .context("Failed to read padded digest from stdin")?;
+
+    if padded_digest.len() != private_key.size() {
+        bail!(
+            "Padded digest size ({}) bytes does not match key size ({})",
+            padded_digest.len(),
+            private_key.size()
+        );
+    }
+
+    // The input is already padded, so perform a raw RSA signing operation.
+    let mut signature = rsa::hazmat::rsa_decrypt_and_check(
+        &private_key,
+        None::<&mut OsRng>,
+        &BigUint::from_bytes_be(&padded_digest),
+    )
+    .context("Failed to sign digest")?
+    .to_bytes_le();
+    signature.resize(private_key.size(), 0);
+    signature.reverse();
+
+    io::stdout()
+        .write_all(&signature)
+        .context("Failed to write signature to stdout")?;
+
+    Ok(())
+}
+
 fn main() -> Result<()> {
+    const ENV_HELPER_MODE: &str = "E2E_HELPER_MODE";
+
+    // Re-invoking ourselves will execute as the helper script instead.
+    if env::var_os(ENV_HELPER_MODE).is_some() {
+        return helper_mode();
+    }
+    env::set_var(ENV_HELPER_MODE, "true");
+
     // Set up a cancel signal so we can properly clean up any temporary files.
     let cancel_signal = Arc::new(AtomicBool::new(false));
     {

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())?;