[ACR] Update manifest digest check (Azure#29207)

Author: Yalin Li

sdk/containerregistry/azure-containerregistry/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/containerregistry/azure-containerregistry",
-  "Tag": "python/containerregistry/azure-containerregistry_a2d3792ebb"
+  "Tag": "python/containerregistry/azure-containerregistry_448865ae2e"
 }

sdk/containerregistry/azure-containerregistry/azure/containerregistry/_container_registry_client.py

@@ -42,9 +42,6 @@
 def _return_response_and_deserialized(pipeline_response, deserialized, _):
     return pipeline_response, deserialized
 
-def _return_deserialized(_, deserialized, __):
-    return deserialized
-
 def _return_response_headers(_, __, response_headers):
     return response_headers
 
@@ -799,15 +796,13 @@ def upload_manifest(
                 cls=_return_response_headers,
                 **kwargs
             )
-
             digest = response_headers['Docker-Content-Digest']
+            if not _validate_digest(data, digest):
+                raise ValueError("The digest in the response does not match the digest of the uploaded manifest.")
         except ValueError:
             if repository is None or manifest is None:
                 raise ValueError("The parameter repository and manifest cannot be None.")
-            if not _validate_digest(data, digest):
-                raise ValueError("The digest in the response does not match the digest of the uploaded manifest.")
             raise
-
         return digest
 
     @distributed_trace
@@ -855,6 +850,8 @@ def download_manifest(self, repository, tag_or_digest, **kwargs):
 
         :param str repository: Name of the repository
         :param str tag_or_digest: The tag or digest of the manifest to download.
+            When digest is provided, will use this digest to compare with the one calculated by the response payload.
+            When tag is provided, will use the digest in response headers to compare.
         :returns: DownloadManifestResult
         :rtype: ~azure.containerregistry.models.DownloadManifestResult
         :raises ValueError: If the parameter repository or tag_or_digest is None.
@@ -872,16 +869,18 @@ def download_manifest(self, repository, tag_or_digest, **kwargs):
                     **kwargs
                 )
             )
-            digest = response.http_response.headers['Docker-Content-Digest']
             manifest = OCIManifest.deserialize(cast(ManifestWrapper, manifest_wrapper).serialize())
             manifest_stream = _serialize_manifest(manifest)
+            if tag_or_digest.startswith("sha256:"):
+                digest = tag_or_digest
+            else:
+                digest = response.http_response.headers['Docker-Content-Digest']
+            if not _validate_digest(manifest_stream, digest):
+                raise ValueError("The requested digest does not match the digest of the received manifest.")
         except ValueError:
             if repository is None or tag_or_digest is None:
                 raise ValueError("The parameter repository and tag_or_digest cannot be None.")
-            if not _validate_digest(manifest_stream, digest):
-                raise ValueError("The requested digest does not match the digest of the received manifest.")
             raise
-
         return DownloadManifestResult(digest=digest, data=manifest_stream, manifest=manifest)
 
     @distributed_trace
@@ -896,9 +895,7 @@ def download_blob(self, repository, digest, **kwargs):
         :raises ValueError: If the parameter repository or digest is None.
         """
         try:
-            deserialized = self._client.container_registry_blob.get_blob( # type: ignore
-                repository, digest, cls=_return_deserialized, **kwargs
-            )
+            deserialized = self._client.container_registry_blob.get_blob(repository, digest, **kwargs)
         except ValueError:
             if repository is None or digest is None:
                 raise ValueError("The parameter repository and digest cannot be None.")

sdk/containerregistry/azure-containerregistry/tests/test_container_registry_client.py

@@ -16,7 +16,7 @@
     ArtifactTagOrder,
     ContainerRegistryClient,
 )
-from azure.containerregistry._helpers import _deserialize_manifest
+from azure.containerregistry._helpers import _deserialize_manifest, _serialize_manifest
 from azure.core.exceptions import ResourceNotFoundError, ClientAuthenticationError
 from azure.core.paging import ItemPaged
 from azure.identity import AzureAuthorityHosts
@@ -461,7 +461,6 @@ def test_upload_oci_manifest(self, containerregistry_endpoint):
 
             # Assert
             response = client.download_manifest(repo, digest)
-            assert response.digest == digest
             assert response.data.tell() == 0
             self.assert_manifest(response.manifest, manifest)
 
@@ -472,9 +471,8 @@ def test_upload_oci_manifest(self, containerregistry_endpoint):
     @recorded_by_proxy
     def test_upload_oci_manifest_stream(self, containerregistry_endpoint):
         repo = self.get_resource_name("repo")
-        base_path = os.path.join(self.get_test_directory(), "data", "oci_artifact")
-        manifest_stream = open(os.path.join(base_path, "manifest.json"), "rb")
-        manifest = _deserialize_manifest(manifest_stream)     
+        manifest = self.create_oci_manifest()
+        manifest_stream = _serialize_manifest(manifest)
         with self.create_registry_client(containerregistry_endpoint) as client:
             self.upload_manifest_prerequisites(repo, client)
 
@@ -483,7 +481,6 @@ def test_upload_oci_manifest_stream(self, containerregistry_endpoint):
 
             # Assert
             response = client.download_manifest(repo, digest)
-            assert response.digest == digest
             assert response.data.tell() == 0
             self.assert_manifest(response.manifest, manifest)
 
@@ -494,18 +491,16 @@ def test_upload_oci_manifest_stream(self, containerregistry_endpoint):
     @recorded_by_proxy
     def test_upload_oci_manifest_with_tag(self, containerregistry_endpoint):
         repo = self.get_resource_name("repo")
+        tag = "v1"
         manifest = self.create_oci_manifest()
         with self.create_registry_client(containerregistry_endpoint) as client:
-            tag = "v1"
-            
             self.upload_manifest_prerequisites(repo, client)
 
             # Act
             digest = client.upload_manifest(repo, manifest, tag=tag)
 
             # Assert
             response = client.download_manifest(repo, digest)
-            assert response.digest == digest
             assert response.data.tell() == 0
             self.assert_manifest(response.manifest, manifest)
 
@@ -525,20 +520,17 @@ def test_upload_oci_manifest_with_tag(self, containerregistry_endpoint):
     @recorded_by_proxy
     def test_upload_oci_manifest_stream_with_tag(self, containerregistry_endpoint):
         repo = self.get_resource_name("repo")
-        base_path = os.path.join(self.get_test_directory(), "data", "oci_artifact")
-        manifest_stream = open(os.path.join(base_path, "manifest.json"), "rb")
-        manifest = _deserialize_manifest(manifest_stream)
+        tag = "v1"
+        manifest = self.create_oci_manifest()
+        manifest_stream = _serialize_manifest(manifest)
         with self.create_registry_client(containerregistry_endpoint) as client:
-            tag = "v1"
-            
             self.upload_manifest_prerequisites(repo, client)
 
             # Act
             digest = client.upload_manifest(repo, manifest_stream, tag=tag)
 
             # Assert
             response = client.download_manifest(repo, digest)
-            assert response.digest == digest
             assert response.data.tell() == 0
             self.assert_manifest(response.manifest, manifest)