amadan

carlos / platform Public

edge: gzip ResponseWriter with a private header map

The wrapper owns its own http.Header: every other wrapper here embeds
the writer below, so Header() is one shared map and maybeStore caches
rec.Header().Clone() — setting Content-Encoding on it would store an
identity body labelled gzip and serve it to every later visitor.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NoqzkUVahW5ykubsBQitxc
Paul Campbell 73ba62789644f12cb9d0811d943acaace061be15 parent 3a5a706
2 files changed, +350 −0
  • internal/edge/compress.go +152 −0
  • internal/edge/compress_test.go +198 −0
diff --git a/internal/edge/compress.go b/internal/edge/compress.go
index 4d3de221a..eb11eaaa5 100644
--- a/internal/edge/compress.go
+++ b/internal/edge/compress.go
@@ -10,9 +10,15 @@
package edge
import (
+ "bufio"
+ "compress/gzip"
+ "fmt"
+ "io"
+ "net"
"net/http"
"strconv"
"strings"
+ "sync"
)
// acceptsGzip reports whether r's Accept-Encoding admits gzip. Parsed as the
@@ -147,3 +153,149 @@ func weakenETag(h http.Header) {
h.Set("ETag", `W/`+et)
}
}
+
+// gzipPool recycles compressors. We compress on every delivery, cache hits
+// included, so a fresh gzip.Writer per response would be a steady stream of
+// garbage on a busy edge for no reason.
+var gzipPool = sync.Pool{New: func() any {
+ w, _ := gzip.NewWriterLevel(io.Discard, gzip.DefaultCompression)
+ return w
+}}
+
+// gzipWriter compresses an eligible response on its way to the client.
+//
+// It owns a PRIVATE header map. Every other ResponseWriter wrapper in this
+// package embeds the writer below, so Header() returns one shared map from the
+// handler all the way down to the connection — and maybeStore caches
+// rec.Header().Clone(). Were the compression headers set on that shared map,
+// the cache would store an identity body labelled Content-Encoding: gzip and
+// hand it to every later visitor. So callers above read and write hdr, and the
+// edits that describe the COMPRESSED representation go on the real writer's
+// map alone.
+//
+// Deliberately does NOT implement io.ReaderFrom: io.Copy (and so
+// http.ServeFile) falls back to Write, which is what compresses.
+type gzipWriter struct {
+ http.ResponseWriter
+ hdr http.Header
+ gz *gzip.Writer // nil unless this response is being compressed
+ wroteHdr bool
+}
+
+// newGzipWriter returns a wrapper for r, or nil when this request can never
+// produce a compressed response — no wrapper at all, so the writer stack is
+// exactly what it is today. The request-side questions are settled here, once:
+// everything after depends only on the response.
+//
+// HEAD is excluded rather than compressed-with-no-body: there is nothing to
+// compress, and a HEAD's truthful Content-Length is the entire value of the
+// response. The accepted cost is that a HEAD advertises the identity length
+// while a GET of the same URL returns gzip with no length at all.
+func newGzipWriter(w http.ResponseWriter, r *http.Request) *gzipWriter {
+ if r.Method == http.MethodHead || !acceptsGzip(r.Header) {
+ return nil
+ }
+ return &gzipWriter{ResponseWriter: w, hdr: http.Header{}}
+}
+
+// Header returns the PRIVATE map. See the type's doc comment: this is what
+// keeps the response cache storing identity bodies.
+func (g *gzipWriter) Header() http.Header { return g.hdr }
+
+// syncHeader replaces the real writer's header block with the private one.
+// Wholesale, not a merge: a handler that DELETED a header between blocks (the
+// shape the purge-on-103 test exercises) must see that deletion reach the
+// client, which a merge would silently undo.
+func (g *gzipWriter) syncHeader() {
+ down := g.ResponseWriter.Header()
+ for k := range down {
+ delete(down, k)
+ }
+ for k, v := range g.hdr {
+ down[k] = v
+ }
+}
+
+func (g *gzipWriter) WriteHeader(status int) {
+ // 1xx blocks (103 Early Hints, forwarded by httputil.ReverseProxy) precede
+ // the final status. Sync and forward without latching — the same shape
+ // recordingResponseWriter.WriteHeader uses, for the same reason: latching
+ // one would decide compression on a status that has no body.
+ if status < 200 {
+ g.syncHeader()
+ g.ResponseWriter.WriteHeader(status)
+ return
+ }
+ if g.wroteHdr {
+ g.ResponseWriter.WriteHeader(status) // superfluous; let the stdlib log it
+ return
+ }
+ g.wroteHdr = true
+ g.syncHeader()
+ if compressible(g.hdr, status) {
+ h := g.ResponseWriter.Header()
+ h.Set("Content-Encoding", "gzip")
+ h.Del("Content-Length") // it described the identity body
+ h.Add("Vary", "Accept-Encoding")
+ weakenETag(h)
+ g.gz = gzipPool.Get().(*gzip.Writer)
+ g.gz.Reset(g.ResponseWriter)
+ }
+ g.ResponseWriter.WriteHeader(status)
+}
+
+func (g *gzipWriter) Write(p []byte) (int, error) {
+ if !g.wroteHdr {
+ g.WriteHeader(http.StatusOK)
+ }
+ if g.gz != nil {
+ return g.gz.Write(p)
+ }
+ return g.ResponseWriter.Write(p)
+}
+
+// finish closes the compressor and returns it to the pool. ServeHTTP defers
+// it, AFTER the usage meter's defer, so LIFO runs this one first and the
+// metered byte count includes gzip's trailer.
+//
+// The no-header case is not a no-op: a handler that returned without ever
+// writing set its headers on the private map, and nothing has copied them
+// down. Without this sync the stdlib emits a bare 200 with none of them.
+func (g *gzipWriter) finish() {
+ if !g.wroteHdr {
+ g.syncHeader()
+ return
+ }
+ if g.gz == nil {
+ return
+ }
+ g.gz.Close()
+ g.gz.Reset(io.Discard) // don't pin this request's writer inside the pool
+ gzipPool.Put(g.gz)
+ g.gz = nil
+}
+
+// Flush pushes the compressor's buffered bytes out and then flushes the writer
+// below — needed so a streamed response reaches the client incrementally
+// instead of buffering until the handler returns.
+func (g *gzipWriter) Flush() {
+ if g.gz != nil {
+ g.gz.Flush()
+ }
+ if f, ok := g.ResponseWriter.(http.Flusher); ok {
+ f.Flush()
+ }
+}
+
+// Hijack forwards to the writer below, when it supports it — the same
+// optional-interface pattern countingWriter and recordingResponseWriter use.
+// ServeHTTP returns on the upgrade path before this wrapper is constructed, so
+// this is not expected on the hot path; it exists so the wrapper is never
+// itself the reason a hijack fails.
+func (g *gzipWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ hj, ok := g.ResponseWriter.(http.Hijacker)
+ if !ok {
+ return nil, nil, fmt.Errorf("edge: underlying ResponseWriter does not support http.Hijacker")
+ }
+ return hj.Hijack()
+}
diff --git a/internal/edge/compress_test.go b/internal/edge/compress_test.go
index e68f92a0c..2e37a3050 100644
--- a/internal/edge/compress_test.go
+++ b/internal/edge/compress_test.go
@@ -1,7 +1,13 @@
package edge
import (
+ "bufio"
+ "compress/gzip"
+ "io"
+ "net"
"net/http"
+ "net/http/httptest"
+ "strings"
"testing"
)
@@ -158,3 +164,195 @@ func TestWeakenETag(t *testing.T) {
}
}
}
+
+// hijackTrackingRecorder is edge_test.go's hijackableRecorder plus a record of
+// whether the call actually arrived — this test needs to prove the wrapper
+// FORWARDED the hijack, not merely that it did not error.
+type hijackTrackingRecorder struct {
+ *httptest.ResponseRecorder
+ hijacked bool
+}
+
+func (h *hijackTrackingRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ h.hijacked = true
+ return nil, nil, nil
+}
+
+// gzipReq builds a request that accepts gzip.
+func gzipReq(method, url string) *http.Request {
+ r := httptest.NewRequest(method, url, nil)
+ r.Header.Set("Accept-Encoding", "gzip")
+ return r
+}
+
+// TestGzipWriterKeepsCompressionHeadersOffTheCallerMap is THE test this design
+// exists for. Every writer wrapper in this package embeds the writer below, so
+// Header() returns one shared map all the way down — and maybeStore caches
+// rec.Header().Clone(). If the compression headers landed on that shared map,
+// the cache would store an IDENTITY body labelled Content-Encoding: gzip and
+// serve it to every later visitor as an undecodable response.
+func TestGzipWriterKeepsCompressionHeadersOffTheCallerMap(t *testing.T) {
+ rec := httptest.NewRecorder()
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/tokens.css"))
+ if g == nil {
+ t.Fatal("newGzipWriter returned nil for a gzip-accepting GET")
+ }
+ g.Header().Set("Content-Type", "text/css")
+ g.Header().Set("ETag", `"abc123"`)
+ g.Header().Set("Content-Length", "4000")
+ g.WriteHeader(http.StatusOK)
+ body := strings.Repeat("a{color:red}\n", 400)
+ io.WriteString(g, body)
+ g.finish()
+
+ if got := g.Header().Get("Content-Encoding"); got != "" {
+ t.Errorf("caller-visible Content-Encoding = %q, want empty — the recorder above would cache an identity body labelled gzip", got)
+ }
+ if got := g.Header().Get("ETag"); got != `"abc123"` {
+ t.Errorf("caller-visible ETag = %q, want the strong tag untouched (the cached body IS the identity one)", got)
+ }
+ if got := g.Header().Get("Content-Length"); got != "4000" {
+ t.Errorf("caller-visible Content-Length = %q, want it untouched", got)
+ }
+
+ if got := rec.Header().Get("Content-Encoding"); got != "gzip" {
+ t.Errorf("client Content-Encoding = %q, want gzip", got)
+ }
+ if got := rec.Header().Get("Content-Length"); got != "" {
+ t.Errorf("client Content-Length = %q, want it dropped — it described the identity body", got)
+ }
+ if got := rec.Header().Get("Vary"); got != "Accept-Encoding" {
+ t.Errorf("client Vary = %q, want Accept-Encoding", got)
+ }
+ if got := rec.Header().Get("ETag"); got != `W/"abc123"` {
+ t.Errorf("client ETag = %q, want the weakened form", got)
+ }
+
+ zr, err := gzip.NewReader(rec.Body)
+ if err != nil {
+ t.Fatalf("client body is not gzip: %v", err)
+ }
+ got, err := io.ReadAll(zr)
+ if err != nil {
+ t.Fatalf("read gzip body: %v", err)
+ }
+ if string(got) != body {
+ t.Fatalf("decompressed body differs from what the handler wrote")
+ }
+ if rec.Body.Len() >= len(body) {
+ t.Fatalf("compressed size %d >= identity size %d", rec.Body.Len(), len(body))
+ }
+}
+
+// TestGzipWriterNilForRequestsThatCanNeverCompress: construction is where the
+// request-side questions are settled, once. A nil wrapper means ServeHTTP
+// leaves the writer stack exactly as it is today.
+func TestGzipWriterNilForRequestsThatCanNeverCompress(t *testing.T) {
+ if g := newGzipWriter(httptest.NewRecorder(), httptest.NewRequest("GET", "http://app.test/", nil)); g != nil {
+ t.Error("no Accept-Encoding must yield a nil wrapper")
+ }
+ if g := newGzipWriter(httptest.NewRecorder(), gzipReq("HEAD", "http://app.test/big.css")); g != nil {
+ t.Error("HEAD must yield a nil wrapper — there is no body to compress, and its truthful Content-Length is the whole value of the response")
+ }
+}
+
+// TestGzipWriterPassesThroughIneligibleResponses: an ineligible response must
+// be byte-identical to the unwrapped one, headers included.
+func TestGzipWriterPassesThroughIneligibleResponses(t *testing.T) {
+ rec := httptest.NewRecorder()
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/logo.png"))
+ g.Header().Set("Content-Type", "image/png")
+ g.Header().Set("Content-Length", "9")
+ g.WriteHeader(http.StatusOK)
+ io.WriteString(g, "PNG-BYTES")
+ g.finish()
+
+ if got := rec.Header().Get("Content-Encoding"); got != "" {
+ t.Errorf("Content-Encoding = %q on an image, want empty", got)
+ }
+ if got := rec.Header().Get("Content-Length"); got != "9" {
+ t.Errorf("Content-Length = %q, want it preserved on a pass-through", got)
+ }
+ if got := rec.Body.String(); got != "PNG-BYTES" {
+ t.Errorf("body = %q, want the identity bytes", got)
+ }
+}
+
+// TestGzipWriterSilentHandlerStillDeliversHeaders: the handler set headers and
+// returned without ever writing. Those headers live in the PRIVATE map, so
+// without a sync in finish() the stdlib would emit a bare 200 with none of
+// them — X-Carlos-Version included.
+func TestGzipWriterSilentHandlerStillDeliversHeaders(t *testing.T) {
+ rec := httptest.NewRecorder()
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/"))
+ g.Header().Set("X-Carlos-Version", "abc123")
+ g.finish()
+
+ if got := rec.Header().Get("X-Carlos-Version"); got != "abc123" {
+ t.Fatalf("X-Carlos-Version = %q, want abc123 — a handler that wrote nothing still lost every header it set", got)
+ }
+}
+
+// TestGzipWriter1xxDoesNotLatch mirrors recordingResponseWriter.WriteHeader: a
+// 103 Early Hints block precedes the final status. Latching it would freeze
+// the compression decision on a status that has no body, and the real response
+// would then be delivered uncompressed — or worse, with the hints lost, since
+// they live in the private map.
+func TestGzipWriter1xxDoesNotLatch(t *testing.T) {
+ rec := httptest.NewRecorder()
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/"))
+ g.Header().Set("Link", "</tokens.css>; rel=preload")
+ g.WriteHeader(http.StatusEarlyHints)
+
+ if got := rec.Header().Get("Link"); got != "</tokens.css>; rel=preload" {
+ t.Errorf("early-hints Link = %q, want it forwarded — it was set on the private map", got)
+ }
+
+ g.Header().Del("Link")
+ g.Header().Set("Content-Type", "text/html")
+ g.WriteHeader(http.StatusOK)
+ io.WriteString(g, strings.Repeat("<p>hi</p>", 200))
+ g.finish()
+
+ if got := rec.Header().Get("Content-Encoding"); got != "gzip" {
+ t.Errorf("Content-Encoding = %q after a 103, want gzip — the 1xx latched the decision", got)
+ }
+ if got := rec.Header().Get("Link"); got != "" {
+ t.Errorf("Link = %q, want empty — the final header block replaces the informational one", got)
+ }
+}
+
+// TestGzipWriterFlushReachesTheClient: a flushed response must actually leave
+// the compressor, not sit in its buffer until the handler returns. SSE never
+// compresses (content type), but a streamed HTML page does.
+func TestGzipWriterFlushReachesTheClient(t *testing.T) {
+ rec := httptest.NewRecorder()
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/stream"))
+ g.Header().Set("Content-Type", "text/html")
+ g.WriteHeader(http.StatusOK)
+ io.WriteString(g, strings.Repeat("<p>chunk</p>", 100))
+ g.Flush()
+
+ if rec.Body.Len() == 0 {
+ t.Fatal("Flush left the bytes inside the compressor — a streaming response would hang until the handler returned")
+ }
+ if !rec.Flushed {
+ t.Fatal("Flush did not reach the writer below")
+ }
+ g.finish()
+}
+
+// TestGzipWriterForwardsHijack: the upgrade path returns before this wrapper
+// is ever built, so this is not a hot path — it exists so the wrapper is never
+// itself the reason a hijack fails, matching the convention countingWriter and
+// recordingResponseWriter already follow.
+func TestGzipWriterForwardsHijack(t *testing.T) {
+ rec := &hijackTrackingRecorder{ResponseRecorder: httptest.NewRecorder()}
+ g := newGzipWriter(rec, gzipReq("GET", "http://app.test/"))
+ if _, _, err := g.Hijack(); err != nil {
+ t.Fatalf("Hijack: %v", err)
+ }
+ if !rec.hijacked {
+ t.Fatal("Hijack did not reach the writer below")
+ }
+}