Add support for gzip VABC algorithm

Older devices, like the Pixel 4a 5G (bramble) use gzip instead of lz4. Fixes: #332 Signed-off-by: Andrew Gunnerson <[email protected]>
chenxiaolong · Aug 18, 2024 · fa81f5a · fa81f5a
1 parent 83ab475
commit fa81f5a
Show file tree

Hide file tree

Showing 7 changed files with 122 additions and 45 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,
+}
 
-    while !raw_data.is_empty() {
-        let n = raw_data.len().min(block_size as usize);
-        let compressed = lz4_flex::block::compress(&raw_data[..n]);
+impl VabcAlgo {
+    pub fn new(name: &str) -> Option<Self> {
+        match name {
+            "lz4" => Some(Self::Lz4),
+            "gz" => Some(Self::Gzip),
+            _ => None,
+        }
+    }
 
-        total += compressed.len().min(n) as u64;
+    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..];
+        }
 
-        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,13 +1060,22 @@ 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)
-    } else {
-        None
+    // lz4_flex and miniz_oxide seem to compress better than the lz4 and zlib
+    // libraries used by update_engine. Add some extra headroom to reduce the
+    // chance of update_engine failing to flash a partition due to the size of
+    // the actual data exceeding the estimated size.
+    let cow_estimate = match vabc_algo {
+        Some(VabcAlgo::Lz4) => {
+            // 1%
+            let fudge = cow_estimate / 100;
+            Some(cow_estimate + fudge)
+        }
+        Some(VabcAlgo::Gzip) => {
+            // 2%
+            let fudge = cow_estimate / 50;
+            Some(cow_estimate + fudge)
+        }
+        None => None,
     };
 
     Ok((partition_info, operations, cow_estimate))

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"
@@ -52,6 +55,9 @@ patched = "a68b3a44c0cf015a225f92837c05fd0dec6e159584c5880eff30d12daa7123ff"
 # 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"
@@ -86,6 +92,9 @@ patched = "a13173a006ad94b25db682bb14fde2e36891417727d6541bdf5f9bca57d26751"
 # 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"
@@ -121,6 +130,9 @@ patched = "49e810ae76154bccf2dc7d810b9eec19d23a5b8fa32381469660488d0a25121b"
 # 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 = "54f0f0f20425288130c8e9935348b368950b6f5967251d67ea824f8572699d09"
+patched = "ea1a33722302da023b8a9b6d8fea097b6bbbc574f3703a301bfe2900ead14193"
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,
                 )