From acf57272a756fcd0ab10dbbc2c7da49c8a55f15e Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Sun, 17 Dec 2023 00:39:04 -0500 Subject: [PATCH] fec: Add support for partially updating FEC data FEC data can be updated efficiently with "round" granularity when the regions where the input file was modified are known. Signed-off-by: Andrew Gunnerson --- README.extra.md | 8 +++ avbroot/src/cli/fec.rs | 40 +++++++++++++ avbroot/src/format/fec.rs | 116 ++++++++++++++++++++++++++++++++++++-- avbroot/src/util.rs | 26 ++++++++- e2e/src/main.rs | 24 +------- 5 files changed, 187 insertions(+), 27 deletions(-) diff --git a/README.extra.md b/README.extra.md index acfb37f..ce59683 100644 --- a/README.extra.md +++ b/README.extra.md @@ -175,6 +175,14 @@ The default behavior is to use 2 bytes of parity information per 253 bytes of in The number of parity bytes (between 2 and 24, inclusive) can be configured using `--parity`. +### Updating FEC data + +```bash +avbroot fec update -i -f [-r ]... +``` + +This will update the FEC data corresponding to the specified regions. This can be significantly faster than generating new FEC data from scratch for large files if the regions where data was modified are known. + ### Verifying a file ```bash diff --git a/avbroot/src/cli/fec.rs b/avbroot/src/cli/fec.rs index 01728d5..846fd1f 100644 --- a/avbroot/src/cli/fec.rs +++ b/avbroot/src/cli/fec.rs @@ -61,6 +61,24 @@ fn generate_subcommand(cli: &GenerateCli, cancel_signal: &AtomicBool) -> Result< Ok(()) } +fn update_subcommand(cli: &UpdateCli, cancel_signal: &AtomicBool) -> Result<()> { + let ranges = cli + .range + .chunks_exact(2) + .map(|w| w[0]..w[1]) + .collect::>(); + + let input = open_input(&cli.input, false)?; + let mut fec = read_fec(&cli.fec)?; + + fec.update(&input, &ranges, cancel_signal) + .context("Failed to update FEC data")?; + + write_fec(&cli.fec, &fec)?; + + Ok(()) +} + fn verify_subcommand(cli: &VerifyCli, cancel_signal: &AtomicBool) -> Result<()> { let input = open_input(&cli.input, false)?; let fec = read_fec(&cli.fec)?; @@ -87,6 +105,7 @@ fn repair_subcommand(cli: &RepairCli, cancel_signal: &AtomicBool) -> Result<()> pub fn fec_main(cli: &FecCli, cancel_signal: &AtomicBool) -> Result<()> { match &cli.command { FecCommand::Generate(c) => generate_subcommand(c, cancel_signal), + FecCommand::Update(c) => update_subcommand(c, cancel_signal), FecCommand::Verify(c) => verify_subcommand(c, cancel_signal), FecCommand::Repair(c) => repair_subcommand(c, cancel_signal), } @@ -108,6 +127,26 @@ struct GenerateCli { parity: u8, } +/// Update FEC data after a file is modified. +#[derive(Debug, Parser)] +struct UpdateCli { + /// Path to input data. + #[arg(short, long, value_name = "FILE", value_parser)] + input: PathBuf, + + /// Path to FEC data. + /// + /// The file will be modified in place. + #[arg(short, long, value_name = "FILE", value_parser)] + fec: PathBuf, + + /// Input file ranges that were updated. + /// + /// This is a half-open range and can be specified multiple times. + #[arg(short, long, value_names = ["START", "END"], num_args = 2)] + range: Vec, +} + /// Verify that a file contains no errors. #[derive(Debug, Parser)] struct VerifyCli { @@ -137,6 +176,7 @@ struct RepairCli { #[derive(Debug, Subcommand)] enum FecCommand { Generate(GenerateCli), + Update(UpdateCli), Verify(VerifyCli), Repair(RepairCli), } diff --git a/avbroot/src/format/fec.rs b/avbroot/src/format/fec.rs index f8816a4..044a56d 100644 --- a/avbroot/src/format/fec.rs +++ b/avbroot/src/format/fec.rs @@ -4,8 +4,10 @@ */ use std::{ + collections::HashSet, fmt, io::{self, Cursor, Read, Seek, SeekFrom, Write}, + ops::Range, sync::atomic::AtomicBool, }; @@ -206,6 +208,34 @@ impl Fec { offset / rs_k + offset % rs_k * self.rounds * u64::from(self.block_size) } + /// Get the rounds that correspond to the specified ranges. + fn rounds_for_ranges(&self, ranges: &[Range]) -> Result> { + let ranges = util::merge_overlapping(ranges); + if let Some(last) = ranges.last() { + if last.end > self.file_size { + return Err(Error::FieldOutOfBounds("ranges")); + } + } + + let block_size = u64::from(self.block_size); + let mut result = HashSet::new(); + + for range in ranges { + let start_block = range.start / block_size; + let end_block = if range.end % block_size == 0 { + range.end / block_size + } else { + util::div_ceil(range.end, block_size) + }; + + for block in start_block..end_block { + result.insert(block % self.rounds); + } + } + + Ok(result) + } + /// Read a raw sequential block from the backing file, starting at offset /// `offset` in the interleaved view. This reads a horizontal block-aligned /// slice in the file offset grid. @@ -433,6 +463,41 @@ impl Fec { Ok(fec) } + /// Update FEC data coreesponding to the specified file ranges. + /// + /// This function is multithreaded and uses rayon's global thread pool. + pub fn update( + &self, + input: &(dyn ReadSeekReopen + Sync), + fec: &mut [u8], + ranges: &[Range], + cancel_signal: &AtomicBool, + ) -> Result<()> { + let fec_size = self.fec_size(); + if fec.len() != fec_size { + return Err(Error::InvalidFecSize { + input: self.file_size, + expected: fec_size, + actual: fec.len(), + }); + } + + let rounds_to_update = self.rounds_for_ranges(ranges)?; + + fec.par_chunks_exact_mut(fec_size / self.rounds as usize) + .enumerate() + .filter(|(round, _)| rounds_to_update.contains(&(*round as u64))) + .map(|(round, buf)| -> Result<()> { + stream::check_cancel(cancel_signal)?; + + let reader = input.reopen_boxed()?; + self.generate_one_round(reader, round as u64, buf) + }) + .collect::>()?; + + Ok(()) + } + /// Verify that the file contains no errors. This is significantly faster /// than [`Self::repair()`] if only error detection, not correction, is /// needed. @@ -555,6 +620,17 @@ impl FecImage { }) } + /// Update FEC data coreesponding to the specified file ranges. + pub fn update( + &mut self, + input: &(dyn ReadSeekReopen + Sync), + ranges: &[Range], + cancel_signal: &AtomicBool, + ) -> Result<()> { + let fec = Fec::new(self.data_size, FEC_BLOCK_SIZE as u32, self.parity)?; + fec.update(input, &mut self.fec, ranges, cancel_signal) + } + /// Check that a file contains no errors. This is significantly faster than /// [`Self::repair()`] if performing a repair is not necessary. pub fn verify( @@ -726,6 +802,31 @@ mod tests { use super::*; + #[test] + fn rounds_for_ranges() { + let size = 2 * 253 * 4096; + let fec = Fec::new(size, 4096, 2).unwrap(); + + assert_eq!(fec.rounds_for_ranges(&[0..0]).unwrap(), HashSet::new(),); + assert_eq!( + fec.rounds_for_ranges(&[0..size]).unwrap(), + HashSet::from([0, 1]), + ); + assert_eq!(fec.rounds_for_ranges(&[0..1]).unwrap(), HashSet::from([0]),); + assert_eq!( + fec.rounds_for_ranges(&[4095..4096]).unwrap(), + HashSet::from([0]), + ); + assert_eq!( + fec.rounds_for_ranges(&[4095..4097]).unwrap(), + HashSet::from([0, 1]), + ); + assert_eq!( + fec.rounds_for_ranges(&[size - 1..size]).unwrap(), + HashSet::from([1]), + ); + } + fn corrupt_byte(file: &mut SharedCursor, offset: u64) { let mut buf = [0u8; 1]; @@ -776,7 +877,9 @@ mod tests { corrupt_byte(&mut file, offset as u64); } - // Verify that all the single-byte errors can be fixed. + // Verify that all the single-byte errors can be fixed. We don't test + // for Error::TooManyErrors because of the chance of false positives due + // to the nature of RS. fec.repair(&file, &file, &fec_data, &cancel_signal).unwrap(); let repaired_digest = { @@ -787,12 +890,17 @@ mod tests { }; assert_eq!(repaired_digest.as_ref(), orig_digest.as_ref()); - // We don't test for Error::TooManyErrors because of the chance of false - // positives due to the nature of RS. + // Intentionally update some data. + corrupt_byte(&mut file, 0); + let mut fec_data_updated = fec_data.clone(); + let fec_data = fec.generate(&file, &cancel_signal).unwrap(); + fec.update(&file, &mut fec_data_updated, &[0..1], &cancel_signal) + .unwrap(); + assert_eq!(fec_data_updated, fec_data); } #[test] - fn generate_verify_repair() { + fn generate_update_verify_repair() { for block_size in [1, 2, 4, 8, 16, 32, 64] { for rs_k in verityrs::FN_ENCODE.keys() { println!("Testing block_size={block_size}, rs_k={rs_k}"); diff --git a/avbroot/src/util.rs b/avbroot/src/util.rs index 3130572..1fa5fe4 100644 --- a/avbroot/src/util.rs +++ b/avbroot/src/util.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: GPL-3.0-only */ -use std::{fmt, path::Path}; +use std::{fmt, ops::Range, path::Path}; use num_traits::PrimInt; @@ -58,3 +58,27 @@ pub fn div_ceil(dividend: T, divisor: T) -> T { T::zero() } } + +/// Sort and merge overlapping intervals. +pub fn merge_overlapping(sections: &[Range]) -> Vec> +where + T: Ord + Clone + Copy, +{ + let mut sections = sections.to_vec(); + sections.sort_by_key(|r| (r.start, r.end)); + + let mut result = Vec::>::new(); + + for section in sections { + if let Some(last) = result.last_mut() { + if section.start <= last.end { + last.end = last.end.max(section.end); + continue; + } + } + + result.push(section); + } + + result +} diff --git a/e2e/src/main.rs b/e2e/src/main.rs index 3935258..451867e 100644 --- a/e2e/src/main.rs +++ b/e2e/src/main.rs @@ -28,6 +28,7 @@ use avbroot::{ format::{ota, payload::PayloadHeader}, protobuf::chromeos_update_engine::install_operation::Type, stream::{self, FromReader, HashingReader, PSeekFile, Reopen, SectionReader}, + util, }; use clap::Parser; use tempfile::TempDir; @@ -42,30 +43,9 @@ const DOWNLOAD_TASKS: usize = 4; const DOWNLOAD_RETRIES: u8 = 3; const DOWNLOAD_PROGRESS_INTERVAL: Duration = Duration::from_millis(50); -/// Sort and merge overlapping intervals. -fn merge_overlapping(sections: &[Range]) -> Vec> { - let mut sections = sections.to_vec(); - sections.sort_by_key(|r| (r.start, r.end)); - - let mut result = Vec::>::new(); - - for section in sections { - if let Some(last) = result.last_mut() { - if section.start <= last.end { - last.end = section.end; - continue; - } - } - - result.push(section); - } - - result -} - /// Convert an exclusion list into an inclusion list in the range [start, end). fn exclusion_to_inclusion(holes: &[Range], file_range: Range) -> Result>> { - let exclusions = merge_overlapping(holes); + let exclusions = util::merge_overlapping(holes); if let (Some(first), Some(last)) = (exclusions.first(), exclusions.last()) { if first.start < file_range.start || last.end > file_range.end {