fix: cache PyPI metadata for filtered versions (#258)

Co-authored-by: dindin <dindin@DMBA.local>

Author: joyheroes

internal/handler/pypi.go

@@ -110,24 +110,14 @@ func (h *PyPIHandler) handleSimplePackage(w http.ResponseWriter, r *http.Request
 // that should be filtered out due to cooldown.
 func (h *PyPIHandler) fetchFilteredVersions(r *http.Request, name string) map[string]bool {
 	jsonURL := fmt.Sprintf("%s/pypi/%s/json", h.upstreamURL, name)
-	req, err := http.NewRequestWithContext(r.Context(), http.MethodGet, jsonURL, nil)
-	if err != nil {
-		return nil
-	}
-	req.Header.Set("Accept", "application/json")
 
-	resp, err := h.proxy.HTTPClient.Do(req)
+	body, _, err := h.proxy.FetchOrCacheMetadata(r.Context(), "pypi", name+"/json", jsonURL)
 	if err != nil {
 		return nil
 	}
-	defer func() { _ = resp.Body.Close() }()
-
-	if resp.StatusCode != http.StatusOK {
-		return nil
-	}
 
 	var metadata map[string]any
-	if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
+	if err := json.Unmarshal(body, &metadata); err != nil {
 		return nil
 	}
 

internal/handler/pypi_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -292,3 +293,60 @@ func TestPyPIDownloadCooldown(t *testing.T) {
 		})
 	}
 }
+
+// TestPyPIDownloadCooldownMetadataCache ensures that repeated downloads that
+// trigger cooldown filtering reuse the cached PyPI JSON metadata instead of
+// fetching it from upstream once per download.
+func TestPyPIDownloadCooldownMetadataCache(t *testing.T) {
+	now := time.Now()
+	releases := `{"releases": {
+		"1.0.0": [{"upload_time_iso_8601": "` + now.Add(-30*24*time.Hour).Format(time.RFC3339) + `"}],
+		"2.0.0": [{"upload_time_iso_8601": "` + now.Add(-1*time.Hour).Format(time.RFC3339) + `"}]
+	}}`
+
+	var metadataRequests atomic.Int64
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/pypi/newpkg/json" {
+			metadataRequests.Add(1)
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = io.WriteString(w, releases)
+			return
+		}
+		w.Header().Set("Content-Type", "application/octet-stream")
+		_, _ = io.WriteString(w, "package data")
+	}))
+	defer upstream.Close()
+
+	proxy, _, _, fetcher := setupTestProxy(t)
+	proxy.HTTPClient = upstream.Client()
+	proxy.CacheMetadata = true
+	proxy.MetadataTTL = time.Hour
+	proxy.Cooldown = &cooldown.Config{Default: "7d"}
+	fetcher.artifact = &fetch.Artifact{
+		Body:        io.NopCloser(strings.NewReader("package data")),
+		ContentType: "application/octet-stream",
+	}
+
+	h := &PyPIHandler{
+		proxy:       proxy,
+		upstreamURL: upstream.URL,
+		proxyURL:    "http://localhost",
+	}
+	srv := httptest.NewServer(h.Routes())
+	defer srv.Close()
+
+	// Two downloads of the same package: one outside the cooldown window
+	// (served) and one inside (withheld). Both go through the download path
+	// that resolves filtered versions.
+	for _, filename := range []string{"newpkg-1.0.0.tar.gz", "newpkg-2.0.0.tar.gz"} {
+		resp, err := http.Get(srv.URL + "/packages/packages/ab/cd/ef0123456789/" + filename)
+		if err != nil {
+			t.Fatalf("request failed: %v", err)
+		}
+		_ = resp.Body.Close()
+	}
+
+	if got := metadataRequests.Load(); got != 1 {
+		t.Errorf("upstream metadata JSON requests = %d, want 1 (repeated downloads should reuse the cached metadata)", got)
+	}
+}