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_proxy_registry_mirror_enable_blob_based_task_id is the default value for enable_blob_based_task_id.
+#[inline]
+fn default_proxy_registry_mirror_enable_blob_based_task_id() -> 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_blob_based_task_id 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_proxy_registry_mirror_enable_blob_based_task_id")]
+    pub enable_blob_based_task_id: bool,
 }
 
 /// RegistryMirror implements Default.
@@ -1128,6 +1142,7 @@ impl Default for RegistryMirror {
         Self {
             addr: default_proxy_registry_mirror_addr(),
             cert: None,
+            enable_blob_based_task_id: default_proxy_registry_mirror_enable_blob_based_task_id(),
         }
     }
 }

dragonfly-client/src/proxy/mod.rs

@@ -1097,6 +1097,31 @@ fn make_registry_mirror_request(
     Ok(request)
 }
 
+/// 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>
+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
+}
+
 /// make_download_task_request makes a download task request by the request.
 fn make_download_task_request(
     config: Arc<Config>,
@@ -1121,6 +1146,21 @@ 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_blob_based_task_id 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_blob_based_task_id {
+                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 +1189,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,
@@ -1281,3 +1321,119 @@ fn empty() -> BoxBody<Bytes, ClientError> {
         .map_err(|never| match never {})
         .boxed()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[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 test_make_download_task_request_with_blob_based_task_id_enabled() {
+        // Create a config with enable_blob_based_task_id = true
+        let mut config = Config::default();
+        config.proxy.registry_mirror.enable_blob_based_task_id = true;
+        let config = Arc::new(config);
+
+        // Create a request with a blob URL
+        let uri = "https://registry.example.com/v2/library/nginx/blobs/sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+            .parse::<hyper::Uri>()
+            .unwrap();
+
+        // Create an empty body that can be used with hyper::body::Incoming
+        // Since we can't directly instantiate Incoming, we create a compatible request
+        // using a type that can be converted. For testing purposes, we'll work around this
+        // by testing the extract_blob_digest_from_url function directly and verifying
+        // the logic path through code inspection.
+
+        // Test the URL extraction directly
+        let path = uri.path();
+        let extracted = extract_blob_digest_from_url(path);
+        assert_eq!(
+            extracted,
+            Some(
+                "sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+                    .to_string()
+            )
+        );
+
+        // Verify the config setting
+        assert!(config.proxy.registry_mirror.enable_blob_based_task_id);
+    }
+
+    #[test]
+    fn test_make_download_task_request_with_blob_based_task_id_disabled() {
+        // Create a config with enable_blob_based_task_id = false
+        let mut config = Config::default();
+        config.proxy.registry_mirror.enable_blob_based_task_id = false;
+        let config = Arc::new(config);
+
+        // Create a request with a blob URL
+        let uri = "https://registry.example.com/v2/library/nginx/blobs/sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+            .parse::<hyper::Uri>()
+            .unwrap();
+
+        // Test that the URL would extract a digest
+        let path = uri.path();
+        let extracted = extract_blob_digest_from_url(path);
+        assert!(extracted.is_some());
+
+        // But verify the config setting is disabled
+        assert!(!config.proxy.registry_mirror.enable_blob_based_task_id);
+    }
+
+    #[test]
+    fn test_make_download_task_request_with_non_blob_url() {
+        // Create a config with enable_blob_based_task_id = true
+        let mut config = Config::default();
+        config.proxy.registry_mirror.enable_blob_based_task_id = true;
+        let config = Arc::new(config);
+
+        // Create a request with a non-blob URL (manifest URL)
+        let uri = "https://registry.example.com/v2/library/nginx/manifests/latest"
+            .parse::<hyper::Uri>()
+            .unwrap();
+
+        // Test the URL extraction directly - should return None for non-blob URLs
+        let path = uri.path();
+        let extracted = extract_blob_digest_from_url(path);
+        assert_eq!(extracted, None);
+
+        // Verify the config setting
+        assert!(config.proxy.registry_mirror.enable_blob_based_task_id);
+    }
+}