Auth: Set the default org after User login (#83918)

* poc

* add logger, skip hook when user is not assigned to default org

* Add tests, move to hook folder

* docs

* Skip for OrgId < 1

* Address feedback

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

* lint

* Move the hook to org_sync.go

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

* Handle the case when GetUserOrgList returns error

---------

Co-authored-by: Christopher Moyer <35463610+chri2547@users.noreply.github.com>
Co-authored-by: Karl Persson <kalle.persson@grafana.com>

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

@@ -3,6 +3,7 @@ package sync
 import (
 	"context"
 	"errors"
+	"fmt"
 	"sort"
 
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -11,16 +12,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 +75,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 +93,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 +103,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 +131,59 @@ 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
+	}
+
+	hasAssignedToOrg, err := s.validateUsingOrg(ctx, userID, s.cfg.LoginDefaultOrgId)
+	if err != nil {
+		ctxLogger.Error("Skipping default org sync, failed to validate user's organizations", "id", currentIdentity.ID, "err", err)
+		return nil
+	}
+
+	if !hasAssignedToOrg {
+		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.Error("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, error) {
+	query := org.GetUserOrgListQuery{UserID: userID}
+
+	result, err := s.orgService.GetUserOrgList(ctx, &query)
+	if err != nil {
+		return false, fmt.Errorf("failed to get user's organizations: %w", err)
+	}
+
+	// validate that the org id in the list
+	for _, other := range result {
+		if other.OrgID == orgID {
+			return true, nil
+		}
+	}
+	return false, nil
+}