Auth: Add OrgRoleMapper service

pkg/login/social/connectors/org_role_mapper.go

@@ -0,0 +1,185 @@
+package connectors
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+const MapperMatchAllOrgID = -1
+
+type OrgRoleMapper struct {
+	cfg        *setting.Cfg
+	logger     log.Logger
+	orgService org.Service
+}
+
+func ProvideOrgRoleMapper(cfg *setting.Cfg, orgService org.Service) *OrgRoleMapper {
+	return &OrgRoleMapper{
+		cfg:        cfg,
+		logger:     log.New("orgrole.mapper"),
+		orgService: orgService,
+	}
+}
+
+func (m *OrgRoleMapper) MapOrgRoles(ctx context.Context, externalOrgs []string, orgMappingSettings []string, directlyMappedRole org.RoleType) (map[int64]org.RoleType, error) {
+	orgMapping := m.splitOrgMappingSettings(ctx, orgMappingSettings)
+
+	userOrgRoles := getMappedOrgRoles(externalOrgs, orgMapping)
+
+	m.handleGlobalOrgMapping(userOrgRoles)
+
+	if len(userOrgRoles) == 0 {
+		userOrgRoles = m.GetDefaultOrgMapping(directlyMappedRole)
+	}
+
+	if directlyMappedRole == "" {
+		m.logger.Debug("No direct role mapping found")
+		return userOrgRoles, nil
+	}
+
+	m.logger.Debug("Direct role mapping found", "role", directlyMappedRole)
+
+	// Merge roles from org mapping `org_mapping` with role from direct mapping
+	for orgID, role := range userOrgRoles {
+		userOrgRoles[orgID] = getTopRole(directlyMappedRole, role)
+	}
+
+	return userOrgRoles, nil
+}
+
+func (m *OrgRoleMapper) GetDefaultOrgMapping(directlyMappedRole org.RoleType) map[int64]org.RoleType {
+	orgRoles := make(map[int64]org.RoleType, 0)
+
+	orgID := int64(1)
+	if m.cfg.AutoAssignOrg && m.cfg.AutoAssignOrgId > 0 {
+		orgID = int64(m.cfg.AutoAssignOrgId)
+	}
+
+	if directlyMappedRole == "" || !directlyMappedRole.IsValid() {
+		orgRoles[orgID] = org.RoleType(m.cfg.AutoAssignOrgRole)
+	} else {
+		orgRoles[orgID] = directlyMappedRole
+	}
+
+	return orgRoles
+}
+
+func (m *OrgRoleMapper) handleGlobalOrgMapping(orgRoles map[int64]org.RoleType) {
+	// No global role mapping => return
+	globalRole, ok := orgRoles[MapperMatchAllOrgID]
+	if !ok {
+		return
+	}
+
+	allOrgIDs, err := m.getAllOrgs()
+	if err != nil {
+		// Prevent resetting assignments
+		orgRoles = map[int64]org.RoleType{}
+		m.logger.Warn("error fetching all orgs, removing org mapping to prevent org sync")
+	}
+
+	// Remove the global role mapping
+	delete(orgRoles, MapperMatchAllOrgID)
+
+	// Global mapping => for all orgs get top role mapping
+	for orgID := range allOrgIDs {
+		orgRoles[orgID] = getTopRole(orgRoles[orgID], globalRole)
+	}
+}
+
+func (m *OrgRoleMapper) splitOrgMappingSettings(ctx context.Context, mappings []string) map[string]map[int64]org.RoleType {
+	res := map[string]map[int64]org.RoleType{}
+
+	for _, v := range mappings {
+		kv := strings.Split(v, ":")
+		if len(kv) > 1 {
+			orgID, err := strconv.Atoi(kv[1])
+			if err != nil && kv[1] != "*" {
+				res, getErr := m.orgService.GetByName(ctx, &org.GetOrgByNameQuery{Name: kv[1]})
+
+				if getErr != nil {
+					// ignore not existing org
+					m.logger.Warn("Could not fetch organization. Skipping.", "mapping", fmt.Sprintf("%v", v))
+					continue
+				}
+				orgID, err = int(res.ID), nil
+			}
+			if kv[1] == "*" {
+				orgID, err = MapperMatchAllOrgID, nil
+			}
+			if err == nil {
+				orga := kv[0]
+				if res[orga] == nil {
+					res[orga] = map[int64]org.RoleType{}
+				}
+
+				if len(kv) > 2 && org.RoleType(kv[2]).IsValid() {
+					res[orga][int64(orgID)] = org.RoleType(kv[2])
+				} else {
+					res[orga][int64(orgID)] = org.RoleViewer
+				}
+			}
+		}
+	}
+
+	return res
+}
+
+func (m *OrgRoleMapper) getAllOrgs() (map[int64]bool, error) {
+	allOrgIDs := map[int64]bool{}
+	allOrgs, err := m.orgService.Search(context.Background(), &org.SearchOrgsQuery{})
+	if err != nil {
+		// In case of error, return no orgs
+		return nil, err
+	}
+
+	for _, org := range allOrgs {
+		allOrgIDs[org.ID] = true
+	}
+	return allOrgIDs, nil
+}
+
+func getMappedOrgRoles(externalOrgs []string, orgMapping map[string]map[int64]org.RoleType) map[int64]org.RoleType {
+	userOrgRoles := map[int64]org.RoleType{}
+
+	if len(orgMapping) == 0 {
+		return nil
+	}
+
+	if orgRoles, ok := orgMapping["*"]; ok {
+		for orgID, role := range orgRoles {
+			userOrgRoles[orgID] = role
+		}
+	}
+
+	for _, org := range externalOrgs {
+		orgRoles, ok := orgMapping[org]
+		if !ok {
+			continue
+		}
+
+		for orgID, role := range orgRoles {
+			userOrgRoles[orgID] = getTopRole(userOrgRoles[orgID], role)
+		}
+	}
+
+	return userOrgRoles
+}
+
+func getTopRole(currRole org.RoleType, otherRole org.RoleType) org.RoleType {
+	if currRole == "" {
+		return otherRole
+	}
+
+	if currRole.Includes(otherRole) {
+		return currRole
+	}
+
+	return otherRole
+}

pkg/login/social/connectors/org_role_mapper_test.go

@@ -0,0 +1,147 @@
+package connectors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/services/org/orgtest"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+func TestExternalOrgRoleMapper_MapOrgRoles(t *testing.T) {
+	testCases := []struct {
+		name               string
+		externalOrgs       []string
+		orgMappingSettings []string
+		directlyMappedRole org.RoleType
+		expected           map[int64]org.RoleType
+	}{
+		{
+			name:               "should return the default mapping when no org mapping settings are provided and directly mapped role is not set",
+			externalOrgs:       []string{},
+			orgMappingSettings: []string{},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{2: org.RoleViewer},
+		},
+		{
+			name:               "should return the default mapping when no org mapping settings are provided",
+			externalOrgs:       []string{},
+			orgMappingSettings: []string{},
+			directlyMappedRole: org.RoleEditor,
+			expected:           map[int64]org.RoleType{2: org.RoleEditor},
+		},
+		{
+			name:               "should return the default mapping when org mapping doesn't match any of the external orgs and no directly mapped role is not provided",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"Second:1:Editor"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{2: org.RoleViewer},
+		},
+		{
+			name:               "should map the higher role if directly mapped role is lower than the role found in the org mapping",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:1:Editor"},
+			directlyMappedRole: org.RoleViewer,
+			expected:           map[int64]org.RoleType{1: org.RoleEditor},
+		},
+		{
+			name:               "should map the directly mapped role if it is higher than the role found in the org mapping",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:1:Viewer"},
+			directlyMappedRole: org.RoleEditor,
+			expected:           map[int64]org.RoleType{1: org.RoleEditor},
+		},
+		{
+			name:               "should map to multiple organizations and set the directly mapped role for those if it is higher than the role found in the org mapping",
+			externalOrgs:       []string{"First", "Second"},
+			orgMappingSettings: []string{"First:1:Viewer", "Second:2:Viewer"},
+			directlyMappedRole: org.RoleEditor,
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleEditor},
+		},
+		{
+			name:               "should map to multiple organizations and set the directly mapped role for those if the directly mapped role is not set",
+			externalOrgs:       []string{"First", "Second"},
+			orgMappingSettings: []string{"First:1:Viewer", "Second:2:Viewer"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{1: org.RoleViewer, 2: org.RoleViewer},
+		},
+		{
+			name:               "should map to all of the organizations if global org mapping is provided and the directly mapped role is set",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:*:Editor"},
+			directlyMappedRole: org.RoleViewer,
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleEditor, 3: org.RoleEditor},
+		},
+		{
+			name:               "should map to all of the organizations if global org mapping is provided and the directly mapped role is not set",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:*:Editor"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleEditor, 3: org.RoleEditor},
+		},
+		{
+			name:               "should map to all of the organizations if global org mapping is provided and the directly mapped role is higher than the role found in the org mappings",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:*:Viewer"},
+			directlyMappedRole: org.RoleEditor,
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleEditor, 3: org.RoleEditor},
+		},
+		{
+			name:               "should map correctly when global org mapping is provided",
+			externalOrgs:       []string{"First", "Second", "Third"},
+			orgMappingSettings: []string{"First:1:Viewer", "*:1:Editor", "Second:2:Viewer"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleViewer},
+		},
+		{
+			name:               "should map correctly and respect wildcard precedence when global org mapping is provided",
+			externalOrgs:       []string{"First", "Second"},
+			orgMappingSettings: []string{"*:1:Editor", "First:1:Viewer", "Second:2:Viewer"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{1: org.RoleEditor, 2: org.RoleViewer},
+		},
+		{
+			name:               "should map directly mapped role when global org mapping is provided and the directly mapped role is higher than the role found in the org mappings",
+			externalOrgs:       []string{"First", "Second", "Third"},
+			orgMappingSettings: []string{"First:1:Viewer", "*:1:Editor", "Second:2:Viewer"},
+			directlyMappedRole: org.RoleAdmin,
+			expected:           map[int64]org.RoleType{1: org.RoleAdmin, 2: org.RoleAdmin},
+		},
+		{
+			name:               "should map correctly and respect the mapping precedence when multiple org mappings are provided for the same org",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:1:Editor", "First:1:Viewer"},
+			directlyMappedRole: "",
+			expected:           map[int64]org.RoleType{1: org.RoleViewer},
+		},
+		{
+			name:               "should map to all organizations when global org mapping is provided and the directly mapped role is higher than the role found in the org mappings",
+			externalOrgs:       []string{"First"},
+			orgMappingSettings: []string{"First:*:Editor"},
+			directlyMappedRole: org.RoleAdmin,
+			expected:           map[int64]org.RoleType{1: org.RoleAdmin, 2: org.RoleAdmin, 3: org.RoleAdmin},
+		},
+	}
+	orgService := orgtest.NewOrgServiceFake()
+	orgService.ExpectedOrgs = []*org.OrgDTO{
+		{Name: "First", ID: 1},
+		{Name: "Second", ID: 2},
+		{Name: "Third", ID: 3},
+	}
+	cfg := setting.NewCfg()
+	cfg.AutoAssignOrg = true
+	cfg.AutoAssignOrgId = 2
+	cfg.AutoAssignOrgRole = string(org.RoleViewer)
+	mapper := ProvideOrgRoleMapper(cfg, orgService)
+
+	for _, tc := range testCases {
+		actual, err := mapper.MapOrgRoles(context.Background(), tc.externalOrgs, tc.orgMappingSettings, tc.directlyMappedRole)
+		require.NoError(t, err)
+
+		assert.EqualValues(t, tc.expected, actual)
+	}
+}