Bmap file integrity check

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<BlockRange>,
 }
 
@@ -92,6 +93,11 @@ impl Bmap {
     pub fn block_map(&self) -> impl ExactSizeIterator + Iterator<Item = &BlockRange> {
         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<u64>,
     checksum_type: Option<HashType>,
     mapped_blocks: Option<u64>,
+    bmap_file_checksum: Option<String>,
     blockmap: Vec<BlockRange>,
 }
 
@@ -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,
         })
     }

bmap-parser/src/bmap/xml.rs

@@ -96,6 +96,7 @@ pub(crate) fn from_xml(xml: &str) -> Result<crate::bmap::Bmap, XmlError> {
         .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 {

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"

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<Response> {
     }
 }
 
+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)