From 9dd98fd7155de2a03891ae873f97acebcacea7c1 Mon Sep 17 00:00:00 2001
From: Andrew Gunnerson <accounts+github@chiller3.com>
Date: Sat, 3 May 2025 22:00:45 -0400
Subject: [PATCH] payload: Add support for uncompressed CoW and fix more
 estimation bugs

No known device uses this, but it's very useful for testing that our
overhead calculation is identical to AOSP's. A few more bugs were found
in our overhead calculation logic:

* The CowHeaderV3 size was missing the original CowHeader (v2) fields
  that are supposed to be included due to inheritance in the C++ class.
* The additional 1% overhead was incorrectly calculated against the
  initial CoW estimate before static overhead for CoW headers was added.
* The V3 num_ops estimation did not set a minimum of 25 to match
  delta_generator.
* The V2 size estimation did not take into account that a cluster of CoW
  operations cannot be truncated. It must be a multiple of cluster_ops
  (200 for avbroot).

With these fixes, the CoW estimation when compression is disabled
matches AOSP exactly. This means all the overhead calculation is now
correct and the only difference when compression is enabled is in the
compression ratios of the various lz4/gz implementations.

Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>
---
 avbroot/src/cli/ota.rs        |   7 +-
 avbroot/src/format/payload.rs | 126 +++++++++++++++++++---------------
 e2e/e2e.toml                  |  32 ++++-----
 3 files changed, 89 insertions(+), 76 deletions(-)

diff --git a/avbroot/src/cli/ota.rs b/avbroot/src/cli/ota.rs
index f4bf0db..3bb8b51 100644
--- a/avbroot/src/cli/ota.rs
+++ b/avbroot/src/cli/ota.rs
@@ -10,6 +10,7 @@ use std::{
     io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write},
     ops::Range,
     path::{Path, PathBuf},
+    str::FromStr,
     sync::{atomic::AtomicBool, Mutex},
 };
 
@@ -17,7 +18,7 @@ use anyhow::{anyhow, bail, Context, Result};
 use bitflags::bitflags;
 use cap_std::{ambient_authority, fs::Dir};
 use cap_tempfile::TempDir;
-use clap::{value_parser, ArgAction, Args, Parser, Subcommand, ValueEnum};
+use clap::{value_parser, ArgAction, Args, Parser, Subcommand};
 use rayon::{iter::IntoParallelRefIterator, prelude::ParallelIterator};
 use tempfile::NamedTempFile;
 use topological_sort::TopologicalSort;
@@ -597,7 +598,7 @@ fn get_vabc_params(header: &PayloadHeader) -> Result<Option<VabcParams>> {
     };
 
     let compression = dpm.vabc_compression_param();
-    let Ok(vabc_algo) = VabcAlgo::from_str(compression, false) else {
+    let Ok(vabc_algo) = VabcAlgo::from_str(compression) else {
         bail!("Unsupported VABC compression: {compression}");
     };
 
@@ -631,7 +632,7 @@ fn set_vabc_algo(header: &mut PayloadHeader, vabc_algo: VabcAlgo) -> Result<bool
     }
 
     let compression = dpm.vabc_compression_param();
-    let Ok(old_vabc_algo) = VabcAlgo::from_str(compression, false) else {
+    let Ok(old_vabc_algo) = VabcAlgo::from_str(compression) else {
         bail!("Unsupported VABC compression: {compression}");
     };
 
diff --git a/avbroot/src/format/payload.rs b/avbroot/src/format/payload.rs
index d68d3bc..3d57360 100644
--- a/avbroot/src/format/payload.rs
+++ b/avbroot/src/format/payload.rs
@@ -7,13 +7,13 @@ use std::{
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     num::NonZeroU32,
     ops::{Add, Range},
+    str::FromStr,
     sync::atomic::AtomicBool,
 };
 
 use base64::engine::general_purpose::STANDARD;
 use base64::Engine;
 use bzip2::write::BzDecoder;
-use clap::ValueEnum;
 use flate2::{write::GzEncoder, Compression};
 use liblzma::{
     stream::{Check, Stream},
@@ -974,7 +974,7 @@ impl CowVersion {
         const NUM_RESUME_POINTS: u64 = 4;
         const SIZEOF_COW_FOOTER_V2: u64 = 84;
         const SIZEOF_COW_HEADER_V2: u64 = 38;
-        const SIZEOF_COW_HEADER_V3: u64 = 40;
+        const SIZEOF_COW_HEADER_V3: u64 = SIZEOF_COW_HEADER_V2 + 40;
         const SIZEOF_COW_OPERATION_V2: u64 = 20;
         const SIZEOF_COW_OPERATION_V3: u64 = 16;
         const SIZEOF_RESUME_POINT_V3: u64 = 16;
@@ -991,25 +991,25 @@ impl CowVersion {
                 // AOSP: CowWriterV2::InitPos()
                 overhead += BUFFER_REGION_DEFAULT_SIZE;
 
-                // Add an operation header (sizeof(CowOperationV2)) for each
-                // kCowReplaceOp for the raw data.
+                // Add all the CoW operation headers:
+                //
+                // - There is a kCowReplaceOp for each compressed chunk.
+                // - There is a kCowLabelOp for each InstallOperation in the
+                //   payload. This is added by delta_generator in CowDryRun().
+                // - There is a kCowClusterOp at the end of each cluster of
+                //   operations (which includes the kCowClusterOp itself). The
+                //   cluster size used to be 200, but was changed to 1024 in
+                //   5e8e488c13cbff9e0a305ce7c22fd6a13aabb886. We'll use the
+                //   smaller value because it's better to overestimate.
+                //
                 // AOSP: CowWriterV2::EmitClusterIfNeeded()
-                overhead += cow_replace_ops * SIZEOF_COW_OPERATION_V2;
+                let cow_label_ops = payload_install_ops;
+                let cow_cluster_ops = (cow_replace_ops + cow_label_ops).div_ceil(CLUSTER_OPS - 1);
 
-                // Add an operation header (sizeof(CowOperationV2)) for each
-                // kCowLabelOp. delta_generator adds one for each install
-                // operation.
-                // AOSP: CowDryRun()
-                overhead += payload_install_ops * SIZEOF_COW_OPERATION_V2;
-
-                // Add an operation header (sizeof(CowOperationV2)) for each
-                // kCowClusterOp, which is emitted for each cluster of cow
-                // operations. The cluster size used to be 200, but changed to
-                // 1024 in 5e8e488c13cbff9e0a305ce7c22fd6a13aabb886. We'll use
-                // the smaller value because it's better to overestimate.
-                // AOSP: CowWriterV2::EmitClusterIfNeeded()
-                let num_ops = (cow_replace_ops + payload_install_ops).div_ceil(CLUSTER_OPS - 1);
-                overhead += num_ops * SIZEOF_COW_OPERATION_V2;
+                // A cluster cannot be truncated, so round up to the nearest
+                // cluster boundary.
+                // AOSP: CowWriterV2::AddOperation()
+                overhead += cow_cluster_ops * CLUSTER_OPS * SIZEOF_COW_OPERATION_V2;
 
                 // sizeof(CowFooter).
                 // AOSP: CowWriterV2::GetCowSizeInfo()
@@ -1082,24 +1082,22 @@ pub struct CowEstimate {
 
 impl CowEstimate {
     /// Add fudge factor to account for overhead.
-    fn fudged(
-        &self,
-        payload_install_ops: u64,
-        cow_version: CowVersion,
-        vabc_algo: VabcAlgo,
-    ) -> Option<Self> {
+    fn fudged(&self, payload_install_ops: u64, cow_version: CowVersion) -> Option<Self> {
         let version_overhead = cow_version.size_overhead(self.num_ops, payload_install_ops);
-        let algo_overhead = vabc_algo.size_overhead(self.size);
 
-        let size = self
-            .size
-            .checked_add(version_overhead)
-            .and_then(|e| e.checked_add(algo_overhead))?;
+        let mut size = self.size.checked_add(version_overhead)?;
 
-        Some(Self {
-            size,
-            num_ops: self.num_ops,
-        })
+        // delta_generator adds 1% overhead to the original CoW size estimate,
+        // even if compression is disabled. We'll do the same too. For the
+        // compressed scenario, we rely on this more because lz4_flex and
+        // zlib-rs usually compress better than the lz4 and zlib implementations
+        // used by libsnapshot_cow.
+        size += size / 100;
+
+        // AOSP: PartitionProcessor::Run()
+        let num_ops = self.num_ops.max(25);
+
+        Some(Self { size, num_ops })
     }
 }
 
@@ -1123,22 +1121,14 @@ impl CheckedAdd for CowEstimate {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, ValueEnum)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize)]
 pub enum VabcAlgo {
+    None,
     Lz4,
     Gz,
 }
 
 impl VabcAlgo {
-    /// Compute the size overhead to account for differences in compression
-    /// level between our compression libraries and AOSP's.
-    fn size_overhead(self, estimate: u64) -> u64 {
-        // lz4_flex and zlib-rs usually compress better than the lz4 and zlib
-        // implementations used by libsnapshot_cow. Make up for this by adding
-        // percentage-based overhead.
-        estimate / 100
-    }
-
     /// Compute the compressed size of the raw data when split into chunks based
     /// on the specified [`ChunkingParams`]. The length of `raw_data` must be a
     /// multiple of the block size or else this will panic. The compressed data
@@ -1157,20 +1147,20 @@ impl VabcAlgo {
 
             // This should match CompressWorker::GetDefaultCompressionLevel() in
             // AOSP's libsnapshot.
-            let compressed = match self {
-                Self::Lz4 => lz4_flex::block::compress(chunk),
-                Self::Gz => {
-                    let mut encoder = GzEncoder::new(Vec::new(), Compression::best());
-                    encoder.write_all(chunk).map_err(Error::GzCompress)?;
-                    encoder.finish().map_err(Error::GzCompress)?
-                }
-            };
-
+            //
             // CoW v3 uses the raw data instead of the compressed data if the
             // raw data is smaller. Because we use a different implementation of
             // the compression algorithms, we don't implement this. It's safer
             // to just overestimate and use the (larger) compressed size.
-            size += compressed.len().min(chunk_size) as u64;
+            size += match self {
+                Self::None => chunk_size as u64,
+                Self::Lz4 => lz4_flex::block::compress(chunk).len() as u64,
+                Self::Gz => {
+                    let mut encoder = GzEncoder::new(Vec::new(), Compression::best());
+                    encoder.write_all(chunk).map_err(Error::GzCompress)?;
+                    encoder.finish().map_err(Error::GzCompress)?.len() as u64
+                }
+            };
 
             num_ops += 1;
             raw_data = remaining;
@@ -1182,7 +1172,29 @@ impl VabcAlgo {
 
 impl fmt::Display for VabcAlgo {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str(self.to_possible_value().ok_or(fmt::Error)?.get_name())
+        let name = match self {
+            Self::None => "none",
+            Self::Lz4 => "lz4",
+            Self::Gz => "gz",
+        };
+        f.write_str(name)
+    }
+}
+
+#[derive(Clone, Debug, Error)]
+#[error("Invalid VABC algorithm: {0} (must be {{none|lz4|gz}}[,<level>])")]
+pub struct InvalidVabcAlgo(String);
+
+impl FromStr for VabcAlgo {
+    type Err = InvalidVabcAlgo;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s {
+            "" | "none" => Ok(Self::None),
+            "lz4" => Ok(Self::Lz4),
+            "gz" => Ok(Self::Gz),
+            a => Err(InvalidVabcAlgo(a.to_owned())),
+        }
     }
 }
 
@@ -1300,7 +1312,7 @@ pub fn compute_cow_estimate(
         })?;
 
     initial_estimate
-        .fudged(payload_install_ops, vabc_params.version, vabc_params.algo)
+        .fudged(payload_install_ops, vabc_params.version)
         .ok_or(Error::IntOverflow("fudged_estimate"))
 }
 
@@ -1441,7 +1453,7 @@ pub fn compress_image(
     let cow_estimate = if let Some(p) = vabc_params {
         Some(
             initial_estimate
-                .fudged(operations.len() as u64, p.version, p.algo)
+                .fudged(operations.len() as u64, p.version)
                 .ok_or(Error::IntOverflow("fudged_estimate"))?,
         )
     } else {
diff --git a/e2e/e2e.toml b/e2e/e2e.toml
index fab4853..70b78bf 100644
--- a/e2e/e2e.toml
+++ b/e2e/e2e.toml
@@ -51,12 +51,12 @@ data.version = "vendor_v4"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v4_gki.hashes_streaming]
-original = "0615b681742f243aae0a022aeba6c180a3920bf6f0175a1a17129f13c7a2a214"
-patched = "8e9506f447585e54b060d8b0ccb6d705a6370ebd2425dffd5bc21a438be5a454"
+original = "6b2a71bb4c2057e7625f97a2ab9832416399d4ce5788db1bab978eb5cf6e8c51"
+patched = "0538bdff77c4ead8b6c7cf70e95a44a3559d8651bfd0afdcd3963d08c65999f6"
 
 [profile.pixel_v4_gki.hashes_seekable]
-original = "dda797dfbf8b2a9dfe103d4050a2b801e9e58c52bdd7606a8b73e8846a5cffec"
-patched = "361f2337db95e29366aa473a618d7839d7c43b02f9a374d223546230dd794096"
+original = "46e88872ab1591f86b98b146278ec6392a97e0258d42ceb84a87d43a316a9edd"
+patched = "00eeb6832c62e872e8a226a27728726a5f7d7f2436c4ddeba143efea6206924e"
 
 # Google Pixel 6a
 # What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4, 2 ramdisks)
@@ -93,12 +93,12 @@ data.version = "vendor_v4"
 data.ramdisks = [["init", "otacerts", "first_stage", "dsu_key_dir"], ["dlkm"]]
 
 [profile.pixel_v4_non_gki.hashes_streaming]
-original = "7b7acb994f938b256f07248f0b05dd160a44ed069c4bf8520db8cdc1a4441d8a"
-patched = "593d7b7969c018575c3d5f9c9764354750d0677abc8340e37ac3c500f9f78532"
+original = "8d76f17b33949c8ca3d4d5872858bb725d755e315f67d61db0e0b7b58cf69685"
+patched = "7c7cb69087d4c8c54a7f4bac6525354b1e9c1550545bc954194aaf6882f7982d"
 
 [profile.pixel_v4_non_gki.hashes_seekable]
-original = "23036d93a2b28aedff59390b6850469eb4ba57cf6bc0f80826cfef290a5e0246"
-patched = "73e3852db9a0f6848f83599802fbf790a0d130068a0fe6341641547fe9c68cfc"
+original = "6825bc02115e0c64f05651571138fbcc5cc04459e5c7e225dcc5e3096051c341"
+patched = "373309f28234866625462253d6f41f5a3b016b3676c7cb87a88bf55099ad6b13"
 
 # Google Pixel 4a 5G
 # What's unique: boot (boot v3) + vendor_boot (vendor v3)
@@ -136,12 +136,12 @@ data.version = "vendor_v3"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v3.hashes_streaming]
-original = "ca09a0158beae3f0ad215f9f381f61b28bc27377507664b79ab106c173cea700"
-patched = "e59884857ac5773f66e666ada6c10382220ce4e4d3f602c3f07e6c570e159a1d"
+original = "65048a5a6d2c90c7e3220b0e3500840613dceb0c167388bc44cb19796f13b89e"
+patched = "191945e73dd8032deb485c2b1e0260c00e56d5b003fa5acf32219a8a4c91bf44"
 
 [profile.pixel_v3.hashes_seekable]
-original = "b10d9708c1d1baade25aa07a045e24846b0cd7520f2d6ede94be32f93b32a5da"
-patched = "1e7cedcd4631327531451b32349eb13d8bb8cbc374f587a80bdac0985dd2812b"
+original = "cf200d0c1c82599cc174afdff585548e89ecc7e4089019083b28456baa757204"
+patched = "673c0b24b6d59426e4c96b820d2d57cc1c83fcefa24e175f4631d8d5dc44a40e"
 
 # Google Pixel 4a
 # What's unique: boot (boot v2)
@@ -173,9 +173,9 @@ data.type = "vbmeta"
 data.deps = ["system"]
 
 [profile.pixel_v2.hashes_streaming]
-original = "6c5f1d534bafb687fd1829c42f12f381fc58d3f6c905b09de2ee9c34aa05bfc5"
-patched = "c738e1ecc5e4a99cdc9387c050254eee35045cf1c800aa04942b213bf77a6eb7"
+original = "7c1f648bef642f794716e9177e69520ad868951bbfa4b9004607f00fdb2ad99c"
+patched = "23adbe34a2272603c81fe6632c5f2ec7fe4f7dacd82c4d6409fc09966d71388a"
 
 [profile.pixel_v2.hashes_seekable]
-original = "bdec93bcffdeb8fa9c0e540aeef1bc0df79bfa65baba76363262c5520e0ce6b5"
-patched = "69d768a14580b789050375709f6eba176e117c81871b5f8ae44cfd491f9f30fa"
+original = "e4e874e28bab5692ab6a385812194f4e139e8e13a186bac6d584ac8be6a25e68"
+patched = "c48a71a4eb6eb1f68efa8447ebe995dc092fa58716963cde4208106a6af1dbbb"