Pull integrity checks (#784)

Signed-off-by: Valentin Delaye <jonesbusy@users.noreply.github.com>

Author: jonesbusy

src/main/java/land/oras/Registry.java

@@ -22,6 +22,7 @@
 
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import java.io.FilterInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
@@ -31,8 +32,10 @@
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
 import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HexFormat;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -1160,24 +1163,11 @@ public byte[] getBlob(ContainerRef containerRef) {
     }
 
     private byte[] getBlobDirect(ContainerRef containerRef) {
-        ContainerRef ref = containerRef.forRegistry(this).checkBlocked(this);
-        if (ref.isInsecure(this) && !this.isInsecure()) {
-            return copyForNewTransport(ref.getRegistry(), true).getBlobDirect(ref);
-        }
-        if (!ref.isInsecure(this) && this.isInsecure()) {
-            return copyForNewTransport(ref.getRegistry(), false).getBlobDirect(ref);
+        try (InputStream is = fetchBlobDirect(containerRef)) {
+            return is.readAllBytes();
+        } catch (IOException e) {
+            throw new OrasException("Failed to get blob", e);
         }
-        URI uri = URI.create("%s://%s".formatted(getScheme(), ref.getBlobsPath(this)));
-        HttpClient.ResponseWrapper<String> response = client.get(
-                uri,
-                Map.of(Const.ACCEPT_HEADER, Const.APPLICATION_OCTET_STREAM_HEADER_VALUE),
-                Scopes.of(ref),
-                authProvider);
-        logResponse(response);
-        handleError(response);
-        byte[] data = response.response().getBytes(StandardCharsets.UTF_8);
-        validateDockerContentDigest(response, data);
-        return data;
     }
 
     @Override
@@ -1207,7 +1197,7 @@ private void fetchBlobDirect(ContainerRef containerRef, Path path) {
                 authProvider);
         logResponse(response);
         handleError(response);
-        validateDockerContentDigest(response, path);
+        verifyBlobDigest(ref, path, response.headers());
     }
 
     @Override
@@ -1231,8 +1221,30 @@ private InputStream fetchBlobDirect(ContainerRef containerRef) {
                 authProvider);
         logResponse(response);
         handleError(response);
-        validateDockerContentDigest(response);
-        return response.response();
+        // Verify the content digest incrementally
+        List<String> expected = expectedBlobDigests(ref, response.headers());
+        if (expected.isEmpty()) {
+            return response.response();
+        }
+        return new DigestVerifyingInputStream(response.response(), expected);
+    }
+
+    /**
+     * Collect the digests a downloaded blob must match: the caller-pinned digest (when it is a supported
+     * digest and not a plain tag) and the server-advertised {@code Docker-Content-Digest} header (when
+     * present). Duplicates are collapsed so the content is hashed once per distinct algorithm.
+     */
+    private static List<String> expectedBlobDigests(ContainerRef ref, Map<String, String> headers) {
+        List<String> expected = new ArrayList<>(2);
+        String pinned = ref.getDigest();
+        if (pinned != null && SupportedAlgorithm.isSupported(pinned)) {
+            expected.add(pinned);
+        }
+        String header = headers.get(Const.DOCKER_CONTENT_DIGEST_HEADER.toLowerCase());
+        if (header != null && !expected.contains(header)) {
+            expected.add(header);
+        }
+        return expected;
     }
 
     @Override
@@ -1304,16 +1316,16 @@ public Descriptor getDescriptor(ContainerRef containerRef) {
         HttpClient.ResponseWrapper<String> response = getManifestResponse(containerRef);
         logResponse(response);
         handleError(response);
+        String json = response.response();
+        // When the reference is pinned to a digest, verify the returned manifest/index bytes against it
+        verifyPinnedDigest(containerRef, json.getBytes(StandardCharsets.UTF_8));
         String size = response.headers().get(Const.CONTENT_LENGTH_HEADER.toLowerCase());
         String contentType = response.headers().get(Const.CONTENT_TYPE_HEADER.toLowerCase());
         return Descriptor.of(
                         validateDockerContentDigest(response),
-                        Long.parseLong(
-                                size == null
-                                        ? String.valueOf(response.response().length())
-                                        : size),
+                        Long.parseLong(size == null ? String.valueOf(json.length()) : size),
                         contentType)
-                .withJson(response.response());
+                .withJson(json);
     }
 
     @Override
@@ -1407,30 +1419,37 @@ private HttpClient.ResponseWrapper<String> getManifestResponseDirect(ContainerRe
         return client.get(uri, Map.of("Accept", Const.MANIFEST_ACCEPT_TYPE), Scopes.of(ref), authProvider);
     }
 
-    private void validateDockerContentDigest(HttpClient.ResponseWrapper<String> response, byte[] data) {
-        String digest = response.headers().get(Const.DOCKER_CONTENT_DIGEST_HEADER.toLowerCase());
-        // This might happen when blob are hosted other storage.
-        // We need a way to propagate the headers like scoped.
-        // For now just skip validation
-        if (digest == null) {
-            LOG.debug("Docker-Content-Digest header not found in response. Skipping validation.");
+    /**
+     * Verify in-memory content against the digest pinned in the reference. No-op when the reference is
+     * not digest-pinned (e.g. pulled by tag), since there is then no trusted value to check against.
+     * @param ref The reference the content was requested for
+     * @param content The bytes returned by the registry
+     */
+    private void verifyPinnedDigest(ContainerRef ref, byte[] content) {
+        String digest = ref.getDigest();
+        if (digest == null || !SupportedAlgorithm.isSupported(digest)) {
             return;
         }
-        String computedDigest = SupportedAlgorithm.fromDigest(digest).digest(data);
-        ensureDigest(digest, computedDigest);
+        ensureDigest(digest, SupportedAlgorithm.fromDigest(digest).digest(content));
     }
 
-    private void validateDockerContentDigest(HttpClient.ResponseWrapper<Path> response, Path path) {
-        String digest = response.headers().get(Const.DOCKER_CONTENT_DIGEST_HEADER.toLowerCase());
-        // This might happen when blob are hosted other storage.
-        // We need a way to propagate the headers like scoped.
-        // For now just skip validation
-        if (digest == null) {
-            LOG.debug("Docker-Content-Digest header not found in response. Skipping validation.");
-            return;
+    /**
+     * Verify a downloaded blob against every digest that applies (pinned or from header)
+     * @param ref The reference the blob was requested for
+     * @param content The path the registry response was written to
+     * @param headers The response headers
+     */
+    private void verifyBlobDigest(ContainerRef ref, Path content, Map<String, String> headers) {
+        String pinned = ref.getDigest();
+        String verified = null;
+        if (pinned != null && SupportedAlgorithm.isSupported(pinned)) {
+            ensureDigest(pinned, SupportedAlgorithm.fromDigest(pinned).digest(content));
+            verified = pinned;
+        }
+        String header = headers.get(Const.DOCKER_CONTENT_DIGEST_HEADER.toLowerCase());
+        if (header != null && !header.equals(verified)) {
+            ensureDigest(header, SupportedAlgorithm.fromDigest(header).digest(content));
         }
-        String computedDigest = SupportedAlgorithm.fromDigest(digest).digest(path);
-        ensureDigest(digest, computedDigest);
     }
 
     private @Nullable String validateDockerContentDigest(HttpClient.ResponseWrapper<?> response) {
@@ -1439,9 +1458,7 @@ private void validateDockerContentDigest(HttpClient.ResponseWrapper<Path> respon
 
     private @Nullable String validateDockerContentDigest(Map<String, String> headers) {
         String digest = headers.get(Const.DOCKER_CONTENT_DIGEST_HEADER.toLowerCase());
-        // This might happen when blob are hosted other storage.
-        // We need a way to propagate the headers like scoped.
-        // For now just skip validation
+        // Not mandatory, but require validation if present
         if (digest == null) {
             LOG.debug("Docker-Content-Digest header not found in response. Skipping validation.");
             return null;
@@ -2001,4 +2018,61 @@ public Registry build() {
             return registry.build();
         }
     }
+
+    /**
+     * A stream that verifies the content digest incrementally as it is read, without buffering to a
+     * temporary file. It updates one {@link MessageDigest} per distinct algorithm as bytes pass through
+     * and, at end-of-stream
+     */
+    private static final class DigestVerifyingInputStream extends FilterInputStream {
+
+        private final List<String> expectedDigests;
+        private final Map<String, MessageDigest> digestsByPrefix = new HashMap<>();
+        private boolean verified;
+
+        private DigestVerifyingInputStream(InputStream in, List<String> expectedDigests) {
+            super(in);
+            this.expectedDigests = expectedDigests;
+            for (String expected : expectedDigests) {
+                SupportedAlgorithm algorithm = SupportedAlgorithm.fromDigest(expected);
+                digestsByPrefix.computeIfAbsent(algorithm.getPrefix(), p -> {
+                    try {
+                        return MessageDigest.getInstance(algorithm.getAlgorithmName());
+                    } catch (NoSuchAlgorithmException e) {
+                        throw new OrasException("Unsupported digest algorithm: " + algorithm.getAlgorithmName(), e);
+                    }
+                });
+            }
+        }
+
+        @Override
+        public int read(byte[] buffer, int off, int len) throws IOException {
+            int read = super.read(buffer, off, len);
+            if (read > 0) {
+                for (MessageDigest digest : digestsByPrefix.values()) {
+                    digest.update(buffer, off, read);
+                }
+            } else if (read < 0) {
+                verify();
+            }
+            return read;
+        }
+
+        private void verify() {
+            if (verified) {
+                return;
+            }
+            verified = true;
+            Map<String, String> actualByPrefix = new HashMap<>();
+            digestsByPrefix.forEach((prefix, digest) ->
+                    actualByPrefix.put(prefix, prefix + ":" + HexFormat.of().formatHex(digest.digest())));
+            for (String expected : expectedDigests) {
+                String prefix = SupportedAlgorithm.fromDigest(expected).getPrefix();
+                String actual = actualByPrefix.get(prefix);
+                if (!expected.equals(actual)) {
+                    throw new OrasException("Digest mismatch: %s != %s".formatted(expected, actual));
+                }
+            }
+        }
+    }
 }

src/test/java/land/oras/RegistryWireMockTest.java

@@ -21,6 +21,7 @@
 package land.oras;
 
 import static com.github.tomakehurst.wiremock.client.WireMock.*;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
@@ -1435,8 +1436,10 @@ void shouldFailCopyWhenIndexNestingExceedsMaxDepth(WireMockRuntimeInfo wmRuntime
         // A malicious source serves an unbounded chain of nested indexes
         int levels = 40;
         String[] digests = new String[levels + 2];
-        for (int i = 0; i < digests.length; i++) {
-            digests[i] = SupportedAlgorithm.SHA256.digest(("deep-chain-" + i).getBytes(StandardCharsets.UTF_8));
+        digests[levels + 1] = SupportedAlgorithm.SHA256.digest("deep-chain-leaf".getBytes(StandardCharsets.UTF_8));
+        for (int i = levels; i >= 0; i--) {
+            digests[i] = SupportedAlgorithm.SHA256.digest(
+                    nestedIndexJson(digests[i + 1]).getBytes(StandardCharsets.UTF_8));
         }
 
         // Deeper levels
@@ -1554,4 +1557,128 @@ private static void stubNestedSourceIndex(
                         .withHeader(Const.DOCKER_CONTENT_DIGEST_HEADER, selfDigest)
                         .withBody(nestedIndexJson(childDigest))));
     }
+
+    @Test
+    void shouldRejectBlobWhenContentDoesNotMatchPinnedDigest(WireMockRuntimeInfo wmRuntimeInfo) {
+        String pinnedDigest = SupportedAlgorithm.SHA256.digest("good-data".getBytes(StandardCharsets.UTF_8));
+        String evil = "evil-data";
+        String selfConsistentHeader = SupportedAlgorithm.SHA256.digest(evil.getBytes(StandardCharsets.UTF_8));
+
+        WireMock wireMock = wmRuntimeInfo.getWireMock();
+        wireMock.register(WireMock.get(WireMock.urlEqualTo("/v2/library/evil-blob/blobs/%s".formatted(pinnedDigest)))
+                .willReturn(WireMock.ok()
+                        .withBody(evil)
+                        .withHeader(Const.DOCKER_CONTENT_DIGEST_HEADER, selfConsistentHeader)));
+
+        Registry registry = Registry.Builder.builder()
+                .withAuthProvider(authProvider)
+                .withInsecure(true)
+                .build();
+        ContainerRef ref = ContainerRef.parse("localhost:%d/library/evil-blob".formatted(wmRuntimeInfo.getHttpPort()))
+                .withDigest(pinnedDigest);
+
+        OrasException ex = assertThrows(OrasException.class, () -> registry.getBlob(ref));
+        assertTrue(ex.getMessage().contains("Digest mismatch"), "Unexpected: " + ex.getMessage());
+    }
+
+    @Test
+    void shouldReturnBinaryBlobWithoutCorruption(WireMockRuntimeInfo wmRuntimeInfo) {
+        byte[] binary = new byte[] {0x00, (byte) 0xC3, 0x28, (byte) 0xFF, (byte) 0x80, 0x7F, (byte) 0xFE};
+        String digest = SupportedAlgorithm.SHA256.digest(binary);
+
+        WireMock wireMock = wmRuntimeInfo.getWireMock();
+        wireMock.register(WireMock.get(WireMock.urlEqualTo("/v2/library/bin-blob/blobs/%s".formatted(digest)))
+                .willReturn(WireMock.ok().withBody(binary).withHeader(Const.DOCKER_CONTENT_DIGEST_HEADER, digest)));
+
+        Registry registry = Registry.Builder.builder()
+                .withAuthProvider(authProvider)
+                .withInsecure(true)
+                .build();
+        ContainerRef ref = ContainerRef.parse("localhost:%d/library/bin-blob".formatted(wmRuntimeInfo.getHttpPort()))
+                .withDigest(digest);
+
+        assertArrayEquals(binary, registry.getBlob(ref));
+    }
+
+    @Test
+    void shouldRejectFetchBlobToPathOnDigestMismatch(WireMockRuntimeInfo wmRuntimeInfo) {
+        String pinnedDigest = SupportedAlgorithm.SHA256.digest("good".getBytes(StandardCharsets.UTF_8));
+
+        WireMock wireMock = wmRuntimeInfo.getWireMock();
+        wireMock.register(WireMock.get(WireMock.urlEqualTo("/v2/library/evil-file/blobs/%s".formatted(pinnedDigest)))
+                .willReturn(WireMock.ok().withBody("tampered")));
+
+        Registry registry = Registry.Builder.builder()
+                .withAuthProvider(authProvider)
+                .withInsecure(true)
+                .build();
+        ContainerRef ref = ContainerRef.parse("localhost:%d/library/evil-file".formatted(wmRuntimeInfo.getHttpPort()))
+                .withDigest(pinnedDigest);
+        Path out = configDir.resolve("blob-out.bin");
+
+        OrasException ex = assertThrows(OrasException.class, () -> registry.fetchBlob(ref, out));
+        assertTrue(ex.getMessage().contains("Digest mismatch"), "Unexpected: " + ex.getMessage());
+    }
+
+    @Test
+    void shouldRejectBlobStreamOnDigestMismatch(WireMockRuntimeInfo wmRuntimeInfo) {
+        String pinnedDigest = SupportedAlgorithm.SHA256.digest("good".getBytes(StandardCharsets.UTF_8));
+
+        WireMock wireMock = wmRuntimeInfo.getWireMock();
+        wireMock.register(WireMock.get(WireMock.urlEqualTo("/v2/library/evil-stream/blobs/%s".formatted(pinnedDigest)))
+                .willReturn(WireMock.ok().withBody("tampered")));
+
+        Registry registry = Registry.Builder.builder()
+                .withAuthProvider(authProvider)
+                .withInsecure(true)
+                .build();
+        ContainerRef ref = ContainerRef.parse("localhost:%d/library/evil-stream".formatted(wmRuntimeInfo.getHttpPort()))
+                .withDigest(pinnedDigest);
+
+        // The digest is verified incrementally, so the mismatch surfaces when the stream is read to EOF.
+        OrasException ex = assertThrows(OrasException.class, () -> {
+            try (InputStream is = registry.getBlobStream(ref)) {
+                is.readAllBytes();
+            }
+        });
+        assertTrue(ex.getMessage().contains("Digest mismatch"), "Unexpected: " + ex.getMessage());
+    }
+
+    @Test
+    void shouldRejectManifestWhoseContentDoesNotMatchPinnedDigest(WireMockRuntimeInfo wmRuntimeInfo) {
+        String realManifest =
+                """
+                {"schemaVersion":2,"mediaType":"application/vnd.oci.image.manifest.v1+json",\
+                "config":{"mediaType":"application/vnd.oci.empty.v1+json",\
+                "digest":"sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a","size":2},\
+                "layers":[]}""";
+        String evilManifest = realManifest.replace("\"layers\":[]", "\"layers\":[],\"annotations\":{\"x\":\"y\"}");
+        String pinnedDigest = SupportedAlgorithm.SHA256.digest(realManifest.getBytes(StandardCharsets.UTF_8));
+        String selfConsistentHeader = SupportedAlgorithm.SHA256.digest(evilManifest.getBytes(StandardCharsets.UTF_8));
+
+        WireMock wireMock = wmRuntimeInfo.getWireMock();
+        String manifestPath = "/v2/library/evil-manifest/manifests/%s".formatted(pinnedDigest);
+        wireMock.register(WireMock.head(WireMock.urlEqualTo(manifestPath))
+                .willReturn(WireMock.aResponse()
+                        .withStatus(200)
+                        .withHeader(Const.CONTENT_TYPE_HEADER, Const.DEFAULT_MANIFEST_MEDIA_TYPE)
+                        .withHeader(Const.DOCKER_CONTENT_DIGEST_HEADER, selfConsistentHeader)));
+        wireMock.register(WireMock.get(WireMock.urlEqualTo(manifestPath))
+                .willReturn(WireMock.aResponse()
+                        .withStatus(200)
+                        .withHeader(Const.CONTENT_TYPE_HEADER, Const.DEFAULT_MANIFEST_MEDIA_TYPE)
+                        .withHeader(Const.DOCKER_CONTENT_DIGEST_HEADER, selfConsistentHeader)
+                        .withBody(evilManifest)));
+
+        Registry registry = Registry.Builder.builder()
+                .withAuthProvider(authProvider)
+                .withInsecure(true)
+                .build();
+        ContainerRef ref = ContainerRef.parse(
+                        "localhost:%d/library/evil-manifest".formatted(wmRuntimeInfo.getHttpPort()))
+                .withDigest(pinnedDigest);
+
+        OrasException ex = assertThrows(OrasException.class, () -> registry.getManifest(ref));
+        assertTrue(ex.getMessage().contains("Digest mismatch"), "Unexpected: " + ex.getMessage());
+    }
 }