diff --git a/bmap-parser/src/bmap.rs b/bmap-parser/src/bmap.rs index a294acc..0ce795a 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, } @@ -100,6 +101,11 @@ impl Bmap { self.blockmap.iter() } + /// checksum of the .bmap xml file + pub fn bmap_file_checksum(&self) -> String { + self.bmap_file_checksum.clone() + } + /// Total mapped size in bytes pub fn total_mapped_size(&self) -> u64 { self.block_size * self.mapped_blocks @@ -118,6 +124,8 @@ pub enum BmapBuilderError { MissingMappedBlocks, #[error("Checksum type missing")] MissingChecksumType, + #[error("Bmap file checksum missing")] + MissingBmapFileChecksum, #[error("No block ranges")] NoBlockRanges, } @@ -129,6 +137,7 @@ pub struct BmapBuilder { blocks: Option, checksum_type: Option, mapped_blocks: Option, + bmap_file_checksum: Option, blockmap: Vec, } @@ -158,6 +167,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"); @@ -186,6 +200,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 { @@ -194,6 +211,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 e0f1d2d..0dd32e2 100644 --- a/bmap-parser/src/bmap/xml.rs +++ b/bmap-parser/src/bmap/xml.rs @@ -108,6 +108,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 8f2dc78..194dc99 100644 --- a/bmap-rs/Cargo.toml +++ b/bmap-rs/Cargo.toml @@ -23,3 +23,5 @@ tokio = { version = "1.21.2", features = ["rt", "macros", "fs", "rt-multi-thread reqwest = { version = "0.12.4", 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 b8cdcc9..5621dc2 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})") @@ -203,6 +222,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)