From 041824555ca2a5967699e0281e48e1e7b65eec93 Mon Sep 17 00:00:00 2001
From: Mikel Rychliski <mikel@mikelr.com>
Date: Fri, 26 Jan 2024 18:56:45 -0500
Subject: [PATCH] Include headers when serving blob through proxy

In commit 17952924f3 we updated ServeBlob() to use an io.MultiWriter to
write simultaneously to the local store and the HTTP response.

However, copyContent was using a type assertion to only add headers if
the io.Writer was a http.ResponseWriter. Therefore, this change caused
us to stop sending the expected headers (i.e. Content-Length, Etag,
etc.) on the first request for a blob.

Resolve the issue by explicitly passing in http.Header and setting it
unconditionally.

Signed-off-by: Mikel Rychliski <mikel@mikelr.com>
---
 registry/proxy/proxyblobstore.go      | 20 +++++++++-----------
 registry/proxy/proxyblobstore_test.go | 12 +++++++++++-
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/registry/proxy/proxyblobstore.go b/registry/proxy/proxyblobstore.go
index 0a082888..7dc3a1c4 100644
--- a/registry/proxy/proxyblobstore.go
+++ b/registry/proxy/proxyblobstore.go
@@ -33,22 +33,20 @@ var inflight = make(map[digest.Digest]struct{})
 // mu protects inflight
 var mu sync.Mutex
 
-func setResponseHeaders(w http.ResponseWriter, length int64, mediaType string, digest digest.Digest) {
-	w.Header().Set("Content-Length", strconv.FormatInt(length, 10))
-	w.Header().Set("Content-Type", mediaType)
-	w.Header().Set("Docker-Content-Digest", digest.String())
-	w.Header().Set("Etag", digest.String())
+func setResponseHeaders(h http.Header, length int64, mediaType string, digest digest.Digest) {
+	h.Set("Content-Length", strconv.FormatInt(length, 10))
+	h.Set("Content-Type", mediaType)
+	h.Set("Docker-Content-Digest", digest.String())
+	h.Set("Etag", digest.String())
 }
 
-func (pbs *proxyBlobStore) copyContent(ctx context.Context, dgst digest.Digest, writer io.Writer) (distribution.Descriptor, error) {
+func (pbs *proxyBlobStore) copyContent(ctx context.Context, dgst digest.Digest, writer io.Writer, h http.Header) (distribution.Descriptor, error) {
 	desc, err := pbs.remoteStore.Stat(ctx, dgst)
 	if err != nil {
 		return distribution.Descriptor{}, err
 	}
 
-	if w, ok := writer.(http.ResponseWriter); ok {
-		setResponseHeaders(w, desc.Size, desc.MediaType, dgst)
-	}
+	setResponseHeaders(h, desc.Size, desc.MediaType, dgst)
 
 	remoteReader, err := pbs.remoteStore.Open(ctx, dgst)
 	if err != nil {
@@ -102,7 +100,7 @@ func (pbs *proxyBlobStore) ServeBlob(ctx context.Context, w http.ResponseWriter,
 		// Will return the blob from the remote store directly.
 		// TODO Maybe we could reuse the these blobs are serving remotely and caching locally.
 		mu.Unlock()
-		_, err := pbs.copyContent(ctx, dgst, w)
+		_, err := pbs.copyContent(ctx, dgst, w, w.Header())
 		return err
 	}
 	inflight[dgst] = struct{}{}
@@ -122,7 +120,7 @@ func (pbs *proxyBlobStore) ServeBlob(ctx context.Context, w http.ResponseWriter,
 	// Serving client and storing locally over same fetching request.
 	// This can prevent a redundant blob fetching.
 	multiWriter := io.MultiWriter(w, bw)
-	desc, err := pbs.copyContent(ctx, dgst, multiWriter)
+	desc, err := pbs.copyContent(ctx, dgst, multiWriter, w.Header())
 	if err != nil {
 		return err
 	}
diff --git a/registry/proxy/proxyblobstore_test.go b/registry/proxy/proxyblobstore_test.go
index e7dccc09..91218a40 100644
--- a/registry/proxy/proxyblobstore_test.go
+++ b/registry/proxy/proxyblobstore_test.go
@@ -448,12 +448,22 @@ func testProxyStoreServe(t *testing.T, te *testEnv, numClients int) {
 					return
 				}
 
-				bodyBytes := w.Body.Bytes()
+				resp := w.Result()
+				bodyBytes, err := io.ReadAll(resp.Body)
+				resp.Body.Close()
+				if err != nil {
+					t.Errorf(err.Error())
+					return
+				}
 				localDigest := digest.FromBytes(bodyBytes)
 				if localDigest != remoteBlob.Digest {
 					t.Errorf("Mismatching blob fetch from proxy")
 					return
 				}
+				if resp.Header.Get("Docker-Content-Digest") != localDigest.String() {
+					t.Errorf("Mismatching digest in response header")
+					return
+				}
 
 				desc, err := te.store.localStore.Stat(te.ctx, remoteBlob.Digest)
 				if err != nil {