Add support for gzip VABC algorithm

Older devices, like the Pixel 4a 5G (bramble) use gzip instead of lz4. This commit also reworks the CoW size estimate calculation to add the same constant headroom that AOSP's delta_generator adds. Previously, avbroot was already adding an additional 1% to account for differences in compression ratios across compression library implementations. This papered over the issue for large partitions, but small partitions could still have a CoW size estimate that's too small. Adding the constant headroom prevents ENOSPC when flashing those partitions. Fixes: #332 Signed-off-by: Andrew Gunnerson <[email protected]>
chenxiaolong · Aug 19, 2024 · 1fa35fd · 1fa35fd
1 parent 83ab475
commit 1fa35fd
Show file tree

Hide file tree

Showing 7 changed files with 131 additions and 49 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/avbroot/Cargo.toml b/avbroot/Cargo.toml
@@ -27,6 +27,7 @@ hex = { version = "0.4.3", features = ["serde"] }
 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"] }

diff --git a/avbroot/src/cli/ota.rs b/avbroot/src/cli/ota.rs
@@ -34,7 +34,7 @@ use crate::{
         avb::{self, Descriptor, Header},
         ota::{self, SigningWriter, ZipEntry},
         padding,
-        payload::{self, PayloadHeader, PayloadWriter},
+        payload::{self, PayloadHeader, PayloadWriter, VabcAlgo},
     },
     patch::{
         boot::{
@@ -676,12 +676,11 @@ pub fn compress_image(
     // 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 {
+    let vabc_algo = if partition.estimate_cow_size.is_some() {
         info!("Needs updated CoW size estimate: {name}");
 
-        // Only CoW v2 + lz4 seems to exist in the wild currently, so that is
-        // all we support.
+        // 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");
         };
@@ -696,13 +695,17 @@ pub fn compress_image(
         }
 
         let compression = dpm.vabc_compression_param();
-        if compression != "lz4" {
+        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, need_cow, cancel_signal)?;
+        payload::compress_image(&*file, &writer, name, block_size, vabc_algo, cancel_signal)?;
 
     partition.new_partition_info = Some(partition_info);
     partition.operations = operations;

diff --git a/avbroot/src/format/payload.rs b/avbroot/src/format/payload.rs
@@ -5,6 +5,7 @@
 
 use std::{
     collections::{HashMap, HashSet},
+    fmt,
     io::{self, Cursor, Read, Seek, SeekFrom, Write},
     ops::Range,
     sync::atomic::AtomicBool,
@@ -887,19 +888,50 @@ 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;
+#[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,
+        }
+    }
 
-    while !raw_data.is_empty() {
-        let n = raw_data.len().min(block_size as usize);
-        let compressed = lz4_flex::block::compress(&raw_data[..n]);
+    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;
+            total += compressed.len().min(n) as u64;
 
-        raw_data = &raw_data[n..];
+            raw_data = &raw_data[n..];
+        }
+
+        total
     }
+}
 
-    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
@@ -910,18 +942,18 @@ fn compress_cow_size(mut raw_data: &[u8], block_size: u32) -> u64 {
 /// 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.
+/// 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,
-    need_cow_estimate: bool,
+    vabc_algo: Option<VabcAlgo>,
     cancel_signal: &AtomicBool,
 ) -> Result<(PartitionInfo, Vec<InstallOperation>, Option<u64>)> {
     const CHUNK_SIZE: u64 = 2 * 1024 * 1024;
@@ -980,8 +1012,8 @@ pub fn compress_image(
             .map(
                 |(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)
+                    let cow_size = if let Some(algo) = vabc_algo {
+                        algo.compressed_size(&raw_data, block_size)
                     } else {
                         0
                     };
@@ -1028,11 +1060,25 @@ pub fn compress_image(
         hash: Some(digest_uncompressed.as_ref().to_vec()),
     };
 
-    let cow_estimate = if need_cow_estimate {
-        // Because lz4_flex compresses better than official lz4.
-        let fudge = cow_estimate / 100;
-
-        Some(cow_estimate + fudge)
+    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
     };

diff --git a/e2e/e2e.toml b/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"
@@ -46,12 +49,15 @@ data.version = "vendor_v4"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v4_gki.hashes]
-original = "6b140c378d21eae2fa4fc581bce13a689b21bd32f5fba865698d1fd322f2f8c6"
-patched = "a68b3a44c0cf015a225f92837c05fd0dec6e159584c5880eff30d12daa7123ff"
+original = "c00f891f941f3dddb28966f7b07f3acea773bee104dace82b37c2d1341f09422"
+patched = "2f14b161bbd087fbcd99ca39d8f42217f956cad03eed785f8b7248e8f740d47c"
 
 # 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"
@@ -80,12 +86,15 @@ data.version = "vendor_v4"
 data.ramdisks = [["init", "otacerts", "first_stage", "dsu_key_dir"], ["dlkm"]]
 
 [profile.pixel_v4_non_gki.hashes]
-original = "31963e6f81986c6686111f50e36b89e4d85ee5c02bc8e5ecd560528bc98d6fe7"
-patched = "a13173a006ad94b25db682bb14fde2e36891417727d6541bdf5f9bca57d26751"
+original = "4d692bc777b568b0626d3c08d2e6f83f1b472db5ad903486daaec6a78d0cc26e"
+patched = "c9342b77c165700b8c13bff16288dad552cae7309127dc04098c17a10b9532ef"
 
 # 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"
@@ -115,12 +124,15 @@ data.version = "vendor_v3"
 data.ramdisks = [["otacerts", "first_stage", "dsu_key_dir"]]
 
 [profile.pixel_v3.hashes]
-original = "e684aacb54464098c1b8e3f499efe35dff10ea792e89d71a83404620d0108b3e"
-patched = "49e810ae76154bccf2dc7d810b9eec19d23a5b8fa32381469660488d0a25121b"
+original = "f432dc7931520feb238474aa707dd5299747562ffe6129f3f763b5f11ac473ab"
+patched = "adfee4740c1914a6fba01929c326438541218551fdfe50136bed499a3a667610"
 
 # 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"
@@ -144,5 +156,5 @@ data.type = "vbmeta"
 data.deps = ["system"]
 
 [profile.pixel_v2.hashes]
-original = "ee9568797d9195985f14753b89949d8ebb08c8863a32eceeeec6e8d94661b1cf"
-patched = "e413f1f87ee5ba53d402edad03b0df8af451b5b0323202a9d32b8327d433d340"
+original = "1b45235b58054009cc496f6c3ee11d3dc16ed5c388c861761e26a6fce83103a0"
+patched = "193b2dc70dd465d686f35c7b7f74d2cc1b06a55e48cf5c2e4df0f667e03032fc"
diff --git a/e2e/src/config.rs b/e2e/src/config.rs
@@ -6,6 +6,7 @@
 use std::{collections::BTreeMap, fs, path::Path};
 
 use anyhow::{Context, Result};
+use avbroot::format::payload::VabcAlgo;
 use serde::{Deserialize, Serialize};
 use toml_edit::DocumentMut;
 
@@ -109,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,
 }

diff --git a/e2e/src/main.rs b/e2e/src/main.rs
@@ -577,6 +577,7 @@ fn create_payload(
     partitions: &BTreeMap<String, Partition>,
     inputs: &BTreeMap<String, PSeekFile>,
     ota_info: &OtaInfo,
+    profile: &Profile,
     key_ota: &RsaSigningKey,
     cancel_signal: &AtomicBool,
 ) -> Result<(String, u64)> {
@@ -594,14 +595,14 @@ fn create_payload(
             .map(PSeekFile::new)
             .with_context(|| format!("Failed to create temp file for: {name}"))?;
 
-        let (partition_info, operations, cow_estimate) = payload::compress_image(
-            file,
-            &writer,
-            name,
-            4096,
-            dynamic_partitions_names.contains(name),
-            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);
 
@@ -645,7 +646,7 @@ fn create_payload(
                 }],
                 snapshot_enabled: Some(true),
                 vabc_enabled: Some(true),
-                vabc_compression_param: Some("lz4".to_owned()),
+                vabc_compression_param: profile.vabc_algo.map(|a| a.to_string()),
                 cow_version: Some(2),
                 vabc_feature_set: None,
             }),
@@ -746,6 +747,7 @@ fn create_ota(
                     &profile.partitions,
                     &inputs,
                     ota_info,
+                    profile,
                     key_ota,
                     cancel_signal,
                 )