diff --git a/bmap-parser/src/bmap.rs b/bmap-parser/src/bmap.rs index 1ce56dd..8393b54 100644 --- a/bmap-parser/src/bmap.rs +++ b/bmap-parser/src/bmap.rs @@ -57,6 +57,7 @@ pub struct Bmap { blocks: u64, mapped_blocks: u64, checksum_type: HashType, + bmap_file_checksum: String, blockmap: Vec, } @@ -92,6 +93,11 @@ impl Bmap { pub fn block_map(&self) -> impl ExactSizeIterator + Iterator { self.blockmap.iter() } + + pub fn bmap_file_checksum(&self) -> String { + self.bmap_file_checksum.clone() + } + pub fn total_mapped_size(&self) -> u64 { self.block_size * self.mapped_blocks } @@ -109,6 +115,8 @@ pub enum BmapBuilderError { MissingMappedBlocks, #[error("Checksum type missing")] MissingChecksumType, + #[error("Bmap file checksum missing")] + MissingBmapFileChecksum, #[error("No block ranges")] NoBlockRanges, } @@ -120,6 +128,7 @@ pub struct BmapBuilder { blocks: Option, checksum_type: Option, mapped_blocks: Option, + bmap_file_checksum: Option, blockmap: Vec, } @@ -149,6 +158,11 @@ impl BmapBuilder { self } + pub fn bmap_file_checksum(&mut self, bmap_file_checksum: String) -> &mut Self { + self.bmap_file_checksum = Some(bmap_file_checksum); + self + } + pub fn add_block_range(&mut self, start: u64, end: u64, checksum: HashValue) -> &mut Self { let bs = self.block_size.expect("Blocksize needs to be set first"); let total = self.image_size.expect("Image size needs to be set first"); @@ -177,6 +191,9 @@ impl BmapBuilder { let checksum_type = self .checksum_type .ok_or(BmapBuilderError::MissingChecksumType)?; + let bmap_file_checksum = self + .bmap_file_checksum + .ok_or(BmapBuilderError::MissingBmapFileChecksum)?; let blockmap = self.blockmap; Ok(Bmap { @@ -185,6 +202,7 @@ impl BmapBuilder { blocks, mapped_blocks, checksum_type, + bmap_file_checksum, blockmap, }) } diff --git a/bmap-parser/src/bmap/xml.rs b/bmap-parser/src/bmap/xml.rs index bf7f3e4..c247c66 100644 --- a/bmap-parser/src/bmap/xml.rs +++ b/bmap-parser/src/bmap/xml.rs @@ -96,6 +96,7 @@ pub(crate) fn from_xml(xml: &str) -> Result { .block_size(b.block_size) .blocks(b.blocks_count) .checksum_type(hash_type) + .bmap_file_checksum(b.bmap_file_checksum) .mapped_blocks(b.mapped_blocks_count); for range in b.block_map.ranges { diff --git a/bmap-rs/Cargo.toml b/bmap-rs/Cargo.toml index ec056ca..b51bb59 100644 --- a/bmap-rs/Cargo.toml +++ b/bmap-rs/Cargo.toml @@ -22,3 +22,5 @@ tokio = { version = "1.21.2", features = ["rt", "macros", "fs", "rt-multi-thread reqwest = { version = "0.11.12", features = ["stream"] } tokio-util = { version = "0.7.4", features = ["compat"] } futures = "0.3.25" +sha2 = { version = "0.10.6", features = [ "asm" ] } +hex = "0.4.3" diff --git a/bmap-rs/src/main.rs b/bmap-rs/src/main.rs index 5d1f3e6..2a9ae71 100644 --- a/bmap-rs/src/main.rs +++ b/bmap-rs/src/main.rs @@ -7,6 +7,7 @@ use futures::TryStreamExt; use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use nix::unistd::ftruncate; use reqwest::{Response, Url}; +use sha2::{Digest, Sha256}; use std::ffi::OsStr; use std::fmt::Write; use std::fs::File; @@ -158,6 +159,24 @@ async fn setup_remote_input(url: Url) -> Result { } } +fn bmap_integrity(checksum: String, xml: String) -> Result<()> { + //Unset the checksum + let mut bmap_hash = Sha256::new(); + let default = "0".repeat(64); + let before_checksum = xml.replace(&checksum, &default); + + //Compare given and created checksum + bmap_hash.update(before_checksum); + let digest = bmap_hash.finalize_reset(); + let new_checksum = hex::encode(digest.as_slice()); + ensure!( + checksum == new_checksum, + "Bmap file doesn't match its checksum. It could be corrupted or compromised." + ); + println!("Bmap integrity checked!"); + Ok(()) +} + fn setup_progress_bar(bmap: &Bmap) -> ProgressBar { let pb = ProgressBar::new(bmap.total_mapped_size()); pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})") @@ -204,6 +223,7 @@ fn copy_local_input(source: PathBuf, destination: PathBuf) -> Result<()> { b.read_to_string(&mut xml)?; let bmap = Bmap::from_xml(&xml)?; + bmap_integrity(bmap.bmap_file_checksum(), xml)?; let output = std::fs::OpenOptions::new() .write(true) .create(true) @@ -229,6 +249,7 @@ async fn copy_remote_input(source: Url, destination: PathBuf) -> Result<()> { println!("Found bmap file: {}", bmap_url); let bmap = Bmap::from_xml(&xml)?; + bmap_integrity(bmap.bmap_file_checksum(), xml)?; let mut output = tokio::fs::OpenOptions::new() .write(true) .create(true)