From 1a36fbd46c25378d8e23cd2e388915162eeda594 Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Sat, 16 Dec 2023 19:57:30 -0500 Subject: [PATCH] Split new partition images into chunks and compress in parallel Previously, all replacement partition images (those that have been patched or `--replace`d) were compressed as a whole, which would be very slow for larger images. Instead, we'll split the images into 2 MiB chunks and compress them in parallel. This more closely matches what AOSP's payload_generator does and scales linearly with the number of CPU cores. The compression is less efficient, but the file size generally only increases by 10s of KiB. This commit also reworks the implementation so that patched images are stored in temp files instead of in memory again. In hindsight, doing everything in memory only made things more complex and causes the memory usage to blow up when doing things like `--replace system `. Signed-off-by: Andrew Gunnerson --- avbroot/src/boot.rs | 3 +- avbroot/src/cli/ota.rs | 148 +++++++++++--------- avbroot/src/format/fec.rs | 11 +- avbroot/src/format/payload.rs | 247 ++++++++++++++++++++-------------- avbroot/src/util.rs | 10 ++ e2e/e2e.toml | 20 +-- 6 files changed, 252 insertions(+), 187 deletions(-) diff --git a/avbroot/src/boot.rs b/avbroot/src/boot.rs index dc3df86..ee5a409 100644 --- a/avbroot/src/boot.rs +++ b/avbroot/src/boot.rs @@ -528,7 +528,8 @@ impl PrepatchedImagePatcher { const MAX_LEVEL: u8 = 2; // We compile without Unicode support so we have to use [0-9] instead of \d. - const VERSION_REGEX: &'static str = r"Linux version ([0-9]+\.[0-9]+).[0-9]+-(android[0-9]+)-([0-9]+)-"; + const VERSION_REGEX: &'static str = + r"Linux version ([0-9]+\.[0-9]+).[0-9]+-(android[0-9]+)-([0-9]+)-"; pub fn new( prepatched: &Path, diff --git a/avbroot/src/cli/ota.rs b/avbroot/src/cli/ota.rs index c34be1b..75cdc40 100644 --- a/avbroot/src/cli/ota.rs +++ b/avbroot/src/cli/ota.rs @@ -9,7 +9,7 @@ use std::{ ffi::{OsStr, OsString}, fmt::Display, fs::{self, File}, - io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write}, + io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}, path::{Path, PathBuf}, sync::{atomic::AtomicBool, Mutex}, time::Instant, @@ -37,14 +37,14 @@ use crate::{ bootimage::BootImage, ota::{self, SigningWriter, ZipEntry}, padding, - payload::{self, CompressedPartitionWriter, PayloadHeader, PayloadWriter}, + payload::{self, PayloadHeader, PayloadWriter}, }, protobuf::{ build::tools::releasetools::OtaMetadata, chromeos_update_engine::DeltaArchiveManifest, }, stream::{ - self, CountingWriter, FromReader, HolePunchingWriter, PSeekFile, ReadSeek, ReadSeekReopen, - Reopen, SectionReader, ToWriter, + self, CountingWriter, FromReader, HolePunchingWriter, PSeekFile, ReadSeekReopen, Reopen, + SectionReader, ToWriter, }, util, }; @@ -150,17 +150,18 @@ pub fn get_required_images( Ok(images) } -/// Open all input streams listed in `required_images`. If an image has a path -/// in `external_images`, the real file on the filesystem is opened. Otherwise, -/// the image is extracted from the payload. -fn open_input_streams( +/// Open all input files listed in `required_images`. If an image has a path +/// in `external_images`, that file is opened. Otherwise, the image is extracted +/// from the payload into a temporary file (that is unnamed if supported by the +/// operating system). +fn open_input_files( payload: &(dyn ReadSeekReopen + Sync), required_images: &HashMap, external_images: &HashMap, header: &PayloadHeader, cancel_signal: &AtomicBool, -) -> Result>> { - let mut input_streams = HashMap::>::new(); +) -> Result> { + let mut input_files = HashMap::::new(); // We always include replacement images that the user specifies, even if // they don't need to be patched. @@ -174,18 +175,23 @@ fn open_input_streams( status!("Opening external image: {name}: {path:?}"); let file = File::open(path) + .map(PSeekFile::new) .with_context(|| format!("Failed to open external image: {path:?}"))?; - input_streams.insert(name.clone(), Box::new(file)); + input_files.insert(name.clone(), file); } else { status!("Extracting from original payload: {name}"); - let stream = payload::extract_image_to_memory(payload, header, name, cancel_signal) + let file = tempfile::tempfile() + .map(PSeekFile::new) + .with_context(|| format!("Failed to create temp file for: {name}"))?; + + payload::extract_image(payload, &file, header, name, cancel_signal) .with_context(|| format!("Failed to extract from original payload: {name}"))?; - input_streams.insert(name.clone(), Box::new(stream)); + input_files.insert(name.clone(), file); } } - Ok(input_streams) + Ok(input_files) } /// Patch the boot images listed in `required_images`. An [`OtaCertPatcher`] is @@ -195,7 +201,7 @@ fn open_input_streams( /// be re-signed with `key_avb`. fn patch_boot_images( required_images: &HashMap, - input_streams: &mut HashMap>, + input_files: &mut HashMap, root_patcher: Option>, key_avb: &RsaPrivateKey, cert_ota: &Certificate, @@ -219,29 +225,31 @@ fn patch_boot_images( joined(sorted(boot_patchers.keys())) ); - // Temporarily take the streams out of input_streams so we can easily - // run the patchers in parallel. + // Temporarily take the files out of input_files so we can easily run the + // patchers in parallel. let patchers_list = boot_patchers .into_iter() - .map(|(n, p)| (n, p, input_streams.remove(n).unwrap())) + .map(|(n, p)| (n, p, input_files.remove(n).unwrap())) .collect::>(); // Patch the boot images. The original readers are dropped. let patched = patchers_list .into_par_iter() - .map(|(n, p, s)| -> Result<(&str, Cursor>)> { - let mut writer = Cursor::new(Vec::new()); + .map(|(n, p, f)| -> Result<(&str, PSeekFile)> { + let mut new_file = tempfile::tempfile() + .map(PSeekFile::new) + .with_context(|| format!("Failed to create temp file for: {n}"))?; - boot::patch_boot(s, &mut writer, key_avb, &p, cancel_signal) + boot::patch_boot(f, &mut new_file, key_avb, &p, cancel_signal) .with_context(|| format!("Failed to patch boot image: {n}"))?; - Ok((n, writer)) + Ok((n, new_file)) }) .collect::>>()?; - // Put the patched images back into input_streams. - for (name, stream) in patched { - input_streams.insert(name.to_owned(), Box::new(stream)); + // Put the patched images back into input_files. + for (name, file) in patched { + input_files.insert(name.to_owned(), file); } Ok(()) @@ -251,7 +259,7 @@ fn patch_boot_images( /// then an error is returned because the vbmeta patching logic only ever writes /// root vbmeta images. fn load_vbmeta_images( - images: &mut HashMap>, + images: &mut HashMap, vbmeta_images: &HashSet, ) -> Result> { let mut result = HashMap::new(); @@ -325,7 +333,7 @@ fn ensure_partitions_protected( /// determine the order to patch the vbmeta images so that it can be done in a /// single pass. fn get_vbmeta_patch_order( - images: &mut HashMap>, + images: &mut HashMap, vbmeta_headers: &HashMap, ) -> Result)>> { let mut dep_graph = HashMap::<&str, HashSet>::new(); @@ -533,7 +541,7 @@ fn update_metadata_descriptors(parent_header: &mut Header, child_header: &Header /// replaced with a new in-memory reader containing the new image. Otherwise, /// the image is removed from `images` entirely to avoid needing to repack it. fn update_vbmeta_headers( - images: &mut HashMap>, + images: &mut HashMap, headers: &mut HashMap, order: &mut [(String, HashSet)], clear_vbmeta_flags: bool, @@ -575,7 +583,9 @@ fn update_vbmeta_headers( .sign(key) .with_context(|| format!("Failed to sign vbmeta header for image: {name}"))?; - let mut writer = Cursor::new(Vec::new()); + let mut writer = tempfile::tempfile() + .map(PSeekFile::new) + .with_context(|| format!("Failed to create temp file for: {name}"))?; parent_header .to_writer(&mut writer) .with_context(|| format!("Failed to write vbmeta image: {name}"))?; @@ -583,7 +593,7 @@ fn update_vbmeta_headers( padding::write_zeros(&mut writer, block_size) .with_context(|| format!("Failed to write vbmeta padding: {name}"))?; - *images.get_mut(name).unwrap() = Box::new(writer); + *images.get_mut(name).unwrap() = writer; } else { unchanged.insert(name.as_str()); } @@ -600,17 +610,19 @@ fn update_vbmeta_headers( /// Compress an image and update the OTA manifest partition entry appropriately. fn compress_image( name: &str, - mut stream: &mut Box, + file: &mut PSeekFile, header: &Mutex, block_size: u32, cancel_signal: &AtomicBool, ) -> Result<()> { - stream.rewind()?; + file.rewind()?; - let writer = Cursor::new(Vec::new()); - let mut compressed = CompressedPartitionWriter::new(writer, block_size)?; + let writer = tempfile::tempfile() + .map(PSeekFile::new) + .with_context(|| format!("Failed to create temp file for: {name}"))?; - stream::copy(&mut stream, &mut compressed, cancel_signal)?; + let (partition_info, operations) = + payload::compress_image(&*file, &writer, name, block_size, cancel_signal)?; let mut header_locked = header.lock().unwrap(); let partition = header_locked @@ -619,9 +631,11 @@ fn compress_image( .iter_mut() .find(|p| p.partition_name == name) .unwrap(); - let writer = compressed.finish(partition)?; - *stream = Box::new(writer); + partition.new_partition_info = Some(partition_info); + partition.operations = operations; + + *file = writer; Ok(()) } @@ -678,11 +692,11 @@ fn patch_ota_payload( .collect::>(); // The set of source images to be inserted into the new payload, replacing - // what was in the original payload. Initially, this refers to either real - // files on the filesystem (--replace option) or in-memory files (extracted - // from the old payload). The values will be replaced later if the images - // need to be patched (eg. boot or vbmeta image). - let mut input_streams = open_input_streams( + // what was in the original payload. Initially, this refers to either user + // specified files (--replace option) or temporary files (extracted from the + // old payload). The values will be replaced later if the images need to be + // patched (eg. boot or vbmeta image). + let mut input_files = open_input_files( payload, &required_images, external_images, @@ -692,18 +706,18 @@ fn patch_ota_payload( patch_boot_images( &required_images, - &mut input_streams, + &mut input_files, root_patcher, key_avb, cert_ota, cancel_signal, )?; - let mut vbmeta_headers = load_vbmeta_images(&mut input_streams, &vbmeta_images)?; + let mut vbmeta_headers = load_vbmeta_images(&mut input_files, &vbmeta_images)?; ensure_partitions_protected(&header_locked.manifest, &vbmeta_headers)?; - let mut vbmeta_order = get_vbmeta_patch_order(&mut input_streams, &vbmeta_headers)?; + let mut vbmeta_order = get_vbmeta_patch_order(&mut input_files, &vbmeta_headers)?; status!( "Patching vbmeta images: {}", @@ -714,12 +728,12 @@ fn patch_ota_payload( for name in &vbmeta_images { // Linear search is fast enough. if !vbmeta_order.iter().any(|v| v.0 == *name) { - input_streams.remove(name); + input_files.remove(name); } } update_vbmeta_headers( - &mut input_streams, + &mut input_files, &mut vbmeta_headers, &mut vbmeta_order, clear_vbmeta_flags, @@ -729,16 +743,16 @@ fn patch_ota_payload( status!( "Compressing replacement images: {}", - joined(sorted(input_streams.keys())), + joined(sorted(input_files.keys())), ); let block_size = header_locked.manifest.block_size(); drop(header_locked); - input_streams + input_files .par_iter_mut() - .map(|(name, stream)| -> Result<()> { - compress_image(name, stream, &header, block_size, cancel_signal) + .map(|(name, file)| -> Result<()> { + compress_image(name, file, &header, block_size, cancel_signal) .with_context(|| format!("Failed to compress image: {name}")) }) .collect::>()?; @@ -762,25 +776,31 @@ fn patch_ota_payload( continue; }; - if let Some(mut reader) = input_streams.remove(&name) { - // Copy from our replacement image. + let pi = payload_writer.partition_index().unwrap(); + let oi = payload_writer.operation_index().unwrap(); + let orig_partition = &header_locked.manifest.partitions[pi]; + let orig_operation = &orig_partition.operations[oi]; + let data_offset = orig_operation + .data_offset + .ok_or_else(|| anyhow!("Missing data_offset in partition #{pi} operation #{oi}"))?; + + if let Some(reader) = input_files.get_mut(&name) { + // Copy from our replacement image. The compressed chunks are laid + // out sequentially and data_offset is set to the offset within that + // file. reader - .rewind() + .seek(SeekFrom::Start(data_offset)) .with_context(|| format!("Failed to seek image: {name}"))?; - stream::copy_n(&mut reader, &mut payload_writer, data_length, cancel_signal) + stream::copy_n(reader, &mut payload_writer, data_length, cancel_signal) .with_context(|| format!("Failed to copy from replacement image: {name}"))?; } else { // Copy from the original payload. - let pi = payload_writer.partition_index().unwrap(); - let oi = payload_writer.operation_index().unwrap(); - let orig_partition = &header_locked.manifest.partitions[pi]; - let orig_operation = &orig_partition.operations[oi]; - - let data_offset = orig_operation - .data_offset - .and_then(|o| o.checked_add(header_locked.blob_offset)) - .ok_or_else(|| anyhow!("Missing data_offset in partition #{pi} operation #{oi}"))?; + let data_offset = data_offset + .checked_add(header_locked.blob_offset) + .ok_or_else(|| { + anyhow!("data_offset overflow in partition #{pi} operation #{oi}") + })?; orig_payload_reader .seek(SeekFrom::Start(data_offset)) diff --git a/avbroot/src/format/fec.rs b/avbroot/src/format/fec.rs index da2af9f..f8816a4 100644 --- a/avbroot/src/format/fec.rs +++ b/avbroot/src/format/fec.rs @@ -20,7 +20,7 @@ use thiserror::Error; use crate::{ format::verityrs, stream::{self, FromReader, ReadSeekReopen, ToWriter, WriteSeekReopen, WriteZerosExt}, - util::NumBytes, + util::{self, NumBytes}, }; // Not to be confused with the 255-byte RS block size. @@ -111,11 +111,6 @@ impl Codeword { } } -/// Since Rust's built-in .div_ceil() is still nightly-only. -fn div_ceil(dividend: u64, divisor: u64) -> u64 { - dividend / divisor + u64::from(dividend % divisor != 0) -} - /// A type for performing FEC generation, verification, and error correction for /// a specific file size and Reed Solomon configuration. The implementation uses /// dm-verity's interleaving access pattern. @@ -169,8 +164,8 @@ impl Fec { return Err(Error::UnsupportedParity(parity)); } - let blocks = div_ceil(file_size, u64::from(block_size)); - let rounds = div_ceil(blocks, u64::from(rs_k)); + let blocks = util::div_ceil(file_size, u64::from(block_size)); + let rounds = util::div_ceil(blocks, u64::from(rs_k)); // Check upfront so we don't need to do checked multiplication later. rounds diff --git a/avbroot/src/format/payload.rs b/avbroot/src/format/payload.rs index 4b5f2db..d164538 100644 --- a/avbroot/src/format/payload.rs +++ b/avbroot/src/format/payload.rs @@ -34,9 +34,10 @@ use crate::{ InstallOperation, PartitionInfo, PartitionUpdate, Signatures, }, stream::{ - self, CountingReader, CountingWriter, FromReader, HashingWriter, ReadDiscardExt, - ReadSeekReopen, Reopen, SharedCursor, WriteSeek, + self, CountingReader, FromReader, HashingWriter, ReadDiscardExt, ReadSeekReopen, WriteSeek, + WriteSeekReopen, }, + util, }; const OTA_MAGIC: &[u8; 4] = b"CrAU"; @@ -535,98 +536,6 @@ impl Write for PayloadWriter { } } -/// A writer to produce a compressed partition image suitable for directly -/// inserting into a `payload.bin`'s blob section. The data is XZ-compressed at -/// a low compression level, primarily to collapse zeros, and the -/// [`PartitionUpdate`] instance is updated with the final size and hash. The -/// entire data will be represented as a single [`InstallOperation`] and -/// [`InstallOperation::data_offset`] will be set to `None`. -pub struct CompressedPartitionWriter { - inner: XzEncoder>>, - block_size: u32, - h_uncompressed: Context, - written: u64, -} - -impl CompressedPartitionWriter { - pub fn new(writer: W, block_size: u32) -> Result { - let counting_writer = CountingWriter::new(writer); - let hashing_writer = - HashingWriter::new(counting_writer, Context::new(&ring::digest::SHA256)); - - // AOSP's payload_consumer does not support CRC during decompression. - // Also, we intentionally pick the lowest compression level since we - // primarily care about squishing zeros. The non-zero portions of boot - // images are usually already-compressed kernels and ramdisks. - let stream = Stream::new_easy_encoder(0, Check::None)?; - let xz_writer = XzEncoder::new_stream(hashing_writer, stream); - - Ok(Self { - inner: xz_writer, - block_size, - h_uncompressed: Context::new(&ring::digest::SHA256), - written: 0, - }) - } - - /// Finish writing and update the [`PartitionUpdate`] instance with the new - /// size, hash, and install operation metadata. - pub fn finish(mut self, partition: &mut PartitionUpdate) -> Result { - if self.written % u64::from(self.block_size) != 0 { - return Err(Error::InvalidPartitionSize { - name: partition.partition_name.clone(), - size: self.written, - block_size: self.block_size, - }); - } - - self.inner.flush()?; - - let hashing_writer = self.inner.finish()?; - let digest_uncompressed = self.h_uncompressed.finish(); - let (counting_writer, context_compressed) = hashing_writer.finish(); - let digest_compressed = context_compressed.finish(); - // XzEncoder::total_out() cannot be used for an exact byte count because - // XzEncoder::finish() writes data. - let (writer, size_compressed) = counting_writer.finish(); - - partition.new_partition_info = Some(PartitionInfo { - size: Some(self.written), - hash: Some(digest_uncompressed.as_ref().to_vec()), - }); - - let extent = Extent { - start_block: Some(0), - num_blocks: Some(self.written / u64::from(self.block_size)), - }; - - // data_offset must be manually updated by the caller. - let mut operation = InstallOperation::default(); - operation.set_type(Type::ReplaceXz); - operation.data_length = Some(size_compressed); - operation.dst_extents.push(extent); - operation.data_sha256_hash = Some(digest_compressed.as_ref().to_vec()); - - partition.operations.clear(); - partition.operations.push(operation); - - Ok(writer) - } -} - -impl Write for CompressedPartitionWriter { - fn write(&mut self, buf: &[u8]) -> io::Result { - let n = self.inner.write(buf)?; - self.h_uncompressed.update(&buf[..n]); - self.written += n as u64; - Ok(n) - } - - fn flush(&mut self) -> io::Result<()> { - self.inner.flush() - } -} - /// Verify the payload signatures using the specified certificate and check that /// the digests in `payload_properties.txt` are correct. pub fn verify_payload( @@ -867,29 +776,29 @@ pub fn apply_operation( Ok(()) } -/// Extract the specified image from the payload into memory. This is done -/// multithreaded and uses rayon's global thread pool. `open_payload` will be -/// called from multiple threads. -pub fn extract_image_to_memory( +/// Extract the specified image from the payload. This is done multithreaded and +/// uses rayon's global thread pool. Both the `payload` and `output` streams +/// will be reopened from multiple threads. +pub fn extract_image( payload: &(dyn ReadSeekReopen + Sync), + output: &(dyn WriteSeekReopen + Sync), header: &PayloadHeader, partition_name: &str, cancel_signal: &AtomicBool, -) -> Result { +) -> Result<()> { let partition = header .manifest .partitions .iter() .find(|p| p.partition_name == partition_name) .ok_or_else(|| Error::MissingPartition(partition_name.to_owned()))?; - let stream = SharedCursor::default(); partition .operations .par_iter() .map(|op| -> Result<()> { let reader = payload.reopen_boxed()?; - let writer = stream.reopen()?; + let writer = output.reopen_boxed()?; apply_operation( reader, @@ -902,9 +811,7 @@ pub fn extract_image_to_memory( Ok(()) }) - .collect::>()?; - - Ok(stream) + .collect::>() } /// Extract the specified partition images from the payload into writers. This @@ -954,3 +861,135 @@ pub fn extract_images<'a>( }) .collect() } + +/// 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. +pub fn compress_image( + input: &(dyn ReadSeekReopen + Sync), + output: &(dyn WriteSeekReopen + Sync), + partition_name: &str, + block_size: u32, + cancel_signal: &AtomicBool, +) -> Result<(PartitionInfo, Vec)> { + const CHUNK_SIZE: u64 = 2 * 1024 * 1024; + const CHUNK_GROUP: u64 = 32; + + let file_size = input.reopen_boxed()?.seek(SeekFrom::End(0))?; + let final_chunk_different = file_size % CHUNK_SIZE != 0; + + if file_size % u64::from(block_size) != 0 || CHUNK_SIZE % u64::from(block_size) != 0 { + return Err(Error::InvalidPartitionSize { + name: partition_name.to_owned(), + size: file_size, + block_size, + }); + } + + let chunks_total = util::div_ceil(file_size, CHUNK_SIZE); + let mut bytes_compressed = 0; + let mut context_uncompressed = Context::new(&ring::digest::SHA256); + let mut operations = vec![]; + + // Read the file one group at a time. This allows for some parallelization + // without reading the entire file into memory. This is necessary because we + // need to compute the checksum of the entire file. + while (operations.len() as u64) < chunks_total { + let chunks_done = operations.len() as u64; + let chunks_group = (chunks_total - chunks_done).min(CHUNK_GROUP); + + let uncompressed_data_group = (chunks_done..chunks_done + chunks_group) + .into_par_iter() + .map(|chunk| -> Result<(u64, Vec)> { + let mut reader = input.reopen_boxed()?; + let offset = reader.seek(SeekFrom::Start(chunk * CHUNK_SIZE))?; + + let chunk_size = if final_chunk_different && chunk == chunks_total - 1 { + file_size % CHUNK_SIZE + } else { + CHUNK_SIZE + }; + let mut data = vec![0u8; chunk_size as usize]; + + stream::check_cancel(cancel_signal)?; + reader.read_exact(&mut data)?; + + Ok((offset, data)) + }) + .collect::>>()?; + + for (_, data) in &uncompressed_data_group { + context_uncompressed.update(data); + } + + let mut compressed_data_group = uncompressed_data_group + .into_par_iter() + .map( + |(raw_offset, raw_data)| -> Result<(Vec, InstallOperation)> { + let reader = Cursor::new(raw_data.as_slice()); + let writer = Cursor::new(Vec::new()); + let hashing_writer = + HashingWriter::new(writer, Context::new(&ring::digest::SHA256)); + + // AOSP's payload_consumer does not support checking CRC during + // decompression. Also, we intentionally pick the lowest + // compression level since we primarily care about squishing + // zeros. The non-zero portions of boot images are usually + // already-compressed kernels and ramdisks. + let stream = Stream::new_easy_encoder(0, Check::None)?; + let mut xz_writer = XzEncoder::new_stream(hashing_writer, stream); + + stream::copy_n(reader, &mut xz_writer, raw_data.len() as u64, cancel_signal)?; + + let hashing_writer = xz_writer.finish()?; + let (writer, context_compressed) = hashing_writer.finish(); + let digest_compressed = context_compressed.finish(); + let data = writer.into_inner(); + + let extent = Extent { + start_block: Some(raw_offset / u64::from(block_size)), + num_blocks: Some(raw_data.len() as u64 / u64::from(block_size)), + }; + + let mut operation = InstallOperation::default(); + operation.set_type(Type::ReplaceXz); + operation.data_length = Some(data.len() as u64); + operation.dst_extents.push(extent); + operation.data_sha256_hash = Some(digest_compressed.as_ref().to_vec()); + + Ok((data, operation)) + }, + ) + .collect::>>()?; + + for (data, operation) in &mut compressed_data_group { + operation.data_offset = Some(bytes_compressed); + bytes_compressed += data.len() as u64; + } + + let group_operations = compressed_data_group + .into_par_iter() + .map(|(data, operation)| -> Result { + let mut writer = output.reopen_boxed()?; + writer.seek(SeekFrom::Start(operation.data_offset.unwrap()))?; + writer.write_all(&data)?; + + Ok(operation) + }) + .collect::>>()?; + + operations.extend(group_operations.into_iter()); + } + + let digest_uncompressed = context_uncompressed.finish(); + let partition_info = PartitionInfo { + size: Some(file_size), + hash: Some(digest_uncompressed.as_ref().to_vec()), + }; + + Ok((partition_info, operations)) +} diff --git a/avbroot/src/util.rs b/avbroot/src/util.rs index bd80432..3130572 100644 --- a/avbroot/src/util.rs +++ b/avbroot/src/util.rs @@ -48,3 +48,13 @@ pub fn parent_path(path: &Path) -> &Path { Path::new(".") } + +/// Since Rust's built-in .div_ceil() is still nightly-only. +pub fn div_ceil(dividend: T, divisor: T) -> T { + dividend / divisor + + if dividend % divisor != T::zero() { + T::one() + } else { + T::zero() + } +} diff --git a/e2e/e2e.toml b/e2e/e2e.toml index 9bffd31..a32bbe2 100644 --- a/e2e/e2e.toml +++ b/e2e/e2e.toml @@ -15,8 +15,8 @@ sections = [ ] hash.original.full = "6b881553f012d582080642d660e1cf5c9e6fe41e9f1c6ab12ae87fab7894e307" hash.original.stripped = "9befd7887a125ebd8e9ae0555469dababe6bc04b0aa41aa2562036782a6d87e0" -hash.patched.full = "91e15447ade648c10bce599e75569ce55edc18798b11a704586acf3b51ae7971" -hash.patched.stripped = "84b069ac7f20115ac0b19d7e319bb86d00b9f3f64e0c7a51612f99d6cae297bb" +hash.patched.full = "8de1892395be6173687ce6333d926394ca486c61856dd7585c6a7fc4feb0624e" +hash.patched.stripped = "d24efed2dd68625fa8f4b44149dd07fcb5597f925e16b62a9e968edd15115fea" hash.avb_images."init_boot.img" = "3bedb41be98c46241f11219021dfbb799a6d5c89e6e00d45a66f9a5b42e7dfbc" hash.avb_images."vbmeta.img" = "e8e6e898ca73807edb43af0a0e86d4a94b14256def89581970287a9b1bf7a3ee" hash.avb_images."vbmeta_system.img" = "dbb63e08f26f46ccda501d99058d513ff71e3d6302c14d587442b666ff08862a" @@ -36,8 +36,8 @@ sections = [ ] hash.original.full = "1f1f0abe67a6f6f47287be6dafec2c12628de6a715b82ca7beddaf67ad22aca5" hash.original.stripped = "38b15f5efdc7e056bc799859ba72ef9a73e93c61292c59f85fb4b9c31acc5f82" -hash.patched.full = "65f7e29591fb48ad9c7c3233df4d76bb6986aef96b94c06b73213477bc7486d8" -hash.patched.stripped = "cf77b5e307ce4d2a62e742cd3a300964de6c693880cbfa90dcd0bfec66bd1976" +hash.patched.full = "6916e91351e464339f63be311f8debc9a54d2e3c5509a6fe9c886f0ee63a556c" +hash.patched.stripped = "be40b24c23cefb5ab5704420a6d93dc4f168d38124aee3c663ebe9e2f9a345f0" hash.avb_images."boot.img" = "a19cb4d4fcc7f3e7d3046c3d19e2f243fb02513ca848ff92ad70d1ada55c4e65" hash.avb_images."vbmeta.img" = "2d817e35f7b6cdc2edce58ef249a966fc677085b70a4aced4c829320aeef0be2" hash.avb_images."vbmeta_system.img" = "98a050f0d53a016fbb78147b1b4a9bca3fde615aa4da34bf62c2e07a395104b5" @@ -57,8 +57,8 @@ sections = [ ] hash.original.full = "6d107ffac1cd3da2c972112acc75957ed725e5c13d57ca724d9bcca5404fcebd" hash.original.stripped = "5b889bdab3bb12ddcd3c243a56e1c58bedada8831069f49d56fe5098fb141e35" -hash.patched.full = "8725e03798539070d7075a07c80fb1403652a2b446d5c460d636a32b7368c9a2" -hash.patched.stripped = "c4b5cfa84dc8c15f3ac9661d768062546fa36ca82cdeb6dd943347be0422903c" +hash.patched.full = "67e1fbc5b1c189cb856d9fd9d846a780ba42e398c5d2e161f4ca89ffccecc32f" +hash.patched.stripped = "37280d66e316b72f223fb06a2a004902bbf8a28c0da3d9115af12b0d8aea5186" hash.avb_images."boot.img" = "8e7278a2e8ae44ffc5475717eb0e1aa56bfb7650aef34375b0fe92f790835f95" hash.avb_images."vbmeta.img" = "b036132b867f52a86eef79716261f35b1eb50e843dc4fe42f72cce67b24ae2db" hash.avb_images."vbmeta_system.img" = "9a7c6fd654e7a92aeffbdbd55ea0d87eee36f4c235e1b505423ad8a13a751a00" @@ -76,8 +76,8 @@ sections = [ ] hash.original.full = "01fd34b206152a3559039161c9874ab03df37da4268b86a9e0be899de5fc0af7" hash.original.stripped = "cc311b5bd46e06cfdefbade794d33aa9bc3ceda4ad4f38bfe9f0dfc17033d207" -hash.patched.full = "f2ac798b31a94dc251ca4ce370ebfb3073170d4a34431829df9ed0742149ffe7" -hash.patched.stripped = "4387a5ba30c925f56c67eeaf6512757d5c1e07dfd2a8a99be14e06db8ba2dde7" +hash.patched.full = "dc56c61170c0de6e2bb2cf8c7c0b6aa810f0aaea1eb045b3d4b1643482bb27fd" +hash.patched.stripped = "1b6b1818ff72d79e9a1a3f376beca5d62ea2715cbd97a3997f58ca636b588bb4" hash.avb_images."boot.img" = "506a955080b6cfa2039ef85923e8a4e717ef6c1dc478538599c7ba26ee21e525" hash.avb_images."vbmeta.img" = "3679c7224e3e3e0793b4d1a031e116098460a6b4c5f3f88d5a3a0a4c65b21582" hash.avb_images."vbmeta_system.img" = "1d3efa00fd1d44a594c7317072468fa95c23d83d2759d6d6e757783ceeabc594" @@ -107,8 +107,8 @@ sections = [ ] hash.original.full = "929f892fbd70699cf7f118a119aac1ae1b86351e1ada17715666fa4401e63472" hash.original.stripped = "4eabaf79b6c2b5df305e3ecdc2b9570c0dd27350b4e8d6434584000c4989ff3d" -hash.patched.full = "ec576b7430e5a2788e89f8bbe37ee81b116386d9b101865d3b75eb9c51a0db4c" -hash.patched.stripped = "ddd0cf1b42564f8c2980320c14f8e0eb125412eb7a13ae8fb52d35424546c235" +hash.patched.full = "fb61ac143024dd9c01d3af45fdeacde078c46c4c7f1d6d2d6da24561d555ad46" +hash.patched.stripped = "a222904445c2d17c203f34a113da1c4657c622b38d32cde21487ab5ad6b8afb3" hash.avb_images."boot.img" = "480d7cf519326fcaa5106ecfbbeb907309068635f7f6a25e5e4571d525c4926a" hash.avb_images."recovery.img" = "eeb0f67e2084174fced510f4b59a82022559ffbe1c20d2c0ea4757fcd989a3af" hash.avb_images."vbmeta.img" = "c022cf79da301a8430af5c49704944c490707fa0306031fe3ea22c39ce4734f6"