Filter gzipped gem compact indexes under cooldown

internal/handler/gem.go

@@ -1,11 +1,13 @@
 package handler
 
 import (
-	"bufio"
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"net/http"
 	"regexp"
+	"slices"
+	"strconv"
 	"strings"
 	"time"
 
@@ -208,48 +210,66 @@ func (h *GemHandler) fetchCompactIndex(r *http.Request, name string) (*http.Resp
 	if err != nil {
 		return nil, err
 	}
-	for _, hdr := range []string{"Accept", headerAcceptEncoding, "If-None-Match", "If-Modified-Since"} {
-		if v := r.Header.Get(hdr); v != "" {
-			req.Header.Set(hdr, v)
-		}
+	// Only Accept is passed on. The client's Accept-Encoding would turn off
+	// the transport's transparent decompression and leave a gzip body to
+	// filter, and its validators belong to the filtered copy it was served,
+	// not to upstream's file.
+	if v := r.Header.Get(headerAccept); v != "" {
+		req.Header.Set(headerAccept, v)
 	}
 	return h.proxy.HTTPClient.Do(req)
 }
 
-// writeFilteredIndex writes the compact index response with cooldown-filtered versions removed.
+// gemFilteredIndexDroppedHeaders describe upstream's file as a whole, so
+// they are wrong for a copy with versions removed. Bundler checks the digest
+// headers against the body it receives.
+var gemFilteredIndexDroppedHeaders = []string{
+	headerContentLength, headerContentEncoding, headerETag, headerLastModified,
+	"Digest", "Repr-Digest", "Content-Digest", "Accept-Ranges", "Content-Range",
+}
+
+// writeFilteredIndex writes the compact index response with cooldown-filtered
+// versions removed. The body is read in full first, so a read error or an
+// oversized index answers 502 instead of a truncated 200.
 func (h *GemHandler) writeFilteredIndex(w http.ResponseWriter, resp *http.Response, name string, filtered map[string]bool) {
-	for k, vv := range resp.Header {
-		if strings.EqualFold(k, headerContentLength) {
-			continue // length will change after filtering
-		}
-		for _, v := range vv {
-			w.Header().Add(k, v)
-		}
+	body, err := h.proxy.ReadMetadata(resp.Body)
+	if err != nil {
+		h.proxy.Logger.Error("failed to read gem compact index", "gem", name, "error", err)
+		http.Error(w, "failed to read upstream response", http.StatusBadGateway)
+		return
 	}
-	w.WriteHeader(http.StatusOK)
-
-	scanner := bufio.NewScanner(resp.Body)
-	for scanner.Scan() {
-		line := scanner.Text()
 
-		if line == "---" {
-			_, _ = fmt.Fprintln(w, line)
-			continue
-		}
+	var out bytes.Buffer
+	out.Grow(len(body))
+	for rest := body; len(rest) > 0; {
+		line, tail, _ := bytes.Cut(rest, []byte("\n"))
+		rest = tail
 
 		version := line
-		if spaceIdx := strings.IndexByte(line, ' '); spaceIdx > 0 {
+		if spaceIdx := bytes.IndexByte(line, ' '); spaceIdx > 0 {
 			version = line[:spaceIdx]
 		}
-
-		if filtered[version] {
+		if filtered[string(version)] {
 			h.proxy.Logger.Info("cooldown: filtering gem version",
-				"gem", name, "version", version)
+				"gem", name, "version", string(version))
 			continue
 		}
 
-		_, _ = fmt.Fprintln(w, line)
+		out.Write(line)
+		out.WriteByte('\n')
+	}
+
+	for k, vv := range resp.Header {
+		if slices.ContainsFunc(gemFilteredIndexDroppedHeaders, func(d string) bool { return strings.EqualFold(k, d) }) {
+			continue
+		}
+		for _, v := range vv {
+			w.Header().Add(k, v)
+		}
 	}
+	w.Header().Set(headerContentLength, strconv.Itoa(out.Len()))
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write(out.Bytes())
 }
 
 // gemVersion represents a version entry from the RubyGems versions API.

internal/handler/gem_test.go

@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"log/slog"
@@ -107,6 +108,117 @@ func TestGemCompactIndexCooldown(t *testing.T) {
 	}
 }
 
+// gemCooldownUpstream serves index as /info/testgem, gzipped when the
+// request accepts gzip as rubygems.org does, with 1.0.0 published a week ago
+// and 2.0.0 an hour ago.
+func gemCooldownUpstream(t *testing.T, index string, headers http.Header) *httptest.Server {
+	t.Helper()
+	now := time.Now()
+	versionsJSON, _ := json.Marshal([]gemVersion{
+		{Number: "1.0.0", Platform: "ruby", CreatedAt: now.Add(-7 * 24 * time.Hour).Format(time.RFC3339)},
+		{Number: "2.0.0", Platform: "ruby", CreatedAt: now.Add(-time.Hour).Format(time.RFC3339)},
+	})
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/info/testgem":
+			for k, vv := range headers {
+				w.Header()[k] = vv
+			}
+			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+			if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+				_, _ = w.Write([]byte(index))
+				return
+			}
+			w.Header().Set("Content-Encoding", "gzip")
+			gz := gzip.NewWriter(w)
+			_, _ = gz.Write([]byte(index))
+			_ = gz.Close()
+		case "/api/v1/versions/testgem.json":
+			_, _ = w.Write(versionsJSON)
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	t.Cleanup(upstream.Close)
+	return upstream
+}
+
+func TestGemCompactIndexCooldownClientAcceptsGzip(t *testing.T) {
+	index := "---\n1.0.0 dep1:>= 1.0|checksum:abc123\n2.0.0 dep1:>= 1.0|checksum:def456\n"
+	upstream := gemCooldownUpstream(t, index, http.Header{
+		"Etag":        {`"upstream"`},
+		"Repr-Digest": {"sha-256=:upstream:"},
+	})
+
+	proxy := testProxy()
+	proxy.Cooldown = &cooldown.Config{Default: "3d"}
+	h := &GemHandler{proxy: proxy, upstreamURL: upstream.URL}
+
+	// Bundler sends Accept-Encoding: gzip by default.
+	req := httptest.NewRequest(http.MethodGet, "/info/testgem", nil)
+	req.Header.Set("Accept-Encoding", "gzip;q=1.0,deflate;q=0.6,identity;q=0.3")
+	req.SetPathValue("name", "testgem")
+	w := httptest.NewRecorder()
+	h.handleCompactIndex(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", w.Code, http.StatusOK)
+	}
+	if ce := w.Header().Get("Content-Encoding"); ce != "" {
+		t.Errorf("Content-Encoding = %q, want none", ce)
+	}
+	want := "---\n1.0.0 dep1:>= 1.0|checksum:abc123\n"
+	if got := w.Body.String(); got != want {
+		t.Errorf("body = %q, want %q", got, want)
+	}
+	for _, hdr := range []string{"ETag", "Repr-Digest"} {
+		if v := w.Header().Get(hdr); v != "" {
+			t.Errorf("%s = %q, upstream's value must not describe the filtered index", hdr, v)
+		}
+	}
+}
+
+func TestGemCompactIndexCooldownLongLine(t *testing.T) {
+	// A version line longer than bufio.Scanner's 64 KB default used to end
+	// the response early with a 200.
+	long := "1.0.0 " + strings.Repeat("dep:>= 1.0,", 10_000) + "|checksum:abc"
+	index := "---\n" + long + "\n2.0.0 |checksum:def\n3.0.0 |checksum:ghi\n"
+	upstream := gemCooldownUpstream(t, index, nil)
+
+	proxy := testProxy()
+	proxy.Cooldown = &cooldown.Config{Default: "3d"}
+	h := &GemHandler{proxy: proxy, upstreamURL: upstream.URL}
+
+	req := httptest.NewRequest(http.MethodGet, "/info/testgem", nil)
+	req.SetPathValue("name", "testgem")
+	w := httptest.NewRecorder()
+	h.handleCompactIndex(w, req)
+
+	want := "---\n" + long + "\n3.0.0 |checksum:ghi\n"
+	if w.Code != http.StatusOK || w.Body.String() != want {
+		t.Fatalf("status = %d, body of %d bytes; want 200 with %d bytes", w.Code, w.Body.Len(), len(want))
+	}
+}
+
+func TestGemCompactIndexCooldownOversizedIndex(t *testing.T) {
+	index := "---\n1.0.0 |checksum:abc\n" + strings.Repeat("9.9.9 |checksum:zzz\n", 100)
+	upstream := gemCooldownUpstream(t, index, nil)
+
+	proxy := testProxy()
+	proxy.Cooldown = &cooldown.Config{Default: "3d"}
+	proxy.MetadataMaxSize = 256
+	h := &GemHandler{proxy: proxy, upstreamURL: upstream.URL}
+
+	req := httptest.NewRequest(http.MethodGet, "/info/testgem", nil)
+	req.SetPathValue("name", "testgem")
+	w := httptest.NewRecorder()
+	h.handleCompactIndex(w, req)
+
+	if w.Code != http.StatusBadGateway {
+		t.Fatalf("status = %d, want %d", w.Code, http.StatusBadGateway)
+	}
+}
+
 func TestGemCompactIndexCooldownWithPlatformVersion(t *testing.T) {
 	now := time.Now()
 	recentTime := now.Add(-1 * time.Hour).Format(time.RFC3339)