From c19f5fbbbb01e3d0e60313773a24d66fe003e9fb Mon Sep 17 00:00:00 2001
From: Tom Wieczorek <twieczorek@mirantis.com>
Date: Thu, 10 Oct 2024 09:09:59 +0200
Subject: [PATCH] Reuse the reader returned by FetchReference

FetchReference already returns a reader to the manifest content.
However, content.Successors operates on a Fetcher. In order to funnel
the already open reader into content.Successors, create a FetcherFunc
that returns it.

This saves one additional GET request.

Before:
    GET /v2/repository/artifact/manifests/latest
    GET /v2/repository/artifact/manifests/sha256:341098d11767212bb4d148f3fe8e82fc2939c3be247b233291a62e39a594ed09
    GET /v2/repository/artifact/blobs/sha256:9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08

After:
    GET /v2/repository/artifact/manifests/latest
    GET /v2/repository/artifact/blobs/sha256:9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08

Signed-off-by: Tom Wieczorek <twieczorek@mirantis.com>
---
 internal/oci/oci.go      | 36 +++++++++++++++++++++++++++++-------
 internal/oci/oci_test.go |  5 +++++
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/internal/oci/oci.go b/internal/oci/oci.go
index 2e003eace085..e27302e64f4d 100644
--- a/internal/oci/oci.go
+++ b/internal/oci/oci.go
@@ -19,10 +19,12 @@ package oci
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"os"
+	"sync/atomic"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2/content"
@@ -91,13 +93,7 @@ func Download(ctx context.Context, url string, target io.Writer, options ...Down
 	}
 
 	tag := imgref.Reference
-	desc, data, err := repo.Manifests().FetchReference(ctx, tag)
-	if err != nil {
-		return fmt.Errorf("failed to fetch manifest: %w", err)
-	}
-	defer data.Close()
-
-	successors, err := content.Successors(ctx, repo, desc)
+	successors, err := fetchSuccessors(ctx, repo.Manifests(), tag)
 	if err != nil {
 		return fmt.Errorf("failed to fetch successors: %w", err)
 	}
@@ -121,6 +117,32 @@ func Download(ctx context.Context, url string, target io.Writer, options ...Down
 	return nil
 }
 
+// Fetches the manifest for the given reference and returns all of its successors.
+func fetchSuccessors(ctx context.Context, repo registry.ReferenceFetcher, reference string) ([]ocispec.Descriptor, error) {
+	var dataConsumed atomic.Bool
+	desc, data, err := repo.FetchReference(ctx, reference)
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch manifest: %w", err)
+	}
+	defer func() {
+		if dataConsumed.Swap(true) {
+			return
+		}
+		if closeErr := data.Close(); closeErr != nil {
+			err = errors.Join(err, closeErr)
+		}
+	}()
+
+	fetcher := content.FetcherFunc(func(ctx context.Context, target ocispec.Descriptor) (io.ReadCloser, error) {
+		if target.Digest == desc.Digest && !dataConsumed.Swap(true) {
+			return data, nil
+		}
+		return nil, errors.ErrUnsupported
+	})
+
+	return content.Successors(ctx, fetcher, desc)
+}
+
 // findArtifactDescriptor filters, out of the provided list of descriptors, the
 // one that matches the given options. If no artifact name is provided, it
 // returns the first descriptor.
diff --git a/internal/oci/oci_test.go b/internal/oci/oci_test.go
index 37f9208b04f8..aecd65393ae5 100644
--- a/internal/oci/oci_test.go
+++ b/internal/oci/oci_test.go
@@ -125,6 +125,11 @@ func startOCIMockServer(t *testing.T, tname string, test testFile) string {
 
 	server := starter(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			t.Log(r.Proto, r.Method, r.RequestURI)
+			if !assert.Equal(t, r.Method, http.MethodGet) {
+				w.WriteHeader(http.StatusMethodNotAllowed)
+				return
+			}
 
 			// this is a request to authenticate.
 			if strings.Contains(r.URL.Path, "/token") {