fix(uhttp): opt-in cache-key headers via CacheOption + WithCacheKeyHeaders

pkg/uhttp/client.go

@@ -104,17 +104,53 @@ func NewClient(ctx context.Context, options ...Option) (*http.Client, error) {
 }
 
 type icache interface {
-	Get(req *http.Request) (*http.Response, error)
-	Set(req *http.Request, value *http.Response) error
+	Get(req *http.Request, opts ...CacheOption) (*http.Response, error)
+	Set(req *http.Request, value *http.Response, opts ...CacheOption) error
 	Clear(ctx context.Context) error
 	Stats(ctx context.Context) CacheStats
 }
 
+type cacheKeyConfig struct {
+	headers []string
+}
+
+// CacheOption configures how CreateCacheKey computes its key, beyond the
+// default set of headers (Accept, Content-Type, Cookie, Range). Kept as an
+// interface so future dimensions (TTL, query-param keying, etc.) can be
+// added without changing CreateCacheKey's or icache's signatures again.
+type CacheOption interface {
+	applyCache(*cacheKeyConfig)
+}
+
+type cacheKeyHeadersOption []string
+
+func (o cacheKeyHeadersOption) applyCache(c *cacheKeyConfig) {
+	c.headers = append(c.headers, o...)
+}
+
+// CacheKeyHeaders returns a CacheOption that folds the named headers into
+// the cache key computed by CreateCacheKey (and by GoCache/DBCache's
+// Get/Set), beyond the default set (Accept, Content-Type, Cookie, Range).
+// The value folded in is always read from req.Header at key-computation
+// time, so the key can never describe a value other than the one actually
+// present on the request. Named headers must therefore be set on the
+// request before it reaches the cache lookup; a header only added by a
+// transport-level RoundTripper or a cookie jar after that point is not
+// seen.
+func CacheKeyHeaders(headers ...string) CacheOption {
+	return cacheKeyHeadersOption(headers)
+}
+
 // CreateCacheKey generates a cache key based on the request URL, query parameters, and headers.
-func CreateCacheKey(req *http.Request) (string, error) {
+func CreateCacheKey(req *http.Request, opts ...CacheOption) (string, error) {
 	if req == nil {
 		return "", fmt.Errorf("request is nil")
 	}
+	var cfg cacheKeyConfig
+	for _, o := range opts {
+		o.applyCache(&cfg)
+	}
+
 	var sortedParams []string
 	// Normalize the URL path
 	path := strings.ToLower(req.URL.Path)
@@ -130,13 +166,33 @@ func CreateCacheKey(req *http.Request) (string, error) {
 	queryString := strings.Join(sortedParams, "&")
 	// Include relevant headers in the cache key
 	var headerParts []string
+	seenHeaders := map[string]bool{
+		"Accept":       true,
+		"Content-Type": true,
+		"Cookie":       true,
+		"Range":        true,
+	}
 	for key, values := range req.Header {
 		for _, value := range values {
-			if key == "Accept" || key == "Content-Type" || key == "Cookie" || key == "Range" {
+			if seenHeaders[key] {
 				headerParts = append(headerParts, fmt.Sprintf("%s=%s", key, value))
 			}
 		}
 	}
+	// Opted-in headers are folded in on top of the default set above.
+	// seenHeaders already marks the default set, and gets marked as each
+	// opted-in header is processed, so a header named in cfg.headers -- by
+	// one CacheOption or by several -- is never folded in more than once.
+	for _, h := range cfg.headers {
+		key := http.CanonicalHeaderKey(h)
+		if seenHeaders[key] {
+			continue
+		}
+		seenHeaders[key] = true
+		for _, value := range req.Header[key] {
+			headerParts = append(headerParts, fmt.Sprintf("%s=%s", key, value))
+		}
+	}
 
 	sort.Strings(headerParts)
 	headersString := strings.Join(headerParts, "&")

pkg/uhttp/client_test.go

@@ -0,0 +1,125 @@
+package uhttp
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func newCacheKeyRequest(t *testing.T, headerKey, headerValue string) *http.Request {
+	t.Helper()
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com/widgets?id=1", nil)
+	require.NoError(t, err)
+	if headerKey != "" {
+		req.Header.Set(headerKey, headerValue)
+	}
+	return req
+}
+
+func TestCreateCacheKey_NilRequest(t *testing.T) {
+	_, err := CreateCacheKey(nil)
+	require.Error(t, err)
+}
+
+func TestCreateCacheKey_IdenticalRequestsMatch(t *testing.T) {
+	req1 := newCacheKeyRequest(t, "Accept", "application/json")
+	req2 := newCacheKeyRequest(t, "Accept", "application/json")
+
+	key1, err := CreateCacheKey(req1)
+	require.NoError(t, err)
+	key2, err := CreateCacheKey(req2)
+	require.NoError(t, err)
+	require.Equal(t, key1, key2)
+}
+
+// TestCreateCacheKey_HeadersOutsideDefaultSetAreIgnoredByDefault documents
+// current, intentional behavior: only the default set affects the key
+// unless a caller opts in via a CacheOption. Folding in every header
+// unconditionally would key the cache on values that have nothing to do
+// with the response (transport-injected headers, tracing IDs, etc.) and
+// silently tank the hit rate for every caller who never asked for that.
+func TestCreateCacheKey_HeadersOutsideDefaultSetAreIgnoredByDefault(t *testing.T) {
+	headers := []string{"Authorization", "X-Api-Version", "X-Tenant-Id", "User-Agent"}
+	for _, header := range headers {
+		t.Run(header, func(t *testing.T) {
+			reqA := newCacheKeyRequest(t, header, "value-a")
+			reqB := newCacheKeyRequest(t, header, "value-b")
+
+			keyA, err := CreateCacheKey(reqA)
+			require.NoError(t, err)
+			keyB, err := CreateCacheKey(reqB)
+			require.NoError(t, err)
+			require.Equal(t, keyA, keyB, "%s is not in the default set and must not affect the key", header)
+		})
+	}
+}
+
+func TestCreateCacheKey_DefaultHeadersStillChangeKey(t *testing.T) {
+	headers := []string{"Accept", "Content-Type", "Cookie", "Range"}
+	for _, header := range headers {
+		t.Run(header, func(t *testing.T) {
+			reqA := newCacheKeyRequest(t, header, "value-a")
+			reqB := newCacheKeyRequest(t, header, "value-b")
+
+			keyA, err := CreateCacheKey(reqA)
+			require.NoError(t, err)
+			keyB, err := CreateCacheKey(reqB)
+			require.NoError(t, err)
+			require.NotEqual(t, keyA, keyB)
+		})
+	}
+}
+
+// TestCreateCacheKey_CacheOptionOptsInAdditionalHeaders is the regression
+// test for CE-1056: a caller that knows a header varies the response (e.g.
+// Authorization scoping the result set) can opt that header into the key
+// via a CacheOption instead of two requests silently colliding. The value
+// folded in is read from req.Header, same as the default set, so it can
+// never describe anything other than what was actually sent.
+func TestCreateCacheKey_CacheOptionOptsInAdditionalHeaders(t *testing.T) {
+	reqA := newCacheKeyRequest(t, "Authorization", "value-a")
+	reqB := newCacheKeyRequest(t, "Authorization", "value-b")
+
+	opt := cacheKeyHeadersOption{"Authorization"}
+	keyA, err := CreateCacheKey(reqA, opt)
+	require.NoError(t, err)
+	keyB, err := CreateCacheKey(reqB, opt)
+	require.NoError(t, err)
+	require.NotEqual(t, keyA, keyB)
+}
+
+// TestCreateCacheKey_CacheOptionOnlyAffectsNamedHeaders confirms opting a
+// header in doesn't widen the key to every header on the request -- a
+// header present on req.Header but not named in the CacheOption still
+// falls back to the default-set rule.
+func TestCreateCacheKey_CacheOptionOnlyAffectsNamedHeaders(t *testing.T) {
+	reqA := newCacheKeyRequest(t, "X-Tenant-Id", "tenant-a")
+	reqA.Header.Set("Authorization", "same-token")
+
+	reqB := newCacheKeyRequest(t, "X-Tenant-Id", "tenant-b")
+	reqB.Header.Set("Authorization", "same-token")
+
+	opt := cacheKeyHeadersOption{"Authorization"}
+	keyA, err := CreateCacheKey(reqA, opt)
+	require.NoError(t, err)
+	keyB, err := CreateCacheKey(reqB, opt)
+	require.NoError(t, err)
+	require.Equal(t, keyA, keyB, "X-Tenant-Id was never opted in, so it must not affect the key")
+}
+
+// TestCreateCacheKey_CacheOptionCanonicalizesNames confirms header names
+// passed via a CacheOption match regardless of casing, since req.Header
+// stores them canonicalized.
+func TestCreateCacheKey_CacheOptionCanonicalizesNames(t *testing.T) {
+	reqA := newCacheKeyRequest(t, "Authorization", "value-a")
+	reqB := newCacheKeyRequest(t, "Authorization", "value-b")
+
+	opt := cacheKeyHeadersOption{"authorization"}
+	keyA, err := CreateCacheKey(reqA, opt)
+	require.NoError(t, err)
+	keyB, err := CreateCacheKey(reqB, opt)
+	require.NoError(t, err)
+	require.NotEqual(t, keyA, keyB)
+}

pkg/uhttp/dbcache.go

@@ -190,12 +190,12 @@ func (d *DBCache) removeDB(ctx context.Context) error {
 }
 
 // Get returns cached response (if exists).
-func (d *DBCache) Get(req *http.Request) (*http.Response, error) {
+func (d *DBCache) Get(req *http.Request, opts ...CacheOption) (*http.Response, error) {
 	var (
 		isFound = false
 		resp    *http.Response
 	)
-	key, err := CreateCacheKey(req)
+	key, err := CreateCacheKey(req, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -250,8 +250,8 @@ func (d *DBCache) pick(ctx context.Context, key string) ([]byte, error) {
 }
 
 // Set stores and save response in the db.
-func (d *DBCache) Set(req *http.Request, value *http.Response) error {
-	key, err := CreateCacheKey(req)
+func (d *DBCache) Set(req *http.Request, value *http.Response, opts ...CacheOption) error {
+	key, err := CreateCacheKey(req, opts...)
 	if err != nil {
 		return err
 	}

pkg/uhttp/gocache.go

@@ -58,13 +58,13 @@ func NewNoopCache(ctx context.Context) *NoopCache {
 	return &NoopCache{}
 }
 
-func (g *NoopCache) Get(req *http.Request) (*http.Response, error) {
+func (g *NoopCache) Get(req *http.Request, opts ...CacheOption) (*http.Response, error) {
 	// This isn't threadsafe but who cares? It's the noop cache.
 	g.counter++
 	return nil, nil
 }
 
-func (n *NoopCache) Set(req *http.Request, value *http.Response) error {
+func (n *NoopCache) Set(req *http.Request, value *http.Response, opts ...CacheOption) error {
 	return nil
 }
 
@@ -219,12 +219,12 @@ func (g *GoCache) Stats(ctx context.Context) CacheStats {
 	}
 }
 
-func (g *GoCache) Get(req *http.Request) (*http.Response, error) {
+func (g *GoCache) Get(req *http.Request, opts ...CacheOption) (*http.Response, error) {
 	if g.rootLibrary == nil {
 		return nil, nil
 	}
 
-	key, err := CreateCacheKey(req)
+	key, err := CreateCacheKey(req, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -247,12 +247,12 @@ func (g *GoCache) Get(req *http.Request) (*http.Response, error) {
 	return resp, nil
 }
 
-func (g *GoCache) Set(req *http.Request, value *http.Response) error {
+func (g *GoCache) Set(req *http.Request, value *http.Response, opts ...CacheOption) error {
 	if g.rootLibrary == nil {
 		return nil
 	}
 
-	key, err := CreateCacheKey(req)
+	key, err := CreateCacheKey(req, opts...)
 	if err != nil {
 		return err
 	}

pkg/uhttp/wrapper.go

@@ -83,6 +83,31 @@ func WithMetricsHandler(handler metrics.Handler) WrapperOption {
 	return metricsHandlerOption{handler: handler}
 }
 
+type cacheKeyHeadersWrapperOption struct {
+	opt CacheOption
+}
+
+func (o cacheKeyHeadersWrapperOption) Apply(c *BaseHttpClient) {
+	c.cacheOptions = append(c.cacheOptions, o.opt)
+}
+
+// WithCacheKeyHeaders returns a WrapperOption that additionally folds the
+// named headers into the HTTP response cache key for every request this
+// client makes, on top of the default set (Accept, Content-Type, Cookie,
+// Range). Use this when requests through this client vary by a header the
+// cache wouldn't otherwise key on -- e.g. a per-call Authorization token or
+// a tenant/version header -- so requests that only differ in that header
+// don't collide in the cache. The value folded in is always read from
+// req.Header at request time, so the key can never describe a value other
+// than the one actually sent.
+//
+// Named headers must be set on the request before it reaches Do; a header
+// only added later by a transport-level RoundTripper or a cookie jar is not
+// seen by the cache lookup and will not be reflected in the key.
+func WithCacheKeyHeaders(headers ...string) WrapperOption {
+	return cacheKeyHeadersWrapperOption{opt: CacheKeyHeaders(headers...)}
+}
+
 type WrapperOption interface {
 	Apply(*BaseHttpClient)
 }
@@ -120,6 +145,7 @@ type (
 		rateLimiter    uRateLimit.Limiter
 		baseHttpCache  icache
 		metricsHandler metrics.Handler
+		cacheOptions   []CacheOption
 	}
 
 	DoOption      func(resp *WrapperResponse) error
@@ -448,7 +474,7 @@ func (c *BaseHttpClient) Do(req *http.Request, options ...DoOption) (*http.Respo
 	}
 
 	if req.Method == http.MethodGet && req.Header.Get("Cache-Control") != "no-cache" {
-		resp, err = c.baseHttpCache.Get(req)
+		resp, err = c.baseHttpCache.Get(req, c.cacheOptions...)
 		if err != nil {
 			return nil, err
 		}
@@ -520,7 +546,7 @@ func (c *BaseHttpClient) Do(req *http.Request, options ...DoOption) (*http.Respo
 	}
 
 	if req.Method == http.MethodGet && resp.StatusCode == http.StatusOK {
-		cacheErr := c.baseHttpCache.Set(req, resp)
+		cacheErr := c.baseHttpCache.Set(req, resp, c.cacheOptions...)
 		if cacheErr != nil {
 			l.Warn("error setting cache", zap.String("url", req.URL.String()), zap.Error(cacheErr))
 		}