Update CoW size estimate when replacing entire dynamic partitions

avbroot/src/cli/ota.rs

@@ -674,12 +674,40 @@ 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 need_cow = partition.estimate_cow_size.is_some();
+    if need_cow {
+        info!("Needs updated CoW size estimate: {name}");
+
+        // Only CoW v2 + lz4 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();
+        if compression != "lz4" {
+            bail!("Unsupported VABC compression: {compression}");
+        }
+    }
+
+    let (partition_info, operations, cow_estimate) =
+        payload::compress_image(&*file, &writer, name, block_size, need_cow, cancel_signal)?;
 
     partition.new_partition_info = Some(partition_info);
     partition.operations = operations;
+    partition.estimate_cow_size = cow_estimate;
 
     *file = writer;
 

avbroot/src/format/payload.rs

@@ -890,20 +890,43 @@ fn compress_chunk(raw_data: &[u8], cancel_signal: &AtomicBool) -> Result<(Vec<u8
     Ok((data, digest_compressed))
 }
 
+fn compress_cow_size(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 = lz4_flex::block::compress(&raw_data[..n]);
+
+        total += compressed.len().min(n) as u64;
+
+        raw_data = &raw_data[n..];
+    }
+
+    total
+}
+
 /// 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
 /// then written in parallel (but in order) to the output. Each chunk will have
 /// 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 `need_cow_estimate` is true, the VABC CoW v2 + lz4 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,
+    need_cow_estimate: bool,
     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;
 
@@ -921,6 +944,7 @@ pub fn compress_image(
     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 +981,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 need_cow_estimate {
+                        compress_cow_size(&raw_data, block_size)
+                    } else {
+                        0
+                    };
 
                     let extent = Extent {
                         start_block: Some(raw_offset / u64::from(block_size)),
@@ -971,19 +1000,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 +1031,16 @@ pub fn compress_image(
         hash: Some(digest_uncompressed.as_ref().to_vec()),
     };
 
-    Ok((partition_info, operations))
+    let cow_estimate = if need_cow_estimate {
+        // Because lz4_flex compresses better than official lz4.
+        let fudge = cow_estimate / 100;
+
+        Some(cow_estimate + fudge)
+    } else {
+        None
+    };
+
+    Ok((partition_info, operations, cow_estimate))
 }
 
 fn extents_sorted(operations: &[InstallOperation]) -> bool {

e2e/e2e.toml

@@ -46,8 +46,8 @@ data.version = "vendor_v4"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v4_gki.hashes]
-original = "24a0a62cc08b96563f4872aee2fdd4a84d1a977b55c326dd9d4ddd92a1d326ea"
-patched = "29889670efea78bace221742b19e8ace88b67137fc2f46dcd9dbdf67e4e42267"
+original = "6b140c378d21eae2fa4fc581bce13a689b21bd32f5fba865698d1fd322f2f8c6"
+patched = "f00e9745f90754be28ce8355501d876759cd8336451e4c3633908fbb4217b422"
 
 # Google Pixel 6a
 # What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4, 2 ramdisks)
@@ -80,8 +80,8 @@ data.version = "vendor_v4"
 data.ramdisks = [["init", "otacerts", "first_stage", "dsu_key_dir"], ["dlkm"]]
 
 [profile.pixel_v4_non_gki.hashes]
-original = "0e7d0924a68d46e00abe96abfa0e5f3a98d5d15a32bef7401b91fca9a19748e8"
-patched = "1a2f53d9ac3a1da75e5e7502110bda437c5a725764f16656c564d42460136279"
+original = "31963e6f81986c6686111f50e36b89e4d85ee5c02bc8e5ecd560528bc98d6fe7"
+patched = "43959409034dbb9aa0a605d7c5c0e7885012bbcd63510ec87d8e97015b37a746"
 
 # Google Pixel 4a 5G
 # What's unique: boot (boot v3) + vendor_boot (vendor v3)
@@ -115,8 +115,8 @@ data.version = "vendor_v3"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v3.hashes]
-original = "533e6f233cb98c98c945044c2ee81a6069e66baee6f7dbcfbf7523795a11215e"
-patched = "1eeae9dba0302c2d469bd03c8bccdc0469c171204dbd676a20cb6620d2d11c5c"
+original = "e684aacb54464098c1b8e3f499efe35dff10ea792e89d71a83404620d0108b3e"
+patched = "08e03ec327bf5bd841b91ad8d53028c3439a1722b423aaaac6cc0ceee4ef66b1"
 
 # Google Pixel 4a
 # What's unique: boot (boot v2)
@@ -144,5 +144,5 @@ data.type = "vbmeta"
 data.deps = ["system"]
 
 [profile.pixel_v2.hashes]
-original = "958dfa428abd2901178b90147903d7857ed78d6c016f6cb3af30d024a22a8f9a"
-patched = "b0d18ef350ca7b6499b7de4b0182f4ac3be7288ad238ce888708643e750f7631"
+original = "ee9568797d9195985f14753b89949d8ebb08c8863a32eceeeec6e8d94661b1cf"
+patched = "5e265094d4164cedde8f483911c58860f6008b314dc8e5ed3b44deb53fbb2f96"

e2e/src/main.rs

@@ -593,8 +593,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 (partition_info, operations, cow_estimate) = payload::compress_image(
+            file,
+            &writer,
+            name,
+            4096,
+            dynamic_partitions_names.contains(name),
+            cancel_signal,
+        )?;
 
         compressed.insert(name, writer);
 
@@ -617,7 +623,7 @@ fn create_payload(
             fec_roots: None,
             version: None,
             merge_operations: vec![],
-            estimate_cow_size: None,
+            estimate_cow_size: cow_estimate,
         });
     }
 
@@ -638,7 +644,7 @@ fn create_payload(
                 }],
                 snapshot_enabled: Some(true),
                 vabc_enabled: Some(true),
-                vabc_compression_param: Some("gz".to_owned()),
+                vabc_compression_param: Some("lz4".to_owned()),
                 cow_version: Some(2),
                 vabc_feature_set: None,
             }),
@@ -960,6 +966,7 @@ impl KeySet {
 fn patch_image(
     input_file: &Path,
     output_file: &Path,
+    system_image_file: &Path,
     extra_args: &[&OsStr],
     keys: &KeySet,
     cancel_signal: &AtomicBool,
@@ -973,6 +980,9 @@ 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(),
         OsStr::new("--pass-avb-file"),
@@ -1119,9 +1129,15 @@ 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,
             cancel_signal,
@@ -1149,6 +1165,7 @@ fn test_subcommand(cli: &TestCli, cancel_signal: &AtomicBool) -> Result<()> {
         patch_image(
             &out_original,
             &out_prepatched,
+            &system_image,
             &args_prepatched,
             &test_keys,
             cancel_signal,