feat: api management (#4557)

* feat: api management

* chore: address review comments and typos

* chore: add sort and created by user object on create

* chore: replace expiresAt with expiresInDays for request body

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

@@ -152,9 +152,10 @@ func (ah *APIHandler) RegisterRoutes(router *mux.Router, am *baseapp.AuthMiddlew
 	router.HandleFunc("/api/v2/metrics/query_range", am.ViewAccess(ah.queryRangeMetricsV2)).Methods(http.MethodPost)
 
 	// PAT APIs
-	router.HandleFunc("/api/v1/pat", am.AdminAccess(ah.createPAT)).Methods(http.MethodPost)
-	router.HandleFunc("/api/v1/pat", am.AdminAccess(ah.getPATs)).Methods(http.MethodGet)
-	router.HandleFunc("/api/v1/pat/{id}", am.AdminAccess(ah.deletePAT)).Methods(http.MethodDelete)
+	router.HandleFunc("/api/v1/pats", am.AdminAccess(ah.createPAT)).Methods(http.MethodPost)
+	router.HandleFunc("/api/v1/pats", am.AdminAccess(ah.getPATs)).Methods(http.MethodGet)
+	router.HandleFunc("/api/v1/pats/{id}", am.AdminAccess(ah.updatePAT)).Methods(http.MethodPut)
+	router.HandleFunc("/api/v1/pats/{id}", am.AdminAccess(ah.revokePAT)).Methods(http.MethodDelete)
 
 	router.HandleFunc("/api/v1/checkout", am.AdminAccess(ah.checkout)).Methods(http.MethodPost)
 	router.HandleFunc("/api/v1/billing", am.AdminAccess(ah.getBilling)).Methods(http.MethodGet)

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

@@ -12,6 +12,7 @@ import (
 	"github.com/gorilla/mux"
 	"go.signoz.io/signoz/ee/query-service/model"
 	"go.signoz.io/signoz/pkg/query-service/auth"
+	baseconstants "go.signoz.io/signoz/pkg/query-service/constants"
 	basemodel "go.signoz.io/signoz/pkg/query-service/model"
 	"go.uber.org/zap"
 )
@@ -28,7 +29,7 @@ func generatePATToken() string {
 func (ah *APIHandler) createPAT(w http.ResponseWriter, r *http.Request) {
 	ctx := context.Background()
 
-	req := model.PAT{}
+	req := model.CreatePATRequestBody{}
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		RespondError(w, model.BadRequest(err), nil)
 		return
@@ -41,30 +42,87 @@ func (ah *APIHandler) createPAT(w http.ResponseWriter, r *http.Request) {
 		}, nil)
 		return
 	}
+	pat := model.PAT{
+		Name: 	req.Name,
+		Role: 	req.Role,
+		ExpiresAt: req.ExpiresInDays,
+	}
+	err = validatePATRequest(pat)
+	if err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
+	}
+
+	// All the PATs are associated with the user creating the PAT.
+	pat.UserID = user.Id
+	pat.CreatedAt = time.Now().Unix()
+	pat.UpdatedAt = time.Now().Unix()
+	pat.LastUsed = 0
+	pat.Token = generatePATToken()
+
+	if pat.ExpiresAt != 0 {
+		// convert expiresAt to unix timestamp from days
+		pat.ExpiresAt = time.Now().Unix() + (pat.ExpiresAt * 24 * 60 * 60)
+	}
+
+	zap.S().Debugf("Got Create PAT request: %+v", pat)
+	var apierr basemodel.BaseApiError
+	if pat, apierr = ah.AppDao().CreatePAT(ctx, pat); apierr != nil {
+		RespondError(w, apierr, nil)
+		return
+	}
+
+	ah.Respond(w, &pat)
+}
+
+func validatePATRequest(req model.PAT) error {
+	if req.Role == "" || (req.Role != baseconstants.ViewerGroup && req.Role != baseconstants.EditorGroup && req.Role != baseconstants.AdminGroup) {
+		return fmt.Errorf("valid role is required")
+	}
+	if req.ExpiresAt < 0 {
+		return fmt.Errorf("valid expiresAt is required")
+	}
+	if req.Name == "" {
+		return fmt.Errorf("valid name is required")
+	}
+	return nil
+}
+
+func (ah *APIHandler) updatePAT(w http.ResponseWriter, r *http.Request) {
+	ctx := context.Background()
 
-	// All the PATs are associated with the user creating the PAT. Hence, the permissions
-	// associated with the PAT is also equivalent to that of the user.
-	req.UserID = user.Id
-	req.CreatedAt = time.Now().Unix()
-	req.Token = generatePATToken()
+	req := model.PAT{}
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
+	}
 
-	// default expiry is 30 days
-	if req.ExpiresAt == 0 {
-		req.ExpiresAt = time.Now().AddDate(0, 0, 30).Unix()
+	user, err := auth.GetUserFromRequest(r)
+	if err != nil {
+		RespondError(w, &model.ApiError{
+			Typ: model.ErrorUnauthorized,
+			Err: err,
+		}, nil)
+		return
 	}
-	// max expiry is 1 year
-	if req.ExpiresAt > time.Now().AddDate(1, 0, 0).Unix() {
-		req.ExpiresAt = time.Now().AddDate(1, 0, 0).Unix()
+
+	err = validatePATRequest(req)
+	if err != nil {
+		RespondError(w, model.BadRequest(err), nil)
+		return
 	}
 
-	zap.S().Debugf("Got PAT request: %+v", req)
+	req.UpdatedByUserID = user.Id
+	id := mux.Vars(r)["id"]
+	req.UpdatedAt = time.Now().Unix()
+	zap.S().Debugf("Got Update PAT request: %+v", req)
 	var apierr basemodel.BaseApiError
-	if req, apierr = ah.AppDao().CreatePAT(ctx, req); apierr != nil {
+	if apierr = ah.AppDao().UpdatePAT(ctx, req, id); apierr != nil {
 		RespondError(w, apierr, nil)
 		return
 	}
 
-	ah.Respond(w, &req)
+	ah.Respond(w, map[string]string{"data": "pat updated successfully"})
 }
 
 func (ah *APIHandler) getPATs(w http.ResponseWriter, r *http.Request) {
@@ -86,7 +144,7 @@ func (ah *APIHandler) getPATs(w http.ResponseWriter, r *http.Request) {
 	ah.Respond(w, pats)
 }
 
-func (ah *APIHandler) deletePAT(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)
@@ -105,14 +163,14 @@ func (ah *APIHandler) deletePAT(w http.ResponseWriter, r *http.Request) {
 	if pat.UserID != user.Id {
 		RespondError(w, &model.ApiError{
 			Typ: model.ErrorUnauthorized,
-			Err: fmt.Errorf("unauthorized PAT delete request"),
+			Err: fmt.Errorf("unauthorized PAT revoke request"),
 		}, nil)
 		return
 	}
-	zap.S().Debugf("Delete PAT with id: %+v", id)
-	if apierr := ah.AppDao().DeletePAT(ctx, id); apierr != nil {
+	zap.S().Debugf("Revoke PAT with id: %+v", id)
+	if apierr := ah.AppDao().RevokePAT(ctx, id, user.Id); apierr != nil {
 		RespondError(w, apierr, nil)
 		return
 	}
-	ah.Respond(w, map[string]string{"data": "pat deleted successfully"})
+	ah.Respond(w, map[string]string{"data": "pat revoked successfully"})
 }

ee/query-service/app/server.go

@@ -20,10 +20,11 @@ import (
 	"github.com/soheilhy/cmux"
 	"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"
 	"go.signoz.io/signoz/ee/query-service/constants"
 	"go.signoz.io/signoz/ee/query-service/dao"
 	"go.signoz.io/signoz/ee/query-service/interfaces"
-	"go.signoz.io/signoz/pkg/query-service/auth"
+	baseauth "go.signoz.io/signoz/pkg/query-service/auth"
 	baseInterface "go.signoz.io/signoz/pkg/query-service/interfaces"
 	v3 "go.signoz.io/signoz/pkg/query-service/model/v3"
 
@@ -37,7 +38,6 @@ import (
 	"go.signoz.io/signoz/pkg/query-service/app/logparsingpipeline"
 	"go.signoz.io/signoz/pkg/query-service/app/opamp"
 	opAmpModel "go.signoz.io/signoz/pkg/query-service/app/opamp/model"
-	baseauth "go.signoz.io/signoz/pkg/query-service/auth"
 	"go.signoz.io/signoz/pkg/query-service/cache"
 	baseconst "go.signoz.io/signoz/pkg/query-service/constants"
 	"go.signoz.io/signoz/pkg/query-service/healthcheck"
@@ -304,25 +304,12 @@ func (s *Server) createPublicServer(apiHandler *api.APIHandler) (*http.Server, e
 
 	r := mux.NewRouter()
 
+	// add auth middleware
 	getUserFromRequest := func(r *http.Request) (*basemodel.UserPayload, error) {
-		patToken := r.Header.Get("SIGNOZ-API-KEY")
-		if len(patToken) > 0 {
-			zap.S().Debugf("Received a non-zero length PAT token")
-			ctx := context.Background()
-			dao := apiHandler.AppDao()
-
-			user, err := dao.GetUserByPAT(ctx, patToken)
-			if err == nil && user != nil {
-				zap.S().Debugf("Found valid PAT user: %+v", user)
-				return user, nil
-			}
-			if err != nil {
-				zap.S().Debugf("Error while getting user for PAT: %+v", err)
-			}
-		}
-		return baseauth.GetUserFromRequest(r)
+		return auth.GetUserFromRequest(r, apiHandler)
 	}
 	am := baseapp.NewAuthMiddleware(getUserFromRequest)
+
 	r.Use(setTimeoutMiddleware)
 	r.Use(s.analyticsMiddleware)
 	r.Use(loggingMiddleware)
@@ -439,7 +426,7 @@ func extractQueryRangeV3Data(path string, r *http.Request) (map[string]interface
 			telemetry.GetInstance().AddActiveLogsUser()
 		}
 		data["dataSources"] = dataSources
-		userEmail, err := auth.GetEmailFromJwt(r.Context())
+		userEmail, err := baseauth.GetEmailFromJwt(r.Context())
 		if err == nil {
 			telemetry.GetInstance().SendEvent(telemetry.TELEMETRY_EVENT_QUERY_RANGE_V3, data, userEmail, true)
 		}
@@ -463,7 +450,7 @@ func getActiveLogs(path string, r *http.Request) {
 
 func (s *Server) analyticsMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := auth.AttachJwtToContext(r.Context(), r)
+		ctx := baseauth.AttachJwtToContext(r.Context(), r)
 		r = r.WithContext(ctx)
 		route := mux.CurrentRoute(r)
 		path, _ := route.GetPathTemplate()
@@ -482,7 +469,7 @@ func (s *Server) analyticsMiddleware(next http.Handler) http.Handler {
 		}
 
 		if _, ok := telemetry.EnabledPaths()[path]; ok {
-			userEmail, err := auth.GetEmailFromJwt(r.Context())
+			userEmail, err := baseauth.GetEmailFromJwt(r.Context())
 			if err == nil {
 				telemetry.GetInstance().SendEvent(telemetry.TELEMETRY_EVENT_PATH, data, userEmail)
 			}

ee/query-service/auth/auth.go

@@ -0,0 +1,56 @@
+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.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 {
+		zap.S().Debugf("Received a non-zero length PAT token")
+		ctx := context.Background()
+		dao := apiHandler.AppDao()
+
+		pat, err := dao.GetPAT(ctx, patToken)
+		if err == nil && pat != nil {
+			zap.S().Debugf("Found valid PAT: %+v", pat)
+			if pat.ExpiresAt < time.Now().Unix() && pat.ExpiresAt != 0 {
+				zap.S().Debugf("PAT has expired: %+v", pat)
+				return nil, fmt.Errorf("PAT has expired")
+			}
+			group, apiErr := dao.GetGroupByName(ctx, pat.Role)
+			if apiErr != nil {
+				zap.S().Debugf("Error while getting group for PAT: %+v", apiErr)
+				return nil, apiErr
+			}
+			user, err := dao.GetUser(ctx, pat.UserID)
+			if err != nil {
+				zap.S().Debugf("Error while getting user for PAT: %+v", err)
+				return nil, err
+			}
+			telemetry.GetInstance().SetPatTokenUser()
+			dao.UpdatePATLastUsed(ctx, patToken, time.Now().Unix())
+			user.User.GroupId = group.Id
+			user.User.Id = pat.Id
+			return &basemodel.UserPayload{
+				User: user.User,
+				Role: pat.Role,
+			}, nil
+		}
+		if err != nil {
+			zap.S().Debugf("Error while getting user for PAT: %+v", err)
+			return nil, err
+		}
+	}
+	return baseauth.GetUserFromRequest(r)
+}

ee/query-service/dao/interface.go

@@ -34,9 +34,11 @@ type ModelDao interface {
 	GetDomainByEmail(ctx context.Context, email string) (*model.OrgDomain, basemodel.BaseApiError)
 
 	CreatePAT(ctx context.Context, p model.PAT) (model.PAT, basemodel.BaseApiError)
+	UpdatePAT(ctx context.Context, p model.PAT, id string) (basemodel.BaseApiError)
 	GetPAT(ctx context.Context, pat string) (*model.PAT, basemodel.BaseApiError)
+	UpdatePATLastUsed(ctx context.Context, pat string, lastUsed int64) basemodel.BaseApiError
 	GetPATByID(ctx context.Context, id string) (*model.PAT, basemodel.BaseApiError)
 	GetUserByPAT(ctx context.Context, token string) (*basemodel.UserPayload, basemodel.BaseApiError)
 	ListPATs(ctx context.Context, userID string) ([]model.PAT, basemodel.BaseApiError)
-	DeletePAT(ctx context.Context, id string) basemodel.BaseApiError
+	RevokePAT(ctx context.Context, id string, userID string) basemodel.BaseApiError
 }