Auth: Set the default org after User login

conf/defaults.ini

@@ -473,6 +473,9 @@ auto_assign_org_role = Viewer
 # Require email validation before sign up completes
 verify_email_enabled = false
 
+# Redirect to default OrgId after login
+login_default_org_id =
+
 # Background text for the user field on the login page
 login_hint = email or username
 password_hint = password

conf/sample.ini

@@ -449,6 +449,9 @@
 # Require email validation before sign up completes
 ;verify_email_enabled = false
 
+# Redirect to default OrgId after login
+;login_default_org_id =
+
 # Background text for the user field on the login page
 ;login_hint = email or username
 ;password_hint = password

docs/sources/setup-grafana/configure-grafana/_index.md

@@ -835,6 +835,10 @@ The available options are `Viewer` (default), `Admin`, `Editor`, and `None`. For
 
 Require email validation before sign up completes or when updating a user email address. Default is `false`.
 
+### login_default_org_id
+
+Set the default organization for users when they sign in. The default is `-1`.
+
 ### login_hint
 
 Text used as placeholder text on login page for login/username input.

pkg/services/authn/authnimpl/service.go

@@ -147,7 +147,7 @@ func ProvideService(
 
 	// FIXME (jguer): move to User package
 	userSyncService := sync.ProvideUserSync(userService, userProtectionService, authInfoService, quotaService)
-	orgUserSyncService := sync.ProvideOrgSync(userService, orgService, accessControlService)
+	orgUserSyncService := sync.ProvideOrgSync(userService, orgService, accessControlService, cfg)
 	s.RegisterPostAuthHook(userSyncService.SyncUserHook, 10)
 	s.RegisterPostAuthHook(userSyncService.EnableUserHook, 20)
 	s.RegisterPostAuthHook(orgUserSyncService.SyncOrgRolesHook, 30)
@@ -162,6 +162,8 @@ func ProvideService(
 
 	s.RegisterPostAuthHook(rbacSync.SyncPermissionsHook, 120)
 
+	s.RegisterPostAuthHook(orgUserSyncService.SetDefaultOrgHook, 130)
+
 	return s
 }
 

pkg/services/authn/authnimpl/sync/org_sync.go

@@ -11,16 +11,18 @@ import (
 	"github.com/grafana/grafana/pkg/services/authn"
 	"github.com/grafana/grafana/pkg/services/org"
 	"github.com/grafana/grafana/pkg/services/user"
+	"github.com/grafana/grafana/pkg/setting"
 )
 
-func ProvideOrgSync(userService user.Service, orgService org.Service, accessControl accesscontrol.Service) *OrgSync {
-	return &OrgSync{userService, orgService, accessControl, log.New("org.sync")}
+func ProvideOrgSync(userService user.Service, orgService org.Service, accessControl accesscontrol.Service, cfg *setting.Cfg) *OrgSync {
+	return &OrgSync{userService, orgService, accessControl, cfg, log.New("org.sync")}
 }
 
 type OrgSync struct {
 	userService   user.Service
 	orgService    org.Service
 	accessControl accesscontrol.Service
+	cfg           *setting.Cfg
 
 	log log.Logger
 }
@@ -72,7 +74,7 @@ func (s *OrgSync) SyncOrgRolesHook(ctx context.Context, id *authn.Identity, _ *a
 			// update role
 			cmd := &org.UpdateOrgUserCommand{OrgID: orga.OrgID, UserID: userID, Role: extRole}
 			if err := s.orgService.UpdateOrgUser(ctx, cmd); err != nil {
-				s.log.FromContext(ctx).Error("Failed to update active org user", "id", id.ID, "error", err)
+				ctxLogger.Error("Failed to update active org user", "id", id.ID, "error", err)
 				return err
 			}
 		}
@@ -90,7 +92,7 @@ func (s *OrgSync) SyncOrgRolesHook(ctx context.Context, id *authn.Identity, _ *a
 		cmd := &org.AddOrgUserCommand{UserID: userID, Role: orgRole, OrgID: orgId}
 		err := s.orgService.AddOrgUser(ctx, cmd)
 		if err != nil && !errors.Is(err, org.ErrOrgNotFound) {
-			s.log.FromContext(ctx).Error("Failed to update active org for user", "id", id.ID, "error", err)
+			ctxLogger.Error("Failed to update active org for user", "id", id.ID, "error", err)
 			return err
 		}
 	}
@@ -100,7 +102,7 @@ func (s *OrgSync) SyncOrgRolesHook(ctx context.Context, id *authn.Identity, _ *a
 		ctxLogger.Debug("Removing user's organization membership as part of syncing with OAuth login", "id", id.ID, "orgId", orgID)
 		cmd := &org.RemoveOrgUserCommand{OrgID: orgID, UserID: userID}
 		if err := s.orgService.RemoveOrgUser(ctx, cmd); err != nil {
-			s.log.FromContext(ctx).Error("Failed to remove user from org", "id", id.ID, "orgId", orgID, "error", err)
+			ctxLogger.Error("Failed to remove user from org", "id", id.ID, "orgId", orgID, "error", err)
 			if errors.Is(err, org.ErrLastOrgAdmin) {
 				continue
 			}
@@ -128,3 +130,55 @@ func (s *OrgSync) SyncOrgRolesHook(ctx context.Context, id *authn.Identity, _ *a
 
 	return nil
 }
+
+func (s *OrgSync) SetDefaultOrgHook(ctx context.Context, currentIdentity *authn.Identity, r *authn.Request) error {
+	if s.cfg.LoginDefaultOrgId < 1 || currentIdentity == nil {
+		return nil
+	}
+
+	ctxLogger := s.log.FromContext(ctx)
+
+	namespace, identifier := currentIdentity.GetNamespacedID()
+	if namespace != identity.NamespaceUser {
+		ctxLogger.Debug("Skipping default org sync, not a user", "namespace", namespace)
+		return nil
+	}
+
+	userID, err := identity.IntIdentifier(namespace, identifier)
+	if err != nil {
+		ctxLogger.Debug("Skipping default org sync, invalid ID for identity", "id", currentIdentity.ID, "namespace", namespace, "err", err)
+		return nil
+	}
+
+	if !s.validateUsingOrg(ctx, userID, s.cfg.LoginDefaultOrgId) {
+		ctxLogger.Debug("Skipping default org sync, user is not assigned to org", "id", currentIdentity.ID, "org", s.cfg.LoginDefaultOrgId)
+		return nil
+	}
+
+	cmd := user.SetUsingOrgCommand{UserID: userID, OrgID: s.cfg.LoginDefaultOrgId}
+	if err := s.userService.SetUsingOrg(ctx, &cmd); err != nil {
+		ctxLogger.Warn("Failed to set default org", "id", currentIdentity.ID, "err", err)
+		return err
+	}
+
+	return nil
+}
+
+func (s *OrgSync) validateUsingOrg(ctx context.Context, userID int64, orgID int64) bool {
+	query := org.GetUserOrgListQuery{UserID: userID}
+
+	result, err := s.orgService.GetUserOrgList(ctx, &query)
+	if err != nil {
+		return false
+	}
+
+	// validate that the org id in the list
+	valid := false
+	for _, other := range result {
+		if other.OrgID == orgID {
+			valid = true
+		}
+	}
+
+	return valid
+}

pkg/services/authn/authnimpl/sync/org_sync_test.go

@@ -2,9 +2,11 @@ package sync
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/models/roletype"
@@ -16,6 +18,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/org/orgtest"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/services/user/usertest"
+	"github.com/grafana/grafana/pkg/setting"
 )
 
 func TestOrgSync_SyncOrgRolesHook(t *testing.T) {
@@ -124,3 +127,88 @@ func TestOrgSync_SyncOrgRolesHook(t *testing.T) {
 		})
 	}
 }
+
+func TestOrgSync_SetDefaultOrgHook(t *testing.T) {
+	testCases := []struct {
+		name              string
+		defaultOrgSetting int64
+		identity          *authn.Identity
+		setupMock         func(*usertest.MockService)
+
+		wantErr bool
+	}{
+		{
+			name:              "should set default org",
+			defaultOrgSetting: 2,
+			identity:          &authn.Identity{ID: "user:1"},
+			setupMock: func(userService *usertest.MockService) {
+				userService.On("SetUsingOrg", mock.Anything, mock.MatchedBy(func(cmd *user.SetUsingOrgCommand) bool {
+					return cmd.UserID == 1 && cmd.OrgID == 2
+				})).Return(nil)
+			},
+		},
+		{
+			name:              "should not set default org when default org is not set",
+			defaultOrgSetting: -1,
+			identity:          &authn.Identity{ID: "user:1"},
+		},
+		{
+			name:              "should not set default org when identity is nil",
+			defaultOrgSetting: -1,
+			identity:          nil,
+		},
+		{
+			name:              "should not set default org when identity is not a user",
+			defaultOrgSetting: 2,
+			identity:          &authn.Identity{ID: "service-account:1"},
+		},
+		{
+			name:              "should not set default org when user id is not valid",
+			defaultOrgSetting: 2,
+			identity:          &authn.Identity{ID: "user:invalid"},
+		},
+		{
+			name:              "should not set default org when user is not allowed to use the configured default org",
+			defaultOrgSetting: 3,
+			identity:          &authn.Identity{ID: "user:1"},
+		},
+		{
+			name:              "should return error when the user org update was unsuccessful",
+			defaultOrgSetting: 2,
+			identity:          &authn.Identity{ID: "user:1"},
+			setupMock: func(userService *usertest.MockService) {
+				userService.On("SetUsingOrg", mock.Anything, mock.Anything).Return(fmt.Errorf("error"))
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg := setting.NewCfg()
+			cfg.LoginDefaultOrgId = tt.defaultOrgSetting
+
+			userService := &usertest.MockService{}
+			defer userService.AssertExpectations(t)
+
+			if tt.setupMock != nil {
+				tt.setupMock(userService)
+			}
+
+			orgService := &orgtest.FakeOrgService{
+				ExpectedUserOrgDTO: []*org.UserOrgDTO{{OrgID: 2}},
+			}
+
+			s := &OrgSync{
+				userService:   userService,
+				orgService:    orgService,
+				accessControl: actest.FakeService{},
+				log:           log.NewNopLogger(),
+				cfg:           cfg,
+			}
+
+			if err := s.SetDefaultOrgHook(context.Background(), tt.identity, nil); (err != nil) != tt.wantErr {
+				t.Errorf("OrgSync.SetDefaultOrgHook() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}

pkg/services/user/user.go

@@ -6,6 +6,7 @@ import (
 	"github.com/grafana/grafana/pkg/registry"
 )
 
+//go:generate mockery --name Service --structname MockService --outpkg usertest --filename mock.go --output ./usertest/
 type Service interface {
 	registry.ProvidesUsageStats
 	Create(context.Context, *CreateUserCommand) (*User, error)