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 <accounts+github@chiller3.com>

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 <input data file> -f <output FEC file> [-r <start> <end>]...
+```
+
+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

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::<Vec<_>>();
+
+    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<u64>,
+}
+
 /// 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),
 }

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<u64>]) -> Result<HashSet<u64>> {
+        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<u64>],
+        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::<Result<()>>()?;
+
+        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<u64>],
+        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}");

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<T: PrimInt>(dividend: T, divisor: T) -> T {
             T::zero()
         }
 }
+
+/// Sort and merge overlapping intervals.
+pub fn merge_overlapping<T>(sections: &[Range<T>]) -> Vec<Range<T>>
+where
+    T: Ord + Clone + Copy,
+{
+    let mut sections = sections.to_vec();
+    sections.sort_by_key(|r| (r.start, r.end));
+
+    let mut result = Vec::<Range<T>>::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
+}

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<u64>]) -> Vec<Range<u64>> {
-    let mut sections = sections.to_vec();
-    sections.sort_by_key(|r| (r.start, r.end));
-
-    let mut result = Vec::<Range<u64>>::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<u64>], file_range: Range<u64>) -> Result<Vec<Range<u64>>> {
-    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 {