Add subcommands for packing and unpacking payload binaries

Some devices have "full" OTAs where the payload is missing the recovery
partition. These subcommands make it possible to manually add back the
missing image. Given the strict requirements for how the OTA zip is laid
out and signed, users can't just replace payload.bin in a zip and call
it a day, but it's sufficient for feeding a modified input to
`avbroot ota patch`.

Issue: #328

Signed-off-by: Andrew Gunnerson <[email protected]>

README.extra.md

@@ -255,6 +255,36 @@ This will check if the input file has any corrupted blocks. Currently, the comma
 
 ## `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/Cargo.toml

@@ -71,6 +71,7 @@ features = ["deflate"]
 rustix = { version = "0.38.9", default-features = false, features = ["process"] }
 
 [build-dependencies]
+constcat = "0.5.0"
 prost-build = "0.13.1"
 protox = "0.7.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

@@ -164,7 +164,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),
+        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/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/ota.rs

@@ -623,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,
@@ -900,7 +900,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")?;
 
@@ -1091,7 +1091,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,
@@ -1442,7 +1442,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,
@@ -1654,7 +1654,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,
@@ -1864,7 +1864,7 @@ pub struct PatchCli {
         value_names = ["PARTITION", "FILE"],
         value_parser = value_parser!(OsString),
         num_args = 2,
-        help_heading = HEADING_PATH,
+        help_heading = HEADING_PATH
     )]
     pub replace: Vec<OsString>,