fix: centralize pass-through response relaying

internal/handler/composer.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 	"path"
 	"strings"
@@ -491,12 +490,5 @@ func (h *ComposerHandler) proxyUpstream(w http.ResponseWriter, r *http.Request)
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	h.proxy.relayResponse(w, r, resp, nil)
 }

internal/handler/conan.go

@@ -2,7 +2,6 @@ package handler
 
 import (
 	"fmt"
-	"io"
 	"net/http"
 	"strings"
 )
@@ -194,13 +193,5 @@ func (h *ConanHandler) proxyUpstream(w http.ResponseWriter, r *http.Request) {
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	// Copy response headers
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	h.proxy.relayResponse(w, r, resp, nil)
 }

internal/handler/conda.go

@@ -156,13 +156,7 @@ func (h *CondaHandler) handleRepodata(w http.ResponseWriter, r *http.Request) {
 	defer func() { _ = resp.Body.Close() }()
 
 	if resp.StatusCode != http.StatusOK {
-		for k, vv := range resp.Header {
-			for _, v := range vv {
-				w.Header().Add(k, v)
-			}
-		}
-		w.WriteHeader(resp.StatusCode)
-		_, _ = io.Copy(w, resp.Body)
+		h.proxy.relayResponse(w, r, resp, nil)
 		return
 	}
 

internal/handler/container.go

@@ -274,16 +274,16 @@ func (h *ContainerHandler) proxyBlobHead(w http.ResponseWriter, r *http.Request,
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for _, header := range []string{headerContentType, headerContentLength, "Docker-Content-Digest", headerETag, headerLastModified} {
-		if v := resp.Header.Get(header); v != "" {
-			w.Header().Set(header, v)
+	h.proxy.relayResponse(w, r, resp, func(dst, src http.Header) {
+		for _, header := range []string{headerContentType, headerContentLength, "Docker-Content-Digest", headerETag, headerLastModified} {
+			if v := src.Get(header); v != "" {
+				dst.Set(header, v)
+			}
 		}
-	}
-	if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices && w.Header().Get("Docker-Content-Digest") == "" {
-		w.Header().Set("Docker-Content-Digest", digest)
-	}
-
-	w.WriteHeader(resp.StatusCode)
+		if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices && dst.Get("Docker-Content-Digest") == "" {
+			dst.Set("Docker-Content-Digest", digest)
+		}
+	})
 }
 
 // registryForName resolves a client-visible OCI repository name to an upstream

internal/handler/container_manifest.go

@@ -5,7 +5,6 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
-	"io"
 	"mime"
 	"net/http"
 	"regexp"
@@ -77,15 +76,12 @@ func (h *ContainerHandler) serveManifest(w http.ResponseWriter, r *http.Request,
 			writeContainerManifest(w, r, cached, true)
 			return
 		}
-		copyContainerManifestHeaders(w.Header(), resp.Header)
-		w.WriteHeader(resp.StatusCode)
-		_, _ = io.Copy(w, resp.Body)
+		h.proxy.relayResponse(w, r, resp, copyContainerManifestHeaders)
 		return
 	}
 
 	if r.Method == http.MethodHead {
-		copyContainerManifestHeaders(w.Header(), resp.Header)
-		w.WriteHeader(http.StatusOK)
+		h.proxy.relayResponse(w, r, resp, copyContainerManifestHeaders)
 		return
 	}
 

internal/handler/container_tags.go

@@ -5,7 +5,6 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"regexp"
@@ -73,9 +72,7 @@ func (h *ContainerHandler) serveTagsList(w http.ResponseWriter, r *http.Request,
 			writeContainerTags(w, cached, true)
 			return
 		}
-		copyContainerTagsHeaders(w.Header(), resp.Header)
-		w.WriteHeader(resp.StatusCode)
-		_, _ = io.Copy(w, resp.Body)
+		h.proxy.relayResponse(w, r, resp, copyContainerTagsHeaders)
 		return
 	}
 

internal/handler/gem.go

@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net/http"
 	"strings"
 	"time"
@@ -117,17 +116,13 @@ func (h *GemHandler) handleCompactIndex(w http.ResponseWriter, r *http.Request)
 	defer func() { _ = indexResp.Body.Close() }()
 
 	if indexResp.StatusCode != http.StatusOK {
-		copyResponseHeaders(w, indexResp.Header)
-		w.WriteHeader(indexResp.StatusCode)
-		_, _ = io.Copy(w, indexResp.Body)
+		h.proxy.relayResponse(w, r, indexResp, nil)
 		return
 	}
 
 	if filteredVersions == nil {
 		h.proxy.Logger.Warn("failed to fetch version timestamps, proxying unfiltered", "name", name)
-		copyResponseHeaders(w, indexResp.Header)
-		w.WriteHeader(http.StatusOK)
-		_, _ = io.Copy(w, indexResp.Body)
+		h.proxy.relayResponse(w, r, indexResp, nil)
 		return
 	}
 
@@ -215,15 +210,6 @@ func (h *GemHandler) writeFilteredIndex(w http.ResponseWriter, resp *http.Respon
 	}
 }
 
-// copyResponseHeaders copies HTTP headers from a response to a writer.
-func copyResponseHeaders(w http.ResponseWriter, headers http.Header) {
-	for k, vv := range headers {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-}
-
 // gemVersion represents a version entry from the RubyGems versions API.
 type gemVersion struct {
 	Number    string `json:"number"`
@@ -314,15 +300,7 @@ func (h *GemHandler) proxyUpstream(w http.ResponseWriter, r *http.Request) {
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	// Copy response headers
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	h.proxy.relayResponse(w, r, resp, nil)
 }
 
 func init() {

internal/handler/handler.go

@@ -768,7 +768,8 @@ func serveArtifact(w http.ResponseWriter, method string, result *CacheResult) {
 
 // ProxyUpstream forwards a request to an upstream URL without caching.
 // It copies the request, forwards specified headers, and streams the response back.
-// If forwardHeaders is nil, all response headers are copied.
+// forwardHeaders controls the request headers sent upstream. End-to-end response
+// headers and trailers are relayed independently of that list.
 func (p *Proxy) ProxyUpstream(w http.ResponseWriter, r *http.Request, upstreamURL string, forwardHeaders []string) {
 	p.Logger.Debug("proxying to upstream", "url", upstreamURL)
 
@@ -794,17 +795,10 @@ func (p *Proxy) ProxyUpstream(w http.ResponseWriter, r *http.Request, upstreamUR
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	p.relayResponse(w, r, resp, nil)
 }
 
-// ProxyFile forwards a file request to upstream, copying all response headers.
+// ProxyFile forwards a file request, relaying end-to-end headers and trailers.
 func (p *Proxy) ProxyFile(w http.ResponseWriter, r *http.Request, upstreamURL string) {
 	req, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, nil)
 	if err != nil {
@@ -820,14 +814,7 @@ func (p *Proxy) ProxyFile(w http.ResponseWriter, r *http.Request, upstreamURL st
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for key, values := range resp.Header {
-		for _, v := range values {
-			w.Header().Add(key, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	p.relayResponse(w, r, resp, nil)
 }
 
 // JSONError writes a JSON error response.
@@ -1273,16 +1260,13 @@ func (p *Proxy) proxyMetadataStream(w http.ResponseWriter, r *http.Request, upst
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for _, header := range []string{headerContentType, headerContentLength, headerContentEncoding, headerLastModified, headerETag} {
-		if v := resp.Header.Get(header); v != "" {
-			w.Header().Set(header, v)
+	p.relayResponse(w, r, resp, func(dst, src http.Header) {
+		for _, header := range []string{headerContentType, headerContentLength, headerContentEncoding, headerLastModified, headerETag} {
+			if v := src.Get(header); v != "" {
+				dst.Set(header, v)
+			}
 		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	if r.Method != http.MethodHead {
-		_, _ = io.Copy(w, resp.Body)
-	}
+	})
 }
 
 func (p *Proxy) applyUpstreamAuth(req *http.Request) {

internal/handler/hex.go

@@ -117,13 +117,7 @@ func (h *HexHandler) handlePackages(w http.ResponseWriter, r *http.Request) {
 	defer func() { _ = protoResp.Body.Close() }()
 
 	if protoResp.StatusCode != http.StatusOK {
-		for k, vv := range protoResp.Header {
-			for _, v := range vv {
-				w.Header().Add(k, v)
-			}
-		}
-		w.WriteHeader(protoResp.StatusCode)
-		_, _ = io.Copy(w, protoResp.Body)
+		h.proxy.relayResponse(w, r, protoResp, nil)
 		return
 	}
 

internal/handler/nuget.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 	"strings"
 )
@@ -222,14 +221,7 @@ func (h *NuGetHandler) proxyUpstream(w http.ResponseWriter, r *http.Request) {
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	h.proxy.relayResponse(w, r, resp, nil)
 }
 
 // buildUpstreamURL constructs the upstream URL for a request.

internal/handler/pypi.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"mime"
 	"net/http"
 	"regexp"
@@ -817,13 +816,8 @@ func (h *PyPIHandler) proxySimple(w http.ResponseWriter, r *http.Request, path s
 	}
 	defer func() { _ = resp.Body.Close() }()
 
-	for k, vv := range resp.Header {
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
-	}
-	ensureVaryAccept(w.Header())
-
-	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+	h.proxy.relayResponse(w, r, resp, func(dst, src http.Header) {
+		copyRelayHeaders(dst, src)
+		ensureVaryAccept(dst)
+	})
 }