CasObject::serialize validates hash (#18)

rajatarya · web-flow · commit 148b7205ed1d · 2024-09-23T11:43:58.000-07:00
* Added tests for CAS hash generation

* Added hash validation to CasObject::serialize

- Removed from LocalClient::put
- Updated unit-tests across LocalClient, StagingClient, CacheClient

Now all paths to creating Xorbs will validate the hash provided matches
the contents being offered.

* Fix clippy
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cas_client/src/caching_client.rs b/cas_client/src/caching_client.rs
@@ -256,10 +256,10 @@ mod tests {
         // put the different value with the same hash
         // this should fail
         assert_eq!(
-            CasClientError::HashMismatch,
+            CasClientError::CasObjectError(cas_object::error::CasObjectError::HashMismatch),
             client
                 .put(
-                    "key",
+                    "hellp",
                     &hello_hash,
                     "hellp world".as_bytes().to_vec(),
                     vec![hello.len() as u64],
diff --git a/cas_client/src/local_client.rs b/cas_client/src/local_client.rs
@@ -2,8 +2,6 @@ use crate::error::{CasClientError, Result};
 use crate::interface::Client;
 use cas::key::Key;
 use cas_object::CasObject;
-use merkledb::prelude::*;
-use merkledb::{Chunk, MerkleMemDB};
 use merklehash::MerkleHash;
 use std::fs::{metadata, File};
 use std::io::{BufReader, BufWriter, Write};
@@ -128,30 +126,6 @@ impl LocalClient {
         let _ = std::fs::remove_file(file_path);
     }
 
-    fn validate_root_hash(data: &[u8], chunk_boundaries: &[u64], hash: &MerkleHash) -> bool {
-        // at least 1 chunk, and last entry in chunk boundary must match the length
-        if chunk_boundaries.is_empty()
-            || chunk_boundaries[chunk_boundaries.len() - 1] as usize != data.len()
-        {
-            return false;
-        }
-
-        let mut chunks: Vec<Chunk> = Vec::new();
-        let mut left_edge: usize = 0;
-        for i in chunk_boundaries {
-            let right_edge = *i as usize;
-            let hash = merklehash::compute_data_hash(&data[left_edge..right_edge]);
-            let length = right_edge - left_edge;
-            chunks.push(Chunk { hash, length });
-            left_edge = right_edge;
-        }
-
-        let mut db = MerkleMemDB::default();
-        let mut staging = db.start_insertion_staging();
-        db.add_file(&mut staging, &chunks);
-        let ret = db.finalize(staging);
-        *ret.hash() == *hash
-    }
 }
 
 /// LocalClient is responsible for writing/reading Xorbs on local disk.
@@ -176,10 +150,7 @@ impl Client for LocalClient {
             return Err(CasClientError::InvalidArguments);
         }
 
-        // validate hash
-        if !Self::validate_root_hash(&data, &chunk_boundaries, hash) {
-            return Err(CasClientError::HashMismatch);
-        }
+        // moved hash validation into [CasObject::serialize], so removed from here.
 
         if let Ok(xorb_size) = self.get_length(prefix, hash).await {
             if xorb_size > 0 {
@@ -315,6 +286,7 @@ impl Client for LocalClient {
 mod tests {
 
     use super::*;
+    use merkledb::{prelude::MerkleDBHighLevelMethodsV1, Chunk, MerkleMemDB};
     use merklehash::{compute_data_hash, DataHash};
     use rand::Rng;
 
@@ -432,10 +404,10 @@ mod tests {
         // put the different value with the same hash
         // this should fail
         assert_eq!(
-            CasClientError::HashMismatch,
+            CasClientError::CasObjectError(cas_object::error::CasObjectError::HashMismatch),
             client
                 .put(
-                    "key",
+                    "hellp",
                     &hello_hash,
                     "hellp world".as_bytes().to_vec(),
                     vec![hello.len() as u64],
@@ -448,7 +420,7 @@ mod tests {
             CasClientError::InvalidArguments,
             client
                 .put(
-                    "key",
+                    "hellp2",
                     &hello_hash,
                     "hellp wod".as_bytes().to_vec(),
                     vec![hello.len() as u64],
@@ -462,7 +434,7 @@ mod tests {
             CasClientError::InvalidArguments,
             client
                 .put(
-                    "key",
+                    "again",
                     &hello_hash,
                     "hello world again".as_bytes().to_vec(),
                     vec![hello.len() as u64],
diff --git a/cas_client/src/staging_client.rs b/cas_client/src/staging_client.rs
@@ -424,10 +424,10 @@ mod tests {
         // put the different value with the same hash
         // this should fail
         assert_eq!(
-            CasClientError::HashMismatch,
+            CasClientError::CasObjectError(cas_object::error::CasObjectError::HashMismatch),
             client
                 .put(
-                    "key",
+                    "hellp",
                     &hello_hash,
                     "hellp world".as_bytes().to_vec(),
                     vec![hello.len() as u64],
diff --git a/cas_object/Cargo.toml b/cas_object/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 anyhow = "1.0.88"
 bincode = "1.3.3"
 http = "1.1.0"
+merkledb = { path = "../merkledb" }
 merklehash = { path = "../merklehash" }
 tempfile = "3.12.0"
 tracing = "0.1.40"
diff --git a/cas_object/src/cas_object_format.rs b/cas_object/src/cas_object_format.rs
@@ -1,4 +1,5 @@
 use bytes::Buf;
+use merkledb::{prelude::MerkleDBHighLevelMethodsV1, Chunk, MerkleMemDB};
 use merklehash::{DataHash, MerkleHash};
 use std::{
     cmp::min,
@@ -454,6 +455,12 @@ impl CasObject {
         chunk_boundaries: &Vec<u32>,
         compression_scheme: CompressionScheme,
     ) -> Result<(Self, usize), CasObjectError> {
+
+        // validate hash against contents
+        if !Self::validate_root_hash(data, chunk_boundaries, hash) {
+            return Err(CasObjectError::HashMismatch);
+        }
+
         let mut cas = CasObject::default();
         cas.info.cashash.copy_from_slice(hash.as_slice());
         cas.info.num_chunks = chunk_boundaries.len() as u32 + 1; // extra entry for dummy, see [chunk_size_info] for details.
@@ -473,7 +480,6 @@ impl CasObject {
             let chunk_size = chunk_boundary - raw_start_idx;
 
             // now serialize chunk directly to writer (since chunks come first!)
-            // TODO: add compression scheme to this call
             let chunk_written_bytes =
                 serialize_chunk(&chunk_raw_bytes, writer, compression_scheme)?;
             total_written_bytes += chunk_written_bytes;
@@ -506,15 +512,40 @@ impl CasObject {
 
         Ok((cas, total_written_bytes))
     }
+
+    pub fn validate_root_hash(data: &[u8], chunk_boundaries: &[u32], hash: &MerkleHash) -> bool {
+        // at least 1 chunk, and last entry in chunk boundary must match the length
+        if chunk_boundaries.is_empty()
+            || chunk_boundaries[chunk_boundaries.len() - 1] as usize != data.len()
+        {
+            return false;
+        }
+
+        let mut chunks: Vec<Chunk> = Vec::new();
+        let mut left_edge: usize = 0;
+        for i in chunk_boundaries {
+            let right_edge = *i as usize;
+            let hash = merklehash::compute_data_hash(&data[left_edge..right_edge]);
+            let length = right_edge - left_edge;
+            chunks.push(Chunk { hash, length });
+            left_edge = right_edge;
+        }
+
+        let mut db = MerkleMemDB::default();
+        let mut staging = db.start_insertion_staging();
+        db.add_file(&mut staging, &chunks);
+        let ret = db.finalize(staging);
+        *ret.hash() == *hash
+    }
+
 }
 
 #[cfg(test)]
 mod tests {
 
-    use crate::cas_chunk_format::serialize_chunk;
-
     use super::*;
-    use merklehash::compute_data_hash;
+    use crate::cas_chunk_format::serialize_chunk;
+    use merkledb::{prelude::MerkleDBHighLevelMethodsV1, Chunk, MerkleMemDB};
     use rand::Rng;
     use std::io::Cursor;
 
@@ -633,10 +664,10 @@ mod tests {
         chunk_size_info.push(chunk_info);
 
         c.info.num_chunks = chunk_size_info.len() as u32;
-
-        c.info.cashash = compute_data_hash(&writer.get_ref());
         c.info.chunk_size_info = chunk_size_info;
 
+        c.info.cashash = gen_hash(&data_contents_raw, &c.get_chunk_boundaries());
+
         // now serialize info to end Xorb length
         let len = c.info.serialize(&mut writer).unwrap();
         c.info_length = len as u32;
@@ -646,6 +677,92 @@ mod tests {
         (c, writer.get_ref().to_vec(), data_contents_raw)
     }
 
+    fn gen_hash(data: &[u8], chunk_boundaries: &[u32]) -> DataHash {
+        let mut chunks: Vec<Chunk> = Vec::new();
+        let mut left_edge: usize = 0;
+        for i in chunk_boundaries {
+            let right_edge = *i as usize;
+            let hash = merklehash::compute_data_hash(&data[left_edge..right_edge]);
+            let length = right_edge - left_edge;
+            chunks.push(Chunk { hash, length });
+            left_edge = right_edge;
+        }
+
+        let mut db = MerkleMemDB::default();
+        let mut staging = db.start_insertion_staging();
+        db.add_file(&mut staging, &chunks);
+        let ret = db.finalize(staging);
+        *ret.hash()
+    }
+
+    #[test]
+    fn test_compress_decompress() {
+        // Arrange
+        let (c, _cas_data, raw_data) = build_cas_object(55, 53212, false, true); 
+        let hash = gen_hash(&&raw_data, &c.get_chunk_boundaries());
+
+        // Act & Assert
+        let mut writer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
+        assert!(CasObject::serialize(
+            &mut writer,
+            &c.info.cashash,
+            &raw_data,
+            &c.get_chunk_boundaries(),
+            CompressionScheme::LZ4
+        )
+        .is_ok());
+
+        let mut reader = writer.clone();
+        reader.set_position(0);
+        let res = CasObject::deserialize(&mut reader);
+        assert!(res.is_ok());
+        let c = res.unwrap();
+
+        let c_bytes = c.get_all_bytes(&mut reader).unwrap();
+        let c_boundaries = c.get_chunk_boundaries();
+        let c_hash = gen_hash(&c_bytes, &c_boundaries);
+
+        let mut writer: Cursor<Vec<u8>> = Cursor::new(Vec::new());
+        assert!(CasObject::serialize(
+            &mut writer,
+            &c_hash,
+            &c_bytes,
+            &c_boundaries,
+            CompressionScheme::None
+        )
+        .is_ok());
+
+        let mut reader = writer.clone();
+        reader.set_position(0);
+        let res = CasObject::deserialize(&mut reader);
+        assert!(res.is_ok());
+        let c2 = res.unwrap();
+
+        assert_eq!(hash, c_hash);
+        assert_eq!(c.info.cashash, hash);
+        assert_eq!(c2.info.cashash, c.info.cashash);
+    }
+
+    #[test]
+    fn test_hash_generation_compression() {
+        // Arrange
+        let (c, cas_data, raw_data) = build_cas_object(55, 53212, false, true);
+        // Act & Assert
+        let mut buf: Cursor<Vec<u8>> = Cursor::new(Vec::new());
+        assert!(CasObject::serialize(
+            &mut buf,
+            &c.info.cashash,
+            &raw_data,
+            &c.get_chunk_boundaries(),
+            CompressionScheme::LZ4
+        )
+        .is_ok());
+
+        assert_eq!(c.info.cashash, gen_hash(&raw_data, &c.get_chunk_boundaries()));
+        assert_eq!(raw_data, c.get_all_bytes(&mut buf).unwrap());
+        assert_eq!(&cas_data, buf.get_ref());
+    }
+
     #[test]
     fn test_basic_serialization_mem() {
         // Arrange
diff --git a/cas_object/src/error.rs b/cas_object/src/error.rs
@@ -15,6 +15,9 @@ pub enum CasObjectError {
     #[error("Format Error: {0}")]
     FormatError(anyhow::Error),
 
+    #[error("Hash Mismatch")]
+    HashMismatch,
+
     #[error("Internal IO Error: {0}")]
     InternalIOError(#[from] std::io::Error),
 
