[management] refactor auth (#3296)

2026-04-18 08:16:39 +00:00 · 2025-02-20 20:24:40 +00:00
parent d7d5b1b1d6
commit 77e40f41f2
64 changed files with 2085 additions and 1937 deletions
--- a/management/server/auth/jwt/extractor.go
+++ b/management/server/auth/jwt/extractor.go
@@ -0,0 +1,144 @@
+package jwt
+
+import (
+	"errors"
+	"net/url"
+	"time"
+
+	"github.com/golang-jwt/jwt"
+	log "github.com/sirupsen/logrus"
+
+	nbcontext "github.com/netbirdio/netbird/management/server/context"
+)
+
+const (
+	// AccountIDSuffix suffix for the account id claim
+	AccountIDSuffix = "wt_account_id"
+	// DomainIDSuffix suffix for the domain id claim
+	DomainIDSuffix = "wt_account_domain"
+	// DomainCategorySuffix suffix for the domain category claim
+	DomainCategorySuffix = "wt_account_domain_category"
+	// UserIDClaim claim for the user id
+	UserIDClaim = "sub"
+	// LastLoginSuffix claim for the last login
+	LastLoginSuffix = "nb_last_login"
+	// Invited claim indicates that an incoming JWT is from a user that just accepted an invitation
+	Invited = "nb_invited"
+)
+
+var (
+	errUserIDClaimEmpty = errors.New("user ID claim token value is empty")
+)
+
+// ClaimsExtractor struct that holds the extract function
+type ClaimsExtractor struct {
+	authAudience string
+	userIDClaim  string
+}
+
+// ClaimsExtractorOption is a function that configures the ClaimsExtractor
+type ClaimsExtractorOption func(*ClaimsExtractor)
+
+// WithAudience sets the audience for the extractor
+func WithAudience(audience string) ClaimsExtractorOption {
+	return func(c *ClaimsExtractor) {
+		c.authAudience = audience
+	}
+}
+
+// WithUserIDClaim sets the user id claim for the extractor
+func WithUserIDClaim(userIDClaim string) ClaimsExtractorOption {
+	return func(c *ClaimsExtractor) {
+		c.userIDClaim = userIDClaim
+	}
+}
+
+// NewClaimsExtractor returns an extractor, and if provided with a function with ExtractClaims signature,
+// then it will use that logic. Uses ExtractClaimsFromRequestContext by default
+func NewClaimsExtractor(options ...ClaimsExtractorOption) *ClaimsExtractor {
+	ce := &ClaimsExtractor{}
+	for _, option := range options {
+		option(ce)
+	}
+
+	if ce.userIDClaim == "" {
+		ce.userIDClaim = UserIDClaim
+	}
+	return ce
+}
+
+func parseTime(timeString string) time.Time {
+	if timeString == "" {
+		return time.Time{}
+	}
+	parsedTime, err := time.Parse(time.RFC3339, timeString)
+	if err != nil {
+		return time.Time{}
+	}
+	return parsedTime
+}
+
+func (c ClaimsExtractor) audienceClaim(claimName string) string {
+	url, err := url.JoinPath(c.authAudience, claimName)
+	if err != nil {
+		return c.authAudience + claimName // as it was previously
+	}
+
+	return url
+}
+
+func (c *ClaimsExtractor) ToUserAuth(token *jwt.Token) (nbcontext.UserAuth, error) {
+	claims := token.Claims.(jwt.MapClaims)
+	userAuth := nbcontext.UserAuth{}
+
+	userID, ok := claims[c.userIDClaim].(string)
+	if !ok {
+		return userAuth, errUserIDClaimEmpty
+	}
+	userAuth.UserId = userID
+
+	if accountIDClaim, ok := claims[c.audienceClaim(AccountIDSuffix)]; ok {
+		userAuth.AccountId = accountIDClaim.(string)
+	}
+
+	if domainClaim, ok := claims[c.audienceClaim(DomainIDSuffix)]; ok {
+		userAuth.Domain = domainClaim.(string)
+	}
+
+	if domainCategoryClaim, ok := claims[c.audienceClaim(DomainCategorySuffix)]; ok {
+		userAuth.DomainCategory = domainCategoryClaim.(string)
+	}
+
+	if lastLoginClaimString, ok := claims[c.audienceClaim(LastLoginSuffix)]; ok {
+		userAuth.LastLogin = parseTime(lastLoginClaimString.(string))
+	}
+
+	if invitedBool, ok := claims[c.audienceClaim(Invited)]; ok {
+		if value, ok := invitedBool.(bool); ok {
+			userAuth.Invited = value
+		}
+	}
+
+	return userAuth, nil
+}
+
+func (c *ClaimsExtractor) ToGroups(token *jwt.Token, claimName string) []string {
+	claims := token.Claims.(jwt.MapClaims)
+	userJWTGroups := make([]string, 0)
+
+	if claim, ok := claims[claimName]; ok {
+		if claimGroups, ok := claim.([]interface{}); ok {
+			for _, g := range claimGroups {
+				if group, ok := g.(string); ok {
+					userJWTGroups = append(userJWTGroups, group)
+				} else {
+					log.Debugf("JWT claim %q contains a non-string group (type: %T): %v", claimName, g, g)
+				}
+			}
+		}
+	} else {
+		log.Debugf("JWT claim %q is not a string array", claimName)
+	}
+
+	return userJWTGroups
+}
--- a/management/server/auth/jwt/validator.go
+++ b/management/server/auth/jwt/validator.go
@@ -0,0 +1,302 @@
+package jwt
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rsa"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"math/big"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/golang-jwt/jwt"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation
+type Jwks struct {
+	Keys          []JSONWebKey `json:"keys"`
+	expiresInTime time.Time
+}
+
+// The supported elliptic curves types
+const (
+	// p256 represents a cryptographic elliptical curve type.
+	p256 = "P-256"
+
+	// p384 represents a cryptographic elliptical curve type.
+	p384 = "P-384"
+
+	// p521 represents a cryptographic elliptical curve type.
+	p521 = "P-521"
+)
+
+// JSONWebKey is a representation of a Jason Web Key
+type JSONWebKey struct {
+	Kty string   `json:"kty"`
+	Kid string   `json:"kid"`
+	Use string   `json:"use"`
+	N   string   `json:"n"`
+	E   string   `json:"e"`
+	Crv string   `json:"crv"`
+	X   string   `json:"x"`
+	Y   string   `json:"y"`
+	X5c []string `json:"x5c"`
+}
+
+type Validator struct {
+	lock                     sync.Mutex
+	issuer                   string
+	audienceList             []string
+	keysLocation             string
+	idpSignkeyRefreshEnabled bool
+	keys                     *Jwks
+}
+
+var (
+	errKeyNotFound     = errors.New("unable to find appropriate key")
+	errInvalidAudience = errors.New("invalid audience")
+	errInvalidIssuer   = errors.New("invalid issuer")
+	errTokenEmpty      = errors.New("required authorization token not found")
+	errTokenInvalid    = errors.New("token is invalid")
+	errTokenParsing    = errors.New("token could not be parsed")
+)
+
+func NewValidator(issuer string, audienceList []string, keysLocation string, idpSignkeyRefreshEnabled bool) *Validator {
+	keys, err := getPemKeys(keysLocation)
+	if err != nil {
+		log.WithField("keysLocation", keysLocation).Errorf("could not get keys from location: %s", err)
+	}
+
+	return &Validator{
+		keys:                     keys,
+		issuer:                   issuer,
+		audienceList:             audienceList,
+		keysLocation:             keysLocation,
+		idpSignkeyRefreshEnabled: idpSignkeyRefreshEnabled,
+	}
+}
+
+func (v *Validator) getKeyFunc(ctx context.Context) jwt.Keyfunc {
+	return func(token *jwt.Token) (interface{}, error) {
+		// Verify 'aud' claim
+		var checkAud bool
+		for _, audience := range v.audienceList {
+			checkAud = token.Claims.(jwt.MapClaims).VerifyAudience(audience, false)
+			if checkAud {
+				break
+			}
+		}
+		if !checkAud {
+			return token, errInvalidAudience
+		}
+
+		// Verify 'issuer' claim
+		checkIss := token.Claims.(jwt.MapClaims).VerifyIssuer(v.issuer, false)
+		if !checkIss {
+			return token, errInvalidIssuer
+		}
+
+		// If keys are rotated, verify the keys prior to token validation
+		if v.idpSignkeyRefreshEnabled {
+			// If the keys are invalid, retrieve new ones
+			// @todo propose a separate go routine to regularly check these to prevent blocking when actually
+			// validating the token
+			if !v.keys.stillValid() {
+				v.lock.Lock()
+				defer v.lock.Unlock()
+
+				refreshedKeys, err := getPemKeys(v.keysLocation)
+				if err != nil {
+					log.WithContext(ctx).Debugf("cannot get JSONWebKey: %v, falling back to old keys", err)
+					refreshedKeys = v.keys
+				}
+
+				log.WithContext(ctx).Debugf("keys refreshed, new UTC expiration time: %s", refreshedKeys.expiresInTime.UTC())
+
+				v.keys = refreshedKeys
+			}
+		}
+
+		publicKey, err := getPublicKey(token, v.keys)
+		if err == nil {
+			return publicKey, nil
+		}
+
+		msg := fmt.Sprintf("getPublicKey error: %s", err)
+		if errors.Is(err, errKeyNotFound) && !v.idpSignkeyRefreshEnabled {
+			msg = fmt.Sprintf("getPublicKey error: %s. You can enable key refresh by setting HttpServerConfig.IdpSignKeyRefreshEnabled to true in your management.json file and restart the service", err)
+		}
+
+		log.WithContext(ctx).Error(msg)
+
+		return nil, err
+	}
+}
+
+// ValidateAndParse validates the token and returns the parsed token
+func (m *Validator) ValidateAndParse(ctx context.Context, token string) (*jwt.Token, error) {
+	// If the token is empty...
+	if token == "" {
+		// If we get here, the required token is missing
+		log.WithContext(ctx).Debugf("  Error: No credentials found (CredentialsOptional=false)")
+		return nil, errTokenEmpty
+	}
+
+	// Now parse the token
+	parsedToken, err := jwt.Parse(token, m.getKeyFunc(ctx))
+
+	// Check if there was an error in parsing...
+	if err != nil {
+		err = fmt.Errorf("%w: %s", errTokenParsing, err)
+		log.WithContext(ctx).Error(err.Error())
+		return nil, err
+	}
+
+	// Check if the parsed token is valid...
+	if !parsedToken.Valid {
+		log.WithContext(ctx).Debug(errTokenInvalid.Error())
+		return nil, errTokenInvalid
+	}
+
+	return parsedToken, nil
+}
+
+// stillValid returns true if the JSONWebKey still valid and have enough time to be used
+func (jwks *Jwks) stillValid() bool {
+	return !jwks.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(jwks.expiresInTime)
+}
+
+func getPemKeys(keysLocation string) (*Jwks, error) {
+	jwks := &Jwks{}
+
+	url, err := url.ParseRequestURI(keysLocation)
+	if err != nil {
+		return jwks, err
+	}
+
+	resp, err := http.Get(url.String())
+	if err != nil {
+		return jwks, err
+	}
+	defer resp.Body.Close()
+
+	err = json.NewDecoder(resp.Body).Decode(jwks)
+	if err != nil {
+		return jwks, err
+	}
+
+	cacheControlHeader := resp.Header.Get("Cache-Control")
+	expiresIn := getMaxAgeFromCacheHeader(cacheControlHeader)
+	jwks.expiresInTime = time.Now().Add(time.Duration(expiresIn) * time.Second)
+
+	return jwks, nil
+}
+
+func getPublicKey(token *jwt.Token, jwks *Jwks) (interface{}, error) {
+	// todo as we load the jkws when the server is starting, we should build a JKS map with the pem cert at the boot time
+	for k := range jwks.Keys {
+		if token.Header["kid"] != jwks.Keys[k].Kid {
+			continue
+		}
+
+		if len(jwks.Keys[k].X5c) != 0 {
+			cert := "-----BEGIN CERTIFICATE-----\n" + jwks.Keys[k].X5c[0] + "\n-----END CERTIFICATE-----"
+			return jwt.ParseRSAPublicKeyFromPEM([]byte(cert))
+		}
+
+		if jwks.Keys[k].Kty == "RSA" {
+			return getPublicKeyFromRSA(jwks.Keys[k])
+		}
+		if jwks.Keys[k].Kty == "EC" {
+			return getPublicKeyFromECDSA(jwks.Keys[k])
+		}
+	}
+
+	return nil, errKeyNotFound
+}
+
+func getPublicKeyFromECDSA(jwk JSONWebKey) (publicKey *ecdsa.PublicKey, err error) {
+	if jwk.X == "" || jwk.Y == "" || jwk.Crv == "" {
+		return nil, fmt.Errorf("ecdsa key incomplete")
+	}
+
+	var xCoordinate []byte
+	if xCoordinate, err = base64.RawURLEncoding.DecodeString(jwk.X); err != nil {
+		return nil, err
+	}
+
+	var yCoordinate []byte
+	if yCoordinate, err = base64.RawURLEncoding.DecodeString(jwk.Y); err != nil {
+		return nil, err
+	}
+
+	publicKey = &ecdsa.PublicKey{}
+
+	var curve elliptic.Curve
+	switch jwk.Crv {
+	case p256:
+		curve = elliptic.P256()
+	case p384:
+		curve = elliptic.P384()
+	case p521:
+		curve = elliptic.P521()
+	}
+
+	publicKey.Curve = curve
+	publicKey.X = big.NewInt(0).SetBytes(xCoordinate)
+	publicKey.Y = big.NewInt(0).SetBytes(yCoordinate)
+
+	return publicKey, nil
+}
+
+func getPublicKeyFromRSA(jwk JSONWebKey) (*rsa.PublicKey, error) {
+	decodedE, err := base64.RawURLEncoding.DecodeString(jwk.E)
+	if err != nil {
+		return nil, err
+	}
+	decodedN, err := base64.RawURLEncoding.DecodeString(jwk.N)
+	if err != nil {
+		return nil, err
+	}
+
+	var n, e big.Int
+	e.SetBytes(decodedE)
+	n.SetBytes(decodedN)
+
+	return &rsa.PublicKey{
+		E: int(e.Int64()),
+		N: &n,
+	}, nil
+}
+
+// getMaxAgeFromCacheHeader extracts max-age directive from the Cache-Control header
+func getMaxAgeFromCacheHeader(cacheControl string) int {
+	// Split into individual directives
+	directives := strings.Split(cacheControl, ",")
+
+	for _, directive := range directives {
+		directive = strings.TrimSpace(directive)
+		if strings.HasPrefix(directive, "max-age=") {
+			// Extract the max-age value
+			maxAgeStr := strings.TrimPrefix(directive, "max-age=")
+			maxAge, err := strconv.Atoi(maxAgeStr)
+			if err != nil {
+				return 0
+			}
+
+			return maxAge
+		}
+	}
+
+	return 0
+}
--- a/management/server/auth/manager.go
+++ b/management/server/auth/manager.go
@@ -0,0 +1,170 @@
+package auth
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"hash/crc32"
+
+	"github.com/golang-jwt/jwt"
+
+	"github.com/netbirdio/netbird/base62"
+	nbjwt "github.com/netbirdio/netbird/management/server/auth/jwt"
+	nbcontext "github.com/netbirdio/netbird/management/server/context"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+var _ Manager = (*manager)(nil)
+
+type Manager interface {
+	ValidateAndParseToken(ctx context.Context, value string) (nbcontext.UserAuth, *jwt.Token, error)
+	EnsureUserAccessByJWTGroups(ctx context.Context, userAuth nbcontext.UserAuth, token *jwt.Token) (nbcontext.UserAuth, error)
+	MarkPATUsed(ctx context.Context, tokenID string) error
+	GetPATInfo(ctx context.Context, token string) (user *types.User, pat *types.PersonalAccessToken, domain string, category string, err error)
+}
+
+type manager struct {
+	store store.Store
+
+	validator *nbjwt.Validator
+	extractor *nbjwt.ClaimsExtractor
+}
+
+func NewManager(store store.Store, issuer, audience, keysLocation, userIdClaim string, allAudiences []string, idpRefreshKeys bool) Manager {
+	// @note if invalid/missing parameters are sent the validator will instantiate
+	// but it will fail when validating and parsing the token
+	jwtValidator := nbjwt.NewValidator(
+		issuer,
+		allAudiences,
+		keysLocation,
+		idpRefreshKeys,
+	)
+
+	claimsExtractor := nbjwt.NewClaimsExtractor(
+		nbjwt.WithAudience(audience),
+		nbjwt.WithUserIDClaim(userIdClaim),
+	)
+
+	return &manager{
+		store: store,
+
+		validator: jwtValidator,
+		extractor: claimsExtractor,
+	}
+}
+
+func (m *manager) ValidateAndParseToken(ctx context.Context, value string) (nbcontext.UserAuth, *jwt.Token, error) {
+	token, err := m.validator.ValidateAndParse(ctx, value)
+	if err != nil {
+		return nbcontext.UserAuth{}, nil, err
+	}
+
+	userAuth, err := m.extractor.ToUserAuth(token)
+	if err != nil {
+		return nbcontext.UserAuth{}, nil, err
+	}
+	return userAuth, token, err
+}
+
+func (m *manager) EnsureUserAccessByJWTGroups(ctx context.Context, userAuth nbcontext.UserAuth, token *jwt.Token) (nbcontext.UserAuth, error) {
+	if userAuth.IsChild || userAuth.IsPAT {
+		return userAuth, nil
+	}
+
+	settings, err := m.store.GetAccountSettings(ctx, store.LockingStrengthShare, userAuth.AccountId)
+	if err != nil {
+		return userAuth, err
+	}
+
+	// Ensures JWT group synchronization to the management is enabled before,
+	// filtering access based on the allowed groups.
+	if settings != nil && settings.JWTGroupsEnabled {
+		userAuth.Groups = m.extractor.ToGroups(token, settings.JWTGroupsClaimName)
+		if allowedGroups := settings.JWTAllowGroups; len(allowedGroups) > 0 {
+			if !userHasAllowedGroup(allowedGroups, userAuth.Groups) {
+				return userAuth, fmt.Errorf("user does not belong to any of the allowed JWT groups")
+			}
+		}
+	}
+
+	return userAuth, nil
+}
+
+// MarkPATUsed marks a personal access token as used
+func (am *manager) MarkPATUsed(ctx context.Context, tokenID string) error {
+	return am.store.MarkPATUsed(ctx, store.LockingStrengthUpdate, tokenID)
+}
+
+// GetPATInfo retrieves user, personal access token, domain, and category details from a personal access token.
+func (am *manager) GetPATInfo(ctx context.Context, token string) (user *types.User, pat *types.PersonalAccessToken, domain string, category string, err error) {
+	user, pat, err = am.extractPATFromToken(ctx, token)
+	if err != nil {
+		return nil, nil, "", "", err
+	}
+
+	domain, category, err = am.store.GetAccountDomainAndCategory(ctx, store.LockingStrengthShare, user.AccountID)
+	if err != nil {
+		return nil, nil, "", "", err
+	}
+
+	return user, pat, domain, category, nil
+}
+
+// extractPATFromToken validates the token structure and retrieves associated User and PAT.
+func (am *manager) extractPATFromToken(ctx context.Context, token string) (*types.User, *types.PersonalAccessToken, error) {
+	if len(token) != types.PATLength {
+		return nil, nil, fmt.Errorf("PAT has incorrect length")
+	}
+
+	prefix := token[:len(types.PATPrefix)]
+	if prefix != types.PATPrefix {
+		return nil, nil, fmt.Errorf("PAT has wrong prefix")
+	}
+	secret := token[len(types.PATPrefix) : len(types.PATPrefix)+types.PATSecretLength]
+	encodedChecksum := token[len(types.PATPrefix)+types.PATSecretLength : len(types.PATPrefix)+types.PATSecretLength+types.PATChecksumLength]
+
+	verificationChecksum, err := base62.Decode(encodedChecksum)
+	if err != nil {
+		return nil, nil, fmt.Errorf("PAT checksum decoding failed: %w", err)
+	}
+
+	secretChecksum := crc32.ChecksumIEEE([]byte(secret))
+	if secretChecksum != verificationChecksum {
+		return nil, nil, fmt.Errorf("PAT checksum does not match")
+	}
+
+	hashedToken := sha256.Sum256([]byte(token))
+	encodedHashedToken := base64.StdEncoding.EncodeToString(hashedToken[:])
+
+	var user *types.User
+	var pat *types.PersonalAccessToken
+
+	err = am.store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
+		pat, err = transaction.GetPATByHashedToken(ctx, store.LockingStrengthShare, encodedHashedToken)
+		if err != nil {
+			return err
+		}
+
+		user, err = transaction.GetUserByPATID(ctx, store.LockingStrengthShare, pat.ID)
+		return err
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return user, pat, nil
+}
+
+// userHasAllowedGroup checks if a user belongs to any of the allowed groups.
+func userHasAllowedGroup(allowedGroups []string, userGroups []string) bool {
+	for _, userGroup := range userGroups {
+		for _, allowedGroup := range allowedGroups {
+			if userGroup == allowedGroup {
+				return true
+			}
+		}
+	}
+	return false
+}
--- a/management/server/auth/manager_mock.go
+++ b/management/server/auth/manager_mock.go
@@ -0,0 +1,54 @@
+package auth
+
+import (
+	"context"
+
+	"github.com/golang-jwt/jwt"
+
+	nbcontext "github.com/netbirdio/netbird/management/server/context"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+var (
+	_ Manager = (*MockManager)(nil)
+)
+
+// @note really dislike this mocking approach but rather than have to do additional test refactoring.
+type MockManager struct {
+	ValidateAndParseTokenFunc       func(ctx context.Context, value string) (nbcontext.UserAuth, *jwt.Token, error)
+	EnsureUserAccessByJWTGroupsFunc func(ctx context.Context, userAuth nbcontext.UserAuth, token *jwt.Token) (nbcontext.UserAuth, error)
+	MarkPATUsedFunc                 func(ctx context.Context, tokenID string) error
+	GetPATInfoFunc                  func(ctx context.Context, token string) (user *types.User, pat *types.PersonalAccessToken, domain string, category string, err error)
+}
+
+// EnsureUserAccessByJWTGroups implements Manager.
+func (m *MockManager) EnsureUserAccessByJWTGroups(ctx context.Context, userAuth nbcontext.UserAuth, token *jwt.Token) (nbcontext.UserAuth, error) {
+	if m.EnsureUserAccessByJWTGroupsFunc != nil {
+		return m.EnsureUserAccessByJWTGroupsFunc(ctx, userAuth, token)
+	}
+	return nbcontext.UserAuth{}, nil
+}
+
+// GetPATInfo implements Manager.
+func (m *MockManager) GetPATInfo(ctx context.Context, token string) (user *types.User, pat *types.PersonalAccessToken, domain string, category string, err error) {
+	if m.GetPATInfoFunc != nil {
+		return m.GetPATInfoFunc(ctx, token)
+	}
+	return &types.User{}, &types.PersonalAccessToken{}, "", "", nil
+}
+
+// MarkPATUsed implements Manager.
+func (m *MockManager) MarkPATUsed(ctx context.Context, tokenID string) error {
+	if m.MarkPATUsedFunc != nil {
+		return m.MarkPATUsedFunc(ctx, tokenID)
+	}
+	return nil
+}
+
+// ValidateAndParseToken implements Manager.
+func (m *MockManager) ValidateAndParseToken(ctx context.Context, value string) (nbcontext.UserAuth, *jwt.Token, error) {
+	if m.ValidateAndParseTokenFunc != nil {
+		return m.ValidateAndParseTokenFunc(ctx, value)
+	}
+	return nbcontext.UserAuth{}, &jwt.Token{}, nil
+}
--- a/management/server/auth/manager_test.go
+++ b/management/server/auth/manager_test.go
@@ -0,0 +1,407 @@
+package auth_test
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/golang-jwt/jwt"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/auth"
+	nbjwt "github.com/netbirdio/netbird/management/server/auth/jwt"
+	nbcontext "github.com/netbirdio/netbird/management/server/context"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+func TestAuthManager_GetAccountInfoFromPAT(t *testing.T) {
+	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	if err != nil {
+		t.Fatalf("Error when creating store: %s", err)
+	}
+	t.Cleanup(cleanup)
+
+	token := "nbp_9999EUDNdkeusjentDLSJEn1902u84390W6W"
+	hashedToken := sha256.Sum256([]byte(token))
+	encodedHashedToken := base64.StdEncoding.EncodeToString(hashedToken[:])
+	account := &types.Account{
+		Id: "account_id",
+		Users: map[string]*types.User{"someUser": {
+			Id: "someUser",
+			PATs: map[string]*types.PersonalAccessToken{
+				"tokenId": {
+					ID:          "tokenId",
+					UserID:      "someUser",
+					HashedToken: encodedHashedToken,
+				},
+			},
+		}},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatalf("Error when saving account: %s", err)
+	}
+
+	manager := auth.NewManager(store, "", "", "", "", []string{}, false)
+
+	user, pat, _, _, err := manager.GetPATInfo(context.Background(), token)
+	if err != nil {
+		t.Fatalf("Error when getting Account from PAT: %s", err)
+	}
+
+	assert.Equal(t, "account_id", user.AccountID)
+	assert.Equal(t, "someUser", user.Id)
+	assert.Equal(t, account.Users["someUser"].PATs["tokenId"].ID, pat.ID)
+}
+
+func TestAuthManager_MarkPATUsed(t *testing.T) {
+	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	if err != nil {
+		t.Fatalf("Error when creating store: %s", err)
+	}
+	t.Cleanup(cleanup)
+
+	token := "nbp_9999EUDNdkeusjentDLSJEn1902u84390W6W"
+	hashedToken := sha256.Sum256([]byte(token))
+	encodedHashedToken := base64.StdEncoding.EncodeToString(hashedToken[:])
+	account := &types.Account{
+		Id: "account_id",
+		Users: map[string]*types.User{"someUser": {
+			Id: "someUser",
+			PATs: map[string]*types.PersonalAccessToken{
+				"tokenId": {
+					ID:          "tokenId",
+					HashedToken: encodedHashedToken,
+				},
+			},
+		}},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatalf("Error when saving account: %s", err)
+	}
+
+	manager := auth.NewManager(store, "", "", "", "", []string{}, false)
+
+	err = manager.MarkPATUsed(context.Background(), "tokenId")
+	if err != nil {
+		t.Fatalf("Error when marking PAT used: %s", err)
+	}
+
+	account, err = store.GetAccount(context.Background(), "account_id")
+	if err != nil {
+		t.Fatalf("Error when getting account: %s", err)
+	}
+	assert.True(t, !account.Users["someUser"].PATs["tokenId"].GetLastUsed().IsZero())
+}
+
+func TestAuthManager_EnsureUserAccessByJWTGroups(t *testing.T) {
+	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	if err != nil {
+		t.Fatalf("Error when creating store: %s", err)
+	}
+	t.Cleanup(cleanup)
+
+	userId := "user-id"
+	domain := "test.domain"
+
+	account := &types.Account{
+		Id:     "account_id",
+		Domain: domain,
+		Users: map[string]*types.User{"someUser": {
+			Id: "someUser",
+		}},
+		Settings: &types.Settings{},
+	}
+
+	err = store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatalf("Error when saving account: %s", err)
+	}
+
+	// this has been validated and parsed by ValidateAndParseToken
+	userAuth := nbcontext.UserAuth{
+		AccountId:      account.Id,
+		Domain:         domain,
+		UserId:         userId,
+		DomainCategory: "test-category",
+		// Groups:         []string{"group1", "group2"},
+	}
+
+	// these tests only assert groups are parsed from token as per account settings
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{"idp-groups": []interface{}{"group1", "group2"}})
+
+	manager := auth.NewManager(store, "", "", "", "", []string{}, false)
+
+	t.Run("JWT groups disabled", func(t *testing.T) {
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+		require.Len(t, userAuth.Groups, 0, "account not enabled to ensure access by groups")
+	})
+
+	t.Run("User impersonated", func(t *testing.T) {
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+		require.Len(t, userAuth.Groups, 0, "account not enabled to ensure access by groups")
+	})
+
+	t.Run("User PAT", func(t *testing.T) {
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+		require.Len(t, userAuth.Groups, 0, "account not enabled to ensure access by groups")
+	})
+
+	t.Run("JWT groups enabled without claim name", func(t *testing.T) {
+		account.Settings.JWTGroupsEnabled = true
+		err := store.SaveAccount(context.Background(), account)
+		require.NoError(t, err, "save account failed")
+
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+		require.Len(t, userAuth.Groups, 0, "account missing groups claim name")
+	})
+
+	t.Run("JWT groups enabled without allowed groups", func(t *testing.T) {
+		account.Settings.JWTGroupsEnabled = true
+		account.Settings.JWTGroupsClaimName = "idp-groups"
+		err := store.SaveAccount(context.Background(), account)
+		require.NoError(t, err, "save account failed")
+
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+		require.Equal(t, []string{"group1", "group2"}, userAuth.Groups, "group parsed do not match")
+	})
+
+	t.Run("User in allowed JWT groups", func(t *testing.T) {
+		account.Settings.JWTGroupsEnabled = true
+		account.Settings.JWTGroupsClaimName = "idp-groups"
+		account.Settings.JWTAllowGroups = []string{"group1"}
+		err := store.SaveAccount(context.Background(), account)
+		require.NoError(t, err, "save account failed")
+
+		userAuth, err := manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.NoError(t, err, "ensure user access by JWT groups failed")
+
+		require.Equal(t, []string{"group1", "group2"}, userAuth.Groups, "group parsed do not match")
+	})
+
+	t.Run("User not in allowed JWT groups", func(t *testing.T) {
+		account.Settings.JWTGroupsEnabled = true
+		account.Settings.JWTGroupsClaimName = "idp-groups"
+		account.Settings.JWTAllowGroups = []string{"not-a-group"}
+		err := store.SaveAccount(context.Background(), account)
+		require.NoError(t, err, "save account failed")
+
+		_, err = manager.EnsureUserAccessByJWTGroups(context.Background(), userAuth, token)
+		require.Error(t, err, "ensure user access is not in allowed groups")
+	})
+}
+
+func TestAuthManager_ValidateAndParseToken(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Add("Cache-Control", "max-age=30") // set a 30s expiry to these keys
+		http.ServeFile(w, r, "test_data/jwks.json")
+	}))
+	defer server.Close()
+
+	issuer := "http://issuer.local"
+	audience := "http://audience.local"
+	userIdClaim := "" // defaults to "sub"
+
+	// we're only testing with RSA256
+	keyData, _ := os.ReadFile("test_data/sample_key")
+	key, _ := jwt.ParseRSAPrivateKeyFromPEM(keyData)
+	keyId := "test-key"
+
+	// note, we can use a nil store because ValidateAndParseToken does not use it in it's flow
+	manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false)
+
+	customClaim := func(name string) string {
+		return fmt.Sprintf("%s/%s", audience, name)
+	}
+
+	lastLogin := time.Date(2025, 2, 12, 14, 25, 26, 0, time.UTC) //"2025-02-12T14:25:26.186Z"
+
+	tests := []struct {
+		name      string
+		tokenFunc func() string
+		expected  *nbcontext.UserAuth // nil indicates expected error
+	}{
+		{
+			name: "Valid with custom claims",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss":                                   issuer,
+					"aud":                                   []string{audience},
+					"iat":                                   time.Now().Unix(),
+					"exp":                                   time.Now().Add(time.Hour * 1).Unix(),
+					"sub":                                   "user-id|123",
+					customClaim(nbjwt.AccountIDSuffix):      "account-id|567",
+					customClaim(nbjwt.DomainIDSuffix):       "http://localhost",
+					customClaim(nbjwt.DomainCategorySuffix): "private",
+					customClaim(nbjwt.LastLoginSuffix):      lastLogin.Format(time.RFC3339),
+					customClaim(nbjwt.Invited):              false,
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+			expected: &nbcontext.UserAuth{
+				UserId:         "user-id|123",
+				AccountId:      "account-id|567",
+				Domain:         "http://localhost",
+				DomainCategory: "private",
+				LastLogin:      lastLogin,
+				Invited:        false,
+			},
+		},
+		{
+			name: "Valid without custom claims",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": issuer,
+					"aud": []string{audience},
+					"iat": time.Now().Unix(),
+					"exp": time.Now().Add(time.Hour).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+			expected: &nbcontext.UserAuth{
+				UserId: "user-id|123",
+			},
+		},
+		{
+			name: "Expired token",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": issuer,
+					"aud": []string{audience},
+					"iat": time.Now().Add(time.Hour * -2).Unix(),
+					"exp": time.Now().Add(time.Hour * -1).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+		},
+		{
+			name: "Not yet valid",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": issuer,
+					"aud": []string{audience},
+					"iat": time.Now().Add(time.Hour).Unix(),
+					"exp": time.Now().Add(time.Hour * 2).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+		},
+		{
+			name: "Invalid signature",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": issuer,
+					"aud": []string{audience},
+					"iat": time.Now().Unix(),
+					"exp": time.Now().Add(time.Hour).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				parts := strings.Split(tokenString, ".")
+				parts[2] = "invalid-signature"
+				return strings.Join(parts, ".")
+			},
+		},
+		{
+			name: "Invalid issuer",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": "not-the-issuer",
+					"aud": []string{audience},
+					"iat": time.Now().Unix(),
+					"exp": time.Now().Add(time.Hour).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+		},
+		{
+			name: "Invalid audience",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss": issuer,
+					"aud": []string{"not-the-audience"},
+					"iat": time.Now().Unix(),
+					"exp": time.Now().Add(time.Hour).Unix(),
+					"sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+		},
+		{
+			name: "Invalid user claim",
+			tokenFunc: func() string {
+				token := jwt.New(jwt.SigningMethodRS256)
+				token.Header["kid"] = keyId
+				token.Claims = jwt.MapClaims{
+					"iss":     issuer,
+					"aud":     []string{audience},
+					"iat":     time.Now().Unix(),
+					"exp":     time.Now().Add(time.Hour).Unix(),
+					"not-sub": "user-id|123",
+				}
+				tokenString, _ := token.SignedString(key)
+				return tokenString
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tokenString := tt.tokenFunc()
+
+			userAuth, token, err := manager.ValidateAndParseToken(context.Background(), tokenString)
+
+			if tt.expected != nil {
+				assert.NoError(t, err)
+				assert.True(t, token.Valid)
+				assert.Equal(t, *tt.expected, userAuth)
+			} else {
+				assert.Error(t, err)
+				assert.Nil(t, token)
+				assert.Empty(t, userAuth)
+			}
+		})
+	}
+
+}
--- a/management/server/auth/test_data/jwks.json
+++ b/management/server/auth/test_data/jwks.json
@@ -0,0 +1,11 @@
+{
+    "keys": [
+        {
+            "kty": "RSA",
+            "kid": "test-key",
+            "use": "sig",
+            "n": "4f5wg5l2hKsTeNem_V41fGnJm6gOdrj8ym3rFkEU_wT8RDtnSgFEZOQpHEgQ7JL38xUfU0Y3g6aYw9QT0hJ7mCpz9Er5qLaMXJwZxzHzAahlfA0icqabvJOMvQtzD6uQv6wPEyZtDTWiQi9AXwBpHssPnpYGIn20ZZuNlX2BrClciHhCPUIIZOQn_MmqTD31jSyjoQoV7MhhMTATKJx2XrHhR-1DcKJzQBSTAGnpYVaqpsARap-nwRipr3nUTuxyGohBTSmjJ2usSeQXHI3bODIRe1AuTyHceAbewn8b462yEWKARdpd9AjQW5SIVPfdsz5B6GlYQ5LdYKtznTuy7w",
+            "e": "AQAB"
+        }
+    ]
+}
--- a/management/server/auth/test_data/sample_key
+++ b/management/server/auth/test_data/sample_key
@@ -0,0 +1,27 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIEowIBAAKCAQEA4f5wg5l2hKsTeNem/V41fGnJm6gOdrj8ym3rFkEU/wT8RDtn
+SgFEZOQpHEgQ7JL38xUfU0Y3g6aYw9QT0hJ7mCpz9Er5qLaMXJwZxzHzAahlfA0i
+cqabvJOMvQtzD6uQv6wPEyZtDTWiQi9AXwBpHssPnpYGIn20ZZuNlX2BrClciHhC
+PUIIZOQn/MmqTD31jSyjoQoV7MhhMTATKJx2XrHhR+1DcKJzQBSTAGnpYVaqpsAR
+ap+nwRipr3nUTuxyGohBTSmjJ2usSeQXHI3bODIRe1AuTyHceAbewn8b462yEWKA
+Rdpd9AjQW5SIVPfdsz5B6GlYQ5LdYKtznTuy7wIDAQABAoIBAQCwia1k7+2oZ2d3
+n6agCAbqIE1QXfCmh41ZqJHbOY3oRQG3X1wpcGH4Gk+O+zDVTV2JszdcOt7E5dAy
+MaomETAhRxB7hlIOnEN7WKm+dGNrKRvV0wDU5ReFMRHg31/Lnu8c+5BvGjZX+ky9
+POIhFFYJqwCRlopGSUIxmVj5rSgtzk3iWOQXr+ah1bjEXvlxDOWkHN6YfpV5ThdE
+KdBIPGEVqa63r9n2h+qazKrtiRqJqGnOrHzOECYbRFYhexsNFz7YT02xdfSHn7gM
+IvabDDP/Qp0PjE1jdouiMaFHYnLBbgvlnZW9yuVf/rpXTUq/njxIXMmvmEyyvSDn
+FcFikB8pAoGBAPF77hK4m3/rdGT7X8a/gwvZ2R121aBcdPwEaUhvj/36dx596zvY
+mEOjrWfZhF083/nYWE2kVquj2wjs+otCLfifEEgXcVPTnEOPO9Zg3uNSL0nNQghj
+FuD3iGLTUBCtM66oTe0jLSslHe8gLGEQqyMzHOzYxNqibxcOZIe8Qt0NAoGBAO+U
+I5+XWjWEgDmvyC3TrOSf/KCGjtu0TSv30ipv27bDLMrpvPmD/5lpptTFwcxvVhCs
+2b+chCjlghFSWFbBULBrfci2FtliClOVMYrlNBdUSJhf3aYSG2Doe6Bgt1n2CpNn
+/iu37Y3NfemZBJA7hNl4dYe+f+uzM87cdQ214+jrAoGAXA0XxX8ll2+ToOLJsaNT
+OvNB9h9Uc5qK5X5w+7G7O998BN2PC/MWp8H+2fVqpXgNENpNXttkRm1hk1dych86
+EunfdPuqsX+as44oCyJGFHVBnWpm33eWQw9YqANRI+pCJzP08I5WK3osnPiwshd+
+hR54yjgfYhBFNI7B95PmEQkCgYBzFSz7h1+s34Ycr8SvxsOBWxymG5zaCsUbPsL0
+4aCgLScCHb9J+E86aVbbVFdglYa5Id7DPTL61ixhl7WZjujspeXZGSbmq0Kcnckb
+mDgqkLECiOJW2NHP/j0McAkDLL4tysF8TLDO8gvuvzNC+WQ6drO2ThrypLVZQ+ry
+eBIPmwKBgEZxhqa0gVvHQG/7Od69KWj4eJP28kq13RhKay8JOoN0vPmspXJo1HY3
+CKuHRG+AP579dncdUnOMvfXOtkdM4vk0+hWASBQzM9xzVcztCa+koAugjVaLS9A+
+9uQoqEeVNTckxx0S2bYevRy7hGQmUJTyQm3j1zEUR5jpdbL83Fbq
+-----END RSA PRIVATE KEY-----
--- a/management/server/auth/test_data/sample_key.pub
+++ b/management/server/auth/test_data/sample_key.pub
@@ -0,0 +1,9 @@
+-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4f5wg5l2hKsTeNem/V41
+fGnJm6gOdrj8ym3rFkEU/wT8RDtnSgFEZOQpHEgQ7JL38xUfU0Y3g6aYw9QT0hJ7
+mCpz9Er5qLaMXJwZxzHzAahlfA0icqabvJOMvQtzD6uQv6wPEyZtDTWiQi9AXwBp
+HssPnpYGIn20ZZuNlX2BrClciHhCPUIIZOQn/MmqTD31jSyjoQoV7MhhMTATKJx2
+XrHhR+1DcKJzQBSTAGnpYVaqpsARap+nwRipr3nUTuxyGohBTSmjJ2usSeQXHI3b
+ODIRe1AuTyHceAbewn8b462yEWKARdpd9AjQW5SIVPfdsz5B6GlYQ5LdYKtznTuy
+7wIDAQAB
+-----END PUBLIC KEY-----