Add content verification

Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com>

Author: judahrand

CHANGELOG.md

@@ -14,6 +14,7 @@ and **Merged pull requests**. Critical items to know are:
 The versions coincide with releases on pip. Only major versions will be released as tags on Github.
 
 ## [0.0.x](https://github.com/oras-project/oras-py/tree/main) (0.0.x)
+ - validate pulled blob size and digest before writing or extracting content (0.2.43)
  - add Layout `copy` for pull_from_registry capability (0.2.42)
  - make `get_manifest()` validation optional, fix `Accept` header join, and expand default `Accept` header types to cover all supported response types for the `/v2/<name>/manifests/<reference>` endpoint (0.2.41)
  - fix preemptive exit in non-empty `auths` lookup when `credsStore` or `credHelpers` is used (0.2.40)

oras/provider.py

@@ -3,14 +3,18 @@
 __license__ = "Apache-2.0"
 
 import copy
+import enum
+import hashlib
+import hmac
 import os
+import re
 import sys
 import urllib
 from contextlib import contextmanager, nullcontext
 from dataclasses import asdict
 from http.cookiejar import DefaultCookiePolicy
 from tempfile import TemporaryDirectory
-from typing import Callable, Generator, List, Optional, Tuple, Union
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 
 import jsonschema
 import requests
@@ -27,6 +31,87 @@
 from oras.types import container_type
 from oras.utils.fileio import PathAndOptionalContent
 
+_DIGEST_PATTERN = re.compile(
+    r"^(?P<algorithm>[a-z0-9]+(?:[+._-][a-z0-9]+)*):" r"(?P<encoded>[a-zA-Z0-9=_-]+)$"
+)
+
+
+class RegisteredDigestAlgorithm(str, enum.Enum):
+    SHA256 = "sha256"
+    SHA512 = "sha512"
+
+    def hasher(self) -> Any:
+        try:
+            return hashlib.new(self.value)
+        except ValueError as error:
+            raise ValueError(f"Unsupported OCI digest algorithm: {self}.") from error
+
+
+def _parse_digest(digest: str) -> Tuple[RegisteredDigestAlgorithm, str]:
+    """Parse and validate an OCI digest's algorithm and encoded value."""
+    if not isinstance(digest, str):
+        raise ValueError(f"Invalid OCI digest: {digest!r}.")
+
+    match = _DIGEST_PATTERN.fullmatch(digest)
+    if not match:
+        raise ValueError(f"Invalid OCI digest: {digest!r}.")
+
+    algorithm = match.group("algorithm")
+    try:
+        registered_algorithm = RegisteredDigestAlgorithm(algorithm)
+    except ValueError as error:
+        raise ValueError(f"Unsupported OCI digest algorithm: {algorithm}.") from error
+    encoded = match.group("encoded")
+    encoded_length = registered_algorithm.hasher().digest_size * 2
+    if encoded_length is None:
+        raise ValueError(f"Unsupported OCI digest algorithm: {algorithm}.")
+
+    if not re.fullmatch(f"[a-f0-9]{{{encoded_length}}}", encoded):
+        raise ValueError(
+            f"Invalid {algorithm} digest encoding: expected {encoded_length} "
+            "lowercase hexadecimal characters."
+        )
+
+    return registered_algorithm, encoded
+
+
+class Digest:
+    def __init__(self, digest: str) -> None:
+        self.algorithm, self.encoded = _parse_digest(digest)
+
+    @property
+    def digest(self) -> str:
+        return f"{self.algorithm.value}:{self.encoded}"
+
+    def __str__(self) -> str:
+        return self.digest
+
+
+def _validate_downloaded_blob(
+    path: str,
+    digest: Digest,
+    expected_size: int,
+) -> None:
+    """Validate a downloaded blob's descriptor size and digest."""
+    actual_size = os.path.getsize(path)
+    if actual_size != expected_size:
+        raise ValueError(
+            f"Downloaded blob size mismatch for {digest}: expected "
+            f"{expected_size} bytes, got {actual_size} bytes."
+        )
+
+    hasher = digest.algorithm.hasher()
+    with open(path, "rb") as blob:
+        for chunk in iter(lambda: blob.read(8192), b""):
+            hasher.update(chunk)
+
+    actual_encoded = hasher.hexdigest()
+    if not hmac.compare_digest(actual_encoded, digest.encoded):
+        raise ValueError(
+            f"Downloaded blob digest mismatch: expected {digest}, got "
+            f"{digest.algorithm.value}:{actual_encoded}."
+        )
+
 
 @contextmanager
 def temporary_empty_config() -> Generator[str, None, None]:
@@ -889,14 +974,14 @@ def pull(
         :type outdir: str
         :param target: target location to pull from
         :type target: str
+        :raises ValueError: if a layer descriptor or downloaded layer is invalid
         """
         container = self.get_container(target)
         self.auth.load_configs(
             container, configs=[config_path] if config_path else None
         )
         manifest = self.get_manifest(container, allowed_media_type)
         outdir = outdir or oras.utils.get_tmpdir()
-        overwrite = overwrite
 
         files = []
         for layer in manifest.get("layers", []):
@@ -917,17 +1002,40 @@ def pull(
                 )
                 continue
 
-            # A directory will need to be uncompressed and moved
-            if layer["mediaType"] == oras.defaults.default_blob_dir_media_type:
-                targz = oras.utils.get_tmpfile(suffix=".tar.gz")
-                self.download_blob(container, layer["digest"], targz)
+            digest = Digest(layer["digest"])
+            expected_size = layer["size"]
+            if (
+                not isinstance(expected_size, int)
+                or isinstance(expected_size, bool)
+                or expected_size < 0
+            ):
+                raise ValueError(
+                    f"Invalid OCI descriptor size for {digest.digest}: {expected_size!r}."
+                )
+
+            outfile_dir = os.path.dirname(outfile)
+            if outfile_dir and not os.path.exists(outfile_dir):
+                oras.utils.mkdir_p(outfile_dir)
+
+            # Keep downloaded content private until its descriptor is verified.
+            with TemporaryDirectory(prefix=".oras-", dir=outfile_dir) as tmpdir:
+                is_directory = (
+                    layer["mediaType"] == oras.defaults.default_blob_dir_media_type
+                )
+                staged = os.path.join(tmpdir, "blob.tar.gz" if is_directory else "blob")
+                self.download_blob(container, digest.digest, staged)
+                _validate_downloaded_blob(
+                    staged,
+                    digest,
+                    expected_size,
+                )
 
-                # The artifact will be extracted to the correct name
-                oras.utils.extract_targz(targz, os.path.dirname(outfile))
+                # A verified directory archive can now be safely consumed.
+                if is_directory:
+                    oras.utils.extract_targz(staged, outfile_dir)
+                else:
+                    os.replace(staged, outfile)
 
-            # Anything else just extracted directly
-            else:
-                self.download_blob(container, layer["digest"], outfile)
             logger.info(f"Successfully pulled {outfile}.")
             files.append(outfile)
         return files

oras/tests/test_provider.py

@@ -2,6 +2,7 @@
 __copyright__ = "Copyright The ORAS Authors."
 __license__ = "Apache-2.0"
 
+import hashlib
 import os
 import subprocess
 from pathlib import Path
@@ -17,6 +18,147 @@
 here = Path(__file__).resolve().parent
 
 
+def make_pull_client(monkeypatch, layer, content):
+    """Create a registry client whose pull inputs do not require a live registry."""
+    client = oras.provider.Registry(insecure=True)
+    monkeypatch.setattr(client, "get_container", lambda target: target)
+    monkeypatch.setattr(client.auth, "load_configs", lambda *args, **kwargs: None)
+    monkeypatch.setattr(
+        client,
+        "get_manifest",
+        lambda container, allowed_media_type: {"layers": [layer]},
+    )
+
+    def download_blob(container, digest, outfile):
+        Path(outfile).write_bytes(content)
+        return outfile
+
+    monkeypatch.setattr(client, "download_blob", download_blob)
+    return client
+
+
+def test_digest_string_round_trip():
+    original = f"sha256:{hashlib.sha256(b'content').hexdigest()}"
+
+    assert str(oras.provider.Digest(original)) == original
+
+
+@pytest.mark.parametrize("algorithm", ["sha256", "sha512"])
+def test_pull_validates_registered_digest(monkeypatch, tmp_path, algorithm):
+    content = b"verified content"
+    digest = f"{algorithm}:{hashlib.new(algorithm, content).hexdigest()}"
+    layer = {
+        "mediaType": oras.defaults.default_blob_media_type,
+        "size": len(content),
+        "digest": digest,
+        "annotations": {oras.defaults.annotation_title: "artifact.txt"},
+    }
+    client = make_pull_client(monkeypatch, layer, content)
+
+    files = client.pull("registry.example/repository:tag", outdir=str(tmp_path))
+
+    outfile = tmp_path / "artifact.txt"
+    assert files == [str(outfile)]
+    assert outfile.read_bytes() == content
+
+
+def test_pull_rejects_digest_mismatch_without_replacing_file(monkeypatch, tmp_path):
+    expected_content = b"expected content"
+    downloaded_content = b"corrupt! content"
+    digest = f"sha256:{hashlib.sha256(expected_content).hexdigest()}"
+    layer = {
+        "mediaType": oras.defaults.default_blob_media_type,
+        "size": len(downloaded_content),
+        "digest": digest,
+        "annotations": {oras.defaults.annotation_title: "artifact.txt"},
+    }
+    client = make_pull_client(monkeypatch, layer, downloaded_content)
+    outfile = tmp_path / "artifact.txt"
+    outfile.write_bytes(b"existing content")
+
+    with pytest.raises(ValueError) as error:
+        client.pull("registry.example/repository:tag", outdir=str(tmp_path))
+
+    actual_digest = f"sha256:{hashlib.sha256(downloaded_content).hexdigest()}"
+    assert str(error.value) == (
+        f"Downloaded blob digest mismatch: expected {digest}, got {actual_digest}."
+    )
+    assert outfile.read_bytes() == b"existing content"
+    assert not list(tmp_path.glob(".oras-*"))
+
+
+def test_pull_rejects_size_mismatch(monkeypatch, tmp_path):
+    content = b"content"
+    digest = f"sha256:{hashlib.sha256(content).hexdigest()}"
+    layer = {
+        "mediaType": oras.defaults.default_blob_media_type,
+        "size": len(content) + 1,
+        "digest": digest,
+        "annotations": {oras.defaults.annotation_title: "artifact.txt"},
+    }
+    client = make_pull_client(monkeypatch, layer, content)
+
+    with pytest.raises(ValueError, match="Downloaded blob size mismatch"):
+        client.pull("registry.example/repository:tag", outdir=str(tmp_path))
+
+    assert not (tmp_path / "artifact.txt").exists()
+    assert not list(tmp_path.glob(".oras-*"))
+
+
+def test_pull_validates_directory_before_extraction(monkeypatch, tmp_path):
+    expected_content = b"expected archive"
+    downloaded_content = b"corrupted archive"
+    digest = f"sha256:{hashlib.sha256(expected_content).hexdigest()}"
+    layer = {
+        "mediaType": oras.defaults.default_blob_dir_media_type,
+        "size": len(downloaded_content),
+        "digest": digest,
+        "annotations": {oras.defaults.annotation_title: "artifact"},
+    }
+    client = make_pull_client(monkeypatch, layer, downloaded_content)
+    extracted = False
+
+    def extract_targz(*args, **kwargs):
+        nonlocal extracted
+        extracted = True
+
+    monkeypatch.setattr(oras.utils, "extract_targz", extract_targz)
+
+    with pytest.raises(ValueError, match="Downloaded blob digest mismatch"):
+        client.pull("registry.example/repository:tag", outdir=str(tmp_path))
+
+    assert not extracted
+    assert not (tmp_path / "artifact").exists()
+    assert not list(tmp_path.glob(".oras-*"))
+
+
+@pytest.mark.parametrize(
+    ("digest", "error"),
+    [
+        ("sha256+b64u:YWJj", "Unsupported OCI digest algorithm"),
+        (f"sha384:{'a' * 96}", "Unsupported OCI digest algorithm"),
+        (f"sha256:{'A' * 64}", "Invalid sha256 digest encoding"),
+        (f"sha256:{'a' * 63}", "Invalid sha256 digest encoding"),
+        ("sha256:not!hex", "Invalid OCI digest"),
+        ("sha256:", "Invalid OCI digest"),
+        (f"SHA256:{'a' * 64}", "Invalid OCI digest"),
+    ],
+)
+def test_pull_rejects_invalid_digest_encoding(monkeypatch, tmp_path, digest, error):
+    layer = {
+        "mediaType": oras.defaults.default_blob_media_type,
+        "size": 0,
+        "digest": digest,
+        "annotations": {oras.defaults.annotation_title: "artifact.txt"},
+    }
+    client = make_pull_client(monkeypatch, layer, b"")
+
+    with pytest.raises(ValueError, match=error):
+        client.pull("registry.example/repository:tag", outdir=str(tmp_path))
+
+    assert not (tmp_path / "artifact.txt").exists()
+
+
 @pytest.mark.with_auth(False)
 def test_annotated_registry_push(tmp_path, registry, credentials, target):
     """

oras/version.py

@@ -2,7 +2,7 @@
 __copyright__ = "Copyright The ORAS Authors."
 __license__ = "Apache-2.0"
 
-__version__ = "0.2.42"
+__version__ = "0.2.43"
 AUTHOR = "Vanessa Sochat"
 EMAIL = "vsoch@users.noreply.github.com"
 NAME = "oras"