fix: refactor auth package (#7110)

* fix: refactor auth package

* fix: minor changes

* fix: refactor jwt

* fix: add tests and address comments

* fix: address comments

* fix: add uncomitted file

* fix: address comments

* fix: update tests

ee/http/middleware/pat.go

@@ -0,0 +1,36 @@
+package middleware
+
+import (
+	"net/http"
+
+	"go.signoz.io/signoz/pkg/types/authtypes"
+)
+
+type Pat struct {
+	uuid    *authtypes.UUID
+	headers []string
+}
+
+func NewPat(headers []string) *Pat {
+	return &Pat{uuid: authtypes.NewUUID(), headers: headers}
+}
+
+func (p *Pat) Wrap(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var values []string
+		for _, header := range p.headers {
+			values = append(values, r.Header.Get(header))
+		}
+
+		ctx, err := p.uuid.ContextFromRequest(r.Context(), values...)
+		if err != nil {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		r = r.WithContext(ctx)
+
+		next.ServeHTTP(w, r)
+	})
+
+}

ee/query-service/app/api/api.go

@@ -20,6 +20,7 @@ import (
 	basemodel "go.signoz.io/signoz/pkg/query-service/model"
 	rules "go.signoz.io/signoz/pkg/query-service/rules"
 	"go.signoz.io/signoz/pkg/query-service/version"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 )
 
 type APIHandlerOptions struct {
@@ -41,6 +42,7 @@ type APIHandlerOptions struct {
 	FluxInterval      time.Duration
 	UseLogsNewSchema  bool
 	UseTraceNewSchema bool
+	JWT               *authtypes.JWT
 }
 
 type APIHandler struct {

ee/query-service/app/api/auth.go

@@ -50,7 +50,7 @@ func (ah *APIHandler) loginUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// if all looks good, call auth
-	resp, err := baseauth.Login(ctx, &req)
+	resp, err := baseauth.Login(ctx, &req, ah.opts.JWT)
 	if ah.HandleError(w, err, http.StatusUnauthorized) {
 		return
 	}
@@ -253,7 +253,7 @@ func (ah *APIHandler) receiveGoogleAuth(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	nextPage, err := ah.AppDao().PrepareSsoRedirect(ctx, redirectUri, identity.Email)
+	nextPage, err := ah.AppDao().PrepareSsoRedirect(ctx, redirectUri, identity.Email, ah.opts.JWT)
 	if err != nil {
 		zap.L().Error("[receiveGoogleAuth] failed to generate redirect URI after successful login ", zap.String("domain", domain.String()), zap.Error(err))
 		handleSsoError(w, r, redirectUri)
@@ -331,7 +331,7 @@ func (ah *APIHandler) receiveSAML(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	nextPage, err := ah.AppDao().PrepareSsoRedirect(ctx, redirectUri, email)
+	nextPage, err := ah.AppDao().PrepareSsoRedirect(ctx, redirectUri, email, ah.opts.JWT)
 	if err != nil {
 		zap.L().Error("[receiveSAML] failed to generate redirect URI after successful login ", zap.String("domain", domain.String()), zap.Error(err))
 		handleSsoError(w, r, redirectUri)

ee/query-service/app/api/cloudIntegrations.go

@@ -37,7 +37,7 @@ func (ah *APIHandler) CloudIntegrationsGenerateConnectionParams(w http.ResponseW
 		return
 	}
 
-	currentUser, err := auth.GetUserFromRequest(r)
+	currentUser, err := auth.GetUserFromReqContext(r.Context())
 	if err != nil {
 		RespondError(w, basemodel.UnauthorizedError(fmt.Errorf(
 			"couldn't deduce current user: %w", err,

ee/query-service/app/api/pat.go

@@ -34,7 +34,7 @@ func (ah *APIHandler) createPAT(w http.ResponseWriter, r *http.Request) {
 		RespondError(w, model.BadRequest(err), nil)
 		return
 	}
-	user, err := auth.GetUserFromRequest(r)
+	user, err := auth.GetUserFromReqContext(r.Context())
 	if err != nil {
 		RespondError(w, &model.ApiError{
 			Typ: model.ErrorUnauthorized,
@@ -97,7 +97,7 @@ func (ah *APIHandler) updatePAT(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	user, err := auth.GetUserFromRequest(r)
+	user, err := auth.GetUserFromReqContext(r.Context())
 	if err != nil {
 		RespondError(w, &model.ApiError{
 			Typ: model.ErrorUnauthorized,
@@ -127,7 +127,7 @@ func (ah *APIHandler) updatePAT(w http.ResponseWriter, r *http.Request) {
 
 func (ah *APIHandler) getPATs(w http.ResponseWriter, r *http.Request) {
 	ctx := context.Background()
-	user, err := auth.GetUserFromRequest(r)
+	user, err := auth.GetUserFromReqContext(r.Context())
 	if err != nil {
 		RespondError(w, &model.ApiError{
 			Typ: model.ErrorUnauthorized,
@@ -147,7 +147,7 @@ func (ah *APIHandler) getPATs(w http.ResponseWriter, r *http.Request) {
 func (ah *APIHandler) revokePAT(w http.ResponseWriter, r *http.Request) {
 	ctx := context.Background()
 	id := mux.Vars(r)["id"]
-	user, err := auth.GetUserFromRequest(r)
+	user, err := auth.GetUserFromReqContext(r.Context())
 	if err != nil {
 		RespondError(w, &model.ApiError{
 			Typ: model.ErrorUnauthorized,

ee/query-service/app/server.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/rs/cors"
 	"github.com/soheilhy/cmux"
+	eemiddleware "go.signoz.io/signoz/ee/http/middleware"
 	"go.signoz.io/signoz/ee/query-service/app/api"
 	"go.signoz.io/signoz/ee/query-service/app/db"
 	"go.signoz.io/signoz/ee/query-service/auth"
@@ -24,6 +25,7 @@ import (
 	"go.signoz.io/signoz/ee/query-service/rules"
 	"go.signoz.io/signoz/pkg/http/middleware"
 	"go.signoz.io/signoz/pkg/signoz"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 	"go.signoz.io/signoz/pkg/web"
 
 	licensepkg "go.signoz.io/signoz/ee/query-service/license"
@@ -72,6 +74,7 @@ type ServerOptions struct {
 	GatewayUrl                 string
 	UseLogsNewSchema           bool
 	UseTraceNewSchema          bool
+	Jwt                        *authtypes.JWT
 }
 
 // Server runs HTTP api service
@@ -261,6 +264,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 		GatewayUrl:                    serverOptions.GatewayUrl,
 		UseLogsNewSchema:              serverOptions.UseLogsNewSchema,
 		UseTraceNewSchema:             serverOptions.UseTraceNewSchema,
+		JWT:                           serverOptions.Jwt,
 	}
 
 	apiHandler, err := api.NewAPIHandler(apiOpts)
@@ -303,6 +307,8 @@ func (s *Server) createPrivateServer(apiHandler *api.APIHandler) (*http.Server,
 
 	r := baseapp.NewRouter()
 
+	r.Use(middleware.NewAuth(zap.L(), s.serverOptions.Jwt, []string{"Authorization", "Sec-WebSocket-Protocol"}).Wrap)
+	r.Use(eemiddleware.NewPat([]string{"SIGNOZ-API-KEY"}).Wrap)
 	r.Use(middleware.NewTimeout(zap.L(),
 		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
 		s.serverOptions.Config.APIServer.Timeout.Default,
@@ -334,8 +340,8 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web web.Web) (*h
 	r := baseapp.NewRouter()
 
 	// add auth middleware
-	getUserFromRequest := func(r *http.Request) (*basemodel.UserPayload, error) {
-		user, err := auth.GetUserFromRequest(r, apiHandler)
+	getUserFromRequest := func(ctx context.Context) (*basemodel.UserPayload, error) {
+		user, err := auth.GetUserFromRequestContext(ctx, apiHandler)
 
 		if err != nil {
 			return nil, err
@@ -349,6 +355,8 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler, web web.Web) (*h
 	}
 	am := baseapp.NewAuthMiddleware(getUserFromRequest)
 
+	r.Use(middleware.NewAuth(zap.L(), s.serverOptions.Jwt, []string{"Authorization", "Sec-WebSocket-Protocol"}).Wrap)
+	r.Use(eemiddleware.NewPat([]string{"SIGNOZ-API-KEY"}).Wrap)
 	r.Use(middleware.NewTimeout(zap.L(),
 		s.serverOptions.Config.APIServer.Timeout.ExcludedRoutes,
 		s.serverOptions.Config.APIServer.Timeout.Default,

ee/query-service/auth/auth.go

@@ -3,20 +3,20 @@ package auth
 import (
 	"context"
 	"fmt"
-	"net/http"
 	"time"
 
 	"go.signoz.io/signoz/ee/query-service/app/api"
 	baseauth "go.signoz.io/signoz/pkg/query-service/auth"
 	basemodel "go.signoz.io/signoz/pkg/query-service/model"
 	"go.signoz.io/signoz/pkg/query-service/telemetry"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 
 	"go.uber.org/zap"
 )
 
-func GetUserFromRequest(r *http.Request, apiHandler *api.APIHandler) (*basemodel.UserPayload, error) {
-	patToken := r.Header.Get("SIGNOZ-API-KEY")
-	if len(patToken) > 0 {
+func GetUserFromRequestContext(ctx context.Context, apiHandler *api.APIHandler) (*basemodel.UserPayload, error) {
+	patToken, ok := authtypes.UUIDFromContext(ctx)
+	if ok && patToken != "" {
 		zap.L().Debug("Received a non-zero length PAT token")
 		ctx := context.Background()
 		dao := apiHandler.AppDao()
@@ -52,5 +52,5 @@ func GetUserFromRequest(r *http.Request, apiHandler *api.APIHandler) (*basemodel
 			return nil, err
 		}
 	}
-	return baseauth.GetUserFromRequest(r)
+	return baseauth.GetUserFromReqContext(ctx)
 }

ee/query-service/dao/interface.go

@@ -10,6 +10,7 @@ import (
 	basedao "go.signoz.io/signoz/pkg/query-service/dao"
 	baseint "go.signoz.io/signoz/pkg/query-service/interfaces"
 	basemodel "go.signoz.io/signoz/pkg/query-service/model"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 )
 
 type ModelDao interface {
@@ -22,7 +23,7 @@ type ModelDao interface {
 
 	// auth methods
 	CanUsePassword(ctx context.Context, email string) (bool, basemodel.BaseApiError)
-	PrepareSsoRedirect(ctx context.Context, redirectUri, email string) (redirectURL string, apierr basemodel.BaseApiError)
+	PrepareSsoRedirect(ctx context.Context, redirectUri, email string, jwt *authtypes.JWT) (redirectURL string, apierr basemodel.BaseApiError)
 	GetDomainFromSsoResponse(ctx context.Context, relayState *url.URL) (*model.OrgDomain, error)
 
 	// org domain (auth domains) CRUD ops

ee/query-service/dao/sqlite/auth.go

@@ -14,6 +14,7 @@ import (
 	baseconst "go.signoz.io/signoz/pkg/query-service/constants"
 	basemodel "go.signoz.io/signoz/pkg/query-service/model"
 	"go.signoz.io/signoz/pkg/query-service/utils"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 	"go.uber.org/zap"
 )
 
@@ -64,7 +65,7 @@ func (m *modelDao) createUserForSAMLRequest(ctx context.Context, email string) (
 
 // PrepareSsoRedirect prepares redirect page link after SSO response
 // is successfully parsed (i.e. valid email is available)
-func (m *modelDao) PrepareSsoRedirect(ctx context.Context, redirectUri, email string) (redirectURL string, apierr basemodel.BaseApiError) {
+func (m *modelDao) PrepareSsoRedirect(ctx context.Context, redirectUri, email string, jwt *authtypes.JWT) (redirectURL string, apierr basemodel.BaseApiError) {
 
 	userPayload, apierr := m.GetUserByEmail(ctx, email)
 	if !apierr.IsNil() {
@@ -85,7 +86,7 @@ func (m *modelDao) PrepareSsoRedirect(ctx context.Context, redirectUri, email st
 		user = &userPayload.User
 	}
 
-	tokenStore, err := baseauth.GenerateJWTForUser(user)
+	tokenStore, err := baseauth.GenerateJWTForUser(user, jwt)
 	if err != nil {
 		zap.L().Error("failed to generate token for SSO login user", zap.Error(err))
 		return "", model.InternalErrorStr("failed to generate token for the user")

ee/query-service/license/manager.go

@@ -10,8 +10,8 @@ import (
 
 	"sync"
 
-	"go.signoz.io/signoz/pkg/query-service/auth"
 	baseconstants "go.signoz.io/signoz/pkg/query-service/constants"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 
 	validate "go.signoz.io/signoz/ee/query-service/integrations/signozio"
 	"go.signoz.io/signoz/ee/query-service/model"
@@ -237,10 +237,10 @@ func (lm *Manager) ValidateV3(ctx context.Context) (reterr error) {
 func (lm *Manager) ActivateV3(ctx context.Context, licenseKey string) (licenseResponse *model.LicenseV3, errResponse *model.ApiError) {
 	defer func() {
 		if errResponse != nil {
-			userEmail, err := auth.GetEmailFromJwt(ctx)
-			if err == nil {
+			claims, ok := authtypes.ClaimsFromContext(ctx)
+			if ok {
 				telemetry.GetInstance().SendEvent(telemetry.TELEMETRY_LICENSE_ACT_FAILED,
-					map[string]interface{}{"err": errResponse.Err.Error()}, userEmail, true, false)
+					map[string]interface{}{"err": errResponse.Err.Error()}, claims.Email, true, false)
 			}
 		}
 	}()

ee/query-service/main.go

@@ -20,6 +20,7 @@ import (
 	baseconst "go.signoz.io/signoz/pkg/query-service/constants"
 	"go.signoz.io/signoz/pkg/query-service/version"
 	"go.signoz.io/signoz/pkg/signoz"
+	"go.signoz.io/signoz/pkg/types/authtypes"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 
@@ -154,6 +155,16 @@ func main() {
 		zap.L().Fatal("Failed to create signoz struct", zap.Error(err))
 	}
 
+	jwtSecret := os.Getenv("SIGNOZ_JWT_SECRET")
+
+	if len(jwtSecret) == 0 {
+		zap.L().Warn("No JWT secret key is specified.")
+	} else {
+		zap.L().Info("JWT secret key set successfully.")
+	}
+
+	jwt := authtypes.NewJWT(jwtSecret, 30*time.Minute, 30*24*time.Hour)
+
 	serverOptions := &app.ServerOptions{
 		Config:                     config,
 		SigNoz:                     signoz,
@@ -171,15 +182,7 @@ func main() {
 		GatewayUrl:                 gatewayUrl,
 		UseLogsNewSchema:           useLogsNewSchema,
 		UseTraceNewSchema:          useTraceNewSchema,
-	}
-
-	// Read the jwt secret key
-	auth.JwtSecret = os.Getenv("SIGNOZ_JWT_SECRET")
-
-	if len(auth.JwtSecret) == 0 {
-		zap.L().Warn("No JWT secret key is specified.")
-	} else {
-		zap.L().Info("JWT secret key set successfully.")
+		Jwt:                        jwt,
 	}
 
 	server, err := app.NewServer(serverOptions)

go.mod

@@ -13,7 +13,6 @@ require (
 	github.com/SigNoz/zap_otlp/zap_otlp_encoder v0.0.0-20230822164844-1b861a431974
 	github.com/SigNoz/zap_otlp/zap_otlp_sync v0.0.0-20230822164844-1b861a431974
 	github.com/antonmedv/expr v1.15.3
-	github.com/auth0/go-jwt-middleware v1.0.1
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/coreos/go-oidc/v3 v3.11.0
 	github.com/dustin/go-humanize v1.0.1
@@ -24,7 +23,7 @@ require (
 	github.com/go-redis/redis/v8 v8.11.5
 	github.com/go-redis/redismock/v8 v8.11.5
 	github.com/go-viper/mapstructure/v2 v2.1.0
-	github.com/golang-jwt/jwt v3.2.2+incompatible
+	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/handlers v1.5.1
 	github.com/gorilla/mux v1.8.1
@@ -112,7 +111,6 @@ require (
 	github.com/expr-lang/expr v1.16.9 // indirect
 	github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
-	github.com/form3tech-oss/jwt-go v3.2.5+incompatible // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/go-faster/city v1.0.1 // indirect
 	github.com/go-faster/errors v0.7.1 // indirect
@@ -132,7 +130,6 @@ require (
 	github.com/goccy/go-json v0.10.3 // indirect
 	github.com/gofrs/uuid v4.4.0+incompatible // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/golang/snappy v0.0.5-0.20220116011046-fa5810519dcb // indirect