diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index 4cae872e2..5a561f379 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -13,6 +13,7 @@ import ( "math" "net" "net/http" + "net/netip" "net/url" "os" "strconv" @@ -25,6 +26,7 @@ import ( log "github.com/sirupsen/logrus" "golang.org/x/oauth2" "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" "github.com/netbirdio/netbird/shared/management/domain" @@ -134,6 +136,10 @@ type ProxyServiceServer struct { // initial snapshot delivery. Configurable via NB_PROXY_SNAPSHOT_BATCH_SIZE. snapshotBatchSize int + authAttemptLimiter *authFailureLimiter + authClientLimiter *authFailureLimiter + authFailureMAC []byte + cancel context.CancelFunc } @@ -204,6 +210,10 @@ func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeT snapshotBatchSize: snapshotBatchSizeFromEnv(), cancel: cancel, } + s.authAttemptLimiter = newAuthFailureLimiter() + s.authClientLimiter = newAuthClientLimiter() + s.authFailureMAC = make([]byte, sha256.Size) + _, _ = rand.Read(s.authFailureMAC) go s.cleanupStaleProxies(ctx) return s } @@ -1172,6 +1182,18 @@ func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.Authen return nil, err } + failureKey := s.authFailureKey(req) + limitFailures := failureKey != "" && s.authAttemptLimiter != nil && len(s.authFailureMAC) > 0 + if limitFailures && s.authAttemptLimiter.isLimited(failureKey) { + return nil, status.Errorf(codes.ResourceExhausted, "too many failed authentication attempts for this credential, please try again later") + } + + clientKey := s.authClientKey(ctx, req.GetId()) + limitClient := clientKey != "" && s.authClientLimiter != nil + if limitClient && s.authClientLimiter.isLimited(clientKey) { + return nil, status.Errorf(codes.ResourceExhausted, "too many failed authentication attempts from this client, please try again later") + } + service, err := s.serviceManager.GetServiceByID(ctx, req.GetAccountId(), req.GetId()) if err != nil { log.WithContext(ctx).Debugf("failed to get service from store: %v", err) @@ -1179,6 +1201,14 @@ func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.Authen } authenticated, userId, method := s.authenticateRequest(ctx, req, service) + if !authenticated { + if limitFailures { + s.authAttemptLimiter.recordFailure(failureKey) + } + if limitClient { + s.authClientLimiter.recordFailure(clientKey) + } + } // Non-OIDC schemes (PIN/Password/Header) authenticate against per-service // secrets and have no user-level group context, so groups stay nil. Email @@ -1194,6 +1224,40 @@ func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.Authen }, nil } +func (s *ProxyServiceServer) authClientKey(ctx context.Context, serviceID string) string { + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return "" + } + values := md.Get(proxyauth.ClientIPMetadataKey) + if len(values) == 0 { + return "" + } + addr, err := netip.ParseAddr(strings.TrimSpace(values[0])) + if err != nil { + return "" + } + return serviceID + "|" + addr.Unmap().String() +} + +func (s *ProxyServiceServer) authFailureKey(req *proto.AuthenticateRequest) string { + var secret string + switch v := req.GetRequest().(type) { + case *proto.AuthenticateRequest_Pin: + secret = "pin|" + v.Pin.GetPin() + case *proto.AuthenticateRequest_Password: + secret = "password|" + v.Password.GetPassword() + case *proto.AuthenticateRequest_HeaderAuth: + secret = "header|" + v.HeaderAuth.GetHeaderName() + "|" + v.HeaderAuth.GetHeaderValue() + default: + return "" + } + + mac := hmac.New(sha256.New, s.authFailureMAC) + mac.Write([]byte(secret)) + return req.GetId() + "|" + hex.EncodeToString(mac.Sum(nil)) +} + func (s *ProxyServiceServer) authenticateRequest(ctx context.Context, req *proto.AuthenticateRequest, service *rpservice.Service) (bool, string, proxyauth.Method) { switch v := req.GetRequest().(type) { case *proto.AuthenticateRequest_Pin: diff --git a/management/internals/shared/grpc/proxy_auth_attempts_test.go b/management/internals/shared/grpc/proxy_auth_attempts_test.go new file mode 100644 index 000000000..7d4f4708f --- /dev/null +++ b/management/internals/shared/grpc/proxy_auth_attempts_test.go @@ -0,0 +1,189 @@ +package grpc + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "fmt" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/time/rate" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + proxyauth "github.com/netbirdio/netbird/proxy/auth" + "github.com/netbirdio/netbird/shared/hash/argon2id" + "github.com/netbirdio/netbird/shared/management/proto" +) + +const authAttemptsHeaderName = "X-API-Key" + +func newAuthAttemptsTestServer(t *testing.T) *ProxyServiceServer { + t.Helper() + + firstHash, err := argon2id.Hash("first-key") + require.NoError(t, err) + secondHash, err := argon2id.Hash("second-key") + require.NoError(t, err) + + svc := &rpservice.Service{ + ID: "svc1", + Domain: "example.com", + Auth: rpservice.AuthConfig{ + HeaderAuths: []*rpservice.HeaderAuthConfig{ + {Enabled: true, Header: authAttemptsHeaderName, Value: firstHash}, + {Enabled: true, Header: authAttemptsHeaderName, Value: secondHash}, + }, + }, + } + + ctrl := gomock.NewController(t) + mgr := rpservice.NewMockManager(ctrl) + mgr.EXPECT().GetServiceByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(svc, nil).AnyTimes() + + limiter := newAuthFailureLimiter() + t.Cleanup(limiter.stop) + clientLimiter := newAuthClientLimiter() + t.Cleanup(clientLimiter.stop) + + mac := make([]byte, sha256.Size) + _, err = rand.Read(mac) + require.NoError(t, err) + + return &ProxyServiceServer{ + serviceManager: mgr, + authAttemptLimiter: limiter, + authClientLimiter: clientLimiter, + authFailureMAC: mac, + } +} + +func clientIPContext(ip string) context.Context { + return metadata.NewIncomingContext(context.Background(), metadata.Pairs(proxyauth.ClientIPMetadataKey, ip)) +} + +func authAttemptsRequest(credential string) *proto.AuthenticateRequest { + return &proto.AuthenticateRequest{ + Id: "svc1", + AccountId: "acc1", + Request: &proto.AuthenticateRequest_HeaderAuth{ + HeaderAuth: &proto.HeaderAuthRequest{ + HeaderName: authAttemptsHeaderName, + HeaderValue: credential, + }, + }, + } +} + +func TestAuthenticate_ValidCredentialIsNeverRateLimited(t *testing.T) { + s := newAuthAttemptsTestServer(t) + + for i := 0; i < proxyAuthFailureBurst*2; i++ { + resp, err := s.Authenticate(context.Background(), authAttemptsRequest("first-key")) + require.NoError(t, err, "a valid credential must never be throttled (attempt %d)", i) + require.True(t, resp.GetSuccess()) + } +} + +func TestAuthenticate_FailedCredentialIsRateLimited(t *testing.T) { + s := newAuthAttemptsTestServer(t) + + for i := 0; i < proxyAuthFailureBurst; i++ { + resp, err := s.Authenticate(context.Background(), authAttemptsRequest("wrong-key")) + require.NoError(t, err, "attempt %d should be within the failure budget", i) + require.False(t, resp.GetSuccess()) + } + + _, err := s.Authenticate(context.Background(), authAttemptsRequest("wrong-key")) + require.Error(t, err) + assert.Equal(t, codes.ResourceExhausted, status.Code(err)) +} + +func TestAuthenticate_ThrottledCredentialDoesNotAffectOthers(t *testing.T) { + s := newAuthAttemptsTestServer(t) + + for i := 0; i < proxyAuthFailureBurst+2; i++ { + _, _ = s.Authenticate(context.Background(), authAttemptsRequest("wrong-key")) + } + + resp, err := s.Authenticate(context.Background(), authAttemptsRequest("first-key")) + require.NoError(t, err, "one throttled credential must not block a valid one") + assert.True(t, resp.GetSuccess()) + + resp, err = s.Authenticate(context.Background(), authAttemptsRequest("second-key")) + require.NoError(t, err) + assert.True(t, resp.GetSuccess()) + + _, err = s.Authenticate(context.Background(), authAttemptsRequest("another-wrong-key")) + require.NoError(t, err, "a different failing credential has its own budget") +} + +func TestAuthenticate_DistinctCredentialsThrottledPerClient(t *testing.T) { + const budget = 3 + + s := newAuthAttemptsTestServer(t) + s.authClientLimiter.stop() + s.authClientLimiter = newAuthLimiter(rate.Every(time.Hour), budget) + t.Cleanup(s.authClientLimiter.stop) + + ctx := clientIPContext("198.51.100.7") + + for i := 0; i < budget; i++ { + resp, err := s.Authenticate(ctx, authAttemptsRequest(fmt.Sprintf("garbage-%d", i))) + require.NoError(t, err, "attempt %d should be within the client budget", i) + require.False(t, resp.GetSuccess()) + } + + _, err := s.Authenticate(ctx, authAttemptsRequest("garbage-final")) + require.Error(t, err, "a client rotating distinct credentials must be throttled") + assert.Equal(t, codes.ResourceExhausted, status.Code(err)) + + other := clientIPContext("198.51.100.8") + resp, err := s.Authenticate(other, authAttemptsRequest("first-key")) + require.NoError(t, err, "a different client must be unaffected") + assert.True(t, resp.GetSuccess()) +} + +func TestAuthenticate_OneStaleCredentialDoesNotExhaustSharedClientBudget(t *testing.T) { + s := newAuthAttemptsTestServer(t) + ctx := clientIPContext("198.51.100.9") + + for i := 0; i < proxyAuthFailureBurst*4; i++ { + _, _ = s.Authenticate(ctx, authAttemptsRequest("stale-key")) + } + + resp, err := s.Authenticate(ctx, authAttemptsRequest("first-key")) + require.NoError(t, err, "one client stuck on a stale key must not block others behind the same NAT") + assert.True(t, resp.GetSuccess()) +} + +func TestAuthenticate_ProxyWithoutClientIPIsNotClientLimited(t *testing.T) { + s := newAuthAttemptsTestServer(t) + + for i := 0; i < proxyAuthFailureBurst*2; i++ { + _, _ = s.Authenticate(context.Background(), authAttemptsRequest(fmt.Sprintf("garbage-%d", i))) + } + + resp, err := s.Authenticate(context.Background(), authAttemptsRequest("first-key")) + require.NoError(t, err, "an old proxy must not have its clients share one budget") + assert.True(t, resp.GetSuccess()) +} + +func TestAuthenticate_MalformedClientIPIsIgnored(t *testing.T) { + s := newAuthAttemptsTestServer(t) + ctx := clientIPContext("not-an-ip") + + for i := 0; i < proxyAuthFailureBurst*2; i++ { + _, _ = s.Authenticate(ctx, authAttemptsRequest(fmt.Sprintf("garbage-%d", i))) + } + + resp, err := s.Authenticate(ctx, authAttemptsRequest("first-key")) + require.NoError(t, err) + assert.True(t, resp.GetSuccess()) +} diff --git a/management/internals/shared/grpc/proxy_auth_ratelimit.go b/management/internals/shared/grpc/proxy_auth_ratelimit.go index 78ab1bd20..99c046834 100644 --- a/management/internals/shared/grpc/proxy_auth_ratelimit.go +++ b/management/internals/shared/grpc/proxy_auth_ratelimit.go @@ -18,12 +18,16 @@ const ( proxyAuthLimiterCleanup = 5 * time.Minute // proxyAuthLimiterTTL is how long a limiter is kept after the last failure. proxyAuthLimiterTTL = 15 * time.Minute + + proxyAuthClientBurst = 30 ) // defaultProxyAuthFailureRate is the token replenishment rate for failed auth attempts. // One token every 12 seconds = 5 per minute. var defaultProxyAuthFailureRate = rate.Every(12 * time.Second) +var defaultProxyAuthClientRate = rate.Limit(1) + // clientIP identifies a client by its IP address for rate limiting purposes. type clientIP = string @@ -37,6 +41,7 @@ type authFailureLimiter struct { mu sync.Mutex limiters map[clientIP]*limiterEntry failureRate rate.Limit + burst int cancel context.CancelFunc } @@ -45,10 +50,19 @@ func newAuthFailureLimiter() *authFailureLimiter { } func newAuthFailureLimiterWithRate(failureRate rate.Limit) *authFailureLimiter { + return newAuthLimiter(failureRate, proxyAuthFailureBurst) +} + +func newAuthClientLimiter() *authFailureLimiter { + return newAuthLimiter(defaultProxyAuthClientRate, proxyAuthClientBurst) +} + +func newAuthLimiter(failureRate rate.Limit, burst int) *authFailureLimiter { ctx, cancel := context.WithCancel(context.Background()) l := &authFailureLimiter{ limiters: make(map[clientIP]*limiterEntry), failureRate: failureRate, + burst: burst, cancel: cancel, } go l.cleanupLoop(ctx) @@ -77,7 +91,7 @@ func (l *authFailureLimiter) recordFailure(ip clientIP) { entry, exists := l.limiters[ip] if !exists { entry = &limiterEntry{ - limiter: rate.NewLimiter(l.failureRate, proxyAuthFailureBurst), + limiter: rate.NewLimiter(l.failureRate, l.burst), } l.limiters[ip] = entry } diff --git a/proxy/auth/auth.go b/proxy/auth/auth.go index 78f0097d5..4e72fbb59 100644 --- a/proxy/auth/auth.go +++ b/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". diff --git a/proxy/internal/auth/header_cache.go b/proxy/internal/auth/header_cache.go new file mode 100644 index 000000000..a37610597 --- /dev/null +++ b/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[:]) +} diff --git a/proxy/internal/auth/header_cache_test.go b/proxy/internal/auth/header_cache_test.go new file mode 100644 index 000000000..e0966d545 --- /dev/null +++ b/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") +} diff --git a/proxy/internal/auth/middleware.go b/proxy/internal/auth/middleware.go index 72630b085..dc29e21d5 100644 --- a/proxy/internal/auth/middleware.go +++ b/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) diff --git a/proxy/internal/auth/tunnel_cache.go b/proxy/internal/auth/tunnel_cache.go index 185c53c62..d59ccb9d8 100644 --- a/proxy/internal/auth/tunnel_cache.go +++ b/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:]...)