add checksum in download request to verify its consistency (#979)

Signed-off-by: cormick <[email protected]>

dragonfly-client-storage/src/metadata.rs

@@ -15,9 +15,10 @@
  */
 
 use chrono::{NaiveDateTime, Utc};
+use crc::*;
 use dragonfly_client_config::dfdaemon::Config;
 use dragonfly_client_core::{Error, Result};
-use dragonfly_client_util::http::headermap_to_hashmap;
+use dragonfly_client_util::{digest, http::headermap_to_hashmap};
 use rayon::prelude::*;
 use reqwest::header::HeaderMap;
 use serde::{Deserialize, Serialize};
@@ -301,6 +302,27 @@ impl Piece {
             None => None,
         }
     }
+
+    /// calculate_digest return the digest of the piece metadata,
+    /// which is different from the digest of the piece content.
+    pub fn calculate_digest(&self) -> String {
+        let crc = Crc::<u32, Table<16>>::new(&CRC_32_ISCSI);
+        let mut crc_digest = crc.digest();
+
+        crc_digest.update(&self.number.to_be_bytes());
+        crc_digest.update(&self.offset.to_be_bytes());
+        crc_digest.update(&self.length.to_be_bytes());
+        crc_digest.update(self.digest.as_bytes());
+
+        if let Some(parent_id) = &self.parent_id {
+            crc_digest.update(parent_id.as_bytes());
+        }
+
+        let digest =
+            digest::Digest::new(digest::Algorithm::Crc32, crc_digest.finalize().to_string());
+
+        digest.to_string()
+    }
 }
 
 /// Metadata manages the metadata of [Task], [Piece] and [PersistentCacheTask].
@@ -928,6 +950,21 @@ mod tests {
     use super::*;
     use tempdir::TempDir;
 
+    #[test]
+    fn test_calculate_digest() {
+        let piece = Piece {
+            number: 1,
+            offset: 0,
+            length: 1024,
+            digest: "crc32:3810626145".to_string(),
+            parent_id: Some("d3c4e940ad06c47fc36ac67801e6f8e36cb4".to_string()),
+            ..Default::default()
+        };
+
+        let digest = piece.calculate_digest();
+        assert_eq!(digest, "crc32:523852508");
+    }
+
     #[test]
     fn should_create_metadata() {
         let dir = TempDir::new("metadata").unwrap();

dragonfly-client-util/src/digest/mod.rs

@@ -131,7 +131,7 @@ impl FromStr for Digest {
 
 /// calculate_file_hash calculates the hash of a file.
 #[instrument(skip_all)]
-pub fn calculate_file_hash(algorithm: Algorithm, path: &Path) -> ClientResult<Digest> {
+pub fn calculate_file_digest(algorithm: Algorithm, path: &Path) -> ClientResult<Digest> {
     let f = std::fs::File::open(path)?;
     let mut reader = std::io::BufReader::new(f);
     match algorithm {
@@ -200,31 +200,31 @@ mod tests {
     }
 
     #[test]
-    fn test_calculate_file_hash() {
+    fn test_calculate_file_digest() {
         let content = b"test content";
         let temp_file = tempfile::NamedTempFile::new().expect("failed to create temp file");
         let path = temp_file.path();
         let mut file = File::create(path).expect("failed to create file");
         file.write_all(content).expect("failed to write to file");
 
         let expected_blake3 = "ead3df8af4aece7792496936f83b6b6d191a7f256585ce6b6028db161278017e";
-        let digest =
-            calculate_file_hash(Algorithm::Blake3, path).expect("failed to calculate Blake3 hash");
+        let digest = calculate_file_digest(Algorithm::Blake3, path)
+            .expect("failed to calculate Blake3 hash");
         assert_eq!(digest.encoded(), expected_blake3);
 
         let expected_sha256 = "6ae8a75555209fd6c44157c0aed8016e763ff435a19cf186f76863140143ff72";
-        let digest =
-            calculate_file_hash(Algorithm::Sha256, path).expect("failed to calculate Sha256 hash");
+        let digest = calculate_file_digest(Algorithm::Sha256, path)
+            .expect("failed to calculate Sha256 hash");
         assert_eq!(digest.encoded(), expected_sha256);
 
         let expected_sha512 = "0cbf4caef38047bba9a24e621a961484e5d2a92176a859e7eb27df343dd34eb98d538a6c5f4da1ce302ec250b821cc001e46cc97a704988297185a4df7e99602";
-        let digest =
-            calculate_file_hash(Algorithm::Sha512, path).expect("failed to calculate Sha512 hash");
+        let digest = calculate_file_digest(Algorithm::Sha512, path)
+            .expect("failed to calculate Sha512 hash");
         assert_eq!(digest.encoded(), expected_sha512);
 
         let expected_crc32 = "422618885";
         let digest =
-            calculate_file_hash(Algorithm::Crc32, path).expect("failed to calculate Sha512 hash");
+            calculate_file_digest(Algorithm::Crc32, path).expect("failed to calculate Sha512 hash");
         assert_eq!(digest.encoded(), expected_crc32);
     }
 }

dragonfly-client/src/grpc/dfdaemon_upload.rs

@@ -873,16 +873,16 @@ impl DfdaemonUpload for DfdaemonUploadServerHandler {
         Ok(Response::new(DownloadPieceResponse {
             piece: Some(Piece {
                 number: piece.number,
-                parent_id: piece.parent_id,
+                parent_id: piece.parent_id.clone(),
                 offset: piece.offset,
                 length: piece.length,
-                digest: piece.digest,
+                digest: piece.digest.clone(),
                 content: Some(content),
                 traffic_type: None,
                 cost: None,
                 created_at: None,
             }),
-            digest: None,
+            digest: Some(piece.calculate_digest()),
         }))
     }
 

dragonfly-client/src/resource/piece_downloader.rs

@@ -18,6 +18,7 @@ use crate::grpc::dfdaemon_upload::DfdaemonUploadClient;
 use dragonfly_api::dfdaemon::v2::{DownloadPersistentCachePieceRequest, DownloadPieceRequest};
 use dragonfly_client_config::dfdaemon::Config;
 use dragonfly_client_core::{Error, Result};
+use dragonfly_client_storage::metadata;
 use std::sync::Arc;
 use tracing::{error, instrument};
 
@@ -120,6 +121,25 @@ impl Downloader for GRPCDownloader {
             return Err(Error::InvalidParameter);
         };
 
+        let piece_metadata = metadata::Piece {
+            number,
+            length: piece.length,
+            offset: piece.offset,
+            digest: piece.digest.clone(),
+            parent_id: piece.parent_id.clone(),
+            ..Default::default()
+        };
+
+        if let Some(expected_digest) = &response.digest {
+            let actual_digest = piece_metadata.calculate_digest();
+            if expected_digest != &actual_digest {
+                return Err(Error::ValidationError(format!(
+                    "checksum mismatch, expected: {}, actual: {}",
+                    expected_digest, actual_digest
+                )));
+            }
+        }
+
         let Some(content) = piece.content else {
             return Err(Error::InvalidParameter);
         };