The file /Users/misha/Documents/GolandProjects/netbird/netbird/management/server/auth/manager_test.go has been updated successfully.

package auth_test

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/types"
	nbauth "github.com/netbirdio/netbird/shared/auth"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

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 := "REDACTED"
	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 := "REDACTED"
	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 := nbauth.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  *nbauth.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: &nbauth.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: &nbauth.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)
			}
		})
	}

}

package auth_test

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/types"
	nbauth "github.com/netbirdio/netbird/shared/auth"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

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 := "REDACTED"
	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, nil)

	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 := "REDACTED"
	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, nil)

	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 := nbauth.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, nil)

	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  *nbauth.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: &nbauth.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: &nbauth.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)
			}
		})
	}

}

-  manager := auth.NewManager(store, "", "", "", "", []string{}, false)

+  manager := auth.NewManager(store, "", "", "", "", []string{}, false, nil)


Edit

serverauth.NewManager(store, "", "", "", "", []string{}, false)

serverauth.NewManager(store, "", "", "", "", []string{}, false, nil)


The file /Users/misha/Documents/GolandProjects/netbird/netbird/management/server/http/testing/testing_tools/channel/channel.go has been updated successfully.

package channel

import (
	"context"
	"errors"
	"net/http"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"

	"go.opentelemetry.io/otel/metric/noop"

	"github.com/netbirdio/management-integrations/integrations"

	accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain/manager"
	proxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy/manager"
	reverseproxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service/manager"
	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"

	zonesManager "github.com/netbirdio/netbird/management/internals/modules/zones/manager"
	recordsManager "github.com/netbirdio/netbird/management/internals/modules/zones/records/manager"
	"github.com/netbirdio/netbird/management/internals/server/config"

	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
	"github.com/netbirdio/netbird/management/internals/modules/peers"
	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
	"github.com/netbirdio/netbird/management/server/job"

	"github.com/netbirdio/netbird/management/server"
	"github.com/netbirdio/netbird/management/server/account"
	"github.com/netbirdio/netbird/management/server/activity"
	serverauth "github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/geolocation"
	"github.com/netbirdio/netbird/management/server/groups"
	http2 "github.com/netbirdio/netbird/management/server/http"
	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools"
	"github.com/netbirdio/netbird/management/server/networks"
	"github.com/netbirdio/netbird/management/server/networks/resources"
	"github.com/netbirdio/netbird/management/server/networks/routers"
	"github.com/netbirdio/netbird/management/server/permissions"
	"github.com/netbirdio/netbird/management/server/settings"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/telemetry"
	"github.com/netbirdio/netbird/management/server/users"
	"github.com/netbirdio/netbird/shared/auth"
)

func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPeerUpdate *network_map.UpdateMessage, validateUpdate bool) (http.Handler, account.Manager, chan struct{}) {
	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), sqlFile, t.TempDir())
	if err != nil {
		t.Fatalf("Failed to create test store: %v", err)
	}
	t.Cleanup(cleanup)

	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
	if err != nil {
		t.Fatalf("Failed to create metrics: %v", err)
	}

	peersUpdateManager := update_channel.NewPeersUpdateManager(nil)
	updMsg := peersUpdateManager.CreateChannel(context.Background(), testing_tools.TestPeerId)
	done := make(chan struct{})
	if validateUpdate {
		go func() {
			if expectedPeerUpdate != nil {
				peerShouldReceiveUpdate(t, updMsg, expectedPeerUpdate)
			} else {
				peerShouldNotReceiveUpdate(t, updMsg)
			}
			close(done)
		}()
	}

	geoMock := &geolocation.Mock{}
	validatorMock := server.MockIntegratedValidator{}
	proxyController := integrations.NewController(store)
	userManager := users.NewManager(store)
	permissionsManager := permissions.NewManager(store)
	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManager, settings.IdpConfig{})
	peersManager := peers.NewManager(store, permissionsManager)

	jobManager := job.NewJobManager(nil, store, peersManager)

	ctx := context.Background()
	requestBuffer := server.NewAccountRequestBuffer(ctx, store)
	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peersManager), &config.Config{})
	am, err := server.BuildManager(ctx, nil, store, networkMapController, jobManager, nil, "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager, false)
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	accessLogsManager := accesslogsmanager.NewManager(store, permissionsManager, nil)
	proxyTokenStore, err := nbgrpc.NewOneTimeTokenStore(ctx, 5*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create proxy token store: %v", err)
	}
	pkceverifierStore, err := nbgrpc.NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create PKCE verifier store: %v", err)
	}
	noopMeter := noop.NewMeterProvider().Meter("")
	proxyMgr, err := proxymanager.NewManager(store, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy manager: %v", err)
	}
	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr)
	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy controller: %v", err)
	}
	serviceManager := reverseproxymanager.NewManager(store, am, permissionsManager, serviceProxyController, proxyMgr, domainManager)
	proxyServiceServer.SetServiceManager(serviceManager)
	am.SetServiceManager(serviceManager)

	// @note this is required so that PAT's validate from store, but JWT's are mocked
	authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false)
	authManagerMock := &serverauth.MockManager{
		ValidateAndParseTokenFunc:       mockValidateAndParseToken,
		EnsureUserAccessByJWTGroupsFunc: authManager.EnsureUserAccessByJWTGroups,
		MarkPATUsedFunc:                 authManager.MarkPATUsed,
		GetPATInfoFunc:                  authManager.GetPATInfo,
	}

	groupsManager := groups.NewManager(store, permissionsManager, am)
	routersManager := routers.NewManager(store, permissionsManager, am)
	resourcesManager := resources.NewManager(store, permissionsManager, groupsManager, am, serviceManager)
	networksManager := networks.NewManager(store, permissionsManager, resourcesManager, routersManager, am)
	customZonesManager := zonesManager.NewManager(store, am, permissionsManager, "")
	zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)

	apiHandler, err := http2.NewAPIHandler(context.Background(), am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManager, peersManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil)
	if err != nil {
		t.Fatalf("Failed to create API handler: %v", err)
	}

	return apiHandler, am, done
}

func peerShouldNotReceiveUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) {
	t.Helper()
	select {
	case msg := <-updateMessage:
		t.Errorf("Unexpected message received: %+v", msg)
	case <-time.After(500 * time.Millisecond):
		return
	}
}

func peerShouldReceiveUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage, expected *network_map.UpdateMessage) {
	t.Helper()

	select {
	case msg := <-updateMessage:
		if msg == nil {
			t.Errorf("Received nil update message, expected valid message")
		}
		assert.Equal(t, expected, msg)
	case <-time.After(500 * time.Millisecond):
		t.Errorf("Timed out waiting for update message")
	}
}

// PeerShouldReceiveAnyUpdate waits for a peer update message and returns it.
// Fails the test if no update is received within timeout.
func PeerShouldReceiveAnyUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) *network_map.UpdateMessage {
	t.Helper()
	select {
	case msg := <-updateMessage:
		if msg == nil {
			t.Errorf("Received nil update message, expected valid message")
		}
		return msg
	case <-time.After(500 * time.Millisecond):
		t.Errorf("Timed out waiting for update message")
		return nil
	}
}

// PeerShouldNotReceiveAnyUpdate verifies no peer update message is received.
func PeerShouldNotReceiveAnyUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) {
	t.Helper()
	peerShouldNotReceiveUpdate(t, updateMessage)
}

// BuildApiBlackBoxWithDBStateAndPeerChannel creates the API handler and returns
// the peer update channel directly so tests can verify updates inline.
func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile string) (http.Handler, account.Manager, <-chan *network_map.UpdateMessage) {
	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), sqlFile, t.TempDir())
	if err != nil {
		t.Fatalf("Failed to create test store: %v", err)
	}
	t.Cleanup(cleanup)

	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
	if err != nil {
		t.Fatalf("Failed to create metrics: %v", err)
	}

	peersUpdateManager := update_channel.NewPeersUpdateManager(nil)
	updMsg := peersUpdateManager.CreateChannel(context.Background(), testing_tools.TestPeerId)

	geoMock := &geolocation.Mock{}
	validatorMock := server.MockIntegratedValidator{}
	proxyController := integrations.NewController(store)
	userManager := users.NewManager(store)
	permissionsManager := permissions.NewManager(store)
	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManager, settings.IdpConfig{})
	peersManager := peers.NewManager(store, permissionsManager)

	jobManager := job.NewJobManager(nil, store, peersManager)

	ctx := context.Background()
	requestBuffer := server.NewAccountRequestBuffer(ctx, store)
	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peersManager), &config.Config{})
	am, err := server.BuildManager(ctx, nil, store, networkMapController, jobManager, nil, "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager, false)
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	accessLogsManager := accesslogsmanager.NewManager(store, permissionsManager, nil)
	proxyTokenStore, err := nbgrpc.NewOneTimeTokenStore(ctx, 5*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create proxy token store: %v", err)
	}
	pkceverifierStore, err := nbgrpc.NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create PKCE verifier store: %v", err)
	}
	noopMeter := noop.NewMeterProvider().Meter("")
	proxyMgr, err := proxymanager.NewManager(store, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy manager: %v", err)
	}
	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr)
	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy controller: %v", err)
	}
	serviceManager := reverseproxymanager.NewManager(store, am, permissionsManager, serviceProxyController, proxyMgr, domainManager)
	proxyServiceServer.SetServiceManager(serviceManager)
	am.SetServiceManager(serviceManager)

	// @note this is required so that PAT's validate from store, but JWT's are mocked
	authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false)
	authManagerMock := &serverauth.MockManager{
		ValidateAndParseTokenFunc:       mockValidateAndParseToken,
		EnsureUserAccessByJWTGroupsFunc: authManager.EnsureUserAccessByJWTGroups,
		MarkPATUsedFunc:                 authManager.MarkPATUsed,
		GetPATInfoFunc:                  authManager.GetPATInfo,
	}

	groupsManager := groups.NewManager(store, permissionsManager, am)
	routersManager := routers.NewManager(store, permissionsManager, am)
	resourcesManager := resources.NewManager(store, permissionsManager, groupsManager, am, serviceManager)
	networksManager := networks.NewManager(store, permissionsManager, resourcesManager, routersManager, am)
	customZonesManager := zonesManager.NewManager(store, am, permissionsManager, "")
	zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)

	apiHandler, err := http2.NewAPIHandler(context.Background(), am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManager, peersManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil)
	if err != nil {
		t.Fatalf("Failed to create API handler: %v", err)
	}

	return apiHandler, am, updMsg
}

func mockValidateAndParseToken(_ context.Context, token string) (auth.UserAuth, *jwt.Token, error) {
	userAuth := auth.UserAuth{}

	switch token {
	case "testUserId", "testAdminId", "testOwnerId", "testServiceUserId", "testServiceAdminId", "blockedUserId":
		userAuth.UserId = token
		userAuth.AccountId = "testAccountId"
		userAuth.Domain = "test.com"
		userAuth.DomainCategory = "private"
	case "otherUserId":
		userAuth.UserId = "otherUserId"
		userAuth.AccountId = "otherAccountId"
		userAuth.Domain = "other.com"
		userAuth.DomainCategory = "private"
	case "invalidToken":
		return userAuth, nil, errors.New("invalid token")
	}

	jwtToken := jwt.New(jwt.SigningMethodHS256)
	return userAuth, jwtToken, nil
}

package channel

import (
	"context"
	"errors"
	"net/http"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"

	"go.opentelemetry.io/otel/metric/noop"

	"github.com/netbirdio/management-integrations/integrations"

	accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain/manager"
	proxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy/manager"
	reverseproxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service/manager"
	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"

	zonesManager "github.com/netbirdio/netbird/management/internals/modules/zones/manager"
	recordsManager "github.com/netbirdio/netbird/management/internals/modules/zones/records/manager"
	"github.com/netbirdio/netbird/management/internals/server/config"

	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
	"github.com/netbirdio/netbird/management/internals/modules/peers"
	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
	"github.com/netbirdio/netbird/management/server/job"

	"github.com/netbirdio/netbird/management/server"
	"github.com/netbirdio/netbird/management/server/account"
	"github.com/netbirdio/netbird/management/server/activity"
	serverauth "github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/geolocation"
	"github.com/netbirdio/netbird/management/server/groups"
	http2 "github.com/netbirdio/netbird/management/server/http"
	"github.com/netbirdio/netbird/management/server/http/testing/testing_tools"
	"github.com/netbirdio/netbird/management/server/networks"
	"github.com/netbirdio/netbird/management/server/networks/resources"
	"github.com/netbirdio/netbird/management/server/networks/routers"
	"github.com/netbirdio/netbird/management/server/permissions"
	"github.com/netbirdio/netbird/management/server/settings"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/telemetry"
	"github.com/netbirdio/netbird/management/server/users"
	"github.com/netbirdio/netbird/shared/auth"
)

func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPeerUpdate *network_map.UpdateMessage, validateUpdate bool) (http.Handler, account.Manager, chan struct{}) {
	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), sqlFile, t.TempDir())
	if err != nil {
		t.Fatalf("Failed to create test store: %v", err)
	}
	t.Cleanup(cleanup)

	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
	if err != nil {
		t.Fatalf("Failed to create metrics: %v", err)
	}

	peersUpdateManager := update_channel.NewPeersUpdateManager(nil)
	updMsg := peersUpdateManager.CreateChannel(context.Background(), testing_tools.TestPeerId)
	done := make(chan struct{})
	if validateUpdate {
		go func() {
			if expectedPeerUpdate != nil {
				peerShouldReceiveUpdate(t, updMsg, expectedPeerUpdate)
			} else {
				peerShouldNotReceiveUpdate(t, updMsg)
			}
			close(done)
		}()
	}

	geoMock := &geolocation.Mock{}
	validatorMock := server.MockIntegratedValidator{}
	proxyController := integrations.NewController(store)
	userManager := users.NewManager(store)
	permissionsManager := permissions.NewManager(store)
	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManager, settings.IdpConfig{})
	peersManager := peers.NewManager(store, permissionsManager)

	jobManager := job.NewJobManager(nil, store, peersManager)

	ctx := context.Background()
	requestBuffer := server.NewAccountRequestBuffer(ctx, store)
	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peersManager), &config.Config{})
	am, err := server.BuildManager(ctx, nil, store, networkMapController, jobManager, nil, "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager, false)
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	accessLogsManager := accesslogsmanager.NewManager(store, permissionsManager, nil)
	proxyTokenStore, err := nbgrpc.NewOneTimeTokenStore(ctx, 5*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create proxy token store: %v", err)
	}
	pkceverifierStore, err := nbgrpc.NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create PKCE verifier store: %v", err)
	}
	noopMeter := noop.NewMeterProvider().Meter("")
	proxyMgr, err := proxymanager.NewManager(store, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy manager: %v", err)
	}
	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr)
	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy controller: %v", err)
	}
	serviceManager := reverseproxymanager.NewManager(store, am, permissionsManager, serviceProxyController, proxyMgr, domainManager)
	proxyServiceServer.SetServiceManager(serviceManager)
	am.SetServiceManager(serviceManager)

	// @note this is required so that PAT's validate from store, but JWT's are mocked
	authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false, nil)
	authManagerMock := &serverauth.MockManager{
		ValidateAndParseTokenFunc:       mockValidateAndParseToken,
		EnsureUserAccessByJWTGroupsFunc: authManager.EnsureUserAccessByJWTGroups,
		MarkPATUsedFunc:                 authManager.MarkPATUsed,
		GetPATInfoFunc:                  authManager.GetPATInfo,
	}

	groupsManager := groups.NewManager(store, permissionsManager, am)
	routersManager := routers.NewManager(store, permissionsManager, am)
	resourcesManager := resources.NewManager(store, permissionsManager, groupsManager, am, serviceManager)
	networksManager := networks.NewManager(store, permissionsManager, resourcesManager, routersManager, am)
	customZonesManager := zonesManager.NewManager(store, am, permissionsManager, "")
	zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)

	apiHandler, err := http2.NewAPIHandler(context.Background(), am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManager, peersManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil)
	if err != nil {
		t.Fatalf("Failed to create API handler: %v", err)
	}

	return apiHandler, am, done
}

func peerShouldNotReceiveUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) {
	t.Helper()
	select {
	case msg := <-updateMessage:
		t.Errorf("Unexpected message received: %+v", msg)
	case <-time.After(500 * time.Millisecond):
		return
	}
}

func peerShouldReceiveUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage, expected *network_map.UpdateMessage) {
	t.Helper()

	select {
	case msg := <-updateMessage:
		if msg == nil {
			t.Errorf("Received nil update message, expected valid message")
		}
		assert.Equal(t, expected, msg)
	case <-time.After(500 * time.Millisecond):
		t.Errorf("Timed out waiting for update message")
	}
}

// PeerShouldReceiveAnyUpdate waits for a peer update message and returns it.
// Fails the test if no update is received within timeout.
func PeerShouldReceiveAnyUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) *network_map.UpdateMessage {
	t.Helper()
	select {
	case msg := <-updateMessage:
		if msg == nil {
			t.Errorf("Received nil update message, expected valid message")
		}
		return msg
	case <-time.After(500 * time.Millisecond):
		t.Errorf("Timed out waiting for update message")
		return nil
	}
}

// PeerShouldNotReceiveAnyUpdate verifies no peer update message is received.
func PeerShouldNotReceiveAnyUpdate(t testing_tools.TB, updateMessage <-chan *network_map.UpdateMessage) {
	t.Helper()
	peerShouldNotReceiveUpdate(t, updateMessage)
}

// BuildApiBlackBoxWithDBStateAndPeerChannel creates the API handler and returns
// the peer update channel directly so tests can verify updates inline.
func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile string) (http.Handler, account.Manager, <-chan *network_map.UpdateMessage) {
	store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), sqlFile, t.TempDir())
	if err != nil {
		t.Fatalf("Failed to create test store: %v", err)
	}
	t.Cleanup(cleanup)

	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
	if err != nil {
		t.Fatalf("Failed to create metrics: %v", err)
	}

	peersUpdateManager := update_channel.NewPeersUpdateManager(nil)
	updMsg := peersUpdateManager.CreateChannel(context.Background(), testing_tools.TestPeerId)

	geoMock := &geolocation.Mock{}
	validatorMock := server.MockIntegratedValidator{}
	proxyController := integrations.NewController(store)
	userManager := users.NewManager(store)
	permissionsManager := permissions.NewManager(store)
	settingsManager := settings.NewManager(store, userManager, integrations.NewManager(&activity.InMemoryEventStore{}), permissionsManager, settings.IdpConfig{})
	peersManager := peers.NewManager(store, permissionsManager)

	jobManager := job.NewJobManager(nil, store, peersManager)

	ctx := context.Background()
	requestBuffer := server.NewAccountRequestBuffer(ctx, store)
	networkMapController := controller.NewController(ctx, store, metrics, peersUpdateManager, requestBuffer, server.MockIntegratedValidator{}, settingsManager, "", port_forwarding.NewControllerMock(), ephemeral_manager.NewEphemeralManager(store, peersManager), &config.Config{})
	am, err := server.BuildManager(ctx, nil, store, networkMapController, jobManager, nil, "", &activity.InMemoryEventStore{}, geoMock, false, validatorMock, metrics, proxyController, settingsManager, permissionsManager, false)
	if err != nil {
		t.Fatalf("Failed to create manager: %v", err)
	}

	accessLogsManager := accesslogsmanager.NewManager(store, permissionsManager, nil)
	proxyTokenStore, err := nbgrpc.NewOneTimeTokenStore(ctx, 5*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create proxy token store: %v", err)
	}
	pkceverifierStore, err := nbgrpc.NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100)
	if err != nil {
		t.Fatalf("Failed to create PKCE verifier store: %v", err)
	}
	noopMeter := noop.NewMeterProvider().Meter("")
	proxyMgr, err := proxymanager.NewManager(store, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy manager: %v", err)
	}
	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr)
	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
	if err != nil {
		t.Fatalf("Failed to create proxy controller: %v", err)
	}
	serviceManager := reverseproxymanager.NewManager(store, am, permissionsManager, serviceProxyController, proxyMgr, domainManager)
	proxyServiceServer.SetServiceManager(serviceManager)
	am.SetServiceManager(serviceManager)

	// @note this is required so that PAT's validate from store, but JWT's are mocked
	authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false, nil)
	authManagerMock := &serverauth.MockManager{
		ValidateAndParseTokenFunc:       mockValidateAndParseToken,
		EnsureUserAccessByJWTGroupsFunc: authManager.EnsureUserAccessByJWTGroups,
		MarkPATUsedFunc:                 authManager.MarkPATUsed,
		GetPATInfoFunc:                  authManager.GetPATInfo,
	}

	groupsManager := groups.NewManager(store, permissionsManager, am)
	routersManager := routers.NewManager(store, permissionsManager, am)
	resourcesManager := resources.NewManager(store, permissionsManager, groupsManager, am, serviceManager)
	networksManager := networks.NewManager(store, permissionsManager, resourcesManager, routersManager, am)
	customZonesManager := zonesManager.NewManager(store, am, permissionsManager, "")
	zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)

	apiHandler, err := http2.NewAPIHandler(context.Background(), am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, validatorMock, proxyController, permissionsManager, peersManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil)
	if err != nil {
		t.Fatalf("Failed to create API handler: %v", err)
	}

	return apiHandler, am, updMsg
}

func mockValidateAndParseToken(_ context.Context, token string) (auth.UserAuth, *jwt.Token, error) {
	userAuth := auth.UserAuth{}

	switch token {
	case "testUserId", "testAdminId", "testOwnerId", "testServiceUserId", "testServiceAdminId", "blockedUserId":
		userAuth.UserId = token
		userAuth.AccountId = "testAccountId"
		userAuth.Domain = "test.com"
		userAuth.DomainCategory = "private"
	case "otherUserId":
		userAuth.UserId = "otherUserId"
		userAuth.AccountId = "otherAccountId"
		userAuth.Domain = "other.com"
		userAuth.DomainCategory = "private"
	case "invalidToken":
		return userAuth, nil, errors.New("invalid token")
	}

	jwtToken := jwt.New(jwt.SigningMethodHS256)
	return userAuth, jwtToken, nil
}

-  authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false)

+  authManager := serverauth.NewManager(store, "", "", "", "", []string{}, false, nil)


Now fix the last caller in manager_test.go with the server.URL argument.


225		keyId := "test-key"
226
227		// note, we can use a nil store because ValidateAndParseToken does not use it in it's flow
228		manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false)
229
230		customClaim := func(name string) string {
231			return fmt.Sprintf("%s/%s", audience, name)
232		}
233
234		lastLogin := time.Date(2025, 2, 12, 14, 25, 26, 0, time.UTC) //"2025-02-12T14:25:26.186Z"

	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"


manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false)

manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false, nil)


package auth_test

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/types"
	nbauth "github.com/netbirdio/netbird/shared/auth"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

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 := "REDACTED"
	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, nil)

	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 := "REDACTED"
	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, nil)

	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 := nbauth.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, nil)

	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  *nbauth.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: &nbauth.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: &nbauth.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)
			}
		})
	}

}

package auth_test

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/types"
	nbauth "github.com/netbirdio/netbird/shared/auth"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)

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 := "REDACTED"
	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, nil)

	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 := "REDACTED"
	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, nil)

	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 := nbauth.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, nil)

	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, nil)

	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  *nbauth.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: &nbauth.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: &nbauth.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)
			}
		})
	}

}

-  manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false)

+  manager := auth.NewManager(nil, issuer, audience, server.URL, userIdClaim, []string{audience}, false, nil)


Now check if there are any references to `expiresInTime` in tests or elsewhere.


management/server/idp/zitadel_test.go:211:			creds.jwtToken.expiresInTime = testCase.inputTime
management/server/idp/zitadel_test.go:274:			creds.jwtToken.expiresInTime = testCase.inputExpireToken
management/server/idp/auth0_test.go:236:			creds.jwtToken.expiresInTime = testCase.inputTime
management/server/idp/auth0_test.go:299:			creds.jwtToken.expiresInTime = testCase.inputExpireToken
management/server/idp/azure_test.go:40:			creds.jwtToken.expiresInTime = testCase.inputTime
management/server/idp/azure_test.go:103:			creds.jwtToken.expiresInTime = testCase.inputExpireToken
management/server/idp/auth0.go:183:	return !c.jwtToken.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(c.jwtToken.expiresInTime)
management/server/idp/auth0.go:245:	jwtToken.expiresInTime = time.Unix(IssuedAt.Exp, 0)
management/server/idp/keycloak_test.go:233:			creds.jwtToken.expiresInTime = testCase.inputTime
management/server/idp/keycloak_test.go:296:			creds.jwtToken.expiresInTime = testCase.inputExpireToken
management/server/idp/zitadel.go:203:	return !zc.jwtToken.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(zc.jwtToken.expiresInTime)
management/server/idp/zitadel.go:268:	jwtToken.expiresInTime = time.Unix(IssuedAt.Exp, 0)
management/server/idp/zitadel.go:282:	tok.expiresInTime = time.Now().Add(time.Duration(tok.ExpiresIn) * time.Second)
management/server/idp/authentik.go:114:	return !ac.jwtToken.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(ac.jwtToken.expiresInTime)
management/server/idp/authentik.go:179:	jwtToken.expiresInTime = time.Unix(IssuedAt.Exp, 0)
management/server/idp/azure.go:111:	return !ac.jwtToken.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(ac.jwtToken.expiresInTime)
management/server/idp/azure.go:183:	jwtToken.expiresInTime = time.Unix(IssuedAt.Exp, 0)
management/server/idp/keycloak.go:110:	return !kc.jwtToken.expiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(kc.jwtToken.expiresInTime)
management/server/idp/keycloak.go:173:	jwtToken.expiresInTime = time.Unix(IssuedAt.Exp, 0)
management/server/idp/authentik_test.go:243:			creds.jwtToken.expiresInTime = testCase.inputTime
management/server/idp/authentik_test.go:306:			creds.jwtToken.expiresInTime = testCase.inputExpireToken
management/server/idp/idp.go:108:	expiresInTime time.Time


These are different `expiresInTime` fields in IDP-specific JWT token structs (not the `Jwks` struct), so they're unaffected by my change.

Now let me also check if the `context` import is needed in validator.go (for `NewValidatorWithKeyFetcher`).


1	package jwt
2
3	import (
4		"context"
5		"crypto/ecdsa"
6		"crypto/elliptic"
7		"crypto/rsa"
8		"encoding/base64"
9		"encoding/json"
10		"errors"
11		"fmt"
12		"math/big"
13		"net/http"
14		"net/url"
15		"strconv"
16		"strings"
17		"sync"
18		"time"
19
20		"github.com/golang-jwt/jwt/v5"
21
22		log "github.com/sirupsen/logrus"
23	)
24
25	// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation

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/v5"

	log "github.com/sirupsen/logrus"
)

// Jwks is a collection of JSONWebKey obtained from Config.HttpServerConfig.AuthKeysLocation


Good, `context` is already imported. Now let me also remove the unused `context` import in controllers.go if needed and check for unused imports.


1	package server
2
3	import (
4		"context"
5
6		log "github.com/sirupsen/logrus"
7
8		"github.com/netbirdio/management-integrations/integrations"
9		"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
10		proxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy/manager"
11
12		"github.com/netbirdio/netbird/management/internals/controllers/network_map"
13		nmapcontroller "github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
14		"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
15		"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
16		"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
17		"github.com/netbirdio/netbird/management/internals/shared/grpc"
18		"github.com/netbirdio/netbird/management/server"
19		"github.com/netbirdio/netbird/management/server/auth"
20		"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
21		"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
22		"github.com/netbirdio/netbird/management/server/job"
23		nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
24	)
25

package server

import (
	"context"

	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/management-integrations/integrations"
	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
	proxymanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy/manager"

	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
	nmapcontroller "github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral"
	"github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
	"github.com/netbirdio/netbird/management/internals/shared/grpc"
	"github.com/netbirdio/netbird/management/server"
	"github.com/netbirdio/netbird/management/server/auth"
	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
	"github.com/netbirdio/netbird/management/server/job"
	nbjwt "github.com/netbirdio/netbird/shared/auth/jwt"
)


Let me try to build and see if there are compilation errors.


Bash

Build jwt package


(Bash completed with no output)


Build dex package


go build ./management/server/auth/... 2>&1

Build auth package


go build ./management/server/idp/... 2>&1

Build idp package


go build ./management/internals/server/... 2>&1

Build internals server package


Build full management package


Everything compiles. Let me run the relevant tests.


go test ./shared/auth/jwt/... -v -count=1 -timeout 60s 2>&1 | tail -30

Run jwt tests


go test ./management/server/auth/... -v -count=1 -timeout 60s 2>&1 | tail -30

Run auth manager tests


    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/extracts_Dex_encoded_user_ID (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/handles_missing_email_claim (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/handles_missing_name_claim (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/handles_missing_both_email_and_name (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/extracts_preferred_username (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/fails_when_user_ID_claim_is_empty (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/uses_custom_user_ID_claim (0.00s)
    --- PASS: TestClaimsExtractor_ToUserAuth_ExtractsEmailAndName/extracts_account_ID_with_audience_prefix (0.00s)
=== RUN   TestClaimsExtractor_ToUserAuth_PreferredUsername
--- PASS: TestClaimsExtractor_ToUserAuth_PreferredUsername (0.00s)
=== RUN   TestClaimsExtractor_ToUserAuth_LastLogin
--- PASS: TestClaimsExtractor_ToUserAuth_LastLogin (0.00s)
=== RUN   TestClaimsExtractor_ToUserAuth_Invited
--- PASS: TestClaimsExtractor_ToUserAuth_Invited (0.00s)
=== RUN   TestClaimsExtractor_ToGroups
=== RUN   TestClaimsExtractor_ToGroups/extracts_groups_from_claim
=== RUN   TestClaimsExtractor_ToGroups/returns_empty_slice_when_claim_missing
=== RUN   TestClaimsExtractor_ToGroups/handles_custom_claim_name
=== RUN   TestClaimsExtractor_ToGroups/filters_non-string_values
--- PASS: TestClaimsExtractor_ToGroups (0.00s)
    --- PASS: TestClaimsExtractor_ToGroups/extracts_groups_from_claim (0.00s)
    --- PASS: TestClaimsExtractor_ToGroups/returns_empty_slice_when_claim_missing (0.00s)
    --- PASS: TestClaimsExtractor_ToGroups/handles_custom_claim_name (0.00s)
    --- PASS: TestClaimsExtractor_ToGroups/filters_non-string_values (0.00s)
=== RUN   TestClaimsExtractor_DefaultUserIDClaim
--- PASS: TestClaimsExtractor_DefaultUserIDClaim (0.00s)
=== RUN   TestClaimsExtractor_DexUserIDFormat
--- PASS: TestClaimsExtractor_DexUserIDFormat (0.00s)
PASS
ok  	github.com/netbirdio/netbird/shared/auth/jwt	0.655s


    --- PASS: TestAuthManager_EnsureUserAccessByJWTGroups/User_PAT (0.00s)
    --- PASS: REDACTED (0.00s)
    --- PASS: REDACTED (0.00s)
    --- PASS: TestAuthManager_EnsureUserAccessByJWTGroups/User_in_allowed_JWT_groups (0.00s)
    --- PASS: TestAuthManager_EnsureUserAccessByJWTGroups/User_not_in_allowed_JWT_groups (0.00s)
=== RUN   TestAuthManager_ValidateAndParseToken
=== RUN   TestAuthManager_ValidateAndParseToken/Valid_with_custom_claims
=== RUN   TestAuthManager_ValidateAndParseToken/Valid_without_custom_claims
=== RUN   TestAuthManager_ValidateAndParseToken/Expired_token
time="2026-03-30T15:25:33+02:00" level=error msg="token could not be parsed: token has invalid claims: token is expired"
=== RUN   TestAuthManager_ValidateAndParseToken/Not_yet_valid
time="2026-03-30T15:25:33+02:00" level=error msg="token could not be parsed: token has invalid claims: token used before issued"
=== RUN   TestAuthManager_ValidateAndParseToken/Invalid_signature
time="2026-03-30T15:25:33+02:00" level=error msg="token could not be parsed: token is malformed: could not base64 decode signature: illegal base64 data at input byte 16"
=== RUN   TestAuthManager_ValidateAndParseToken/Invalid_issuer
time="2026-03-30T15:25:33+02:00" level=error msg="token could not be parsed: token has invalid claims: token has invalid issuer"
=== RUN   TestAuthManager_ValidateAndParseToken/Invalid_audience
time="2026-03-30T15:25:33+02:00" level=error msg="token could not be parsed: token has invalid claims: token has invalid audience"
=== RUN   TestAuthManager_ValidateAndParseToken/Invalid_user_claim
--- PASS: TestAuthManager_ValidateAndParseToken (0.01s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Valid_with_custom_claims (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Valid_without_custom_claims (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Expired_token (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Not_yet_valid (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Invalid_signature (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Invalid_issuer (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Invalid_audience (0.00s)
    --- PASS: TestAuthManager_ValidateAndParseToken/Invalid_user_claim (0.00s)
PASS
ok  	github.com/netbirdio/netbird/management/server/auth	0.746s


All tests pass. Here's a summary of the changes:

### What was done

**1. `shared/auth/jwt/validator.go`**
- Exported `Jwks.ExpiresInTime` (was `expiresInTime`) so external packages can set it
- Added `KeyFetcher` type — `func(ctx context.Context) (*Jwks, error)` — for direct key fetching
- Added `keyFetcher` field to `Validator` struct
- Added `NewValidatorWithKeyFetcher` constructor that uses a `KeyFetcher` instead of HTTP
- Added `fetchKeys` helper that delegates to `keyFetcher` if set, otherwise falls back to `getPemKeys` (HTTP)
- `refreshKeys` and `forceRefreshKeys` now use `fetchKeys` instead of `getPemKeys` directly

**2. `idp/dex/provider.go`**
- Added `GetJWKS(ctx) (*Jwks, error)` method on `Provider` — reads signing keys directly from Dex storage via `storage.GetKeys()`, converts `jose.JSONWebKey` (RSA/ECDSA) to the `nbjwt.JSONWebKey` format
- Added `joseKeyToJSONWebKey` helper for the conversion

**3. `management/server/idp/embedded.go`**
- Added `GetKeyFetcher() nbjwt.KeyFetcher` to `OAuthConfigProvider` interface
- Implemented it on `EmbeddedIdPManager` — returns `provider.GetJWKS`

**4. `management/server/auth/manager.go`**
- `NewManager` now accepts an optional `KeyFetcher` parameter; when non-nil, creates the validator with `NewValidatorWithKeyFetcher`

**5. `management/internals/server/controllers.go`**
- Wired the key fetcher from `OAuthConfigProvider` into `auth.NewManager`
- Falls back to `GetLocalKeysLocation()` (HTTP) if `GetKeyFetcher()` returns nil

**6. Test files** — updated `NewManager` call sites to pass `nil` for the new parameter.


make sure the verification of the keys from db is done according to the original function. e.g., expiration checks etc


Let me review how the original `stillValid()` and `getPemKeys` handle expiration, and how the new `GetJWKS` compares.


108		}
109	}
110
111	// forcedRefreshCooldown is the minimum time between forced key refreshes
112	// to prevent abuse from invalid tokens with fake kid values
113	const forcedRefreshCooldown = 30 * time.Second
114
115	// fetchKeys retrieves keys using the keyFetcher if available, otherwise falls back to HTTP.
116	func (v *Validator) fetchKeys(ctx context.Context) (*Jwks, error) {
117		if v.keyFetcher != nil {
118			return v.keyFetcher(ctx)
119		}
120		return getPemKeys(v.keysLocation)
121	}
122
123	func (v *Validator) getKeyFunc(ctx context.Context) jwt.Keyfunc {
124		return func(token *jwt.Token) (interface{}, error) {
125			// If keys are rotated, verify the keys prior to token validation
126			if v.idpSignkeyRefreshEnabled {
127				// If the keys are invalid, retrieve new ones
128				if !v.keys.stillValid() {
129					v.refreshKeys(ctx)
130				}
131			}
132
133			publicKey, err := getPublicKey(token, v.keys)
134			if err == nil {
135				return publicKey, nil
136			}
137

	}
}

// forcedRefreshCooldown is the minimum time between forced key refreshes
// to prevent abuse from invalid tokens with fake kid values
const forcedRefreshCooldown = 30 * time.Second

// fetchKeys retrieves keys using the keyFetcher if available, otherwise falls back to HTTP.
func (v *Validator) fetchKeys(ctx context.Context) (*Jwks, error) {
	if v.keyFetcher != nil {
		return v.keyFetcher(ctx)
	}
	return getPemKeys(v.keysLocation)
}

func (v *Validator) getKeyFunc(ctx context.Context) jwt.Keyfunc {
	return func(token *jwt.Token) (interface{}, error) {
		// If keys are rotated, verify the keys prior to token validation
		if v.idpSignkeyRefreshEnabled {
			// If the keys are invalid, retrieve new ones
			if !v.keys.stillValid() {
				v.refreshKeys(ctx)
			}
		}

		publicKey, err := getPublicKey(token, v.keys)
		if err == nil {
			return publicKey, nil
		}


220		if err != nil {
221			err = fmt.Errorf("%w: %s", errTokenParsing, err)
222			log.WithContext(ctx).Error(err.Error())
223			return nil, err
224		}
225
226		// Check if the parsed token is valid...
227		if !parsedToken.Valid {
228			log.WithContext(ctx).Debug(errTokenInvalid.Error())
229			return nil, errTokenInvalid
230		}
231
232		return parsedToken, nil
233	}
234
235	// stillValid returns true if the JSONWebKey still valid and have enough time to be used
236	func (jwks *Jwks) stillValid() bool {
237		return !jwks.ExpiresInTime.IsZero() && time.Now().Add(5*time.Second).Before(jwks.ExpiresInTime)
238	}
239
240	func getPemKeys(keysLocation string) (*Jwks, error) {
241		jwks := &Jwks{}
242
243		requestURI, err := url.ParseRequestURI(keysLocation)
244		if err != nil {
245			return jwks, err
246		}
247
248		resp, err := http.Get(requestURI.String())
249		if err != nil {

	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{}

	requestURI, err := url.ParseRequestURI(keysLocation)
	if err != nil {
		return jwks, err
	}

	resp, err := http.Get(requestURI.String())
	if err != nil {


Now let me look at the current `GetJWKS` implementation.