feat(authn): clean the rotate handling

ee/query-service/app/server.go

@@ -216,8 +216,7 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web web.Web) (*h
 		}),
 		otelmux.WithPublicEndpoint(),
 	))
-	r.Use(middleware.NewAuthN([]string{"Authorization", "Sec-WebSocket-Protocol"}, s.signoz.Sharder, s.signoz.Tokenizer, s.signoz.Instrumentation.Logger()).Wrap)
-	r.Use(middleware.NewAPIKey(s.signoz.SQLStore, []string{"SIGNOZ-API-KEY"}, s.signoz.Instrumentation.Logger(), s.signoz.Sharder).Wrap)
+	r.Use(middleware.NewIdentN(s.signoz.IdentNResolver, s.signoz.Sharder, s.signoz.Instrumentation.Logger()).Wrap)
 	r.Use(middleware.NewTimeout(s.signoz.Instrumentation.Logger(),
 		s.config.APIServer.Timeout.ExcludedRoutes,
 		s.config.APIServer.Timeout.Default,

pkg/apiserver/signozapiserver/provider.go

@@ -23,7 +23,7 @@ import (
 	"github.com/SigNoz/signoz/pkg/modules/user"
 	"github.com/SigNoz/signoz/pkg/querier"
 	"github.com/SigNoz/signoz/pkg/types"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
 	"github.com/SigNoz/signoz/pkg/zeus"
 	"github.com/gorilla/mux"
 )
@@ -238,13 +238,13 @@ func (provider *provider) AddToRouter(router *mux.Router) error {
 
 func newSecuritySchemes(role types.Role) []handler.OpenAPISecurityScheme {
 	return []handler.OpenAPISecurityScheme{
-		{Name: ctxtypes.AuthTypeAPIKey.StringValue(), Scopes: []string{role.String()}},
-		{Name: ctxtypes.AuthTypeTokenizer.StringValue(), Scopes: []string{role.String()}},
+		{Name: authtypes.IdentNProviderAPIkey.StringValue(), Scopes: []string{role.String()}},
+		{Name: authtypes.IdentNProviderTokenizer.StringValue(), Scopes: []string{role.String()}},
 	}
 }
 
 func newAnonymousSecuritySchemes(scopes []string) []handler.OpenAPISecurityScheme {
 	return []handler.OpenAPISecurityScheme{
-		{Name: ctxtypes.AuthTypeAnonymous.StringValue(), Scopes: scopes},
+		{Name: authtypes.IdentNProviderAnonymous.StringValue(), Scopes: scopes},
 	}
 }

pkg/apiserver/signozapiserver/session.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
 	"github.com/gorilla/mux"
 )
 
@@ -73,7 +72,7 @@ func (provider *provider) addSessionRoutes(router *mux.Router) error {
 		SuccessStatusCode:   http.StatusNoContent,
 		ErrorStatusCodes:    []int{http.StatusBadRequest},
 		Deprecated:          false,
-		SecuritySchemes:     []handler.OpenAPISecurityScheme{{Name: ctxtypes.AuthTypeTokenizer.StringValue()}},
+		SecuritySchemes:     []handler.OpenAPISecurityScheme{{Name: authtypes.IdentNProviderTokenizer.StringValue()}},
 	})).Methods(http.MethodDelete).GetError(); err != nil {
 		return err
 	}

pkg/apiserver/signozapiserver/user.go

@@ -5,7 +5,7 @@ import (
 
 	"github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/types"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
 	"github.com/gorilla/mux"
 )
 
@@ -208,7 +208,7 @@ func (provider *provider) addUserRoutes(router *mux.Router) error {
 		SuccessStatusCode:   http.StatusOK,
 		ErrorStatusCodes:    []int{},
 		Deprecated:          false,
-		SecuritySchemes:     []handler.OpenAPISecurityScheme{{Name: ctxtypes.AuthTypeTokenizer.StringValue()}},
+		SecuritySchemes:     []handler.OpenAPISecurityScheme{{Name: authtypes.IdentNProviderTokenizer.StringValue()}},
 	})).Methods(http.MethodGet).GetError(); err != nil {
 		return err
 	}

pkg/authn/passwordauthn/emailpasswordauthn/authn.go

@@ -30,5 +30,5 @@ func (a *AuthN) Authenticate(ctx context.Context, email string, password string,
 		return nil, errors.New(errors.TypeUnauthenticated, types.ErrCodeIncorrectPassword, "invalid email or password")
 	}
 
-	return authtypes.NewIdentity(user.ID, orgID, user.Email, user.Role), nil
+	return authtypes.NewIdentity(user.ID, orgID, user.Email, user.Role, authtypes.IdentNProviderTokenizer), nil
 }

pkg/http/middleware/api_key.go

@@ -1,143 +0,0 @@
-package middleware
-
-import (
-	"context"
-	"log/slog"
-	"net/http"
-	"time"
-
-	"github.com/SigNoz/signoz/pkg/sharder"
-	"github.com/SigNoz/signoz/pkg/sqlstore"
-	"github.com/SigNoz/signoz/pkg/types"
-	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
-	"github.com/SigNoz/signoz/pkg/valuer"
-	"golang.org/x/sync/singleflight"
-)
-
-const (
-	apiKeyCrossOrgMessage string = "::API-KEY-CROSS-ORG::"
-)
-
-type APIKey struct {
-	store   sqlstore.SQLStore
-	uuid    *authtypes.UUID
-	headers []string
-	logger  *slog.Logger
-	sharder sharder.Sharder
-	sfGroup *singleflight.Group
-}
-
-func NewAPIKey(store sqlstore.SQLStore, headers []string, logger *slog.Logger, sharder sharder.Sharder) *APIKey {
-	return &APIKey{
-		store:   store,
-		uuid:    authtypes.NewUUID(),
-		headers: headers,
-		logger:  logger,
-		sharder: sharder,
-		sfGroup: &singleflight.Group{},
-	}
-}
-
-func (a *APIKey) Wrap(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		var values []string
-		var apiKeyToken string
-		var apiKey types.StorableAPIKey
-
-		for _, header := range a.headers {
-			values = append(values, r.Header.Get(header))
-		}
-
-		ctx, err := a.uuid.ContextFromRequest(r.Context(), values...)
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		apiKeyToken, ok := authtypes.UUIDFromContext(ctx)
-		if !ok {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		err = a.
-			store.
-			BunDB().
-			NewSelect().
-			Model(&apiKey).
-			Where("token = ?", apiKeyToken).
-			Scan(r.Context())
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		// allow the APIKey if expires_at is not set
-		if apiKey.ExpiresAt.Before(time.Now()) && !apiKey.ExpiresAt.Equal(types.NEVER_EXPIRES) {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		// get user from db
-		user := types.User{}
-		err = a.store.BunDB().NewSelect().Model(&user).Where("id = ?", apiKey.UserID).Scan(r.Context())
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		jwt := authtypes.Claims{
-			UserID: user.ID.String(),
-			Role:   apiKey.Role,
-			Email:  user.Email.String(),
-			OrgID:  user.OrgID.String(),
-		}
-
-		ctx = authtypes.NewContextWithClaims(ctx, jwt)
-
-		claims, err := authtypes.ClaimsFromContext(ctx)
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		if err := a.sharder.IsMyOwnedKey(r.Context(), types.NewOrganizationKey(valuer.MustNewUUID(claims.OrgID))); err != nil {
-			a.logger.ErrorContext(r.Context(), apiKeyCrossOrgMessage, "claims", claims, "error", err)
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		ctx = ctxtypes.SetAuthType(ctx, ctxtypes.AuthTypeAPIKey)
-
-		comment := ctxtypes.CommentFromContext(ctx)
-		comment.Set("auth_type", ctxtypes.AuthTypeAPIKey.StringValue())
-		comment.Set("user_id", claims.UserID)
-		comment.Set("org_id", claims.OrgID)
-
-		r = r.WithContext(ctxtypes.NewContextWithComment(ctx, comment))
-
-		next.ServeHTTP(w, r)
-
-		lastUsedCtx := context.WithoutCancel(r.Context())
-		_, _, _ = a.sfGroup.Do(apiKey.ID.StringValue(), func() (any, error) {
-			apiKey.LastUsed = time.Now()
-			_, err = a.
-				store.
-				BunDB().
-				NewUpdate().
-				Model(&apiKey).
-				Column("last_used").
-				Where("token = ?", apiKeyToken).
-				Where("revoked = false").
-				Exec(lastUsedCtx)
-			if err != nil {
-				a.logger.ErrorContext(lastUsedCtx, "failed to update last used of api key", "error", err)
-			}
-
-			return true, nil
-		})
-
-	})
-
-}

pkg/http/middleware/authn.go

@@ -1,150 +0,0 @@
-package middleware
-
-import (
-	"context"
-	"log/slog"
-	"net/http"
-	"strings"
-	"time"
-
-	"github.com/SigNoz/signoz/pkg/errors"
-	"github.com/SigNoz/signoz/pkg/sharder"
-	"github.com/SigNoz/signoz/pkg/tokenizer"
-	"github.com/SigNoz/signoz/pkg/types"
-	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
-	"github.com/SigNoz/signoz/pkg/valuer"
-	"golang.org/x/sync/singleflight"
-)
-
-const (
-	authCrossOrgMessage string = "::AUTH-CROSS-ORG::"
-)
-
-type AuthN struct {
-	tokenizer tokenizer.Tokenizer
-	headers   []string
-	sharder   sharder.Sharder
-	logger    *slog.Logger
-	sfGroup   *singleflight.Group
-}
-
-func NewAuthN(headers []string, sharder sharder.Sharder, tokenizer tokenizer.Tokenizer, logger *slog.Logger) *AuthN {
-	return &AuthN{
-		headers:   headers,
-		sharder:   sharder,
-		tokenizer: tokenizer,
-		logger:    logger,
-		sfGroup:   &singleflight.Group{},
-	}
-}
-
-func (a *AuthN) Wrap(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		var values []string
-		for _, header := range a.headers {
-			values = append(values, r.Header.Get(header))
-		}
-
-		ctx, err := a.contextFromRequest(r.Context(), values...)
-		if err != nil {
-			r = r.WithContext(ctx)
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		r = r.WithContext(ctx)
-
-		claims, err := authtypes.ClaimsFromContext(ctx)
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		if err := a.sharder.IsMyOwnedKey(r.Context(), types.NewOrganizationKey(valuer.MustNewUUID(claims.OrgID))); err != nil {
-			a.logger.ErrorContext(r.Context(), authCrossOrgMessage, "claims", claims, "error", err)
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		ctx = ctxtypes.SetAuthType(ctx, ctxtypes.AuthTypeTokenizer)
-
-		comment := ctxtypes.CommentFromContext(ctx)
-		comment.Set("auth_type", ctxtypes.AuthTypeTokenizer.StringValue())
-		comment.Set("tokenizer_provider", a.tokenizer.Config().Provider)
-		comment.Set("user_id", claims.UserID)
-		comment.Set("org_id", claims.OrgID)
-
-		r = r.WithContext(ctxtypes.NewContextWithComment(ctx, comment))
-
-		next.ServeHTTP(w, r)
-
-		accessToken, err := authtypes.AccessTokenFromContext(r.Context())
-		if err != nil {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		lastObservedAtCtx := context.WithoutCancel(r.Context())
-		_, _, _ = a.sfGroup.Do(accessToken, func() (any, error) {
-			if err := a.tokenizer.SetLastObservedAt(lastObservedAtCtx, accessToken, time.Now()); err != nil {
-				a.logger.ErrorContext(lastObservedAtCtx, "failed to set last observed at", "error", err)
-				return false, err
-			}
-
-			return true, nil
-		})
-	})
-}
-
-func (a *AuthN) contextFromRequest(ctx context.Context, values ...string) (context.Context, error) {
-	ctx, err := a.contextFromAccessToken(ctx, values...)
-	if err != nil {
-		return ctx, err
-	}
-
-	accessToken, err := authtypes.AccessTokenFromContext(ctx)
-	if err != nil {
-		return ctx, err
-	}
-
-	authenticatedUser, err := a.tokenizer.GetIdentity(ctx, accessToken)
-	if err != nil {
-		return ctx, err
-	}
-
-	return authtypes.NewContextWithClaims(ctx, authenticatedUser.ToClaims()), nil
-}
-
-func (a *AuthN) contextFromAccessToken(ctx context.Context, values ...string) (context.Context, error) {
-	var value string
-	for _, v := range values {
-		if v != "" {
-			value = v
-			break
-		}
-	}
-
-	if value == "" {
-		return ctx, errors.New(errors.TypeUnauthenticated, errors.CodeUnauthenticated, "missing authorization header")
-	}
-
-	// parse from
-	bearerToken, ok := parseBearerAuth(value)
-	if !ok {
-		// this will take care that if the value is not of type bearer token, directly use it
-		bearerToken = value
-	}
-
-	return authtypes.NewContextWithAccessToken(ctx, bearerToken), nil
-}
-
-func parseBearerAuth(auth string) (string, bool) {
-	const prefix = "Bearer "
-	// Case insensitive prefix match
-	if len(auth) < len(prefix) || !strings.EqualFold(auth[:len(prefix)], prefix) {
-		return "", false
-	}
-
-	return auth[len(prefix):], true
-}

pkg/http/middleware/authz.go

@@ -44,7 +44,7 @@ func (middleware *AuthZ) ViewAccess(next http.HandlerFunc) http.HandlerFunc {
 
 		commentCtx := ctxtypes.CommentFromContext(ctx)
 		authtype, ok := commentCtx.Map()["auth_type"]
-		if ok && authtype == ctxtypes.AuthTypeAPIKey.StringValue() {
+		if ok && (authtype == authtypes.IdentNProviderAPIkey.StringValue()) {
 			if err := claims.IsViewer(); err != nil {
 				middleware.logger.WarnContext(ctx, authzDeniedMessage, "claims", claims)
 				render.Error(rw, err)
@@ -96,7 +96,7 @@ func (middleware *AuthZ) EditAccess(next http.HandlerFunc) http.HandlerFunc {
 
 		commentCtx := ctxtypes.CommentFromContext(ctx)
 		authtype, ok := commentCtx.Map()["auth_type"]
-		if ok && authtype == ctxtypes.AuthTypeAPIKey.StringValue() {
+		if ok && (authtype == authtypes.IdentNProviderAPIkey.StringValue()) {
 			if err := claims.IsEditor(); err != nil {
 				middleware.logger.WarnContext(ctx, authzDeniedMessage, "claims", claims)
 				render.Error(rw, err)
@@ -147,7 +147,7 @@ func (middleware *AuthZ) AdminAccess(next http.HandlerFunc) http.HandlerFunc {
 
 		commentCtx := ctxtypes.CommentFromContext(ctx)
 		authtype, ok := commentCtx.Map()["auth_type"]
-		if ok && authtype == ctxtypes.AuthTypeAPIKey.StringValue() {
+		if ok && (authtype == authtypes.IdentNProviderAPIkey.StringValue()) {
 			if err := claims.IsAdmin(); err != nil {
 				middleware.logger.WarnContext(ctx, authzDeniedMessage, "claims", claims)
 				render.Error(rw, err)

pkg/http/middleware/identn.go

@@ -0,0 +1,75 @@
+package middleware
+
+import (
+	"context"
+	"log/slog"
+	"net/http"
+
+	"github.com/SigNoz/signoz/pkg/identn"
+	"github.com/SigNoz/signoz/pkg/sharder"
+	"github.com/SigNoz/signoz/pkg/types"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
+	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
+	"github.com/SigNoz/signoz/pkg/valuer"
+)
+
+const (
+	identityCrossOrgMessage string = "::IDENTITY-CROSS-ORG::"
+)
+
+type IdentN struct {
+	resolver identn.IdentNResolver
+	sharder  sharder.Sharder
+	logger   *slog.Logger
+}
+
+func NewIdentN(resolver identn.IdentNResolver, sharder sharder.Sharder, logger *slog.Logger) *IdentN {
+	return &IdentN{
+		resolver: resolver,
+		sharder:  sharder,
+		logger:   logger,
+	}
+}
+
+func (m *IdentN) Wrap(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		idn := m.resolver.GetIdentN(r)
+		if idn == nil {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		if pre, ok := idn.(identn.IdentNWithPreHook); ok {
+			r = pre.Pre(r)
+		}
+
+		identity, err := idn.GetIdentity(r)
+		if err != nil {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		ctx := r.Context()
+		claims := identity.ToClaims()
+		if err := m.sharder.IsMyOwnedKey(ctx, types.NewOrganizationKey(valuer.MustNewUUID(claims.OrgID))); err != nil {
+			m.logger.ErrorContext(ctx, identityCrossOrgMessage, "claims", claims, "error", err)
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		ctx = authtypes.NewContextWithClaims(ctx, claims)
+
+		comment := ctxtypes.CommentFromContext(ctx)
+		comment.Set("identn_provider", claims.IdentNProvider)
+		comment.Set("user_id", claims.UserID)
+		comment.Set("org_id", claims.OrgID)
+		ctx = ctxtypes.NewContextWithComment(ctx, comment)
+
+		r = r.WithContext(ctx)
+		next.ServeHTTP(w, r)
+
+		if hook, ok := idn.(identn.IdentNWithPostHook); ok {
+			hook.Post(context.WithoutCancel(r.Context()), r, claims)
+		}
+	})
+}

pkg/identn/apikeyidentn/identn.go

@@ -0,0 +1,131 @@
+package apikeyidentn
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/SigNoz/signoz/pkg/factory"
+	"github.com/SigNoz/signoz/pkg/identn"
+	"github.com/SigNoz/signoz/pkg/sqlstore"
+	"github.com/SigNoz/signoz/pkg/types"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
+	"golang.org/x/sync/singleflight"
+)
+
+// todo: will move this in types layer with service account integration
+type apiKeyTokenKey struct{}
+
+type resolver struct {
+	store    sqlstore.SQLStore
+	headers  []string
+	settings factory.ScopedProviderSettings
+	sfGroup  *singleflight.Group
+}
+
+func New(providerSettings factory.ProviderSettings, store sqlstore.SQLStore, headers []string) identn.IdentN {
+	return &resolver{
+		store:    store,
+		headers:  headers,
+		settings: factory.NewScopedProviderSettings(providerSettings, "github.com/SigNoz/signoz/pkg/identn/apikeyidentn"),
+		sfGroup:  &singleflight.Group{},
+	}
+}
+
+func (r *resolver) Name() authtypes.IdentNProvider {
+	return authtypes.IdentNProviderAPIkey
+}
+
+func (r *resolver) Test(req *http.Request) bool {
+	for _, header := range r.headers {
+		if req.Header.Get(header) != "" {
+			return true
+		}
+	}
+	return false
+}
+
+func (r *resolver) Pre(req *http.Request) *http.Request {
+	token := r.extractToken(req)
+	if token == "" {
+		return req
+	}
+
+	ctx := context.WithValue(req.Context(), apiKeyTokenKey{}, token)
+	return req.WithContext(ctx)
+}
+
+func (r *resolver) GetIdentity(req *http.Request) (*authtypes.Identity, error) {
+	ctx := req.Context()
+
+	apiKeyToken, ok := ctx.Value(apiKeyTokenKey{}).(string)
+	if !ok || apiKeyToken == "" {
+		return nil, errors.New(errors.TypeUnauthenticated, errors.CodeUnauthenticated, "missing api key")
+	}
+
+	var apiKey types.StorableAPIKey
+	err := r.store.
+		BunDB().
+		NewSelect().
+		Model(&apiKey).
+		Where("token = ?", apiKeyToken).
+		Scan(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if apiKey.ExpiresAt.Before(time.Now()) && !apiKey.ExpiresAt.Equal(types.NEVER_EXPIRES) {
+		return nil, errors.New(errors.TypeUnauthenticated, errors.CodeUnauthenticated, "api key has expired")
+	}
+
+	var user types.User
+	err = r.store.
+		BunDB().
+		NewSelect().
+		Model(&user).
+		Where("id = ?", apiKey.UserID).
+		Scan(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	identity := authtypes.Identity{
+		UserID: user.ID,
+		Role:   apiKey.Role,
+		Email:  user.Email,
+		OrgID:  user.OrgID,
+	}
+	return &identity, nil
+}
+
+func (r *resolver) Post(ctx context.Context, _ *http.Request, _ authtypes.Claims) {
+	apiKeyToken, ok := ctx.Value(apiKeyTokenKey{}).(string)
+	if !ok || apiKeyToken == "" {
+		return
+	}
+
+	_, _, _ = r.sfGroup.Do(apiKeyToken, func() (any, error) {
+		_, err := r.store.
+			BunDB().
+			NewUpdate().
+			Model(new(types.StorableAPIKey)).
+			Set("last_used = ?", time.Now()).
+			Where("token = ?", apiKeyToken).
+			Where("revoked = false").
+			Exec(ctx)
+		if err != nil {
+			r.settings.Logger().ErrorContext(ctx, "failed to update last used of api key", "error", err)
+		}
+		return true, nil
+	})
+}
+
+func (r *resolver) extractToken(req *http.Request) string {
+	for _, header := range r.headers {
+		if v := req.Header.Get(header); v != "" {
+			return v
+		}
+	}
+	return ""
+}

pkg/identn/identn.go

@@ -0,0 +1,43 @@
+package identn
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
+)
+
+type IdentNResolver interface {
+	// GetIdentN returns the first IdentN whose Test() returns true for the request.
+	// Returns nil if no resolver matched.
+	GetIdentN(r *http.Request) IdentN
+}
+
+type IdentN interface {
+	// Test checks if this identn can handle the request.
+	// This should be a cheap check (e.g., header presence) with no I/O.
+	Test(r *http.Request) bool
+
+	// GetIdentity returns the resolved identity.
+	// Only called when Test() returns true.
+	GetIdentity(r *http.Request) (*authtypes.Identity, error)
+
+	Name() authtypes.IdentNProvider
+}
+
+// IdentNWithPreHook is optionally implemented by resolvers that need to
+// enrich the request before authentication (e.g., storing the access token
+// in context so downstream handlers can use it even on auth failure).
+type IdentNWithPreHook interface {
+	IdentN
+
+	Pre(r *http.Request) *http.Request
+}
+
+// IdentNWithPostHook is optionally implemented by resolvers that need
+// post-response side-effects (e.g., updating last_observed_at).
+type IdentNWithPostHook interface {
+	IdentN
+
+	Post(ctx context.Context, r *http.Request, claims authtypes.Claims)
+}

pkg/identn/resolver.go

@@ -0,0 +1,31 @@
+package identn
+
+import (
+	"net/http"
+
+	"github.com/SigNoz/signoz/pkg/factory"
+)
+
+type identNResolver struct {
+	identNs  []IdentN
+	settings factory.ScopedProviderSettings
+}
+
+func NewIdentNResolver(providerSettings factory.ProviderSettings, identNs ...IdentN) IdentNResolver {
+	return &identNResolver{
+		identNs:  identNs,
+		settings: factory.NewScopedProviderSettings(providerSettings, "github.com/SigNoz/signoz/pkg/identn"),
+	}
+}
+
+// GetIdentN returns the first IdentN whose Test() returns true.
+// Returns nil if no resolver matched.
+func (c *identNResolver) GetIdentN(r *http.Request) IdentN {
+	for _, idn := range c.identNs {
+		if idn.Test(r) {
+			c.settings.Logger().DebugContext(r.Context(), "identn matched", "provider", idn.Name())
+			return idn
+		}
+	}
+	return nil
+}

pkg/identn/tokenizeridentn/identn.go

@@ -0,0 +1,103 @@
+package tokenizeridentn
+
+import (
+	"context"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/SigNoz/signoz/pkg/factory"
+	"github.com/SigNoz/signoz/pkg/identn"
+	"github.com/SigNoz/signoz/pkg/tokenizer"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
+	"golang.org/x/sync/singleflight"
+)
+
+type resolver struct {
+	tokenizer tokenizer.Tokenizer
+	headers   []string
+	settings  factory.ScopedProviderSettings
+	sfGroup   *singleflight.Group
+}
+
+func New(providerSettings factory.ProviderSettings, tokenizer tokenizer.Tokenizer, headers []string) identn.IdentN {
+	return &resolver{
+		tokenizer: tokenizer,
+		headers:   headers,
+		settings:  factory.NewScopedProviderSettings(providerSettings, "github.com/SigNoz/signoz/pkg/identn/tokenizeridentn"),
+		sfGroup:   &singleflight.Group{},
+	}
+}
+
+func (r *resolver) Name() authtypes.IdentNProvider {
+	return authtypes.IdentNProviderTokenizer
+}
+
+func (r *resolver) Test(req *http.Request) bool {
+	for _, header := range r.headers {
+		if req.Header.Get(header) != "" {
+			return true
+		}
+	}
+	return false
+}
+
+func (r *resolver) Pre(req *http.Request) *http.Request {
+	accessToken := r.extractToken(req)
+	if accessToken == "" {
+		return req
+	}
+
+	ctx := authtypes.NewContextWithAccessToken(req.Context(), accessToken)
+	return req.WithContext(ctx)
+}
+
+func (r *resolver) GetIdentity(req *http.Request) (*authtypes.Identity, error) {
+	ctx := req.Context()
+
+	accessToken, err := authtypes.AccessTokenFromContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return r.tokenizer.GetIdentity(ctx, accessToken)
+}
+
+func (r *resolver) Post(ctx context.Context, _ *http.Request, _ authtypes.Claims) {
+	accessToken, err := authtypes.AccessTokenFromContext(ctx)
+	if err != nil {
+		return
+	}
+
+	_, _, _ = r.sfGroup.Do(accessToken, func() (any, error) {
+		if err := r.tokenizer.SetLastObservedAt(ctx, accessToken, time.Now()); err != nil {
+			r.settings.Logger().ErrorContext(ctx, "failed to set last observed at", "error", err)
+			return false, err
+		}
+		return true, nil
+	})
+}
+
+func (r *resolver) extractToken(req *http.Request) string {
+	var value string
+	for _, header := range r.headers {
+		if v := req.Header.Get(header); v != "" {
+			value = v
+			break
+		}
+	}
+
+	accessToken, ok := r.parseBearerAuth(value)
+	if !ok {
+		return value
+	}
+	return accessToken
+}
+
+func (r *resolver) parseBearerAuth(auth string) (string, bool) {
+	const prefix = "Bearer "
+	if len(auth) < len(prefix) || !strings.EqualFold(auth[:len(prefix)], prefix) {
+		return "", false
+	}
+	return auth[len(prefix):], true
+}

pkg/modules/dashboard/impldashboard/handler.go

@@ -15,7 +15,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/transition"
 	"github.com/SigNoz/signoz/pkg/types"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
 	"github.com/SigNoz/signoz/pkg/types/dashboardtypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
@@ -109,7 +108,7 @@ func (handler *handler) Update(rw http.ResponseWriter, r *http.Request) {
 
 	diff := 0
 	// Allow multiple deletions for API key requests; enforce for others
-	if authType, ok := ctxtypes.AuthTypeFromContext(ctx); ok && authType == ctxtypes.AuthTypeTokenizer {
+	if claims.IdentNProvider == authtypes.IdentNProviderTokenizer.StringValue() {
 		diff = 1
 	}
 

pkg/modules/session/implsession/module.go

@@ -158,7 +158,7 @@ func (module *module) CreateCallbackAuthNSession(ctx context.Context, authNProvi
 		return "", errors.WithAdditionalf(err, "root user can only authenticate via password")
 	}
 
-	token, err := module.tokenizer.CreateToken(ctx, authtypes.NewIdentity(user.ID, user.OrgID, user.Email, user.Role), map[string]string{})
+	token, err := module.tokenizer.CreateToken(ctx, authtypes.NewIdentity(user.ID, user.OrgID, user.Email, user.Role, authtypes.IdentNProviderTokenizer), map[string]string{})
 	if err != nil {
 		return "", err
 	}

pkg/query-service/app/server.go

@@ -195,13 +195,12 @@ func (s *Server) createPublicServer(api *APIHandler, web web.Web) (*http.Server,
 		}),
 		otelmux.WithPublicEndpoint(),
 	))
-	r.Use(middleware.NewAuthN([]string{"Authorization", "Sec-WebSocket-Protocol"}, s.signoz.Sharder, s.signoz.Tokenizer, s.signoz.Instrumentation.Logger()).Wrap)
+	r.Use(middleware.NewIdentN(s.signoz.IdentNResolver, s.signoz.Sharder, s.signoz.Instrumentation.Logger()).Wrap)
 	r.Use(middleware.NewTimeout(s.signoz.Instrumentation.Logger(),
 		s.config.APIServer.Timeout.ExcludedRoutes,
 		s.config.APIServer.Timeout.Default,
 		s.config.APIServer.Timeout.Max,
 	).Wrap)
-	r.Use(middleware.NewAPIKey(s.signoz.SQLStore, []string{"SIGNOZ-API-KEY"}, s.signoz.Instrumentation.Logger(), s.signoz.Sharder).Wrap)
 	r.Use(middleware.NewLogging(s.signoz.Instrumentation.Logger(), s.config.APIServer.Logging.ExcludedRoutes).Wrap)
 	r.Use(middleware.NewComment().Wrap)
 

pkg/signoz/openapi.go

@@ -26,7 +26,7 @@ import (
 	"github.com/SigNoz/signoz/pkg/modules/session"
 	"github.com/SigNoz/signoz/pkg/modules/user"
 	"github.com/SigNoz/signoz/pkg/querier"
-	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
+	"github.com/SigNoz/signoz/pkg/types/authtypes"
 	"github.com/SigNoz/signoz/pkg/zeus"
 	"github.com/swaggest/jsonschema-go"
 	"github.com/swaggest/openapi-go"
@@ -82,8 +82,8 @@ func NewOpenAPI(ctx context.Context, instrumentation instrumentation.Instrumenta
 
 	reflector.SpecSchema().SetTitle("SigNoz")
 	reflector.SpecSchema().SetDescription("OpenTelemetry-Native Logs, Metrics and Traces in a single pane")
-	reflector.SpecSchema().SetAPIKeySecurity(ctxtypes.AuthTypeAPIKey.StringValue(), "SigNoz-Api-Key", openapi.InHeader, "API Keys")
-	reflector.SpecSchema().SetHTTPBearerTokenSecurity(ctxtypes.AuthTypeTokenizer.StringValue(), "Tokenizer", "Tokens generated by the tokenizer")
+	reflector.SpecSchema().SetAPIKeySecurity(authtypes.IdentNProviderAPIkey.StringValue(), "SigNoz-Api-Key", openapi.InHeader, "API Keys")
+	reflector.SpecSchema().SetHTTPBearerTokenSecurity(authtypes.IdentNProviderTokenizer.StringValue(), "Tokenizer", "Tokens generated by the tokenizer")
 
 	collector := handler.NewOpenAPICollector(reflector)
 

pkg/signoz/signoz.go

@@ -16,6 +16,9 @@ import (
 	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/flagger"
 	"github.com/SigNoz/signoz/pkg/gateway"
+	"github.com/SigNoz/signoz/pkg/identn"
+	"github.com/SigNoz/signoz/pkg/identn/apikeyidentn"
+	"github.com/SigNoz/signoz/pkg/identn/tokenizeridentn"
 	"github.com/SigNoz/signoz/pkg/instrumentation"
 	"github.com/SigNoz/signoz/pkg/licensing"
 	"github.com/SigNoz/signoz/pkg/modules/dashboard"
@@ -65,6 +68,7 @@ type SigNoz struct {
 	Sharder                sharder.Sharder
 	StatsReporter          statsreporter.StatsReporter
 	Tokenizer              pkgtokenizer.Tokenizer
+	IdentNResolver         identn.IdentNResolver
 	Authz                  authz.AuthZ
 	Modules                Modules
 	Handlers               Handlers
@@ -390,6 +394,11 @@ func New(
 	// Initialize all modules
 	modules := NewModules(sqlstore, tokenizer, emailing, providerSettings, orgGetter, alertmanager, analytics, querier, telemetrystore, telemetryMetadataStore, authNs, authz, cache, queryParser, config, dashboard, userGetter)
 
+	// Initialize identN resolver
+	tokenizerIdentn := tokenizeridentn.New(providerSettings, tokenizer, []string{"Authorization", "Sec-WebSocket-Protocol"})
+	apikeyIdentn := apikeyidentn.New(providerSettings, sqlstore, []string{"SIGNOZ-API-KEY"})
+	identNResolver := identn.NewIdentNResolver(providerSettings, tokenizerIdentn, apikeyIdentn)
+
 	userService := impluser.NewService(providerSettings, impluser.NewStore(sqlstore, providerSettings), modules.User, orgGetter, authz, config.User.Root)
 
 	// Initialize the querier handler via callback (allows EE to decorate with anomaly detection)
@@ -468,6 +477,7 @@ func New(
 		Emailing:               emailing,
 		Sharder:                sharder,
 		Tokenizer:              tokenizer,
+		IdentNResolver:         identNResolver,
 		Authz:                  authz,
 		Modules:                modules,
 		Handlers:               handlers,

pkg/tokenizer/jwttokenizer/provider.go

@@ -125,7 +125,7 @@ func (provider *provider) GetIdentity(ctx context.Context, accessToken string) (
 		return nil, errors.Newf(errors.TypeUnauthenticated, errors.CodeUnauthenticated, "claim role mismatch")
 	}
 
-	return authtypes.NewIdentity(valuer.MustNewUUID(claims.UserID), valuer.MustNewUUID(claims.OrgID), valuer.MustNewEmail(claims.Email), claims.Role), nil
+	return authtypes.NewIdentity(valuer.MustNewUUID(claims.UserID), valuer.MustNewUUID(claims.OrgID), valuer.MustNewEmail(claims.Email), claims.Role, authtypes.IdentNProviderTokenizer), nil
 }
 
 func (provider *provider) DeleteToken(ctx context.Context, accessToken string) error {

pkg/tokenizer/tokenizerstore/sqltokenizerstore/store.go

@@ -47,7 +47,7 @@ func (store *store) GetIdentityByUserID(ctx context.Context, userID valuer.UUID)
 		return nil, store.sqlstore.WrapNotFoundErrf(err, types.ErrCodeUserNotFound, "user with id: %s does not exist", userID)
 	}
 
-	return authtypes.NewIdentity(userID, user.OrgID, user.Email, types.Role(user.Role)), nil
+	return authtypes.NewIdentity(userID, user.OrgID, user.Email, types.Role(user.Role), authtypes.IdentNProviderTokenizer), nil
 }
 
 func (store *store) GetByAccessToken(ctx context.Context, accessToken string) (*authtypes.StorableToken, error) {

pkg/types/authtypes/authn.go

@@ -25,10 +25,11 @@ var (
 type AuthNProvider struct{ valuer.String }
 
 type Identity struct {
-	UserID valuer.UUID  `json:"userId"`
-	OrgID  valuer.UUID  `json:"orgId"`
-	Email  valuer.Email `json:"email"`
-	Role   types.Role   `json:"role"`
+	UserID        valuer.UUID    `json:"userId"`
+	OrgID         valuer.UUID    `json:"orgId"`
+	IdenNProvider IdentNProvider `json:"identNProvider"`
+	Email         valuer.Email   `json:"email"`
+	Role          types.Role     `json:"role"`
 }
 
 type CallbackIdentity struct {
@@ -78,12 +79,13 @@ func NewStateFromString(state string) (State, error) {
 	}, nil
 }
 
-func NewIdentity(userID valuer.UUID, orgID valuer.UUID, email valuer.Email, role types.Role) *Identity {
+func NewIdentity(userID valuer.UUID, orgID valuer.UUID, email valuer.Email, role types.Role, identNProvider IdentNProvider) *Identity {
 	return &Identity{
-		UserID: userID,
-		OrgID:  orgID,
-		Email:  email,
-		Role:   role,
+		UserID:        userID,
+		OrgID:         orgID,
+		Email:         email,
+		Role:          role,
+		IdenNProvider: identNProvider,
 	}
 }
 
@@ -116,10 +118,11 @@ func (typ *Identity) UnmarshalBinary(data []byte) error {
 
 func (typ *Identity) ToClaims() Claims {
 	return Claims{
-		UserID: typ.UserID.String(),
-		Email:  typ.Email.String(),
-		Role:   typ.Role,
-		OrgID:  typ.OrgID.String(),
+		UserID:         typ.UserID.String(),
+		Email:          typ.Email.String(),
+		Role:           typ.Role,
+		OrgID:          typ.OrgID.String(),
+		IdentNProvider: typ.IdenNProvider.StringValue(),
 	}
 }
 

pkg/types/authtypes/claims.go

@@ -13,10 +13,11 @@ type claimsKey struct{}
 type accessTokenKey struct{}
 
 type Claims struct {
-	UserID string
-	Email  string
-	Role   types.Role
-	OrgID  string
+	UserID         string
+	Email          string
+	Role           types.Role
+	OrgID          string
+	IdentNProvider string
 }
 
 // NewContextWithClaims attaches individual claims to the context.
@@ -53,6 +54,7 @@ func (c *Claims) LogValue() slog.Value {
 		slog.String("email", c.Email),
 		slog.String("role", c.Role.String()),
 		slog.String("org_id", c.OrgID),
+		slog.String("identn_provider", c.IdentNProvider),
 	)
 }
 

pkg/types/authtypes/identn.go

@@ -0,0 +1,11 @@
+package authtypes
+
+import "github.com/SigNoz/signoz/pkg/valuer"
+
+var (
+	IdentNProviderTokenizer = IdentNProvider{valuer.NewString("tokenizer")}
+	IdentNProviderAPIkey    = IdentNProvider{valuer.NewString("api_key")}
+	IdentNProviderAnonymous = IdentNProvider{valuer.NewString("anonymous")}
+)
+
+type IdentNProvider struct{ valuer.String }

pkg/types/authtypes/uuid.go

@@ -1,41 +0,0 @@
-package authtypes
-
-import (
-	"context"
-
-	"github.com/SigNoz/signoz/pkg/errors"
-)
-
-type uuidKey struct{}
-
-type UUID struct {
-}
-
-func NewUUID() *UUID {
-	return &UUID{}
-}
-
-func (u *UUID) ContextFromRequest(ctx context.Context, values ...string) (context.Context, error) {
-	var value string
-	for _, v := range values {
-		if v != "" {
-			value = v
-			break
-		}
-	}
-
-	if value == "" {
-		return ctx, errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput, "missing Authorization header")
-	}
-
-	return NewContextWithUUID(ctx, value), nil
-}
-
-func NewContextWithUUID(ctx context.Context, uuid string) context.Context {
-	return context.WithValue(ctx, uuidKey{}, uuid)
-}
-
-func UUIDFromContext(ctx context.Context) (string, bool) {
-	uuid, ok := ctx.Value(uuidKey{}).(string)
-	return uuid, ok
-}

pkg/types/ctxtypes/auth.go

@@ -1,31 +0,0 @@
-package ctxtypes
-
-import (
-	"context"
-
-	"github.com/SigNoz/signoz/pkg/valuer"
-)
-
-type AuthType struct {
-	valuer.String
-}
-
-var (
-	AuthTypeTokenizer = AuthType{valuer.NewString("tokenizer")}
-	AuthTypeAPIKey    = AuthType{valuer.NewString("api_key")}
-	AuthTypeInternal  = AuthType{valuer.NewString("internal")}
-	AuthTypeAnonymous = AuthType{valuer.NewString("anonymous")}
-)
-
-type authTypeKey struct{}
-
-// SetAuthType stores the auth type (e.g., AuthTypeJWT, AuthTypeAPIKey, AuthTypeInternal) in context.
-func SetAuthType(ctx context.Context, authType AuthType) context.Context {
-	return context.WithValue(ctx, authTypeKey{}, authType)
-}
-
-// AuthTypeFromContext retrieves the auth type from context if set.
-func AuthTypeFromContext(ctx context.Context) (AuthType, bool) {
-	v, ok := ctx.Value(authTypeKey{}).(AuthType)
-	return v, ok
-}