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 = -1
+
 # 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. Default is `false`.
 
+### login_default_org_id
+
+Set the default organization for users when they log in. The default is `-1`.
+
 ### login_hint
 
 Text used as placeholder text on login page for login/username input.

pkg/services/authn/authnimpl/hook/default_org_hook.go

@@ -0,0 +1,74 @@
+package hook
+
+import (
+	"context"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/auth/identity"
+	"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"
+)
+
+type DefaultOrgHook struct {
+	cfg         *setting.Cfg
+	logger      log.Logger
+	userService user.Service
+	orgService  org.Service
+}
+
+// ProvideDefaultOrgHook sets the default org for a user after login.
+func ProvideDefaultOrgHook(cfg *setting.Cfg, userService user.Service, orgService org.Service) *DefaultOrgHook {
+	return &DefaultOrgHook{
+		cfg:         cfg,
+		logger:      log.New("authn.default_org_hook"),
+		userService: userService,
+		orgService:  orgService,
+	}
+}
+
+func (h *DefaultOrgHook) SetDefaultOrg(ctx context.Context,
+	currentIdentity *authn.Identity, r *authn.Request, err error) {
+	if err != nil || h.cfg.LoginDefaultOrgId < 1 || currentIdentity == nil {
+		return
+	}
+
+	namespace, identifier := currentIdentity.GetNamespacedID()
+	if namespace != identity.NamespaceUser {
+		return
+	}
+
+	userID, err := identity.IntIdentifier(namespace, identifier)
+	if err != nil {
+		return
+	}
+
+	if !h.validateUsingOrg(ctx, userID, h.cfg.LoginDefaultOrgId) {
+		return
+	}
+
+	cmd := user.SetUsingOrgCommand{UserID: userID, OrgID: h.cfg.LoginDefaultOrgId}
+	if err := h.userService.SetUsingOrg(ctx, &cmd); err != nil {
+		h.logger.Warn("failed to set default org", "err", err)
+	}
+}
+
+func (h *DefaultOrgHook) validateUsingOrg(ctx context.Context, userID int64, orgID int64) bool {
+	query := org.GetUserOrgListQuery{UserID: userID}
+
+	result, err := h.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/hook/default_org_hook_test.go

@@ -0,0 +1,101 @@
+package hook
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/org"
+	"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"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestDefaultOrgHook_SetDefaultOrg(t *testing.T) {
+	testCases := []struct {
+		name              string
+		defaultOrgSetting int64
+		inputErr          error
+		identity          *authn.Identity
+
+		expectSetUsingOrgCalled bool
+	}{
+		{
+			name:                    "should set default org",
+			defaultOrgSetting:       2,
+			identity:                &authn.Identity{ID: "user:1"},
+			expectSetUsingOrgCalled: true,
+		},
+		{
+			name:                    "should not set default org when default org is not set",
+			defaultOrgSetting:       -1,
+			identity:                &authn.Identity{ID: "user:1"},
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			name:                    "should not set default org when identity is nil",
+			defaultOrgSetting:       -1,
+			identity:                nil,
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			name:                    "should not set default org when error is not nil",
+			defaultOrgSetting:       1,
+			identity:                &authn.Identity{ID: "user:1"},
+			inputErr:                fmt.Errorf("error"),
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			name:                    "should not set default org when identity is not a user",
+			defaultOrgSetting:       2,
+			identity:                &authn.Identity{ID: "service-account:1"},
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			name:                    "should not set default org when error is not nil",
+			defaultOrgSetting:       2,
+			identity:                &authn.Identity{ID: "user:1"},
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			name:                    "should not set default org when user id is not valid",
+			defaultOrgSetting:       2,
+			identity:                &authn.Identity{ID: "user:invalid"},
+			expectSetUsingOrgCalled: false,
+		},
+		{
+			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"},
+			expectSetUsingOrgCalled: false,
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg := setting.NewCfg()
+			cfg.LoginDefaultOrgId = tt.defaultOrgSetting
+
+			userService := &usertest.MockService{}
+			userService.On("SetUsingOrg", mock.Anything, mock.MatchedBy(func(cmd *user.SetUsingOrgCommand) bool {
+				return cmd.UserID == 1 && cmd.OrgID == 2
+			})).Return(nil)
+
+			orgService := &orgtest.FakeOrgService{
+				ExpectedUserOrgDTO: []*org.UserOrgDTO{{OrgID: 2}},
+			}
+
+			hook := ProvideDefaultOrgHook(cfg, userService, orgService)
+
+			hook.SetDefaultOrg(context.Background(), tt.identity, nil, tt.inputErr)
+
+			if tt.expectSetUsingOrgCalled {
+				userService.AssertExpectations(t)
+			} else {
+				userService.AssertNotCalled(t, "SetUsingOrg")
+			}
+		})
+	}
+}

pkg/services/authn/authnimpl/service.go

@@ -22,6 +22,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/auth/identity"
 	"github.com/grafana/grafana/pkg/services/authn"
+	"github.com/grafana/grafana/pkg/services/authn/authnimpl/hook"
 	"github.com/grafana/grafana/pkg/services/authn/authnimpl/sync"
 	"github.com/grafana/grafana/pkg/services/authn/clients"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
@@ -162,6 +163,8 @@ func ProvideService(
 
 	s.RegisterPostAuthHook(rbacSync.SyncPermissionsHook, 120)
 
+	s.RegisterPostLoginHook(hook.ProvideDefaultOrgHook(cfg, userService, orgService).SetDefaultOrg, 10)
+
 	return s
 }
 

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)