feat: added license manager and feature flags

pkg/query-service/app/server.go

@@ -95,7 +95,7 @@ func NewServer(serverOptions *ServerOptions) (*Server, error) {
 	}
 
 	telemetry.GetInstance().SetReader(reader)
-	apiHandler, err := NewAPIHandler(&reader, dao.DB(), rm)
+	apiHandler, err := NewAPIHandler(reader, dao.DB(), rm)
 	if err != nil {
 		return nil, err
 	}

pkg/query-service/auth/auth.go

@@ -299,39 +299,10 @@ func Login(ctx context.Context, request *model.LoginRequest) (*model.LoginRespon
 		return nil, err
 	}
 
-	accessJwtExpiry := time.Now().Add(JwtExpiry).Unix()
-
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
-		"id":    user.Id,
-		"gid":   user.GroupId,
-		"email": user.Email,
-		"exp":   accessJwtExpiry,
-	})
-
-	accessJwt, err := token.SignedString([]byte(JwtSecret))
-	if err != nil {
-		return nil, errors.Errorf("failed to encode jwt: %v", err)
-	}
-
-	refreshJwtExpiry := time.Now().Add(JwtRefresh).Unix()
-	token = jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
-		"id":    user.Id,
-		"gid":   user.GroupId,
-		"email": user.Email,
-		"exp":   refreshJwtExpiry,
-	})
-
-	refreshJwt, err := token.SignedString([]byte(JwtSecret))
-	if err != nil {
-		return nil, errors.Errorf("failed to encode jwt: %v", err)
-	}
+	userjwt, err := GenerateJWTForUser(&user.User)
 
 	return &model.LoginResponse{
-		AccessJwt:        accessJwt,
-		AccessJwtExpiry:  accessJwtExpiry,
-		RefreshJwt:       refreshJwt,
-		RefreshJwtExpiry: refreshJwtExpiry,
-		UserId:           user.Id,
+		userjwt,
 	}, nil
 }
 
@@ -375,3 +346,35 @@ func passwordMatch(hash, password string) bool {
 	}
 	return true
 }
+
+func GenerateJWTForUser(user *model.User) (model.UserJwtObject, error) {
+	j := model.UserJwtObject{}
+	var err error
+	j.AccessJwtExpiry = time.Now().Add(JwtExpiry).Unix()
+
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"id":    user.Id,
+		"gid":   user.GroupId,
+		"email": user.Email,
+		"exp":   j.AccessJwtExpiry,
+	})
+
+	j.AccessJwt, err = token.SignedString([]byte(JwtSecret))
+	if err != nil {
+		return j, errors.Errorf("failed to encode jwt: %v", err)
+	}
+
+	j.RefreshJwtExpiry = time.Now().Add(JwtRefresh).Unix()
+	token = jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"id":    user.Id,
+		"gid":   user.GroupId,
+		"email": user.Email,
+		"exp":   j.RefreshJwtExpiry,
+	})
+
+	j.RefreshJwt, err = token.SignedString([]byte(JwtSecret))
+	if err != nil {
+		return j, errors.Errorf("failed to encode jwt: %v", err)
+	}
+	return j, nil
+}

pkg/query-service/dao/factory.go

@@ -24,6 +24,11 @@ func InitDao(engine, path string) error {
 	return nil
 }
 
+// SetDB is used by ee for setting modelDAO
+func SetDB(m ModelDao) {
+	db = m
+}
+
 func DB() ModelDao {
 	if db == nil {
 		// Should never reach here

pkg/query-service/dao/sqlite/connection.go

@@ -88,6 +88,11 @@ func InitDB(dataSourceName string) (*ModelDaoSqlite, error) {
 	return mds, nil
 }
 
+// DB returns database connection
+func (mds *ModelDaoSqlite) DB() *sqlx.DB {
+	return mds.db
+}
+
 // initializeOrgPreferences initializes in-memory telemetry settings. It is planned to have
 // multiple orgs in the system. In case of multiple orgs, there will be separate instance
 // of in-memory telemetry for each of the org, having their own settings. As of now, we only

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

@@ -0,0 +1,80 @@
+package api
+
+import (
+	"github.com/gorilla/mux"
+	baseapp "go.signoz.io/query-service/app"
+	"go.signoz.io/query-service/ee/dao"
+	"go.signoz.io/query-service/ee/interfaces"
+	"go.signoz.io/query-service/ee/license"
+	baseint "go.signoz.io/query-service/interfaces"
+	rules "go.signoz.io/query-service/rules"
+	"go.signoz.io/query-service/version"
+	"net/http"
+)
+
+type APIHandlerOptions struct {
+	QueryBackend   interfaces.QueryBackend
+	AppDB          dao.ModelDao
+	RulesManager   *rules.Manager
+	FeatureFlags   baseint.FeatureLookup
+	LicenseManager *license.Manager
+}
+
+type APIHandler struct {
+	opts APIHandlerOptions
+	baseapp.APIHandler
+}
+
+// NewAPIHandler returns an APIHandler
+func NewAPIHandler(opts APIHandlerOptions) (*APIHandler, error) {
+	baseHandler, err := baseapp.NewAPIHandler(opts.QueryBackend, opts.AppDB, opts.RulesManager)
+	if err != nil {
+		return nil, err
+	}
+
+	ah := &APIHandler{
+		opts:       opts,
+		APIHandler: *baseHandler,
+	}
+	return ah, nil
+}
+
+func (ah *APIHandler) FF() baseint.FeatureLookup {
+	return ah.opts.FeatureFlags
+}
+
+func (ah *APIHandler) RM() *rules.Manager {
+	return ah.opts.RulesManager
+}
+
+func (ah *APIHandler) LM() *license.Manager {
+	return ah.opts.LicenseManager
+}
+
+func (ah *APIHandler) AppDB() dao.ModelDao {
+	return ah.opts.AppDB
+}
+
+// RegisterRoutes registers routes for this handler on the given router
+func (ah *APIHandler) RegisterRoutes(router *mux.Router) {
+	// note: add ee override methods first
+
+	// routes available only in ee version
+	router.HandleFunc("/api/v1/apply/license", baseapp.AdminAccess(ah.applyLicense)).Methods(http.MethodPost)
+	router.HandleFunc("/api/v1/featureFlags", baseapp.OpenAccess(ah.getFeatureFlags)).Methods(http.MethodGet)
+	router.HandleFunc("/api/v1/loginPrecheck", baseapp.OpenAccess(ah.precheckLogin)).Methods(http.MethodGet)
+
+	// paid plans specific routes
+	router.HandleFunc("/api/v1/organization/{org_id}/complete/saml", baseapp.OpenAccess(ah.ReceiveSAML)).Methods(http.MethodPost)
+
+	// base overrides
+	router.HandleFunc("/api/v1/version", baseapp.OpenAccess(ah.getVersion)).Methods(http.MethodGet)
+
+	ah.APIHandler.RegisterRoutes(router)
+
+}
+
+func (ah *APIHandler) getVersion(w http.ResponseWriter, r *http.Request) {
+	version := version.GetVersion()
+	ah.WriteJSON(w, r, map[string]string{"version": version, "eeAvailable": "Y"})
+}

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

@@ -0,0 +1,138 @@
+package api
+
+import (
+	"context"
+	"fmt"
+	"github.com/gorilla/mux"
+	baseauth "go.signoz.io/query-service/auth"
+	"go.signoz.io/query-service/ee/constants"
+	"go.signoz.io/query-service/ee/model"
+	"go.signoz.io/query-service/ee/saml"
+	"go.uber.org/zap"
+	"net/http"
+)
+
+// methods that use user authentication
+// precheckLogin checks if SSO or SAML is available, the check happens
+// when user enters email address in login screen.
+func (aH *APIHandler) precheckLogin(w http.ResponseWriter, r *http.Request) {
+
+	// todo(amol): validate email with org domain
+	// email := r.URL.Query().Get("email")
+
+	// path := r.URL.Query().Get("path")
+
+	// type precheckLoginResponse struct {
+	// 	SSOEnabled   bool   `json:"ssoEnabled"`
+	// 	SAMLEnabled  bool   `json:"samlEnabled"`
+	// 	SAMLLoginUrl string `json:"samlLoginUrl"`
+	// }
+
+	// var org *model.Organization
+	// var apiError *model.ApiError
+
+	// if !aH.IsMultiOrgAvailable {
+	// 	org, apiError = aH.relationalDB.GetSingleOrg(context.Background())
+	// 	if apiError != nil {
+	// 		zap.S().Debugf("[precheckLogin] failed to fetch organization: %v", apiError)
+	// 		RespondError(w, apiError, nil)
+	// 		return
+	// 	}
+	// } else {
+	// 	// todo(amol): read email address from request and determine org using domain
+	// }
+
+	// // todo(amol) just responding dummy data for now
+	// precheckResp := precheckLoginResponse{
+	// 	SSOEnabled:  org.IsSSOEnabled(),
+	// 	SAMLEnabled: org.IsSAMLEnabled(),
+	// }
+
+	// if org.IsSAMLAvailable() {
+	// 	loginURL, err := saml.BuildLoginURLWithOrg(org, path)
+	// 	if err != nil {
+	// 		RespondError(w, &model.ApiError{
+	// 			Typ: model.ErrorInternal,
+	// 			Err: err,
+	// 		}, nil)
+	// 		return
+	// 	}
+	// 	precheckResp.SAMLLoginUrl = loginURL
+	// }
+
+	// aH.WriteJSON(w, r, precheckResp)
+}
+
+func (ah *APIHandler) ReceiveSAML(w http.ResponseWriter, r *http.Request) {
+
+	domainID := mux.Vars(r)["domain_id"]
+	redirectUri := constants.GetSAMLRedirectURL()
+
+	// get org
+	domain, apiError := ah.AppDB().GetOrgDomain(context.Background(), domainID)
+	if apiError != nil {
+		zap.S().Errorf("[ReceiveSAML] failed to fetch organization (%s): %v", domainID, apiError)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "failed to identify user organization, please contact your administrator"), 301)
+		return
+	}
+
+	if err := ah.CheckFeature(model.SSO); err != nil {
+		zap.S().Errorf("[ReceiveSAML] feature unavailable %s in org %s", model.SAML, domainID)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "feature unavailable, please upgrade your billing plan to access this feature"), 301)
+		return
+	}
+
+	sp, err := saml.PrepRequestWithOrg(domain, "")
+	if err != nil {
+		zap.S().Errorf("[ReceiveSAML] failed to prepare saml request for organization (%s): %v", domainID, err)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "failed to send request to SSO, please contact your administrator"), 301)
+		return
+	}
+
+	err = r.ParseForm()
+	if err != nil {
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "failed to authenticate with the SSO provider"), 301)
+		return
+	}
+
+	assertionInfo, err := sp.RetrieveAssertionInfo(r.FormValue("SAMLResponse"))
+	if err != nil {
+		zap.S().Errorf("[ReceiveSAML] failed to retrieve assertion info from  saml response for organization (%s): %v", domainID, err)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "user not found, please contact your administrator"), 301)
+		return
+	}
+
+	if assertionInfo.WarningInfo.InvalidTime {
+		zap.S().Errorf("[ReceiveSAML] expired saml response for organization (%s): %v", domainID, err)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "saml response expired, please contact your administrator"), 301)
+		return
+	}
+
+	if assertionInfo.WarningInfo.NotInAudience {
+		zap.S().Errorf("[ReceiveSAML] NotInAudience error for orgID: %s", domainID)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "this app does not have accesss to SSO provider login"), 301)
+		return
+	}
+
+	email := assertionInfo.NameID
+	firstName := assertionInfo.Values.Get("FirstName")
+	lastName := assertionInfo.Values.Get("LastName")
+
+	userPayload, err := ah.AppDB().FetchOrRegisterSAMLUser(email, firstName, lastName)
+	if err != nil {
+		zap.S().Errorf("[ReceiveSAML] failed to find or register a new user for email %s and org %s", email, domainID)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "failed to authenticate, please contact your administrator"), 301)
+		return
+	}
+
+	tokenStore, err := baseauth.GenerateJWTForUser(&userPayload.User)
+	if err != nil {
+		zap.S().Errorf("[ReceiveSAML] failed to generate access token for email %s and org %s", email, domainID)
+		http.Redirect(w, r, fmt.Sprintf("%s?ssoerror=%s", redirectUri, "failed to login, please contact your administrator"), 301)
+		return
+	}
+	userID := userPayload.User.Id
+	nextPage := fmt.Sprintf("%s?jwt=%s&usr=%s&refreshjwt=%s", redirectUri, tokenStore.AccessJwt, userID, tokenStore.RefreshJwt)
+
+	http.Redirect(w, r, nextPage, 301)
+}

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

@@ -0,0 +1,17 @@
+package api
+
+import (
+	"net/http"
+)
+
+// methods that use user session or jwt to return
+// user specific data
+
+func (ah *APIHandler) getFeatureFlags(w http.ResponseWriter, r *http.Request) {
+	featureSet := ah.LM().GetFeatureFlags()
+	ah.WriteJSON(w, r, featureSet)
+}
+
+func (ah *APIHandler) CheckFeature(f string) error {
+	return ah.FF().CheckFeature(f)
+}

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

@@ -0,0 +1,32 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"go.signoz.io/query-service/ee/model"
+	"net/http"
+)
+
+func (ah *APIHandler) applyLicense(w http.ResponseWriter, r *http.Request) {
+	ctx := context.Background()
+	var l model.License
+
+	if err := json.NewDecoder(r.Body).Decode(&l); err != nil {
+		RespondError(w, model.NewBadRequestError(err), nil)
+		return
+	}
+
+	if l.Key == "" {
+		RespondError(w, model.NewBadRequestError(fmt.Errorf("license key is required")), nil)
+		return
+	}
+
+	license, apiError := ah.LM().Activate(ctx, l.Key)
+	if apiError != nil {
+		RespondError(w, apiError, nil)
+		return
+	}
+
+	ah.Respond(w, license)
+}

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

@@ -0,0 +1,11 @@
+package api
+
+import (
+	baseapp "go.signoz.io/query-service/app"
+	"go.signoz.io/query-service/ee/model"
+	"net/http"
+)
+
+func RespondError(w http.ResponseWriter, apiErr *model.ApiError, data interface{}) {
+	baseapp.RespondError(w, &apiErr.ApiError, data)
+}