JWT: Add support for assigning org roles (#54277)

* feat: allow jwt role to be set

* chore: update documentation

* fix: cr suggestions

* fix: lint issues

* respect org auto assign and default org ID

* add server admin to devenv

Co-authored-by: jguer <joao.guerreiro@grafana.com>

pkg/services/contexthandler/auth_jwt.go

@@ -2,15 +2,21 @@ package contexthandler
 
 import (
 	"errors"
+	"fmt"
 	"net/http"
 
 	"github.com/grafana/grafana/pkg/login"
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/org"
 	"github.com/grafana/grafana/pkg/services/user"
+	"github.com/jmespath/go-jmespath"
 )
 
-const InvalidJWT = "Invalid JWT"
-const UserNotFound = "User not found"
+const (
+	InvalidJWT   = "Invalid JWT"
+	InvalidRole  = "Invalid Role"
+	UserNotFound = "User not found"
+)
 
 func (h *ContextHandler) initContextWithJWT(ctx *models.ReqContext, orgId int64) bool {
 	if !h.Cfg.JWTAuthEnabled || h.Cfg.JWTAuthHeaderName == "" {
@@ -45,6 +51,7 @@ func (h *ContextHandler) initContextWithJWT(ctx *models.ReqContext, orgId int64)
 	extUser := &models.ExternalUserInfo{
 		AuthModule: "jwt",
 		AuthId:     sub,
+		OrgRoles:   map[int64]org.RoleType{},
 	}
 
 	if key := h.Cfg.JWTAuthUsernameClaim; key != "" {
@@ -60,6 +67,31 @@ func (h *ContextHandler) initContextWithJWT(ctx *models.ReqContext, orgId int64)
 		extUser.Name = name
 	}
 
+	role, grafanaAdmin := h.extractJWTRoleAndAdmin(claims)
+	if h.Cfg.JWTAuthRoleAttributeStrict && !role.IsValid() {
+		ctx.Logger.Debug("Extracted Role is invalid")
+		ctx.JsonApiErr(http.StatusForbidden, InvalidRole, nil)
+		return true
+	}
+
+	if role.IsValid() {
+		var orgID int64
+		if h.Cfg.AutoAssignOrg && h.Cfg.AutoAssignOrgId > 0 {
+			orgID = int64(h.Cfg.AutoAssignOrgId)
+			ctx.Logger.Debug("The user has a role assignment and organization membership is auto-assigned",
+				"role", role, "orgId", orgID)
+		} else {
+			orgID = int64(1)
+			ctx.Logger.Debug("The user has a role assignment and organization membership is not auto-assigned",
+				"role", role, "orgId", orgID)
+		}
+
+		extUser.OrgRoles[orgID] = role
+		if h.Cfg.JWTAuthAllowAssignGrafanaAdmin {
+			extUser.IsGrafanaAdmin = &grafanaAdmin
+		}
+	}
+
 	if query.Login == "" && query.Email == "" {
 		ctx.Logger.Debug("Failed to get an authentication claim from JWT")
 		ctx.JsonApiErr(http.StatusUnauthorized, InvalidJWT, err)
@@ -105,3 +137,52 @@ func (h *ContextHandler) initContextWithJWT(ctx *models.ReqContext, orgId int64)
 
 	return true
 }
+
+const roleGrafanaAdmin = "GrafanaAdmin"
+
+func (h *ContextHandler) extractJWTRoleAndAdmin(claims map[string]interface{}) (org.RoleType, bool) {
+	if h.Cfg.JWTAuthRoleAttributePath == "" {
+		return "", false
+	}
+
+	role, err := searchClaimsForStringAttr(h.Cfg.JWTAuthRoleAttributePath, claims)
+	if err != nil || role == "" {
+		return "", false
+	}
+
+	if role == roleGrafanaAdmin {
+		return org.RoleAdmin, true
+	}
+	return org.RoleType(role), false
+}
+
+func searchClaimsForAttr(attributePath string, claims map[string]interface{}) (interface{}, error) {
+	if attributePath == "" {
+		return "", errors.New("no attribute path specified")
+	}
+
+	if len(claims) == 0 {
+		return "", errors.New("empty claims provided")
+	}
+
+	val, err := jmespath.Search(attributePath, claims)
+	if err != nil {
+		return "", fmt.Errorf("failed to search claims with provided path: %q: %w", attributePath, err)
+	}
+
+	return val, nil
+}
+
+func searchClaimsForStringAttr(attributePath string, claims map[string]interface{}) (string, error) {
+	val, err := searchClaimsForAttr(attributePath, claims)
+	if err != nil {
+		return "", err
+	}
+
+	strVal, ok := val.(string)
+	if ok {
+		return strVal, nil
+	}
+
+	return "", nil
+}