feat: verify digest when file is downloaded (#1208)

Signed-off-by: Gaius <gaius.qi@gmail.com>

Cargo.lock

@@ -939,12 +939,12 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-api"
-version = "2.1.39"
+version = "2.1.40"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ef3a36f55cedea2a004d17cff39bcfe906fc94579cb0b440cf185a0663b645d"
+checksum = "36c49e3e1d805c3efa1834af4c0d549589e2e83acb0a1065b261619de23990aa"
 dependencies = [
  "prost 0.13.5",
- "prost-types",
+ "prost-types 0.14.1",
  "prost-wkt-types",
  "serde",
  "tokio",
@@ -954,7 +954,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "anyhow",
  "bytes",
@@ -1026,7 +1026,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-backend"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "dragonfly-api",
  "dragonfly-client-core",
@@ -1057,7 +1057,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-config"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "bytesize",
  "bytesize-serde",
@@ -1087,7 +1087,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-core"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "headers 0.4.1",
  "hyper 1.6.0",
@@ -1105,7 +1105,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-init"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "anyhow",
  "clap",
@@ -1123,7 +1123,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-storage"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "bincode",
  "bytes",
@@ -1152,7 +1152,7 @@ dependencies = [
 
 [[package]]
 name = "dragonfly-client-util"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "base64 0.22.1",
  "bytesize",
@@ -1561,7 +1561,7 @@ dependencies = [
 
 [[package]]
 name = "hdfs"
-version = "0.2.41"
+version = "0.2.42"
 dependencies = [
  "dragonfly-client-backend",
  "dragonfly-client-core",
@@ -3445,6 +3445,16 @@ dependencies = [
  "prost-derive 0.13.5",
 ]
 
+[[package]]
+name = "prost"
+version = "0.14.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7231bd9b3d3d33c86b58adbac74b5ec0ad9f496b19d22801d773636feaa95f3d"
+dependencies = [
+ "bytes",
+ "prost-derive 0.14.1",
+]
+
 [[package]]
 name = "prost-build"
 version = "0.13.1"
@@ -3460,7 +3470,7 @@ dependencies = [
  "petgraph",
  "prettyplease",
  "prost 0.13.5",
- "prost-types",
+ "prost-types 0.13.5",
  "regex",
  "syn 2.0.90",
  "tempfile",
@@ -3492,6 +3502,19 @@ dependencies = [
  "syn 2.0.90",
 ]
 
+[[package]]
+name = "prost-derive"
+version = "0.14.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425"
+dependencies = [
+ "anyhow",
+ "itertools",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.90",
+]
+
 [[package]]
 name = "prost-types"
 version = "0.13.5"
@@ -3501,6 +3524,15 @@ dependencies = [
  "prost 0.13.5",
 ]
 
+[[package]]
+name = "prost-types"
+version = "0.14.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9b4db3d6da204ed77bb26ba83b6122a73aeb2e87e25fbf7ad2e84c4ccbf8f72"
+dependencies = [
+ "prost 0.14.1",
+]
+
 [[package]]
 name = "prost-wkt"
 version = "0.6.1"
@@ -3525,7 +3557,7 @@ dependencies = [
  "heck",
  "prost 0.13.5",
  "prost-build",
- "prost-types",
+ "prost-types 0.13.5",
  "quote",
 ]
 
@@ -3538,7 +3570,7 @@ dependencies = [
  "chrono",
  "prost 0.13.5",
  "prost-build",
- "prost-types",
+ "prost-types 0.13.5",
  "prost-wkt",
  "prost-wkt-build",
  "regex",
@@ -5036,7 +5068,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "878d81f52e7fcfd80026b7fdb6a9b578b3c3653ba987f87f0dce4b64043cba27"
 dependencies = [
  "prost 0.13.5",
- "prost-types",
+ "prost-types 0.13.5",
  "tokio",
  "tokio-stream",
  "tonic",

Cargo.toml

@@ -12,7 +12,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.2.41"
+version = "0.2.42"
 authors = ["The Dragonfly Developers"]
 homepage = "https://d7y.io/"
 repository = "https://github.com/dragonflyoss/client.git"
@@ -22,14 +22,14 @@ readme = "README.md"
 edition = "2021"
 
 [workspace.dependencies]
-dragonfly-client = { path = "dragonfly-client", version = "0.2.41" }
+dragonfly-client = { path = "dragonfly-client", version = "0.2.42" }
-dragonfly-client-core = { path = "dragonfly-client-core", version = "0.2.41" }
+dragonfly-client-core = { path = "dragonfly-client-core", version = "0.2.42" }
-dragonfly-client-config = { path = "dragonfly-client-config", version = "0.2.41" }
+dragonfly-client-config = { path = "dragonfly-client-config", version = "0.2.42" }
-dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.2.41" }
+dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.2.42" }
-dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.2.41" }
+dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.2.42" }
-dragonfly-client-util = { path = "dragonfly-client-util", version = "0.2.41" }
+dragonfly-client-util = { path = "dragonfly-client-util", version = "0.2.42" }
-dragonfly-client-init = { path = "dragonfly-client-init", version = "0.2.41" }
+dragonfly-client-init = { path = "dragonfly-client-init", version = "0.2.42" }
-dragonfly-api = "=2.1.39"
+dragonfly-api = "=2.1.40"
 thiserror = "2.0"
 futures = "0.3.31"
 reqwest = { version = "0.12.4", features = [

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

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-use dragonfly_client_core::Result as ClientResult;
+use dragonfly_client_core::{Error as ClientError, Result as ClientResult};
 use sha2::Digest as Sha2Digest;
 use std::fmt;
 use std::io::{self, Read};
@@ -112,9 +112,36 @@ impl FromStr for Digest {
         }
 
         let algorithm = match parts[0] {
-            "crc32" => Algorithm::Crc32,
+            "crc32" => {
-            "sha256" => Algorithm::Sha256,
+                if parts[1].len() != 10 {
-            "sha512" => Algorithm::Sha512,
+                    return Err(format!(
+                        "invalid crc32 digest length: {}, expected 10",
+                        parts[1].len()
+                    ));
+                }
+
+                Algorithm::Crc32
+            }
+            "sha256" => {
+                if parts[1].len() != 64 {
+                    return Err(format!(
+                        "invalid sha256 digest length: {}, expected 64",
+                        parts[1].len()
+                    ));
+                }
+
+                Algorithm::Sha256
+            }
+            "sha512" => {
+                if parts[1].len() != 128 {
+                    return Err(format!(
+                        "invalid sha512 digest length: {}, expected 128",
+                        parts[1].len()
+                    ));
+                }
+
+                Algorithm::Sha512
+            }
             _ => return Err(format!("invalid digest algorithm: {}", parts[0])),
         };
 
@@ -155,6 +182,25 @@ pub fn calculate_file_digest(algorithm: Algorithm, path: &Path) -> ClientResult<
     }
 }
 
+/// verify_file_digest verifies the digest of a file against an expected digest.
+pub fn verify_file_digest(expected_digest: Digest, file_path: &Path) -> ClientResult<()> {
+    let digest = match calculate_file_digest(expected_digest.algorithm(), file_path) {
+        Ok(digest) => digest,
+        Err(err) => {
+            return Err(err);
+        }
+    };
+
+    if digest.to_string() != expected_digest.to_string() {
+        return Err(ClientError::DigestMismatch(
+            expected_digest.to_string(),
+            digest.to_string(),
+        ));
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -205,4 +251,28 @@ mod tests {
             calculate_file_digest(Algorithm::Crc32, path).expect("failed to calculate Crc32 hash");
         assert_eq!(digest.encoded(), expected_crc32);
     }
+
+    #[test]
+    fn test_verify_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_sha256_digest = Digest::new(
+            Algorithm::Sha256,
+            "6ae8a75555209fd6c44157c0aed8016e763ff435a19cf186f76863140143ff72".to_string(),
+        );
+        assert!(verify_file_digest(expected_sha256_digest, path).is_ok());
+
+        let expected_sha512_digest = Digest::new(
+            Algorithm::Sha512,
+            "0cbf4caef38047bba9a24e621a961484e5d2a92176a859e7eb27df343dd34eb98d538a6c5f4da1ce302ec250b821cc001e46cc97a704988297185a4df7e99602".to_string(),
+        );
+        assert!(verify_file_digest(expected_sha512_digest, path).is_ok());
+
+        let expected_crc32_digest = Digest::new(Algorithm::Crc32, "1475635037".to_string());
+        assert!(verify_file_digest(expected_crc32_digest, path).is_ok());
+    }
 }

dragonfly-client/src/bin/dfcache/export.rs

@@ -85,6 +85,13 @@ pub struct ExportCommand {
     )]
     timeout: Duration,
 
+    #[arg(
+        long = "digest",
+        required = false,
+        help = "Verify the integrity of the downloaded file using the specified digest, support sha256, sha512, crc32. If the digest is not specified, the downloaded file will not be verified. Format: <algorithm>:<digest>, e.g. sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, crc32:12345678"
+    )]
+    digest: Option<String>,
+
     #[arg(
         short = 'e',
         long = "endpoint",
@@ -457,6 +464,7 @@ impl ExportCommand {
                 ),
                 need_piece_content,
                 force_hard_link: self.force_hard_link,
+                digest: self.digest.clone(),
             })
             .await
             .inspect_err(|err| {

dragonfly-client/src/bin/dfget/main.rs

@@ -141,12 +141,11 @@ struct Args {
     timeout: Duration,
 
     #[arg(
-        short = 'd',
         long = "digest",
-        default_value = "",
+        required = false,
-        help = "Verify the integrity of the downloaded file using the specified digest, e.g. md5:86d3f3a95c324c9479bd8986968f4327"
+        help = "Verify the integrity of the downloaded file using the specified digest, support sha256, sha512, crc32. If the digest is not specified, the downloaded file will not be verified. Format: <algorithm>:<digest>, e.g. sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, crc32:12345678"
     )]
-    digest: String,
+    digest: Option<String>,
 
     #[arg(
         short = 'p',
@@ -759,7 +758,7 @@ async fn download(
         .download_task(DownloadTaskRequest {
             download: Some(Download {
                 url: args.url.to_string(),
-                digest: Some(args.digest),
+                digest: args.digest,
                 // NOTE: Dfget does not support range download.
                 range: None,
                 r#type: TaskType::Standard as i32,

dragonfly-client/src/grpc/dfdaemon_download.rs

@@ -42,6 +42,7 @@ use dragonfly_client_core::{
     Error as ClientError, Result as ClientResult,
 };
 use dragonfly_client_util::{
+    digest::{verify_file_digest, Digest},
     http::{get_range, hashmap_to_headermap, headermap_to_hashmap},
     id_generator::{PersistentCacheTaskIDParameter, TaskIDParameter},
 };
@@ -488,22 +489,48 @@ impl DfdaemonDownload for DfdaemonDownloadServerHandler {
                                                     )),
                                                 )
                                                 .await;
+                                                return;
                                             }
                                             Err(err) => {
                                                 error!("check output path: {}", err);
                                                 handle_error(&out_stream_tx, err).await;
+                                                return;
                                             }
                                         }
-
+                                    } else if let Err(err) = task_manager_clone
-                                        return;
-                                    }
-
-                                    if let Err(err) = task_manager_clone
                                         .copy_task(task_clone.id.as_str(), output_path)
                                         .await
                                     {
                                         error!("copy task: {}", err);
                                         handle_error(&out_stream_tx, err).await;
+                                        return;
+                                    }
+                                }
+
+                                // Verify the file digest if it is provided.
+                                if let Some(raw_digest) = &download_clone.digest {
+                                    let digest = match raw_digest.parse::<Digest>() {
+                                        Ok(digest) => digest,
+                                        Err(err) => {
+                                            error!("parse digest: {}", err);
+                                            handle_error(
+                                                &out_stream_tx,
+                                                Status::invalid_argument(format!(
+                                                    "invalid digest({}): {}",
+                                                    raw_digest, err
+                                                )),
+                                            )
+                                            .await;
+                                            return;
+                                        }
+                                    };
+
+                                    if let Err(err) =
+                                        verify_file_digest(digest, Path::new(output_path.as_str()))
+                                    {
+                                        error!("verify file digest: {}", err);
+                                        handle_error(&out_stream_tx, err).await;
+                                        return;
                                     }
                                 }
                             }
@@ -928,22 +955,48 @@ impl DfdaemonDownload for DfdaemonDownloadServerHandler {
                                                 )),
                                             )
                                             .await;
+                                            return;
                                         }
                                         Err(err) => {
                                             error!("check output path: {}", err);
                                             handle_error(&out_stream_tx, err).await;
+                                            return;
                                         }
                                     }
-
+                                } else if let Err(err) = task_manager_clone
-                                    return;
-                                }
-
-                                if let Err(err) = task_manager_clone
                                     .copy_task(task_clone.id.as_str(), output_path)
                                     .await
                                 {
                                     error!("copy task: {}", err);
                                     handle_error(&out_stream_tx, err).await;
+                                    return;
+                                }
+                            }
+
+                            // Verify the file digest if it is provided.
+                            if let Some(raw_digest) = &request_clone.digest {
+                                let digest = match raw_digest.parse::<Digest>() {
+                                    Ok(digest) => digest,
+                                    Err(err) => {
+                                        error!("parse digest: {}", err);
+                                        handle_error(
+                                            &out_stream_tx,
+                                            Status::invalid_argument(format!(
+                                                "invalid digest({}): {}",
+                                                raw_digest, err
+                                            )),
+                                        )
+                                        .await;
+                                        return;
+                                    }
+                                };
+
+                                if let Err(err) =
+                                    verify_file_digest(digest, Path::new(output_path.as_str()))
+                                {
+                                    error!("verify file digest: {}", err);
+                                    handle_error(&out_stream_tx, err).await;
+                                    return;
                                 }
                             }
                         }