klauspost for compression

Author: jcalabro

go.mod

@@ -2,6 +2,9 @@ module github.com/jcalabro/jttp
 
 go 1.26
 
-require golang.org/x/net v0.53.0
+require (
+	github.com/klauspost/compress v1.18.6
+	golang.org/x/net v0.53.0
+)
 
 require golang.org/x/text v0.36.0 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao=
+github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
 golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA=
 golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs=
 golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg=

guard.go

@@ -1,13 +1,14 @@
 package jttp
 
 import (
-	"compress/gzip"
 	"context"
 	"fmt"
 	"io"
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/klauspost/compress/gzip"
 )
 
 // idleWatchdog cancels a context with a given cause if Reset is not called
@@ -91,6 +92,21 @@ type guardedBodyConfig struct {
 	minRateWindow  time.Duration // window over which to average
 	decompressGzip bool          // wrap inner in gzip.Reader
 	maxRatio       float64       // 0 disables the ratio guard
+	bodyObserver   func(BodyObservation)
+	observeURL     string
+	statusCode     int
+	observeStarted time.Time
+}
+
+// BodyObservation reports response-body byte counts when a guarded response
+// body is closed. It is intentionally generic: callers that need
+// request-specific labels can derive them from URL.
+type BodyObservation struct {
+	URL               string
+	StatusCode        int
+	CompressedBytes   int64
+	UncompressedBytes int64
+	Duration          time.Duration
 }
 
 // guardedBody wraps an http.Response.Body with robustness protections.
@@ -179,6 +195,19 @@ func (g *guardedBody) Close() error {
 			}
 		}
 		g.closeErr = g.inner.Close()
+		if g.cfg.bodyObserver != nil {
+			compressed := g.readBytes
+			if g.compressed != nil {
+				compressed = g.compressed.total
+			}
+			g.cfg.bodyObserver(BodyObservation{
+				URL:               g.cfg.observeURL,
+				StatusCode:        g.cfg.statusCode,
+				CompressedBytes:   compressed,
+				UncompressedBytes: g.readBytes,
+				Duration:          time.Since(g.cfg.observeStarted),
+			})
+		}
 	})
 	return g.closeErr
 }

guard_test.go

@@ -461,6 +461,92 @@ func BenchmarkRateTrackerObserveFullWindow(b *testing.B) {
 	}
 }
 
+func BenchmarkGuardedBodyGzipDecode(b *testing.B) {
+	for _, tc := range []struct {
+		name    string
+		payload []byte
+	}{
+		{name: "repoish", payload: benchmarkRepoishPayload(32 << 20)},
+		{name: "high_entropy", payload: benchmarkHighEntropyPayload(32 << 20)},
+	} {
+		b.Run(tc.name, func(b *testing.B) {
+			var gzipped bytes.Buffer
+			zw, err := gzip.NewWriterLevel(&gzipped, gzip.DefaultCompression)
+			if err != nil {
+				b.Fatal(err)
+			}
+			if _, err := zw.Write(tc.payload); err != nil {
+				b.Fatal(err)
+			}
+			if err := zw.Close(); err != nil {
+				b.Fatal(err)
+			}
+
+			b.SetBytes(int64(len(tc.payload)))
+			b.ReportMetric(float64(len(tc.payload))/float64(gzipped.Len()), "ratio")
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				ctx, cancel := context.WithCancelCause(context.Background())
+				gb, err := newGuardedBody(io.NopCloser(bytes.NewReader(gzipped.Bytes())), guardedBodyConfig{
+					ctx:            ctx,
+					cancel:         cancel,
+					decompressGzip: true,
+					maxRatio:       1000,
+				})
+				if err != nil {
+					cancel(nil)
+					b.Fatal(err)
+				}
+				if _, err := io.CopyBuffer(io.Discard, gb, make([]byte, 128<<10)); err != nil {
+					_ = gb.Close()
+					cancel(nil)
+					b.Fatal(err)
+				}
+				if err := gb.Close(); err != nil {
+					cancel(nil)
+					b.Fatal(err)
+				}
+				cancel(nil)
+			}
+		})
+	}
+}
+
+func benchmarkRepoishPayload(n int) []byte {
+	payload := make([]byte, 0, n)
+	var x uint64 = 0x123456789abcdef0
+	templates := [][]byte{
+		[]byte(`{"$type":"app.bsky.feed.post","text":"hello from a busy repo","createdAt":"2026-06-16T18:00:00.000Z"}`),
+		[]byte(`app.bsky.feed.like/3lq2example cid bafyreigenericrecord did:plc:example`),
+		[]byte(`app.bsky.graph.follow/3lq2example subject did:plc:subject collection rkey`),
+		[]byte(`com.atproto.repo.strongRef uri at://did:plc:author/app.bsky.feed.post/3lq2example`),
+	}
+	for len(payload) < n {
+		tpl := templates[len(payload)/4096%len(templates)]
+		payload = append(payload, tpl...)
+		for range 64 {
+			x ^= x << 13
+			x ^= x >> 7
+			x ^= x << 17
+			payload = append(payload, byte(x))
+		}
+	}
+	return payload[:n]
+}
+
+func benchmarkHighEntropyPayload(n int) []byte {
+	payload := make([]byte, n)
+	var x uint64 = 0x123456789abcdef0
+	for i := range payload {
+		x ^= x << 13
+		x ^= x >> 7
+		x ^= x << 17
+		payload[i] = byte(x) ^ byte(i/4096)
+	}
+	return payload
+}
+
 func gzipZeros(t *testing.T, n int) []byte {
 	t.Helper()
 	var buf bytes.Buffer

jttp.go

@@ -119,6 +119,7 @@ type config struct {
 	// Response size / decompression
 	maxResponseBodyBytes int64
 	maxCompressionRatio  float64
+	bodyObserver         func(BodyObservation)
 
 	// Redirect safety
 	allowSchemeDowngrade  bool
@@ -301,6 +302,7 @@ func New(opts ...Option) *http.Client {
 			minRate:            cfg.minRate,
 			minRateWindow:      cfg.minRateWindow,
 			maxRatio:           cfg.maxCompressionRatio,
+			bodyObserver:       cfg.bodyObserver,
 			strictSSRFInitial:  cfg.strictSSRFInitial,
 			redirectGuard:      rGuard,
 		},
@@ -540,6 +542,13 @@ func WithDisableCompression() Option {
 	return func(c *config) { c.disableCompression = true }
 }
 
+// WithBodyObserver registers a callback invoked when a guarded response body
+// is closed. This is intended for temporary diagnostics where callers need
+// precise response byte counts without changing higher-level APIs.
+func WithBodyObserver(fn func(BodyObservation)) Option {
+	return func(c *config) { c.bodyObserver = fn }
+}
+
 // WithForceHTTP2 controls whether HTTP/2 is attempted when a custom TLS
 // config is set. Default: true.
 func WithForceHTTP2(force bool) Option {

safety.go

@@ -19,6 +19,7 @@ type safetyConfig struct {
 	minRate       int64
 	minRateWindow time.Duration
 	maxRatio      float64
+	bodyObserver  func(BodyObservation)
 
 	// SSRF: if true, the safety transport blocks the initial request too
 	// (the redirectGuard handles subsequent ones via http.Client.CheckRedirect).
@@ -37,6 +38,8 @@ type safetyTransport struct {
 }
 
 func (t *safetyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	observeStarted := time.Now()
+
 	// Strict SSRF: reject the initial request if target resolves to blocked.
 	if t.cfg.strictSSRFInitial && t.cfg.redirectGuard != nil {
 		if err := t.cfg.redirectGuard.checkIPPolicy(req.Context(), req.URL.Hostname()); err != nil {
@@ -106,6 +109,10 @@ func (t *safetyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		minRateWindow:  t.cfg.minRateWindow,
 		decompressGzip: decompressGzip,
 		maxRatio:       t.cfg.maxRatio,
+		bodyObserver:   t.cfg.bodyObserver,
+		observeURL:     req.URL.String(),
+		statusCode:     resp.StatusCode,
+		observeStarted: observeStarted,
 	})
 	if gerr != nil {
 		// Close the raw body and tear down the watchdogs — otherwise a

safety_test.go

@@ -214,3 +214,56 @@ func TestSafetyTransportDecodesGzipResponseAndStripsHeader(t *testing.T) {
 		t.Errorf("decoded = %q, want hello", got)
 	}
 }
+
+func TestSafetyTransportObservesResponseBodyBytesOnClose(t *testing.T) {
+	// A valid gzip stream for "hello".
+	gzipHello := []byte{
+		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff,
+		0xca, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x04, 0x00, 0x00, 0xff, 0xff,
+		0x86, 0xa6, 0x10, 0x36, 0x05, 0x00, 0x00, 0x00,
+	}
+	s := &stubTransport{
+		body: string(gzipHello),
+		header: http.Header{
+			"Content-Encoding": []string{"gzip"},
+		},
+	}
+
+	var got BodyObservation
+	st := &safetyTransport{
+		next: s,
+		cfg: safetyConfig{
+			compressionEnabled: true,
+			bodyObserver: func(obs BodyObservation) {
+				got = obs
+			},
+		},
+	}
+	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://x/path", http.NoBody)
+	resp, err := st.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := io.ReadAll(resp.Body); err != nil {
+		t.Fatal(err)
+	}
+	if err := resp.Body.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	if got.URL != "http://x/path" {
+		t.Errorf("URL = %q, want request URL", got.URL)
+	}
+	if got.StatusCode != http.StatusOK {
+		t.Errorf("StatusCode = %d, want 200", got.StatusCode)
+	}
+	if got.CompressedBytes != int64(len(gzipHello)) {
+		t.Errorf("CompressedBytes = %d, want %d", got.CompressedBytes, len(gzipHello))
+	}
+	if got.UncompressedBytes != int64(len("hello")) {
+		t.Errorf("UncompressedBytes = %d, want %d", got.UncompressedBytes, len("hello"))
+	}
+	if got.Duration <= 0 {
+		t.Errorf("Duration = %s, want positive", got.Duration)
+	}
+}