fix: task ID of image blobs

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

Author: wuliang

dragonfly-client-config/src/dfdaemon.rs

@@ -1105,6 +1105,12 @@ impl Default for Rule {
     }
 }
 
+/// default_enable_task_id_based_blob_digest is the default value for enable_task_id_based_blob_digest.
+#[inline]
+fn default_enable_task_id_based_blob_digest() -> bool {
+    true
+}
+
 /// RegistryMirror is the registry mirror configuration.
 #[derive(Debug, Clone, Validate, Deserialize)]
 #[serde(default, rename_all = "camelCase")]
@@ -1120,6 +1126,14 @@ pub struct RegistryMirror {
     /// If registry use self-signed cert, the client should set the
     /// cert for the registry mirror.
     pub cert: Option<PathBuf>,
+
+    /// enable_task_id_based_blob_digest indicates whether to calculate the task ID based on the blob's SHA256 digest
+    /// for OCI registry blob download URLs. When enabled, if the download URL is for an image blob
+    /// (e.g., /v2/<name>/blobs/sha256:<digest>), the task ID will be calculated based on the blob's digest
+    /// instead of the full URL. This allows the same blob from different registries to share the same task ID,
+    /// avoiding redundant downloads and storage.
+    #[serde(default = "default_enable_task_id_based_blob_digest")]
+    pub enable_task_id_based_blob_digest: bool,
 }
 
 /// RegistryMirror implements Default.
@@ -1128,6 +1142,7 @@ impl Default for RegistryMirror {
         Self {
             addr: default_proxy_registry_mirror_addr(),
             cert: None,
+            enable_task_id_based_blob_digest: default_enable_task_id_based_blob_digest(),
         }
     }
 }

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

@@ -31,6 +31,31 @@ const SEED_PEER_SUFFIX: &str = "seed";
 /// PERSISTENT_CACHE_TASK_SUFFIX is the suffix of the persistent cache task.
 const PERSISTENT_CACHE_TASK_SUFFIX: &str = "persistent-cache-task";
 
+/// extract_blob_digest_from_url extracts the blob digest from a registry blob URL.
+/// Returns the digest if the URL is a blob URL, otherwise returns None.
+///
+/// Example blob URLs:
+/// - /v2/<name>/blobs/sha256:<digest>
+/// - /v2/<namespace>/<repo>/blobs/sha256:<digest>
+pub fn extract_blob_digest_from_url(path: &str) -> Option<String> {
+    // Check if the path contains /blobs/sha256:
+    if let Some(blobs_idx) = path.find("/blobs/sha256:") {
+        // Extract everything after "/blobs/sha256:"
+        let after_blobs = &path[blobs_idx + "/blobs/sha256:".len()..];
+
+        // The digest should be 64 hex characters (SHA256)
+        // It might be followed by query parameters or nothing
+        let digest = after_blobs.split(&['?', '#'][..]).next().unwrap_or("");
+
+        // Validate that it looks like a SHA256 hash (64 hex characters)
+        if digest.len() == 64 && digest.chars().all(|c| c.is_ascii_hexdigit()) {
+            return Some(format!("sha256:{}", digest));
+        }
+    }
+
+    None
+}
+
 /// TaskIDParameter is the parameter of the task id.
 pub enum TaskIDParameter {
     /// Content uses the content to generate the task id.
@@ -44,6 +69,11 @@ pub enum TaskIDParameter {
         application: Option<String>,
         filtered_query_params: Vec<String>,
     },
+    /// BlobDigestBased uses the blob digest to generate the task id.
+    /// The digest can use other algorithms (like sha256, sha512, etc.),
+    /// but the task ID in Dragonfly must be SHA256.
+    /// Task ID needs to compute a SHA256 hash based on the digest content.
+    BlobDigestBased(String),
 }
 
 /// PersistentCacheTaskIDParameter is the parameter of the persistent cache task id.
@@ -153,6 +183,12 @@ impl IDGenerator {
                 // Generate the task id.
                 Ok(hex::encode(hasher.finalize()))
             }
+            TaskIDParameter::BlobDigestBased(digest) => {
+                // The digest can use other algorithms (sha256, sha512, etc.),
+                // but the task ID in Dragonfly must be SHA256.
+                // Compute SHA256 hash based on the digest content.
+                Ok(hex::encode(Sha256::digest(digest.as_bytes())))
+            }
         }
     }
 
@@ -335,6 +371,14 @@ mod tests {
                 TaskIDParameter::Content("This is a test file".to_string()),
                 "e2d0fe1585a63ec6009c8016ff8dda8b17719a637405a4e23c0ff81339148249",
             ),
+            (
+                IDGenerator::new("127.0.0.1".to_string(), "localhost".to_string(), false),
+                TaskIDParameter::BlobDigestBased(
+                    "sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+                        .to_string(),
+                ),
+                "719498689c2f5bd76140f3bd2b03bcbc3134890e72b4d5b788d8b41ec0cd0f93",
+            ),
         ];
 
         for (generator, parameter, expected_id) in test_cases {
@@ -343,6 +387,42 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_extract_blob_digest_from_url() {
+        // Test cases for valid blob URLs
+        let test_cases = vec![
+            (
+                "/v2/library/nginx/blobs/sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef",
+                Some("sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef".to_string()),
+            ),
+            (
+                "/v2/myorg/myrepo/blobs/sha256:abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890",
+                Some("sha256:abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890".to_string()),
+            ),
+            (
+                "/v2/namespace/repo/blobs/sha256:0000000000000000000000000000000000000000000000000000000000000000?query=param",
+                Some("sha256:0000000000000000000000000000000000000000000000000000000000000000".to_string()),
+            ),
+        ];
+
+        for (input, expected) in test_cases {
+            assert_eq!(extract_blob_digest_from_url(input), expected);
+        }
+
+        // Test cases for invalid URLs (should return None)
+        let invalid_cases = vec![
+            "/v2/library/nginx/manifests/latest",
+            "/v2/library/nginx/blobs/sha256:short",
+            "/v2/library/nginx/blobs/sha256:xyz", // Not hex
+            "/api/v1/some/other/path",
+            "",
+        ];
+
+        for input in invalid_cases {
+            assert_eq!(extract_blob_digest_from_url(input), None);
+        }
+    }
+
     #[test]
     fn should_generate_persistent_cache_task_id() {
         let dir = tempdir().unwrap();

dragonfly-client/src/grpc/dfdaemon_download.rs

@@ -263,7 +263,33 @@ impl DfdaemonDownload for DfdaemonDownloadServerHandler {
             .task
             .id_generator
             .task_id(match download.content_for_calculating_task_id.clone() {
-                Some(content) => TaskIDParameter::Content(content),
+                Some(content) => {
+                    // Check if the content matches OCI digest format: algorithm:encoded
+                    // See: https://github.com/opencontainers/image-spec/blob/main/descriptor.md#digests
+                    // Format: algorithm can be [a-z0-9+._-]+, encoded can be [a-zA-Z0-9=_-]+
+                    // If it's a digest, use BlobDigestBased to ensure SHA256 hash is calculated
+                    // from the digest content, regardless of the digest algorithm used.
+                    let is_digest = content.split_once(':').is_some_and(|(alg, enc)| {
+                        // Validate algorithm: [a-z0-9+._-]+
+                        !alg.is_empty()
+                            && alg.chars().all(|c| {
+                                c.is_ascii_lowercase()
+                                    || c.is_ascii_digit()
+                                    || matches!(c, '+' | '.' | '_' | '-')
+                            })
+                            // Validate encoded: [a-zA-Z0-9=_-]+ and minimum length
+                            && enc.len() >= 32
+                            && enc.chars().all(|c| {
+                                c.is_ascii_alphanumeric() || matches!(c, '=' | '_' | '-')
+                            })
+                    });
+
+                    if is_digest {
+                        TaskIDParameter::BlobDigestBased(content)
+                    } else {
+                        TaskIDParameter::Content(content)
+                    }
+                }
                 None => TaskIDParameter::URLBased {
                     url: download.url.clone(),
                     piece_length: download.piece_length,

dragonfly-client/src/grpc/dfdaemon_upload.rs

@@ -260,7 +260,33 @@ impl DfdaemonUpload for DfdaemonUploadServerHandler {
             .task
             .id_generator
             .task_id(match download.content_for_calculating_task_id.clone() {
-                Some(content) => TaskIDParameter::Content(content),
+                Some(content) => {
+                    // Check if the content matches OCI digest format: algorithm:encoded
+                    // See: https://github.com/opencontainers/image-spec/blob/main/descriptor.md#digests
+                    // Format: algorithm can be [a-z0-9+._-]+, encoded can be [a-zA-Z0-9=_-]+
+                    // If it's a digest, use BlobDigestBased to ensure SHA256 hash is calculated
+                    // from the digest content, regardless of the digest algorithm used.
+                    let is_digest = content.split_once(':').is_some_and(|(alg, enc)| {
+                        // Validate algorithm: [a-z0-9+._-]+
+                        !alg.is_empty()
+                            && alg.chars().all(|c| {
+                                c.is_ascii_lowercase()
+                                    || c.is_ascii_digit()
+                                    || matches!(c, '+' | '.' | '_' | '-')
+                            })
+                            // Validate encoded: [a-zA-Z0-9=_-]+ and minimum length
+                            && enc.len() >= 32
+                            && enc.chars().all(|c| {
+                                c.is_ascii_alphanumeric() || matches!(c, '=' | '_' | '-')
+                            })
+                    });
+
+                    if is_digest {
+                        TaskIDParameter::BlobDigestBased(content)
+                    } else {
+                        TaskIDParameter::Content(content)
+                    }
+                }
                 None => TaskIDParameter::URLBased {
                     url: download.url.clone(),
                     piece_length: download.piece_length,

dragonfly-client/src/proxy/mod.rs

@@ -31,6 +31,7 @@ use dragonfly_client_metric::{
 };
 use dragonfly_client_util::{
     http::{hashmap_to_headermap, headermap_to_hashmap},
+    id_generator::extract_blob_digest_from_url,
     shutdown,
     tls::{generate_self_signed_certs_by_ca_cert, generate_simple_self_signed_certs, NoVerifier},
 };
@@ -1121,6 +1122,25 @@ fn make_download_task_request(
         }
     }
 
+    // Determine the content for calculating task ID.
+    // Priority:
+    // 1. Explicit header value (X-Dragonfly-Content-For-Calculating-Task-ID)
+    // 2. Blob digest from URL (if enable_task_id_based_blob_digest is true)
+    // 3. None (will use URL-based calculation)
+    let content_for_calculating_task_id = header::get_content_for_calculating_task_id(&header)
+        .or_else(|| {
+            if config
+                .proxy
+                .registry_mirror
+                .enable_task_id_based_blob_digest
+            {
+                let path = request.uri().path();
+                extract_blob_digest_from_url(path)
+            } else {
+                None
+            }
+        });
+
     Ok(DownloadTaskRequest {
         download: Some(Download {
             url: make_download_url(request.uri(), rule.use_tls, rule.redirect.clone())?,
@@ -1149,7 +1169,7 @@ fn make_download_task_request(
             is_prefetch: false,
             need_piece_content: false,
             force_hard_link: header::get_force_hard_link(&header),
-            content_for_calculating_task_id: header::get_content_for_calculating_task_id(&header),
+            content_for_calculating_task_id,
             remote_ip: Some(remote_ip.to_string()),
             concurrent_piece_count: Some(config.download.concurrent_piece_count),
             overwrite: false,