From 7ea8e8fa9e8f8dca935d4d1b615d21e98538e83a Mon Sep 17 00:00:00 2001 From: Andrew Gunnerson Date: Sun, 17 Dec 2023 20:36:51 -0500 Subject: [PATCH] Split dm-verify hash tree logic out of avb module * Refactor hash tree computation to work on a preallocated hash tree buffer. This makes it possible to partially update a hash tree, which is now supported. * Add new subcommands for working with hash trees. There's no standard header format for dm-verity information, so these commands write hash tree files with a custom header. The commands are not really useful outside of debugging avbroot's hash tree implementation. Using AVB was considered, but it has no support for the hash tree data living in a separate file from the input. If other parties agree on a standard header in the future, avbroot will switch to that format. * Add tests for the hash tree implementation. Signed-off-by: Andrew Gunnerson --- README.extra.md | 48 ++- avbroot/src/cli/args.rs | 4 +- avbroot/src/cli/avb.rs | 5 +- avbroot/src/cli/hashtree.rs | 176 ++++++++ avbroot/src/cli/mod.rs | 1 + avbroot/src/format/avb.rs | 219 +--------- avbroot/src/format/fec.rs | 2 +- avbroot/src/format/hashtree.rs | 767 +++++++++++++++++++++++++++++++++ avbroot/src/format/mod.rs | 1 + 9 files changed, 1020 insertions(+), 203 deletions(-) create mode 100644 avbroot/src/cli/hashtree.rs create mode 100644 avbroot/src/format/hashtree.rs diff --git a/README.extra.md b/README.extra.md index ce59683..e3a7169 100644 --- a/README.extra.md +++ b/README.extra.md @@ -178,7 +178,7 @@ The number of parity bytes (between 2 and 24, inclusive) can be configured using ### Updating FEC data ```bash -avbroot fec update -i -f [-r ]... +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. @@ -202,3 +202,49 @@ avbroot fec repair -i -f This will repair the file in place. As described above, in each column, up to `parity / 2` bytes can be corrected. Note that FEC is **not** a replacement for checksums, like SHA-256. When there are too many errors, the file can potentially be "successfully repaired" to some incorrect data. + +## `avbroot hash-tree` + +This set of commands is for working with dm-verity hash tree data. They are not especially useful outside of debugging avbroot itself because the output format is custom. There is a custom header that sits in front of the standard dm-verity hash tree data. + +| Offsets | Type | Description | +|------------|--------|--------------------------------| +| 0..16 | ASCII | `avbroot!hashtree` magic bytes | +| 16..18 | U16LE | Version (currently 1) | +| 18..26 | U64LE | Image size | +| 26..30 | U32LE | Block size | +| 30..46 | ASCII | Hash algorithm | +| 46..48 | U16LE | Salt size | +| 48..50 | U16LE | Root digest size | +| 50..54 | U32LE | Hash tree size | +| (Variable) | BINARY | Salt | +| (Variable) | BINARY | Root digest | +| (Variable) | BINARY | Hash tree | + +For more information on the hash tree data, see the [Linux kernel documentation](https://docs.kernel.org/admin-guide/device-mapper/verity.html#hash-tree) or avbroot's implementation in [`hashtree.rs`](./avbroot/src/format/hashtree.rs). + +### Generating hash tree + +```bash +avbroot hash-tree generate -i -H +``` + +The default behavior is to use a block size of 4096, the `sha256` algorithm, and an empty salt. These can be changed with the `-b`, `-a`, and `-s` options, respectively. + +All parameters needed for verification are included in the hash tree file's header. + +### Updating hash tree + +```bash +avbroot hash-tree update -i -H [-r ]... +``` + +This will update the hash tree data corresponding to the specified regions. This can be significantly faster than generating new hash tree data from scratch for large files if the regions where data was modified are known. + +### Verifying a file + +```bash +avbroot hash-tree verify -i -H +``` + +This will check if the input file has any corrupted blocks. Currently, the command cannot report which specific blocks are corrupted, only whether the file is valid. diff --git a/avbroot/src/cli/args.rs b/avbroot/src/cli/args.rs index 45b7701..90acb83 100644 --- a/avbroot/src/cli/args.rs +++ b/avbroot/src/cli/args.rs @@ -8,7 +8,7 @@ use std::sync::atomic::AtomicBool; use anyhow::Result; use clap::{Parser, Subcommand}; -use crate::cli::{avb, boot, completion, cpio, fec, key, ota}; +use crate::cli::{avb, boot, completion, cpio, fec, hashtree, key, ota}; #[allow(clippy::large_enum_variant)] #[derive(Debug, Subcommand)] @@ -18,6 +18,7 @@ pub enum Command { Completion(completion::CompletionCli), Cpio(cpio::CpioCli), Fec(fec::FecCli), + HashTree(hashtree::HashTreeCli), Key(key::KeyCli), Ota(ota::OtaCli), /// (Deprecated: Use `avbroot ota patch` instead.) @@ -44,6 +45,7 @@ pub fn main(cancel_signal: &AtomicBool) -> Result<()> { Command::Completion(c) => completion::completion_main(&c), Command::Cpio(c) => cpio::cpio_main(&c, cancel_signal), Command::Fec(c) => fec::fec_main(&c, cancel_signal), + Command::HashTree(c) => hashtree::hash_tree_main(&c, cancel_signal), Command::Key(c) => key::key_main(&c), Command::Ota(c) => ota::ota_main(&c, cancel_signal), // Deprecated aliases. diff --git a/avbroot/src/cli/avb.rs b/avbroot/src/cli/avb.rs index 0358378..120b6a0 100644 --- a/avbroot/src/cli/avb.rs +++ b/avbroot/src/cli/avb.rs @@ -512,10 +512,7 @@ fn verify_and_repair( status!("Verifying hash tree descriptor{suffix}"); match d.verify(&file, cancel_signal) { - Err( - e @ avb::Error::InvalidRootDigest { .. } - | e @ avb::Error::InvalidHashTree { .. }, - ) if repair => { + Err(e @ avb::Error::HashTree(_)) if repair => { warning!("Failed to verify hash tree descriptor{suffix}: {e}"); warning!("Attempting to repair using FEC data{suffix}"); diff --git a/avbroot/src/cli/hashtree.rs b/avbroot/src/cli/hashtree.rs new file mode 100644 index 0000000..e4b782a --- /dev/null +++ b/avbroot/src/cli/hashtree.rs @@ -0,0 +1,176 @@ +/* + * SPDX-FileCopyrightText: 2023 Andrew Gunnerson + * SPDX-License-Identifier: GPL-3.0-only + */ + +use std::{ + fs::{File, OpenOptions}, + io::{BufReader, BufWriter, Write}, + path::{Path, PathBuf}, + sync::atomic::AtomicBool, +}; + +use anyhow::{Context, Result}; +use clap::{Parser, Subcommand}; + +use crate::{ + format::hashtree::HashTreeImage, + stream::{FromReader, PSeekFile, ToWriter}, +}; + +fn open_input(path: &Path, rw: bool) -> Result { + OpenOptions::new() + .read(true) + .write(rw) + .open(path) + .map(PSeekFile::new) + .with_context(|| format!("Failed to open file: {path:?}")) +} + +fn read_hash_tree(path: &Path) -> Result { + let reader = File::open(path) + .map(BufReader::new) + .with_context(|| format!("Failed to open for reading: {path:?}"))?; + let hash_tree = HashTreeImage::from_reader(reader) + .with_context(|| format!("Failed to read hash tree data: {path:?}"))?; + + Ok(hash_tree) +} + +fn write_hash_tree(path: &Path, hash_tree: &HashTreeImage) -> Result<()> { + let mut writer = File::create(path) + .map(BufWriter::new) + .with_context(|| format!("Failed to open for writing: {path:?}"))?; + hash_tree + .to_writer(&mut writer) + .with_context(|| format!("Failed to write hash tree data: {path:?}"))?; + writer + .flush() + .with_context(|| format!("Failed to flush hash tree data: {path:?}"))?; + + Ok(()) +} + +fn generate_subcommand(cli: &GenerateCli, cancel_signal: &AtomicBool) -> Result<()> { + let salt = hex::decode(&cli.salt).context("Invalid salt")?; + let input = open_input(&cli.input, false)?; + + let hash_tree = + HashTreeImage::generate(&input, cli.block_size, &cli.algorithm, &salt, cancel_signal) + .context("Failed to generate hash tree data")?; + + write_hash_tree(&cli.hash_tree, &hash_tree)?; + + 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 hash_tree = read_hash_tree(&cli.hash_tree)?; + + hash_tree + .update(&input, &ranges, cancel_signal) + .context("Failed to update hash tree data")?; + + write_hash_tree(&cli.hash_tree, &hash_tree)?; + + Ok(()) +} + +fn verify_subcommand(cli: &VerifyCli, cancel_signal: &AtomicBool) -> Result<()> { + let input = open_input(&cli.input, false)?; + let hash_tree = read_hash_tree(&cli.hash_tree)?; + + hash_tree + .verify(&input, cancel_signal) + .context("Failed to verify data")?; + + Ok(()) +} + +pub fn hash_tree_main(cli: &HashTreeCli, cancel_signal: &AtomicBool) -> Result<()> { + match &cli.command { + HashTreeCommand::Generate(c) => generate_subcommand(c, cancel_signal), + HashTreeCommand::Update(c) => update_subcommand(c, cancel_signal), + HashTreeCommand::Verify(c) => verify_subcommand(c, cancel_signal), + } +} + +/// Generate hash tree data for a file. +#[derive(Debug, Parser)] +struct GenerateCli { + /// Path to input data. + #[arg(short, long, value_name = "FILE", value_parser)] + input: PathBuf, + + /// Path to output hash tree data. + #[arg(short = 'H', long, value_name = "FILE", value_parser)] + hash_tree: PathBuf, + + /// Block size. + #[arg(short, long, value_name = "BYTES", default_value = "4096")] + block_size: u32, + + /// Hash algorithm. + #[arg(short, long, value_name = "NAME", default_value = "sha256")] + algorithm: String, + + /// Salt (in hex). + #[arg(short, long, value_name = "HEX", default_value = "")] + salt: String, +} + +/// Update hash tree 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 hash tree data. + /// + /// The file will be modified in place. + #[arg(short = 'H', long, value_name = "FILE", value_parser)] + hash_tree: 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 { + /// Path to input data. + #[arg(short, long, value_name = "FILE", value_parser)] + input: PathBuf, + + /// Path to input hash tree data. + #[arg(short = 'H', long, value_name = "FILE", value_parser)] + hash_tree: PathBuf, +} + +#[derive(Debug, Subcommand)] +enum HashTreeCommand { + Generate(GenerateCli), + Update(UpdateCli), + Verify(VerifyCli), +} + +/// Generate dm-verity hash tree data and verify files. +/// +/// These commands operate on a standard hash tree data prepended by a custom +/// header. +#[derive(Debug, Parser)] +pub struct HashTreeCli { + #[command(subcommand)] + command: HashTreeCommand, +} diff --git a/avbroot/src/cli/mod.rs b/avbroot/src/cli/mod.rs index 20fa26b..c11cb23 100644 --- a/avbroot/src/cli/mod.rs +++ b/avbroot/src/cli/mod.rs @@ -9,6 +9,7 @@ pub mod boot; pub mod completion; pub mod cpio; pub mod fec; +pub mod hashtree; pub mod key; pub mod ota; diff --git a/avbroot/src/format/avb.rs b/avbroot/src/format/avb.rs index ae1c025..b39920e 100644 --- a/avbroot/src/format/avb.rs +++ b/avbroot/src/format/avb.rs @@ -14,7 +14,6 @@ use bstr::ByteSlice; use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; use num_bigint_dig::{ModInverse, ToBigInt}; use num_traits::{Pow, ToPrimitive}; -use rayon::prelude::{IntoParallelIterator, ParallelIterator}; use ring::digest::{Algorithm, Context}; use rsa::{traits::PublicKeyParts, BigUint, Pkcs1v15Sign, RsaPrivateKey, RsaPublicKey}; use serde::{Deserialize, Serialize}; @@ -25,6 +24,7 @@ use crate::{ escape, format::{ fec::{self, Fec}, + hashtree::{self, HashTree}, padding, }, stream::{ @@ -92,10 +92,6 @@ pub enum Error { }, #[error("RSA key size (0) is not compatible with any AVB signing algorithm")] UnsupportedKey(usize), - #[error("Expected root digest {expected}, but have {actual}")] - InvalidRootDigest { expected: String, actual: String }, - #[error("Expected hash tree {expected}, but have {actual}")] - InvalidHashTree { expected: String, actual: String }, #[error("Hash tree does not immediately follow image data")] HashTreeGap, #[error("FEC data does not immediately follow hash tree")] @@ -112,6 +108,8 @@ pub enum Error { RsaVerify(#[source] rsa::Error), #[error("{0} byte image size is too small to fit header or footer")] ImageSizeTooSmall(u64), + #[error("Hash tree error")] + HashTree(#[from] hashtree::Error), #[error("FEC error")] Fec(#[from] fec::Error), #[error("I/O error")] @@ -120,7 +118,7 @@ pub enum Error { type Result = std::result::Result; -fn ring_algorithm(name: &str, for_verify: bool) -> Result<&'static Algorithm> { +pub(crate) fn ring_algorithm(name: &str, for_verify: bool) -> Result<&'static Algorithm> { match name { "sha1" if for_verify => Ok(&ring::digest::SHA1_FOR_LEGACY_USE_ONLY), "sha256" => Ok(&ring::digest::SHA256), @@ -382,156 +380,6 @@ impl HashTreeDescriptor { pub const FLAG_DO_NOT_USE_AB: u32 = 1 << 0; pub const FLAG_CHECK_AT_MOST_ONCE: u32 = 1 << 1; - /// Calculate the hash tree digests for a single level of the tree. If the - /// reader's position is block-aligned and `image_size` is a multiple of the - /// block size, then this function can also be used to calculate the digests - /// for a portion of a level. - /// - /// NOTE: The result is **not** padded to the block size. - fn hash_one_level( - mut reader: impl Read, - mut image_size: u64, - block_size: u32, - algorithm: &'static Algorithm, - salt: &[u8], - cancel_signal: &AtomicBool, - ) -> io::Result> { - // Each digest must be a power of 2. - let digest_padding = algorithm.output_len().next_power_of_two() - algorithm.output_len(); - let mut buf = vec![0u8; block_size as usize]; - let mut result = vec![]; - - while image_size > 0 { - stream::check_cancel(cancel_signal)?; - - let n = image_size.min(buf.len() as u64) as usize; - reader.read_exact(&mut buf[..n])?; - - // For undersized blocks, we still hash the whole buffer, except - // with padding. - buf[n..].fill(0); - - let mut context = Context::new(algorithm); - context.update(salt); - context.update(&buf); - - // Add the digest to the tree level. Each tree node must be a power - // of two. - let digest = context.finish(); - result.extend(digest.as_ref()); - result.resize(result.len() + digest_padding, 0); - - image_size -= n as u64; - } - - Ok(result) - } - - /// Calls [`Self::hash_one_level()`] in parallel. - /// - /// NOTE: The result is **not** padded to the block size. - fn hash_one_level_parallel( - input: &(dyn ReadSeekReopen + Sync), - image_size: u64, - block_size: u32, - algorithm: &'static Algorithm, - salt: &[u8], - cancel_signal: &AtomicBool, - ) -> io::Result> { - assert!( - image_size > block_size as u64, - "Images smaller than block size must use a normal hash", - ); - - // Parallelize in 16 MiB chunks to avoid too much seek thrashing. - let chunk_size = padding::round(16 * 1024 * 1024, u64::from(block_size)).unwrap(); - let chunk_count = image_size / chunk_size + u64::from(image_size % chunk_size != 0); - - let pieces = (0..chunk_count) - .into_par_iter() - .map(|c| -> io::Result> { - let start = c * chunk_size; - let size = chunk_size.min(image_size - start); - - let mut reader = input.reopen_boxed()?; - reader.seek(SeekFrom::Start(start))?; - - Self::hash_one_level(reader, size, block_size, algorithm, salt, cancel_signal) - }) - .collect::>>()?; - - Ok(pieces.into_iter().flatten().collect()) - } - - /// Calculate the hash tree for the given input in parallel. - fn calculate_hash_tree( - input: &(dyn ReadSeekReopen + Sync), - image_size: u64, - block_size: u32, - algorithm: &'static Algorithm, - salt: &[u8], - cancel_signal: &AtomicBool, - ) -> io::Result<(Vec, Vec)> { - // Small files are hashed directly, exactly like a hash descriptor. - if image_size <= u64::from(block_size) { - let mut reader = input.reopen_boxed()?; - let mut buf = vec![0u8; image_size as usize]; - reader.read_exact(&mut buf)?; - - let mut context = Context::new(algorithm); - context.update(salt); - context.update(&buf); - let digest = context.finish(); - - return Ok((digest.as_ref().to_vec(), vec![])); - } - - // Large files use the hash tree. - let mut levels = Vec::>::new(); - let mut level_size = image_size; - - while level_size > u64::from(block_size) { - let mut level = if let Some(prev_level) = levels.last() { - // Hash the previous level. - Self::hash_one_level( - Cursor::new(prev_level), - level_size, - block_size, - algorithm, - salt, - cancel_signal, - )? - } else { - // Initially read from file. - Self::hash_one_level_parallel( - input, - level_size, - block_size, - algorithm, - salt, - cancel_signal, - )? - }; - - // Pad to the block size. - level.resize(padding::round(level.len(), block_size as usize).unwrap(), 0); - - level_size = level.len() as u64; - levels.push(level); - } - - // Calculate the root hash. - let mut context = Context::new(algorithm); - context.update(salt); - context.update(levels.last().unwrap()); - let root_hash = context.finish().as_ref().to_vec(); - - // The tree is oriented such that the leaves are at the end. - let hash_tree = levels.into_iter().rev().flatten().collect(); - - Ok((root_hash, hash_tree)) - } - /// Ensure that the image data is immediately followed by the hash tree and /// then the FEC data. fn check_offsets(&self) -> Result<()> { @@ -599,25 +447,20 @@ impl HashTreeDescriptor { cancel_signal: &AtomicBool, ) -> Result<()> { let algorithm = ring_algorithm(&self.hash_algorithm, false)?; - let (root_digest, hash_tree) = Self::calculate_hash_tree( - input, - self.image_size, - self.data_block_size, - algorithm, - &self.salt, - cancel_signal, - )?; + let hash_tree = HashTree::new(self.data_block_size, algorithm, &self.salt); + let (root_digest, hash_tree_data) = + hash_tree.generate(input, self.image_size, cancel_signal)?; - if hash_tree.len() > HASH_TREE_MAX_SIZE as usize { + if hash_tree_data.len() > HASH_TREE_MAX_SIZE as usize { return Err(Error::FieldOutOfBounds("tree_size")); } - let tree_size = hash_tree.len() as u64; + let tree_size = hash_tree_data.len() as u64; let mut writer = output.reopen_boxed()?; writer.seek(SeekFrom::Start(self.image_size))?; writer - .write_all(&hash_tree) + .write_all(&hash_tree_data) .map_err(|e| Error::WriteFieldError("hash_tree", e))?; // The FEC data section is optional. @@ -675,40 +518,23 @@ impl HashTreeDescriptor { return Err(Error::FieldOutOfBounds("tree_size")); } - let (actual_root_digest, actual_hash_tree) = Self::calculate_hash_tree( - input, - self.image_size, - self.data_block_size, - algorithm, - &self.salt, - cancel_signal, - )?; - - if self.root_digest != actual_root_digest { - return Err(Error::InvalidRootDigest { - expected: hex::encode(&self.root_digest), - actual: hex::encode(actual_root_digest), - }); - } - let mut reader = input.reopen_boxed()?; reader.seek(SeekFrom::Start(self.tree_offset))?; - let mut hash_tree = vec![0u8; self.tree_size as usize]; + let mut hash_tree_data = vec![0u8; self.tree_size as usize]; reader - .read_exact(&mut hash_tree) + .read_exact(&mut hash_tree_data) .map_err(|e| Error::ReadFieldError("hash_tree", e))?; - if hash_tree != actual_hash_tree { - // These are multiple megabytes, so only report the hashes. - let expected = ring::digest::digest(algorithm, &hash_tree); - let actual = ring::digest::digest(algorithm, &actual_hash_tree); + let hash_tree = HashTree::new(self.data_block_size, algorithm, &self.salt); - return Err(Error::InvalidHashTree { - expected: hex::encode(expected), - actual: hex::encode(actual), - }); - } + hash_tree.verify( + input, + self.image_size, + &self.root_digest, + &hash_tree_data, + cancel_signal, + )?; // The FEC data section is optional. if self.fec_num_roots != 0 { @@ -945,10 +771,11 @@ impl HashDescriptor { let digest = self.calculate(reader, true, cancel_signal)?; if self.root_digest != digest.as_ref() { - return Err(Error::InvalidRootDigest { + return Err(hashtree::Error::InvalidRootDigest { expected: hex::encode(&self.root_digest), actual: hex::encode(digest), - }); + } + .into()); } Ok(()) diff --git a/avbroot/src/format/fec.rs b/avbroot/src/format/fec.rs index 044a56d..7724d93 100644 --- a/avbroot/src/format/fec.rs +++ b/avbroot/src/format/fec.rs @@ -590,7 +590,7 @@ pub struct FecImage { impl fmt::Debug for FecImage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Fec") + f.debug_struct("FecImage") .field("fec", &NumBytes(self.fec.len())) .field("data_size", &self.data_size) .field("parity", &self.parity) diff --git a/avbroot/src/format/hashtree.rs b/avbroot/src/format/hashtree.rs new file mode 100644 index 0000000..dc4ce7d --- /dev/null +++ b/avbroot/src/format/hashtree.rs @@ -0,0 +1,767 @@ +/* + * SPDX-FileCopyrightText: 2023 Andrew Gunnerson + * SPDX-License-Identifier: GPL-3.0-only + */ + +use std::{ + cmp::Ordering, + fmt, + io::{self, Cursor, Read, SeekFrom, Write}, + ops::Range, + sync::atomic::AtomicBool, +}; + +use bstr::ByteSlice; +use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; +use num_traits::ToPrimitive; +use rayon::{ + iter::{IndexedParallelIterator, ParallelIterator}, + slice::ParallelSliceMut, +}; +use ring::digest::{Algorithm, Context}; +use thiserror::Error; + +use crate::{ + format::{avb, padding}, + stream::{self, FromReader, ReadSeekReopen, ReadStringExt, ToWriter, WriteStringExt}, + util::{self, NumBytes}, +}; + +// TODO: io::Error +#[derive(Debug, Error)] +pub enum Error { + #[error("Hash tree should have size {expected} for input size {input}, but has size {actual}")] + InvalidHashTreeSize { + input: u64, + expected: usize, + actual: usize, + }, + #[error("Expected root digest {expected}, but have {actual}")] + InvalidRootDigest { expected: String, actual: String }, + #[error("Expected hash tree {expected}, but have {actual}")] + InvalidHashTree { expected: String, actual: String }, + #[error("Invalid hash tree header magic: {:?}", .0.as_bstr())] + InvalidHeaderMagic([u8; 16]), + #[error("Invalid hash tree header version: {0}")] + InvalidHeaderVersion(u16), + #[error("Hashing algorithm not supported: {0:?}")] + UnsupportedHashAlgorithm(String), + #[error("{0:?} field is out of bounds")] + FieldOutOfBounds(&'static str), + #[error("I/O error")] + Io(#[from] io::Error), +} + +type Result = std::result::Result; + +pub struct HashTree<'a> { + block_size: u32, + algorithm: &'static Algorithm, + salt: &'a [u8], +} + +impl<'a> HashTree<'a> { + pub fn new(block_size: u32, algorithm: &'static Algorithm, salt: &'a [u8]) -> Self { + Self { + block_size, + algorithm, + salt, + } + } + + /// Compute the list of offset ranges that each level occupies in the hash + /// tree data. The items are returned with the bottom level's offsets first + /// in the list. Note that the bottom level is stored at the end of the hash + /// tree data. + fn compute_level_offsets(&self, image_size: u64) -> Result>> { + let digest_size = self.algorithm.output_len().next_power_of_two(); + let mut ranges = vec![]; + let mut level_size = image_size; + + while level_size > u64::from(self.block_size) { + let blocks = util::div_ceil(level_size, u64::from(self.block_size)); + level_size = blocks + .checked_mul(digest_size as u64) + .and_then(|s| padding::round(s, u64::from(self.block_size))) + .ok_or_else(|| Error::FieldOutOfBounds("level_size"))?; + + // Depending on the chosen block size, the original file size could + // overflow a usize without the first level's size doing the same. + let level_size_usize = level_size + .to_usize() + .ok_or_else(|| Error::FieldOutOfBounds("level_size"))?; + + ranges.push(0..level_size_usize); + } + + // The hash tree puts the leaves at the end. + let mut offset = 0; + for range in ranges.iter_mut().rev() { + let level_size = range.end - range.start; + range.start += offset; + range.end += offset; + offset += level_size; + } + + Ok(ranges) + } + + /// Check if a sorted list of block ranges contains a block. + fn block_range_contains(ranges: &[Range], block: u64) -> bool { + ranges + .binary_search_by(|range| { + if range.start > block { + Ordering::Greater + } else if range.end <= block { + Ordering::Less + } else { + Ordering::Equal + } + }) + .is_ok() + } + + /// Convert a list of ranges of byte offsets to a sorted, non-overlapping + /// list of block ranges. + fn blocks_for_ranges(&self, image_size: u64, ranges: &[Range]) -> Result>> { + let ranges = util::merge_overlapping(ranges); + if let Some(last) = ranges.last() { + if last.end > image_size { + return Err(Error::FieldOutOfBounds("ranges")); + } + } + + let block_size = u64::from(self.block_size); + let mut result = Vec::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) + }; + + result.push(start_block..end_block); + } + + Ok(util::merge_overlapping(&result)) + } + + /// Calculate the hash tree digests for a single level of the tree. If the + /// reader's position is block-aligned and `image_size` is a multiple of the + /// block size, then this function can also be used to calculate the digests + /// for a portion of a level. + fn hash_partial_level( + &self, + mut reader: impl Read, + mut size: u64, + mut level_data: &mut [u8], + cancel_signal: &AtomicBool, + ) -> io::Result<()> { + // Each digest must be a power of 2. + let digest_padding = + self.algorithm.output_len().next_power_of_two() - self.algorithm.output_len(); + let mut buf = vec![0u8; self.block_size as usize]; + + while size > 0 { + stream::check_cancel(cancel_signal)?; + + let n = size.min(buf.len() as u64) as usize; + reader.read_exact(&mut buf[..n])?; + + // For undersized blocks, we still hash the whole buffer, except + // with padding. + buf[n..].fill(0); + + let mut context = Context::new(self.algorithm); + context.update(self.salt); + context.update(&buf); + + // Add the digest to the tree level. Each tree node must be a power + // of two. + let digest = context.finish(); + + level_data[..digest.as_ref().len()].copy_from_slice(digest.as_ref()); + level_data = &mut level_data[digest.as_ref().len()..]; + + level_data[..digest_padding].fill(0); + level_data = &mut level_data[digest_padding..]; + + size -= n as u64; + } + + Ok(()) + } + + /// Hash one full level in parallel. + fn hash_one_level_parallel( + &self, + input: &(dyn ReadSeekReopen + Sync), + size: u64, + level_data: &mut [u8], + cancel_signal: &AtomicBool, + ) -> io::Result<()> { + assert!( + size > self.block_size as u64, + "Images smaller than block size must use a normal hash", + ); + + // Parallelize in larger chunks to avoid too much seek thrashing. + let digest_size = self.algorithm.output_len().next_power_of_two(); + let multiplier = 1024u64; + + level_data + .par_chunks_mut(digest_size * multiplier as usize) + .enumerate() + .map(|(chunk, out_data)| -> io::Result<()> { + let digests = out_data.len() / digest_size; + let in_start = (chunk as u64) * multiplier * u64::from(self.block_size); + let in_size = ((digests as u64) * u64::from(self.block_size)).min(size - in_start); + + let mut reader = input.reopen_boxed()?; + reader.seek(SeekFrom::Start(in_start))?; + + self.hash_partial_level(reader, in_size, out_data, cancel_signal) + }) + .collect::>()?; + + Ok(()) + } + + /// Update parts of the hash tree level corresponding to the specified + /// blocks. + fn hash_partial_level_parallel( + &self, + input: &(dyn ReadSeekReopen + Sync), + size: u64, + block_ranges: &[Range], + level_data: &mut [u8], + cancel_signal: &AtomicBool, + ) -> io::Result<()> { + let digest_size = self.algorithm.output_len().next_power_of_two(); + + level_data + .par_chunks_exact_mut(digest_size) + .enumerate() + .filter(|(chunk, _)| Self::block_range_contains(block_ranges, *chunk as u64)) + .map(|(chunk, out_data)| -> io::Result<()> { + let in_start = (chunk as u64) * u64::from(self.block_size); + let in_size = u64::from(self.block_size).min(size - in_start); + + let mut reader = input.reopen_boxed()?; + reader.seek(SeekFrom::Start(in_start))?; + + self.hash_partial_level(reader, in_size, out_data, cancel_signal) + }) + .collect::>()?; + + Ok(()) + } + + /// Compute the hash tree and return the root digest. If `ranges` is + /// specified, then only the input file blocks containing those ranges are + /// recomputed. + /// + /// `hash_tree_data` must match `level_offsets`. In other words, the ending + /// offset of the leaf layer of the tree must equal `hash_tree_data`'s size. + fn calculate( + &self, + input: &(dyn ReadSeekReopen + Sync), + image_size: u64, + ranges: Option<&[Range]>, + level_offsets: &[Range], + hash_tree_data: &mut [u8], + cancel_signal: &AtomicBool, + ) -> Result> { + // Small files are hashed directly. + if image_size <= u64::from(self.block_size) { + let mut reader = input.reopen_boxed()?; + let mut buf = vec![0u8; image_size as usize]; + reader.read_exact(&mut buf)?; + + let mut context = Context::new(self.algorithm); + context.update(self.salt); + context.update(&buf); + let digest = context.finish(); + + return Ok(digest.as_ref().to_vec()); + } + + // Large files use the hash tree. + for (i, level_range) in level_offsets.iter().enumerate() { + let (front, back) = hash_tree_data.split_at_mut(level_range.end); + let level_data = &mut front[level_range.clone()]; + + if i > 0 { + // Hash the previous level. + let prev_range = level_offsets[i - 1].clone(); + let prev_size = prev_range.end - prev_range.start; + let prev_data = &back[..prev_size]; + + self.hash_partial_level( + Cursor::new(prev_data), + prev_size as u64, + level_data, + cancel_signal, + )?; + } else if let Some(r) = ranges { + // Read partial blocks from file. + let block_ranges = self.blocks_for_ranges(image_size, r)?; + + self.hash_partial_level_parallel( + input, + image_size, + &block_ranges, + level_data, + cancel_signal, + )?; + } else { + // Read entire file. + self.hash_one_level_parallel(input, image_size, level_data, cancel_signal)?; + } + + // No need to explicitly ensure the level is padded to the block + // size since the tree is initialized with zeros. + } + + // Calculate the root hash. + let mut context = Context::new(self.algorithm); + context.update(self.salt); + context.update(&hash_tree_data[level_offsets.last().unwrap().clone()]); + let root_hash = context.finish().as_ref().to_vec(); + + Ok(root_hash) + } + + /// Generate hash tree data for the file. Returns the root digest and the + /// hash tree data. + pub fn generate( + &self, + input: &(dyn ReadSeekReopen + Sync), + image_size: u64, + cancel_signal: &AtomicBool, + ) -> Result<(Vec, Vec)> { + let offsets = self.compute_level_offsets(image_size)?; + let hash_tree_size = offsets.get(0).map(|r| r.end).unwrap_or(0); + let mut hash_tree_data = vec![0u8; hash_tree_size]; + + let root_digest = self.calculate( + input, + image_size, + None, + &offsets, + &mut hash_tree_data, + cancel_signal, + )?; + + Ok((root_digest, hash_tree_data)) + } + + /// Update hash tree data corresponding to the specified file ranges. + /// Returns the new root digest. + pub fn update( + &self, + input: &(dyn ReadSeekReopen + Sync), + image_size: u64, + ranges: &[Range], + hash_tree_data: &mut [u8], + cancel_signal: &AtomicBool, + ) -> Result> { + let offsets = self.compute_level_offsets(image_size)?; + let hash_tree_size = offsets.get(0).map(|r| r.end).unwrap_or(0); + if hash_tree_data.len() != hash_tree_size { + return Err(Error::InvalidHashTreeSize { + input: image_size, + expected: hash_tree_size, + actual: hash_tree_data.len(), + }); + } + + self.calculate( + input, + image_size, + Some(ranges), + &offsets, + hash_tree_data, + cancel_signal, + ) + } + + /// Verify that the file contains no errors. + pub fn verify( + &self, + input: &(dyn ReadSeekReopen + Sync), + image_size: u64, + root_digest: &[u8], + hash_tree_data: &[u8], + cancel_signal: &AtomicBool, + ) -> Result<()> { + let offsets = self.compute_level_offsets(image_size)?; + let hash_tree_size = offsets.get(0).map(|r| r.end).unwrap_or(0); + if hash_tree_data.len() != hash_tree_size { + return Err(Error::InvalidHashTreeSize { + input: image_size, + expected: hash_tree_size, + actual: hash_tree_data.len(), + }); + } + + let (actual_root_digest, actual_hash_tree_data) = + self.generate(input, image_size, cancel_signal)?; + + if root_digest != actual_root_digest { + return Err(Error::InvalidRootDigest { + expected: hex::encode(root_digest), + actual: hex::encode(&actual_root_digest), + }); + } + + if hash_tree_data != actual_hash_tree_data { + // These are multiple megabytes, so only report the hashes. + let expected = ring::digest::digest(self.algorithm, hash_tree_data); + let actual = ring::digest::digest(self.algorithm, &actual_hash_tree_data); + + return Err(Error::InvalidHashTree { + expected: hex::encode(expected), + actual: hex::encode(actual), + }); + } + + Ok(()) + } +} + +/// A type for reading and writing a custom hash tree image format. +/// +/// File format: +/// - [0 .. 16] - ASCII - "avbroot!hashtree" +/// - [16 .. 18] - U16LE - Version +/// - [18 .. 26] - U64LE - Image size +/// - [26 .. 30] - U32LE - Block size +/// - [30 .. 46] - ASCII - Hash algorithm +/// - [46 .. 48] - U16LE - Salt size +/// - [48 .. 50] - U16LE - Root digest size +/// - [50 .. 54] - U32LE - Hash tree size +/// - [] - BINARY - Salt +/// - [] - BINARY - Root digest +/// - [] - BINARY - Hash tree +#[derive(Clone, PartialEq, Eq)] +pub struct HashTreeImage { + pub image_size: u64, + pub block_size: u32, + pub algorithm: String, + pub salt: Vec, + pub root_digest: Vec, + pub hash_tree: Vec, +} + +impl fmt::Debug for HashTreeImage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HashTreeImage") + .field("block_size", &self.block_size) + .field("algorithm", &self.algorithm) + .field("salt", &hex::encode(&self.salt)) + .field("root_digest", &hex::encode(&self.root_digest)) + .field("hash_tree", &NumBytes(self.hash_tree.len())) + .finish() + } +} + +impl HashTreeImage { + const MAGIC: &'static [u8; 16] = b"avbroot!hashtree"; + const VERSION: u16 = 1; + + pub fn ring_algorithm(name: &str) -> Result<&'static Algorithm> { + avb::ring_algorithm(name, false) + .map_err(|_| Error::UnsupportedHashAlgorithm(name.to_owned())) + } + + /// Generate hash tree data for a file. + pub fn generate( + input: &(dyn ReadSeekReopen + Sync), + block_size: u32, + algorithm: &str, + salt: &[u8], + cancel_signal: &AtomicBool, + ) -> Result { + let image_size = { + let mut file = input.reopen_boxed()?; + file.seek(SeekFrom::End(0))? + }; + let ring_algorithm = Self::ring_algorithm(algorithm)?; + let hash_tree = HashTree::new(block_size, ring_algorithm, salt); + let (root_digest, hash_tree_data) = hash_tree.generate(input, image_size, cancel_signal)?; + + Ok(Self { + image_size, + block_size, + algorithm: algorithm.to_owned(), + salt: salt.to_vec(), + root_digest, + hash_tree: hash_tree_data, + }) + } + + /// Update hash tree data coreesponding to the specified file ranges. + pub fn update( + &mut self, + input: &(dyn ReadSeekReopen + Sync), + ranges: &[Range], + cancel_signal: &AtomicBool, + ) -> Result<()> { + let ring_algorithm = Self::ring_algorithm(&self.algorithm)?; + let hash_tree = HashTree::new(self.block_size, ring_algorithm, &self.salt); + + self.root_digest = hash_tree.update( + input, + self.image_size, + ranges, + &mut self.hash_tree, + cancel_signal, + )?; + + Ok(()) + } + + /// Check that a file contains no errors. + pub fn verify( + &self, + input: &(dyn ReadSeekReopen + Sync), + cancel_signal: &AtomicBool, + ) -> Result<()> { + let ring_algorithm = Self::ring_algorithm(&self.algorithm)?; + let hash_tree = HashTree::new(self.block_size, ring_algorithm, &self.salt); + + hash_tree.verify( + input, + self.image_size, + &self.root_digest, + &self.hash_tree, + cancel_signal, + ) + } +} + +impl FromReader for HashTreeImage { + type Error = Error; + + fn from_reader(mut reader: R) -> Result { + let mut magic = [0u8; 16]; + reader.read_exact(&mut magic)?; + if magic != *Self::MAGIC { + return Err(Error::InvalidHeaderMagic(magic)); + } + + let version = reader.read_u16::()?; + if version != Self::VERSION { + return Err(Error::InvalidHeaderVersion(version)); + } + + let image_size = reader.read_u64::()?; + let block_size = reader.read_u32::()?; + let algorithm = reader.read_string_padded(16)?; + let salt_size = reader.read_u16::()?; + let root_digest_size = reader.read_u16::()?; + let hash_tree_size = reader + .read_u32::()? + .to_usize() + .ok_or_else(|| Error::FieldOutOfBounds("hash_tree_size"))?; + + let mut salt = vec![0u8; usize::from(salt_size)]; + reader.read_exact(&mut salt)?; + + let mut root_digest = vec![0u8; usize::from(root_digest_size)]; + reader.read_exact(&mut root_digest)?; + + let mut hash_tree = vec![0u8; hash_tree_size]; + reader.read_exact(&mut hash_tree)?; + + Ok(Self { + image_size, + block_size, + algorithm, + salt, + root_digest, + hash_tree, + }) + } +} + +impl ToWriter for HashTreeImage { + type Error = Error; + + fn to_writer(&self, mut writer: W) -> Result<()> { + let salt_size = self + .salt + .len() + .to_u16() + .ok_or_else(|| Error::FieldOutOfBounds("salt_size"))?; + let root_digest_size = self + .root_digest + .len() + .to_u16() + .ok_or_else(|| Error::FieldOutOfBounds("root_digest_size"))?; + let hash_tree_size = self + .hash_tree + .len() + .to_u32() + .ok_or_else(|| Error::FieldOutOfBounds("hash_tree_size"))?; + + writer.write_all(Self::MAGIC)?; + writer.write_u16::(Self::VERSION)?; + writer.write_u64::(self.image_size)?; + writer.write_u32::(self.block_size)?; + writer.write_string_padded(&self.algorithm, 16)?; + writer.write_u16::(salt_size)?; + writer.write_u16::(root_digest_size)?; + writer.write_u32::(hash_tree_size)?; + writer.write_all(&self.salt)?; + writer.write_all(&self.root_digest)?; + writer.write_all(&self.hash_tree)?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::io::{Seek, Write}; + + use assert_matches::assert_matches; + + use crate::stream::SharedCursor; + + use super::*; + + #[test] + fn calculate_level_ranges() { + let hash_tree = HashTree::new(4096, &ring::digest::SHA256, &[]); + assert_eq!( + hash_tree.compute_level_offsets(0).unwrap(), + &[] as &[Range], + ); + assert_eq!( + hash_tree.compute_level_offsets(1024 * 1024 * 1024).unwrap(), + &[69632..8458240, 4096..69632, 0..4096], + ) + } + + #[test] + fn block_range_contains() { + assert_eq!(HashTree::block_range_contains(&[0..4], 0), true); + assert_eq!(HashTree::block_range_contains(&[0..4], 4), false); + assert_eq!(HashTree::block_range_contains(&[0..4, 5..8], 4), false); + assert_eq!(HashTree::block_range_contains(&[0..4, 5..8], 6), true); + } + + #[test] + fn blocks_for_ranges() { + let hash_tree = HashTree::new(4096, &ring::digest::SHA256, b"Salt"); + assert_eq!( + hash_tree.blocks_for_ranges(16384, &[0..16384]).unwrap(), + &[0..4], + ); + assert_eq!(hash_tree.blocks_for_ranges(16384, &[0..0]).unwrap(), &[],); + assert_eq!( + hash_tree + .blocks_for_ranges(16384, &[12287..12289, 0..1, 5000..5001]) + .unwrap(), + &[0..4], + ); + assert_matches!( + hash_tree.blocks_for_ranges(16384, &[0..16385]), + Err(Error::FieldOutOfBounds(_)) + ); + } + + #[test] + fn generate_update_verify() { + let cancel_signal = AtomicBool::new(false); + let hash_tree = HashTree::new(64, &ring::digest::SHA256, b"Salt"); + let mut input = SharedCursor::new(); + + // Try input smaller than one block. + let (root_digest, hash_tree_data) = hash_tree.generate(&input, 0, &cancel_signal).unwrap(); + + assert_eq!( + root_digest, + &[ + 0x15, 0x0f, 0xe5, 0x51, 0x40, 0x30, 0xb1, 0x43, 0x4a, 0x5d, 0xea, 0xf4, 0x91, 0xec, + 0xe9, 0x2c, 0x0e, 0x64, 0x97, 0x44, 0x7d, 0x6d, 0xe7, 0xbd, 0x6b, 0xa8, 0x5e, 0x8c, + 0xae, 0x1e, 0x00, 0xa3 + ], + ); + assert_eq!(hash_tree_data, &[]); + + // Try larger input that spans multiple blocks are results in an actual + // hash tree being created. + input.write_all(&b"Data".repeat(25)).unwrap(); + + let (root_digest, mut hash_tree_data) = + hash_tree.generate(&input, 100, &cancel_signal).unwrap(); + assert_eq!( + root_digest, + &[ + 0x92, 0xc3, 0xd7, 0x4a, 0x64, 0x03, 0x4b, 0xcc, 0xa9, 0x9a, 0x44, 0xf6, 0x81, 0xa2, + 0x4d, 0xdd, 0x97, 0xd3, 0xda, 0x84, 0xdc, 0xe2, 0x1b, 0x83, 0xd1, 0x7b, 0xab, 0x60, + 0x59, 0xe8, 0x45, 0x59 + ], + ); + assert_eq!( + hash_tree_data, + &[ + 0x7e, 0x33, 0x47, 0xb6, 0xf3, 0x7c, 0xde, 0x0e, 0xe2, 0x8d, 0x9e, 0x49, 0x8e, 0xd4, + 0xbd, 0x53, 0x3a, 0xa1, 0xff, 0xeb, 0x4f, 0x6d, 0x5a, 0x5f, 0x55, 0x28, 0x37, 0x79, + 0xd0, 0x25, 0x07, 0xd5, 0xb7, 0x7f, 0x1a, 0x48, 0x92, 0x12, 0x91, 0xdb, 0x92, 0x04, + 0x74, 0xf6, 0x86, 0x31, 0xfc, 0x64, 0xb6, 0xc8, 0x72, 0xb0, 0xf7, 0x7d, 0x24, 0xa4, + 0x3c, 0x87, 0x1f, 0xc9, 0xd8, 0x17, 0x8a, 0xd9 + ], + ); + + // Change some data and update the hash tree. + input.rewind().unwrap(); + input.write_all(b"Changed").unwrap(); + + let root_digest = hash_tree + .update(&input, 100, &[0..7], &mut hash_tree_data, &cancel_signal) + .unwrap(); + assert_eq!( + root_digest, + &[ + 0x8d, 0x03, 0xad, 0x18, 0xf2, 0x53, 0x13, 0x59, 0xf5, 0xbf, 0x68, 0x0e, 0x0c, 0x4a, + 0x86, 0xe2, 0x6e, 0xaa, 0x3d, 0x4b, 0x0f, 0x1b, 0x57, 0xad, 0x92, 0xe7, 0xbf, 0x3e, + 0xa6, 0xb1, 0x2e, 0xcc + ], + ); + assert_eq!( + hash_tree_data, + &[ + 0xfe, 0x46, 0xf7, 0x8c, 0xa1, 0xd9, 0xc8, 0xdd, 0x47, 0x9e, 0x6c, 0x32, 0x7c, 0x38, + 0x7f, 0x09, 0xe1, 0x58, 0x92, 0xa3, 0xb6, 0xbd, 0x96, 0xef, 0x10, 0xe8, 0x30, 0xb0, + 0x37, 0x8d, 0xef, 0x9a, 0xb7, 0x7f, 0x1a, 0x48, 0x92, 0x12, 0x91, 0xdb, 0x92, 0x04, + 0x74, 0xf6, 0x86, 0x31, 0xfc, 0x64, 0xb6, 0xc8, 0x72, 0xb0, 0xf7, 0x7d, 0x24, 0xa4, + 0x3c, 0x87, 0x1f, 0xc9, 0xd8, 0x17, 0x8a, 0xd9 + ], + ); + + // Updated hash tree should match newly generated tree. + let (new_root_digest, new_hash_tree_data) = + hash_tree.generate(&input, 100, &cancel_signal).unwrap(); + assert_eq!(new_root_digest, root_digest); + assert_eq!(new_hash_tree_data, hash_tree_data); + + // Data should validate successfully. + hash_tree + .verify(&input, 100, &root_digest, &hash_tree_data, &cancel_signal) + .unwrap(); + + // But not if the data is corrupted. + input.rewind().unwrap(); + input.write_all(b"Bad").unwrap(); + + hash_tree + .verify(&input, 100, &root_digest, &hash_tree_data, &cancel_signal) + .unwrap_err(); + } +} diff --git a/avbroot/src/format/mod.rs b/avbroot/src/format/mod.rs index a14b650..177b8fe 100644 --- a/avbroot/src/format/mod.rs +++ b/avbroot/src/format/mod.rs @@ -8,6 +8,7 @@ pub mod bootimage; pub mod compression; pub mod cpio; pub mod fec; +pub mod hashtree; pub mod ota; pub mod padding; pub mod payload;