fix(apps): select layer digest for the correct source repository

apps/docker_store.py

@@ -91,6 +91,7 @@ class Image:
 
         def __init__(self, image_ref, data_root, image_conf_hash):
             self._image_ref = image_ref
+            self._image_src_rep = get_source_repository(image_ref)
             self._data_root = data_root
             self._layer_digests = []
             if not image_conf_hash.startswith(self._SUPPORTED_HASH_TYPE):
@@ -129,25 +130,41 @@ def get_layer_digest(self, diff_id, idx):
             # `docker push`.
             digest_file_path = os.path.join(self._data_root, self._DISTRIBUTION_DIGEST_PATH,
                                             diff_id[len(self._SUPPORTED_HASH_TYPE):])
-            if os.path.exists(digest_file_path):
-                with open(digest_file_path) as f:
-                    digests = json.load(f)
-                    return digests[0]["Digest"]
-            else:
-                print(f"Image layer diff ID to digest mapping is not found in: {digest_file_path}, "
-                      f"fetching image manifest to get its layer digests; uri: {self._image_ref}...")
-                output = subprocess.check_output(
-                    ["skopeo", "inspect", f"docker://{self._image_ref}"])
-                image_desc = json.loads(output)
-                for layer in image_desc["Layers"]:
-                    self._layer_digests.append(layer)
-
-                if len(self._layer_digests) <= idx:
-                    raise Exception("the number of image layer diffIDs and layer digests does not"
-                                    f" match; digests number: {len(self._layer_digests)},"
-                                    f" diffID index: {idx}, image: {self._image_ref}")
-
-                return self._layer_digests[idx]
+            if self._image_src_rep and os.path.exists(digest_file_path):
+                try:
+                    with open(digest_file_path) as f:
+                        digests = json.load(f)
+
+                    for mapping in digests:
+                        if mapping.get("SourceRepository") == self._image_src_rep:
+                            digest = mapping.get("Digest")
+                            if digest:
+                                return digest
+
+                except (OSError, json.JSONDecodeError, TypeError) as err:
+                    print(
+                        "Failed to get layer digest from the diff-ID-to-digest mapping; "
+                        f"image: {self._image_ref}, diff_id: {diff_id}, err: {err}"
+                    )
+
+            # If the image source repo is unknown or if the diffid to digest mapping file does
+            # not exist or the digest is not found for the given diff_id and image source repo,
+            # then fallback to fetching an image manifest containing a list of layer digests
+            # and get the digest for the given diff_if from it by using its index.
+            print(f"Image layer diff ID to digest mapping is not found in: {digest_file_path}, "
+                  f"fetching image manifest to get its layer digests; uri: {self._image_ref}...")
+            output = subprocess.check_output(
+                ["skopeo", "inspect", f"docker://{self._image_ref}"])
+            image_desc = json.loads(output)
+            for layer in image_desc["Layers"]:
+                self._layer_digests.append(layer)
+
+            if len(self._layer_digests) <= idx:
+                raise Exception("the number of image layer diffIDs and layer digests does not"
+                                f" match; digests number: {len(self._layer_digests)},"
+                                f" diffID index: {idx}, image: {self._image_ref}")
+
+            return self._layer_digests[idx]
 
     def __init__(self, data_root="/var/lib/docker"):
         self.data_root = data_root
@@ -169,3 +186,16 @@ def _parse_repositories(self):
                 if image_conf_hash not in self._cfg_to_image:
                     self._cfg_to_image[image_conf_hash] = self.Image(ref, self.data_root, image_conf_hash)
                 self.images_by_ref[ref] = self._cfg_to_image[image_conf_hash]
+
+
+def get_source_repository(image_ref) -> str:
+    if not isinstance(image_ref, str) or not image_ref:
+        return ""
+
+    repository = image_ref.split("@", 1)[0]
+
+    last_part = repository.rsplit("/", 1)[-1]
+    if ":" in last_part:
+        repository = repository.rsplit(":", 1)[0]
+
+    return repository

tests/test_docker_store.py

@@ -0,0 +1,147 @@
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+from apps.docker_store import get_source_repository, DockerStore
+
+
+class GetSourceRepositoryTests(unittest.TestCase):
+    def test_reference_with_tag(self):
+        image_ref = (
+            "ghcr.io/linux-validation/lava/"
+            "lava-server-amd64:2026.05-qualcomm-2026.05-22-d5988428c"
+        )
+
+        self.assertEqual(
+            get_source_repository(image_ref),
+            "ghcr.io/linux-validation/lava/lava-server-amd64",
+        )
+
+    def test_reference_with_digest(self):
+        image_ref = (
+            "ghcr.io/linux-validation/lava/"
+            "lava-server-amd64@sha256:0123456789abcdef"
+        )
+
+        self.assertEqual(
+            get_source_repository(image_ref),
+            "ghcr.io/linux-validation/lava/lava-server-amd64",
+        )
+
+    def test_reference_with_tag_and_digest(self):
+        image_ref = (
+            "ghcr.io/linux-validation/lava/"
+            "lava-server-amd64:latest@sha256:0123456789abcdef"
+        )
+
+        self.assertEqual(
+            get_source_repository(image_ref),
+            "ghcr.io/linux-validation/lava/lava-server-amd64",
+        )
+
+    def test_registry_with_port(self):
+        self.assertEqual(
+            get_source_repository("localhost:5000/example/image:latest"),
+            "localhost:5000/example/image",
+        )
+
+    def test_reference_without_tag_or_digest(self):
+        self.assertEqual(
+            get_source_repository("ghcr.io/example/image"),
+            "ghcr.io/example/image",
+        )
+
+    def test_short_image_name_with_tag(self):
+        self.assertEqual(
+            get_source_repository("image:latest"),
+            "image",
+        )
+
+    def test_empty_string(self):
+        self.assertEqual(get_source_repository(""), "")
+
+    def test_none(self):
+        self.assertEqual(get_source_repository(None), "")
+
+    def test_non_string(self):
+        self.assertEqual(get_source_repository(123), "")
+
+
+class GetLayerDigestTests(unittest.TestCase):
+    def test_get_layer_digest_from_local_mapping(self):
+        diff_id = "sha256:diff-id"
+        expected_digest = "sha256:layer-digest"
+        source_repository = "ghcr.io/example/image"
+
+        with tempfile.TemporaryDirectory() as data_root:
+            mapping_file = (
+                Path(data_root)
+                / DockerStore.Image._DISTRIBUTION_DIGEST_PATH
+                / "diff-id"
+            )
+            mapping_file.parent.mkdir(parents=True)
+            mapping_file.write_text(
+                json.dumps(
+                    [
+                        {
+                            "SourceRepository": source_repository,
+                            "Digest": expected_digest,
+                        }
+                    ]
+                )
+            )
+
+            image = DockerStore.Image.__new__(DockerStore.Image)
+            image._data_root = data_root
+            image._image_src_rep = source_repository
+            image._image_ref = source_repository + ":latest"
+            image._layer_digests = []
+
+            self.assertEqual(
+                image.get_layer_digest(diff_id, 0),
+                expected_digest,
+            )
+
+    def test_get_layer_digest_selects_matching_source_repository(self):
+        diff_id = "sha256:diff-id"
+        source_repository = "ghcr.io/example/image"
+        expected_digest = "sha256:digest-from-example-repo"
+
+        with tempfile.TemporaryDirectory() as data_root:
+            mapping_file = (
+                    Path(data_root)
+                    / DockerStore.Image._DISTRIBUTION_DIGEST_PATH
+                    / "diff-id"
+            )
+            mapping_file.parent.mkdir(parents=True)
+            mapping_file.write_text(
+                json.dumps(
+                    [
+                        {
+                            "SourceRepository": "ghcr.io/other/image",
+                            "Digest": "sha256:digest-from-other-repo",
+                        },
+                        {
+                            "SourceRepository": source_repository,
+                            "Digest": expected_digest,
+                        },
+                    ]
+                ),
+                encoding="utf-8",
+            )
+
+            image = DockerStore.Image.__new__(DockerStore.Image)
+            image._data_root = data_root
+            image._image_src_rep = source_repository
+            image._image_ref = source_repository + ":latest"
+            image._layer_digests = []
+
+            self.assertEqual(
+                image.get_layer_digest(diff_id, 0),
+                expected_digest,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()