Bmap file integrity check

Before using a Bmap file checks if its checksum is correct for the
current bmap file.
Bmap checksum is the application of Sha256 to the file data. When the
bmap file is created, the value of the checksum has to be zero (all ASCII
"0" symbols). Once calculated, zeros are replaced by the checksum, notice
this modifies the file itself.
In order to calculate the checksum before using it and compare it with
the original, we need to set the field as all "0" before applying Sha256.

Closes: #50

Signed-off-by: Rafael Garcia Ruiz <[email protected]>

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;
@@ -150,6 +151,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})")
@@ -184,6 +203,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)

bmap/src/bmap.rs

@@ -0,0 +1,223 @@
+use strum::{Display, EnumDiscriminants, EnumString};
+use thiserror::Error;
+mod xml;
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, EnumString, Display)]
+#[strum(serialize_all = "lowercase")]
+#[non_exhaustive]
+pub enum HashType {
+    Sha256,
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, EnumDiscriminants)]
+#[non_exhaustive]
+pub enum HashValue {
+    Sha256([u8; 32]),
+}
+
+impl HashValue {
+    pub fn to_type(&self) -> HashType {
+        match self {
+            HashValue::Sha256(_) => HashType::Sha256,
+        }
+    }
+
+    pub fn as_slice(&self) -> &[u8] {
+        match self {
+            HashValue::Sha256(v) => v,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct BlockRange {
+    offset: u64,
+    length: u64,
+    checksum: HashValue,
+}
+
+impl BlockRange {
+    pub fn checksum(&self) -> HashValue {
+        self.checksum
+    }
+
+    pub fn offset(&self) -> u64 {
+        self.offset
+    }
+
+    pub fn length(&self) -> u64 {
+        self.length
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Bmap {
+    image_size: u64,
+    block_size: u64,
+    blocks: u64,
+    mapped_blocks: u64,
+    checksum_type: HashType,
+    bmap_file_checksum: String,
+    blockmap: Vec<BlockRange>,
+}
+
+impl Bmap {
+    pub fn builder() -> BmapBuilder {
+        BmapBuilder::default()
+    }
+
+    pub fn from_xml(xml: &str) -> Result<Self, xml::XmlError> {
+        xml::from_xml(xml)
+    }
+
+    pub fn image_size(&self) -> u64 {
+        self.image_size
+    }
+
+    pub const fn block_size(&self) -> u64 {
+        self.block_size
+    }
+
+    pub fn blocks(&self) -> u64 {
+        self.blocks
+    }
+
+    pub fn mapped_blocks(&self) -> u64 {
+        self.mapped_blocks
+    }
+
+    pub fn checksum_type(&self) -> HashType {
+        self.checksum_type
+    }
+
+    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
+    }
+}
+
+#[derive(Clone, Debug, Error)]
+pub enum BmapBuilderError {
+    #[error("Image size missing")]
+    MissingImageSize,
+    #[error("Block size missing")]
+    MissingBlockSize,
+    #[error("Blocks missing")]
+    MissingBlocks,
+    #[error("Mapped blocks missing")]
+    MissingMappedBlocks,
+    #[error("Checksum type missing")]
+    MissingChecksumType,
+    #[error("Bmap file checksum missing")]
+    MissingBmapFileChecksum,
+    #[error("No block ranges")]
+    NoBlockRanges,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct BmapBuilder {
+    image_size: Option<u64>,
+    block_size: Option<u64>,
+    blocks: Option<u64>,
+    checksum_type: Option<HashType>,
+    mapped_blocks: Option<u64>,
+    bmap_file_checksum: Option<String>,
+    blockmap: Vec<BlockRange>,
+}
+
+impl BmapBuilder {
+    pub fn image_size(&mut self, size: u64) -> &mut Self {
+        self.image_size = Some(size);
+        self
+    }
+
+    pub fn block_size(&mut self, block_size: u64) -> &mut Self {
+        self.block_size = Some(block_size);
+        self
+    }
+
+    pub fn blocks(&mut self, blocks: u64) -> &mut Self {
+        self.blocks = Some(blocks);
+        self
+    }
+
+    pub fn mapped_blocks(&mut self, blocks: u64) -> &mut Self {
+        self.mapped_blocks = Some(blocks);
+        self
+    }
+
+    pub fn checksum_type(&mut self, checksum_type: HashType) -> &mut Self {
+        self.checksum_type = Some(checksum_type);
+        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");
+        let offset = start * bs;
+        let length = (total - offset).min((end - start + 1) * bs);
+        self.add_byte_range(offset, length, checksum)
+    }
+
+    pub fn add_byte_range(&mut self, offset: u64, length: u64, checksum: HashValue) -> &mut Self {
+        let range = BlockRange {
+            offset,
+            length,
+            checksum,
+        };
+        self.blockmap.push(range);
+        self
+    }
+
+    pub fn build(self) -> Result<Bmap, BmapBuilderError> {
+        let image_size = self.image_size.ok_or(BmapBuilderError::MissingImageSize)?;
+        let block_size = self.block_size.ok_or(BmapBuilderError::MissingBlockSize)?;
+        let blocks = self.blocks.ok_or(BmapBuilderError::MissingBlocks)?;
+        let mapped_blocks = self
+            .mapped_blocks
+            .ok_or(BmapBuilderError::MissingMappedBlocks)?;
+        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 {
+            image_size,
+            block_size,
+            blocks,
+            mapped_blocks,
+            checksum_type,
+            bmap_file_checksum,
+            blockmap,
+        })
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::str::FromStr;
+
+    #[test]
+    fn hashes() {
+        assert_eq!("sha256", &HashType::Sha256.to_string());
+        assert_eq!(HashType::Sha256, HashType::from_str("sha256").unwrap());
+        let h = HashValue::Sha256([0; 32]);
+        assert_eq!(HashType::Sha256, h.to_type());
+    }
+}