From 2e811c5438ba59ff375d6924b8eb06da0a48fe60 Mon Sep 17 00:00:00 2001
From: Misi <mgyongyosi@users.noreply.github.com>
Date: Tue, 11 Jun 2024 14:53:05 +0200
Subject: [PATCH] Chore: Use OrgRoleMapper in Grafana.com client (#89013)

* Use OrgRoleMapper in Grafana.com client

* Clean up
---
 .../social/connectors/grafana_com_oauth.go    | 11 +++---
 .../connectors/grafana_com_oauth_test.go      | 36 ++++++++++++-------
 pkg/services/authn/clients/oauth.go           | 12 -------
 3 files changed, 27 insertions(+), 32 deletions(-)

diff --git a/pkg/login/social/connectors/grafana_com_oauth.go b/pkg/login/social/connectors/grafana_com_oauth.go
index 51964891823..7371b5a2118 100644
--- a/pkg/login/social/connectors/grafana_com_oauth.go
+++ b/pkg/login/social/connectors/grafana_com_oauth.go
@@ -12,7 +12,6 @@ import (
 	"github.com/grafana/grafana/pkg/models/roletype"
 	"github.com/grafana/grafana/pkg/services/auth/identity"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
-	"github.com/grafana/grafana/pkg/services/org"
 	"github.com/grafana/grafana/pkg/services/ssosettings"
 	ssoModels "github.com/grafana/grafana/pkg/services/ssosettings/models"
 	"github.com/grafana/grafana/pkg/services/ssosettings/validation"
@@ -140,17 +139,15 @@ func (s *SocialGrafanaCom) UserInfo(ctx context.Context, client *http.Client, _
 		return nil, fmt.Errorf("Error getting user info: %s", err)
 	}
 
-	// on login we do not want to display the role from the external provider
-	var role roletype.RoleType
-	if !s.info.SkipOrgRoleSync {
-		role = org.RoleType(data.Role)
-	}
 	userInfo := &social.BasicUserInfo{
 		Id:    fmt.Sprintf("%d", data.Id),
 		Name:  data.Name,
 		Login: data.Login,
 		Email: data.Email,
-		Role:  role,
+	}
+
+	if !s.info.SkipOrgRoleSync {
+		userInfo.OrgRoles = s.orgRoleMapper.MapOrgRoles(&MappingConfiguration{strictRoleMapping: false}, nil, roletype.RoleType(data.Role))
 	}
 
 	if !s.isOrganizationMember(data.Orgs) {
diff --git a/pkg/login/social/connectors/grafana_com_oauth_test.go b/pkg/login/social/connectors/grafana_com_oauth_test.go
index b6b082ba43a..a4fdb275e8c 100644
--- a/pkg/login/social/connectors/grafana_com_oauth_test.go
+++ b/pkg/login/social/connectors/grafana_com_oauth_test.go
@@ -12,6 +12,8 @@ import (
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/services/auth/identity"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/services/org/orgtest"
 	ssoModels "github.com/grafana/grafana/pkg/services/ssosettings/models"
 	"github.com/grafana/grafana/pkg/services/ssosettings/ssosettingstests"
 	"github.com/grafana/grafana/pkg/services/user"
@@ -32,7 +34,15 @@ const (
 )
 
 func TestSocialGrafanaCom_UserInfo(t *testing.T) {
-	provider := NewGrafanaComProvider(social.NewOAuthInfo(), &setting.Cfg{}, nil, &ssosettingstests.MockService{}, featuremgmt.WithFeatures())
+	cfg := &setting.Cfg{
+		AutoAssignOrgRole: "Viewer",
+		AutoAssignOrgId:   2,
+	}
+	provider := NewGrafanaComProvider(social.NewOAuthInfo(),
+		cfg,
+		ProvideOrgRoleMapper(cfg, &orgtest.FakeOrgService{}),
+		&ssosettingstests.MockService{},
+		featuremgmt.WithFeatures())
 
 	type conf struct {
 		skipOrgRoleSync bool
@@ -46,27 +56,27 @@ func TestSocialGrafanaCom_UserInfo(t *testing.T) {
 		ExpectedError error
 	}{
 		{
-			Name:         "should return empty role as userInfo when Skip Org Role Sync Enabled",
+			Name:         "should return empty OrgRoles when skip org role sync is enabled",
 			userInfoResp: userResponse,
 			Cfg:          conf{skipOrgRoleSync: true},
 			want: &social.BasicUserInfo{
-				Id:    "1",
-				Name:  "Eric Leijonmarck",
-				Email: "octocat@github.com",
-				Login: "octocat",
-				Role:  "",
+				Id:       "1",
+				Name:     "Eric Leijonmarck",
+				Email:    "octocat@github.com",
+				Login:    "octocat",
+				OrgRoles: map[int64]org.RoleType{},
 			},
 		},
 		{
-			Name:         "should return role as userInfo when Skip Org Role Sync Enabled",
+			Name:         "should return OrgRoles when skip org role sync is disabled",
 			userInfoResp: userResponse,
 			Cfg:          conf{skipOrgRoleSync: false},
 			want: &social.BasicUserInfo{
-				Id:    "1",
-				Name:  "Eric Leijonmarck",
-				Email: "octocat@github.com",
-				Login: "octocat",
-				Role:  "Admin",
+				Id:       "1",
+				Name:     "Eric Leijonmarck",
+				Email:    "octocat@github.com",
+				Login:    "octocat",
+				OrgRoles: map[int64]org.RoleType{2: org.RoleAdmin},
 			},
 		},
 	}
diff --git a/pkg/services/authn/clients/oauth.go b/pkg/services/authn/clients/oauth.go
index bb18f19972d..bad7d9ddc8e 100644
--- a/pkg/services/authn/clients/oauth.go
+++ b/pkg/services/authn/clients/oauth.go
@@ -21,7 +21,6 @@ import (
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/login"
 	"github.com/grafana/grafana/pkg/services/oauthtoken"
-	"github.com/grafana/grafana/pkg/services/org"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/util/errutil"
 )
@@ -166,17 +165,6 @@ func (c *OAuth) Authenticate(ctx context.Context, r *authn.Request) (*authn.Iden
 		return nil, errOAuthEmailNotAllowed.Errorf("provided email is not allowed")
 	}
 
-	// This is required to implement OrgRole mapping for OAuth providers step by step
-	switch c.providerName {
-	case social.GenericOAuthProviderName, social.GitHubProviderName, social.GitlabProviderName,
-		social.OktaProviderName, social.GoogleProviderName, social.AzureADProviderName:
-		// Do nothing, these providers already supports OrgRole mapping
-	default:
-		userInfo.OrgRoles, userInfo.IsGrafanaAdmin, _ = getRoles(c.cfg, func() (org.RoleType, *bool, error) {
-			return userInfo.Role, userInfo.IsGrafanaAdmin, nil
-		})
-	}
-
 	lookupParams := login.UserLookupParams{}
 	allowInsecureEmailLookup := c.settingsProviderSvc.KeyValue("auth", "oauth_allow_insecure_email_lookup").MustBool(false)
 	if allowInsecureEmailLookup {