add header auth cache to proxy

proxy/auth/auth.go

@@ -30,6 +30,8 @@ const (
 	SessionJWTIssuer     = "netbird-management"
 )
 
+const ClientIPMetadataKey = "nb-client-ip"
+
 // ResolveProto determines the protocol scheme based on the forwarded proto
 // configuration. When set to "http" or "https" the value is used directly.
 // Otherwise TLS state is used: if conn is non-nil "https" is returned, else "http".

proxy/internal/auth/header_cache.go

@@ -0,0 +1,188 @@
+package auth
+
+import (
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/sync/singleflight"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+const headerAuthCacheTTL = 60 * time.Second
+
+const envHeaderAuthCacheTTL = "NB_PROXY_HEADER_AUTH_CACHE_TTL"
+
+const headerAuthCachePerService = 1024
+
+const headerAuthCacheSkew = 30 * time.Second
+
+const headerAuthRPCTimeout = 10 * time.Second
+
+type headerCacheKey struct {
+	serviceID  types.ServiceID
+	headerName string
+	credential [sha256.Size]byte
+}
+
+type headerCacheEntry struct {
+	token     string
+	expiresAt time.Time
+}
+
+type headerAuthCache struct {
+	mu      sync.Mutex
+	entries map[types.ServiceID]*headerServiceBucket
+	flight  singleflight.Group
+	ttl     time.Duration
+	maxSize int
+	macKey  []byte
+	now     func() time.Time
+}
+
+type headerServiceBucket struct {
+	items map[headerCacheKey]headerCacheEntry
+	order []headerCacheKey
+}
+
+func newHeaderAuthCache() *headerAuthCache {
+	macKey := make([]byte, sha256.Size)
+	_, _ = rand.Read(macKey)
+
+	return &headerAuthCache{
+		entries: make(map[types.ServiceID]*headerServiceBucket),
+		ttl:     headerAuthCacheTTLFromEnv(),
+		maxSize: headerAuthCachePerService,
+		macKey:  macKey,
+		now:     time.Now,
+	}
+}
+
+func headerAuthCacheTTLFromEnv() time.Duration {
+	raw := strings.TrimSpace(os.Getenv(envHeaderAuthCacheTTL))
+	if raw == "" {
+		return headerAuthCacheTTL
+	}
+	d, err := time.ParseDuration(raw)
+	if err != nil || d <= 0 {
+		log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
+			envHeaderAuthCacheTTL, raw, headerAuthCacheTTL)
+		return headerAuthCacheTTL
+	}
+	return d
+}
+
+func (c *headerAuthCache) key(serviceID types.ServiceID, headerName, credential string) headerCacheKey {
+	mac := hmac.New(sha256.New, c.macKey)
+	mac.Write([]byte(credential))
+
+	key := headerCacheKey{serviceID: serviceID, headerName: headerName}
+	copy(key.credential[:], mac.Sum(nil))
+	return key
+}
+
+func (c *headerAuthCache) get(key headerCacheKey) string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	bucket, ok := c.entries[key.serviceID]
+	if !ok {
+		return ""
+	}
+	entry, ok := bucket.items[key]
+	if !ok {
+		return ""
+	}
+	if !c.now().Before(entry.expiresAt) {
+		delete(bucket.items, key)
+		bucket.order = removeKey(bucket.order, key)
+		return ""
+	}
+	return entry.token
+}
+
+func (c *headerAuthCache) put(key headerCacheKey, token string, sessionExpiration time.Duration) {
+	lifetime := c.ttl
+	if sessionExpiration > 0 && sessionExpiration-headerAuthCacheSkew < lifetime {
+		lifetime = sessionExpiration - headerAuthCacheSkew
+	}
+	if lifetime <= 0 {
+		return
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	bucket, ok := c.entries[key.serviceID]
+	if !ok {
+		bucket = &headerServiceBucket{items: make(map[headerCacheKey]headerCacheEntry)}
+		c.entries[key.serviceID] = bucket
+	}
+	if _, exists := bucket.items[key]; !exists {
+		bucket.order = append(bucket.order, key)
+	}
+	bucket.items[key] = headerCacheEntry{token: token, expiresAt: c.now().Add(lifetime)}
+
+	for len(bucket.order) > c.maxSize {
+		oldest := bucket.order[0]
+		bucket.order = bucket.order[1:]
+		delete(bucket.items, oldest)
+	}
+}
+
+func (c *headerAuthCache) invalidate(key headerCacheKey) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	bucket, ok := c.entries[key.serviceID]
+	if !ok {
+		return
+	}
+	delete(bucket.items, key)
+	bucket.order = removeKey(bucket.order, key)
+}
+
+func (c *headerAuthCache) invalidateService(serviceID types.ServiceID) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	delete(c.entries, serviceID)
+}
+
+type authenticateHeaderFn func() (string, error)
+
+func (c *headerAuthCache) fetch(key headerCacheKey, sessionExpiration time.Duration, authenticate authenticateHeaderFn) (string, bool, error) {
+	if token := c.get(key); token != "" {
+		return token, true, nil
+	}
+
+	res, err, _ := c.flight.Do(headerFlightKey(key), func() (any, error) {
+		if token := c.get(key); token != "" {
+			return token, nil
+		}
+		token, err := authenticate()
+		if err != nil {
+			return "", err
+		}
+		if token != "" {
+			c.put(key, token, sessionExpiration)
+		}
+		return token, nil
+	})
+	if err != nil {
+		return "", false, err
+	}
+
+	token, _ := res.(string)
+	return token, false, nil
+}
+
+func headerFlightKey(key headerCacheKey) string {
+	return string(key.serviceID) + "|" + key.headerName + "|" + string(key.credential[:])
+}

proxy/internal/auth/header_cache_test.go

@@ -0,0 +1,244 @@
+package auth
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+func newCountingHeaderScheme(t *testing.T, kp *sessionkey.KeyPair, headerName, expectedValue string, calls *atomic.Int32) Header {
+	t.Helper()
+	token, err := sessionkey.SignToken(kp.PrivateKey, "header-user", "", "example.com", auth.MethodHeader, nil, nil, time.Hour)
+	require.NoError(t, err)
+
+	mock := &mockAuthenticator{fn: func(_ context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
+		calls.Add(1)
+		ha := req.GetHeaderAuth()
+		if ha != nil && ha.GetHeaderValue() == expectedValue {
+			return &proto.AuthenticateResponse{Success: true, SessionToken: token}, nil
+		}
+		return &proto.AuthenticateResponse{Success: false}, nil
+	}}
+	return NewHeader(mock, "svc1", "acc1", headerName)
+}
+
+func doHeaderRequest(t *testing.T, mw *Middleware, credential string) *httptest.ResponseRecorder {
+	t.Helper()
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/path", nil)
+	req.Header.Set("X-API-Key", credential)
+	req = req.WithContext(proxy.WithCapturedData(req.Context(), proxy.NewCapturedData("")))
+
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+	return rec
+}
+
+func TestProtect_HeaderAuth_ReusesSessionTokenAcrossRequests(t *testing.T) {
+	var calls atomic.Int32
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+	hdr := newCountingHeaderScheme(t, kp, "X-API-Key", "secret-key", &calls)
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	for i := 0; i < 25; i++ {
+		rec := doHeaderRequest(t, mw, "secret-key")
+		require.Equal(t, http.StatusOK, rec.Code)
+	}
+
+	assert.Equal(t, int32(1), calls.Load(), "a repeated credential must be verified once")
+}
+
+func TestProtect_HeaderAuth_DoesNotCacheFailures(t *testing.T) {
+	var calls atomic.Int32
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+	hdr := newCountingHeaderScheme(t, kp, "X-API-Key", "secret-key", &calls)
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	for i := 0; i < 3; i++ {
+		rec := doHeaderRequest(t, mw, "wrong-key")
+		require.Equal(t, http.StatusUnauthorized, rec.Code)
+	}
+
+	assert.Equal(t, int32(3), calls.Load(), "rejected credentials must not be cached")
+}
+
+func TestProtect_HeaderAuth_MissingHeaderSkipsRPC(t *testing.T) {
+	var calls atomic.Int32
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+	hdr := newCountingHeaderScheme(t, kp, "X-API-Key", "secret-key", &calls)
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	rec := doHeaderRequest(t, mw, "")
+	assert.NotEqual(t, http.StatusOK, rec.Code)
+	assert.Zero(t, calls.Load(), "an absent header must not reach management")
+}
+
+func TestHeaderAuthCache_EvictsExpiredEntries(t *testing.T) {
+	c := newHeaderAuthCache()
+	now := time.Now()
+	c.now = func() time.Time { return now }
+
+	key := c.key("svc1", "X-API-Key", "secret")
+	c.put(key, "token", time.Hour)
+	require.Equal(t, "token", c.get(key))
+
+	now = now.Add(c.ttl + time.Second)
+	assert.Empty(t, c.get(key))
+}
+
+func TestHeaderAuthCache_SkipsCacheWhenSessionExpiresWithinSkew(t *testing.T) {
+	c := newHeaderAuthCache()
+
+	key := c.key("svc1", "X-API-Key", "secret")
+	c.put(key, "token", headerAuthCacheSkew)
+
+	assert.Empty(t, c.get(key), "a token must never outlive the session it was minted for")
+}
+
+func TestHeaderAuthCache_SessionExpirationShortensTTL(t *testing.T) {
+	c := newHeaderAuthCache()
+	now := time.Now()
+	c.now = func() time.Time { return now }
+
+	key := c.key("svc1", "X-API-Key", "secret")
+	c.put(key, "token", headerAuthCacheSkew+10*time.Second)
+	require.Equal(t, "token", c.get(key))
+
+	now = now.Add(11 * time.Second)
+	assert.Empty(t, c.get(key))
+}
+
+func TestHeaderAuthCache_BoundsEntriesPerService(t *testing.T) {
+	c := newHeaderAuthCache()
+	c.maxSize = 4
+
+	var first headerCacheKey
+	for i := 0; i < 10; i++ {
+		key := c.key("svc1", "X-API-Key", fmt.Sprintf("secret-%d", i))
+		if i == 0 {
+			first = key
+		}
+		c.put(key, "token", time.Hour)
+	}
+
+	assert.Len(t, c.entries["svc1"].items, 4)
+	assert.Empty(t, c.get(first), "the oldest entry must be evicted")
+}
+
+func TestHeaderAuthCache_DistinguishesCredentials(t *testing.T) {
+	c := newHeaderAuthCache()
+
+	good := c.key("svc1", "X-API-Key", "good")
+	other := c.key("svc1", "X-API-Key", "other")
+	c.put(good, "token", time.Hour)
+
+	assert.Equal(t, "token", c.get(good))
+	assert.Empty(t, c.get(other))
+}
+
+func TestProtect_HeaderAuth_MappingUpdateInvalidatesCache(t *testing.T) {
+	var calls atomic.Int32
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+	hdr := newCountingHeaderScheme(t, kp, "X-API-Key", "secret-key", &calls)
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	require.Equal(t, http.StatusOK, doHeaderRequest(t, mw, "secret-key").Code)
+	require.Equal(t, int32(1), calls.Load())
+
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	require.Equal(t, http.StatusOK, doHeaderRequest(t, mw, "secret-key").Code)
+	assert.Equal(t, int32(2), calls.Load(), "a mapping update must drop the service's cached credentials")
+}
+
+func TestHeaderAuthCache_InvalidateService(t *testing.T) {
+	c := newHeaderAuthCache()
+
+	key := c.key("svc1", "X-API-Key", "secret")
+	other := c.key("svc2", "X-API-Key", "secret")
+	c.put(key, "token", time.Hour)
+	c.put(other, "token", time.Hour)
+
+	c.invalidateService("svc1")
+
+	assert.Empty(t, c.get(key))
+	assert.Equal(t, "token", c.get(other), "other services must be untouched")
+}
+
+func TestHeaderAuthCache_Invalidate(t *testing.T) {
+	c := newHeaderAuthCache()
+
+	key := c.key("svc1", "X-API-Key", "secret")
+	other := c.key("svc1", "X-API-Key", "second")
+	c.put(key, "token", time.Hour)
+	c.put(other, "token", time.Hour)
+
+	c.invalidate(key)
+
+	assert.Empty(t, c.get(key))
+	assert.Equal(t, "token", c.get(other))
+}
+
+func TestProtect_HeaderAuth_RevalidatesWhenCachedTokenRejected(t *testing.T) {
+	var calls atomic.Int32
+	mw := NewMiddleware(log.StandardLogger(), nil, nil)
+	kp := generateTestKeyPair(t)
+	hdr := newCountingHeaderScheme(t, kp, "X-API-Key", "secret-key", &calls)
+	require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
+
+	key := mw.headerCache.key("svc1", "X-API-Key", "secret-key")
+	mw.headerCache.put(key, "not-a-valid-token", time.Hour)
+
+	rec := doHeaderRequest(t, mw, "secret-key")
+
+	assert.Equal(t, http.StatusOK, rec.Code, "an unusable cached token must not fail the request")
+	assert.Equal(t, int32(1), calls.Load(), "the credential must be re-verified once")
+}
+
+func TestHeaderAuthCache_CollapsesConcurrentMisses(t *testing.T) {
+	c := newHeaderAuthCache()
+	key := c.key("svc1", "X-API-Key", "secret")
+
+	var calls atomic.Int32
+	release := make(chan struct{})
+	var wg sync.WaitGroup
+
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, _, _ = c.fetch(key, time.Hour, func() (string, error) {
+				calls.Add(1)
+				<-release
+				return "token", nil
+			})
+		}()
+	}
+
+	time.Sleep(50 * time.Millisecond)
+	close(release)
+	wg.Wait()
+
+	assert.Equal(t, int32(1), calls.Load(), "a burst of cold requests must collapse into one RPC")
+}

proxy/internal/auth/middleware.go

@@ -16,6 +16,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/proxy/internal/proxy"
@@ -82,6 +83,7 @@ type Middleware struct {
 	sessionValidator SessionValidator
 	geo              restrict.GeoResolver
 	tunnelCache      *tunnelValidationCache
+	headerCache      *headerAuthCache
 }
 
 // NewMiddleware creates a new authentication middleware. The sessionValidator is
@@ -96,6 +98,7 @@ func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator, geo re
 		sessionValidator: sessionValidator,
 		geo:              geo,
 		tunnelCache:      newTunnelValidationCache(),
+		headerCache:      newHeaderAuthCache(),
 	}
 }
 
@@ -452,7 +455,23 @@ func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Reque
 }
 
 func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, hdr Header, next http.Handler) bool {
-	token, _, err := hdr.Authenticate(r)
+	credential := r.Header.Get(hdr.headerName)
+	if credential == "" {
+		return false
+	}
+
+	key := mw.headerCache.key(hdr.id, hdr.headerName, credential)
+	authenticate := func() (string, error) {
+		ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), headerAuthRPCTimeout)
+		defer cancel()
+		if clientIP := mw.resolveClientIP(r); clientIP.IsValid() {
+			ctx = metadata.AppendToOutgoingContext(ctx, auth.ClientIPMetadataKey, clientIP.String())
+		}
+		token, _, err := hdr.Authenticate(r.WithContext(ctx))
+		return token, err
+	}
+
+	token, cached, err := mw.headerCache.fetch(key, config.SessionExpiration, authenticate)
 	if err != nil {
 		return mw.handleHeaderAuthError(w, r, err)
 	}
@@ -461,6 +480,17 @@ func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, ho
 	}
 
 	result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
+	if err != nil && cached {
+		mw.headerCache.invalidate(key)
+		if token, err = authenticate(); err != nil {
+			return mw.handleHeaderAuthError(w, r, err)
+		}
+		if token == "" {
+			return false
+		}
+		mw.headerCache.put(key, token, config.SessionExpiration)
+		result, err = mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
+	}
 	if err != nil {
 		setHeaderCapturedData(r.Context(), "", "", nil, nil)
 		status := http.StatusBadRequest
@@ -645,6 +675,8 @@ func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
 // AddDomain registers authentication schemes for the given domain. With schemes a valid session public key is required.
 // private=true forces ValidateTunnelPeer enforcement (403 on failure) regardless of the schemes list.
 func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter, private bool) error {
+	mw.headerCache.invalidateService(serviceID)
+
 	if len(schemes) == 0 {
 		mw.domainsMux.Lock()
 		defer mw.domainsMux.Unlock()
@@ -681,6 +713,10 @@ func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 st
 
 // RemoveDomain unregisters authentication for the given domain.
 func (mw *Middleware) RemoveDomain(domain string) {
+	if config, exists := mw.getDomainConfig(domain); exists {
+		mw.headerCache.invalidateService(config.ServiceID)
+	}
+
 	mw.domainsMux.Lock()
 	defer mw.domainsMux.Unlock()
 	delete(mw.domains, domain)

proxy/internal/auth/tunnel_cache.go

@@ -146,7 +146,7 @@ func (c *tunnelValidationCache) put(key tunnelCacheKey, resp *proto.ValidateTunn
 
 // removeKey drops the first occurrence of needle from order. The cache
 // uses small slices so a linear scan is cheaper than a map+slice combo.
-func removeKey(order []tunnelCacheKey, needle tunnelCacheKey) []tunnelCacheKey {
+func removeKey[T comparable](order []T, needle T) []T {
 	for i, k := range order {
 		if k == needle {
 			return append(order[:i], order[i+1:]...)