Fetch rewritten mutable responses on HEAD to match GET length

e2e/head_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -98,3 +99,70 @@ func TestProxyHeadGitHubRPMRepodata(t *testing.T) {
 	}
 	assertHeadMatchesGet(t, url)
 }
+
+func TestProxyHeadNpmUncachedPackument(t *testing.T) {
+	const abbreviated = "application/vnd.npm.install-v1+json"
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		doc := map[string]any{"name": "left-pad", "versions": map[string]any{"1.3.0": map[string]any{}}}
+		if r.Header.Get("Accept") != abbreviated {
+			doc["readme"] = strings.Repeat("full packument ", 200)
+		}
+		json.NewEncoder(w).Encode(doc)
+	}))
+	defer upstream.Close()
+
+	createRemote(t, mustFormat(`{"name": "head-npm", "package_type": "npm", "base_url": %q}`, upstream.URL))
+	defer deleteRemote(t, "head-npm")
+
+	url := apiURL("/api/v1/remote/head-npm/left-pad")
+	do := func(method string) (*http.Response, []byte) {
+		req, _ := http.NewRequest(method, url, nil)
+		req.Header.Set("Accept", abbreviated)
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("%s: %v", method, err)
+		}
+		defer resp.Body.Close()
+		b, _ := io.ReadAll(resp.Body)
+		return resp, b
+	}
+	head, _ := do(http.MethodHead)
+	get, body := do(http.MethodGet)
+	if head.StatusCode != http.StatusOK || get.StatusCode != http.StatusOK {
+		t.Fatalf("HEAD %d, GET %d", head.StatusCode, get.StatusCode)
+	}
+	if head.ContentLength != int64(len(body)) {
+		t.Errorf("uncached HEAD Content-Length %d, GET body %d bytes", head.ContentLength, len(body))
+	}
+}
+
+func TestProxyHeadNonRewritingUsesUpstreamHead(t *testing.T) {
+	methods := make(chan string, 4)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		methods <- r.Method
+		w.Header().Set("Content-Type", "text/xml")
+		w.Write([]byte("<repomd/>"))
+	}))
+	defer upstream.Close()
+
+	createRemote(t, mustFormat(`{"name": "head-rpm", "package_type": "rpm", "base_url": %q}`, upstream.URL))
+	defer deleteRemote(t, "head-rpm")
+
+	resp, err := http.Head(apiURL("/api/v1/remote/head-rpm/repodata/repomd.xml"))
+	if err != nil {
+		t.Fatalf("HEAD: %v", err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != http.StatusOK || resp.Header.Get("X-Artifact-Source") != "remote" {
+		t.Fatalf("HEAD %d source %q", resp.StatusCode, resp.Header.Get("X-Artifact-Source"))
+	}
+	close(methods)
+	var seen []string
+	for m := range methods {
+		seen = append(seen, m)
+	}
+	if len(seen) != 1 || seen[0] != http.MethodHead {
+		t.Errorf("upstream saw %v, want a single HEAD", seen)
+	}
+}

internal/api/v1/proxy.go

@@ -141,6 +141,12 @@ func (h *ProxyHandler) serveFetch(w http.ResponseWriter, r *http.Request, remote
 	io.Copy(w, body)
 }
 
+func rewritesMutable(prov provider.Provider, path string) bool {
+	_, sr := prov.(provider.ServeRewriter)
+	_, rr := prov.(provider.ResponseRewriter)
+	return (sr || rr) && prov.Classify(path) == provider.Mutable
+}
+
 // handleProxyHead answers HEAD like GET without a body. Responses a provider
 // synthesizes or rewrites per request go through the GET path so their headers
 // match; everything else is answered from store metadata or an upstream HEAD.
@@ -153,7 +159,7 @@ func (h *ProxyHandler) handleProxyHead(w http.ResponseWriter, r *http.Request) {
 	if h.serveRemote(w, r, remote, prov, path) {
 		return
 	}
-	if _, ok := prov.(provider.ServeRewriter); ok && prov.Classify(path) == provider.Mutable {
+	if rewritesMutable(prov, path) {
 		h.serveFetch(w, r, remote, prov, path)
 		return
 	}

internal/api/v1/routes_test.go

@@ -5,6 +5,11 @@ import (
 	"testing"
 
 	"github.com/go-chi/chi/v5"
+
+	"git.unkin.net/unkin/artifactapi/internal/provider"
+	"git.unkin.net/unkin/artifactapi/internal/provider/npm"
+	"git.unkin.net/unkin/artifactapi/internal/provider/nuget"
+	"git.unkin.net/unkin/artifactapi/internal/provider/rpm"
 )
 
 func TestRoutesRemoteHead(t *testing.T) {
@@ -15,3 +20,22 @@ func TestRoutesRemoteHead(t *testing.T) {
 		}
 	}
 }
+
+func TestRewritesMutable(t *testing.T) {
+	tests := []struct {
+		name string
+		prov provider.Provider
+		path string
+		want bool
+	}{
+		{"npm packument", &npm.Provider{}, "left-pad", true},
+		{"npm tarball", &npm.Provider{}, "left-pad/-/left-pad-1.3.0.tgz", false},
+		{"nuget index", &nuget.Provider{}, "index.json", true},
+		{"rpm repodata", &rpm.Provider{}, "repodata/repomd.xml", false},
+	}
+	for _, tc := range tests {
+		if got := rewritesMutable(tc.prov, tc.path); got != tc.want {
+			t.Errorf("%s: rewritesMutable = %v, want %v", tc.name, got, tc.want)
+		}
+	}
+}

internal/provider/helm/helm.go

@@ -39,6 +39,8 @@ func (p *Provider) UpstreamURL(remote models.Remote, path string) string {
 	return strings.TrimRight(remote.BaseURL, "/") + "/" + strings.TrimLeft(path, "/")
 }
 
+func (p *Provider) RewritesResponse() {}
+
 func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) {
 	if proxyBaseURL == "" {
 		return nil, nil

internal/provider/npm/npm.go

@@ -37,6 +37,8 @@ func (p *Provider) UpstreamURL(remote models.Remote, path string) string {
 	return strings.TrimRight(remote.BaseURL, "/") + "/" + strings.TrimLeft(path, "/")
 }
 
+func (p *Provider) RewritesResponse() {}
+
 func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) {
 	if proxyBaseURL == "" || !json.Valid(body) {
 		return nil, nil

internal/provider/provider.go

@@ -82,6 +82,12 @@ type ServeRewriter interface {
 	RewriteServed(body []byte, remote models.Remote, proxyBaseURL string) []byte
 }
 
+// ResponseRewriter marks providers whose RewriteResponse changes mutable
+// bodies, so the served length is only known after a fetch.
+type ResponseRewriter interface {
+	RewritesResponse()
+}
+
 // RemoteMetadataStore is the persistence surface a RemoteServer needs to cache
 // and read the metadata it derives per upstream asset. *database.DB satisfies it.
 type RemoteMetadataStore interface {

internal/provider/puppet/puppet.go

@@ -39,6 +39,8 @@ func (p *Provider) UpstreamURL(remote models.Remote, path string) string {
 	return strings.TrimRight(remote.BaseURL, "/") + "/" + strings.TrimLeft(path, "/")
 }
 
+func (p *Provider) RewritesResponse() {}
+
 func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) {
 	if proxyBaseURL == "" {
 		return nil, nil

internal/provider/pypi/pypi.go

@@ -52,6 +52,8 @@ func (p *Provider) UpstreamURL(remote models.Remote, path string) string {
 	return strings.TrimRight(remote.BaseURL, "/") + "/" + strings.TrimLeft(path, "/")
 }
 
+func (p *Provider) RewritesResponse() {}
+
 func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) {
 	if proxyBaseURL == "" {
 		return nil, nil

internal/provider/terraform/terraform.go

@@ -73,6 +73,8 @@ func (p *Provider) UpstreamURL(remote models.Remote, path string) string {
 	return strings.TrimRight(remote.BaseURL, "/") + "/v1/providers/" + strings.TrimLeft(path, "/")
 }
 
+func (p *Provider) RewritesResponse() {}
+
 func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) {
 	if remote.ReleasesRemote == "" {
 		return nil, nil