Fix various bad tenant isolation

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/server/api/console/v1/resolver.go

@@ -410,6 +410,10 @@ func (r *Resolver) AuthzService(ctx context.Context, tenantID gid.TenantID) *aut
 	return GetTenantAuthzService(ctx, r.authzSvc, tenantID)
 }
 
+func (r *Resolver) AuthService(ctx context.Context, tenantID gid.TenantID) *auth.TenantAuthService {
+	return GetTenantAuthService(ctx, r.authSvc, tenantID)
+}
+
 func UnwrapOmittable[T any](field graphql.Omittable[T]) *T {
 	if !field.IsSet() {
 		return nil
@@ -428,6 +432,11 @@ func GetTenantAuthzService(ctx context.Context, authzSvc *authz.Service, tenantI
 	return authzSvc.WithTenant(tenantID)
 }
 
+func GetTenantAuthService(ctx context.Context, authSvc *auth.Service, tenantID gid.TenantID) *auth.TenantAuthService {
+	validateTenantAccess(ctx, tenantID)
+	return authSvc.WithTenant(tenantID)
+}
+
 func validateTenantAccess(ctx context.Context, tenantID gid.TenantID) {
 	access, _ := ctx.Value(userTenantContextKey).(*userTenantAccess)
 

pkg/server/api/console/v1/schema.graphql

@@ -2402,15 +2402,6 @@ type Viewer {
     before: CursorKey
     orderBy: OrganizationOrder
   ): OrganizationConnection! @goField(forceResolver: true)
-
-  invitations(
-    first: Int
-    after: CursorKey
-    last: Int
-    before: CursorKey
-    orderBy: InvitationOrder
-    filter: InvitationFilter
-  ): InvitationConnection! @goField(forceResolver: true)
 }
 
 # Connection Types

pkg/server/api/console/v1/schema/schema.go

@@ -1749,7 +1749,6 @@ type ComplexityRoot struct {
 
 	Viewer struct {
 		ID            func(childComplexity int) int
-		Invitations   func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.InvitationOrder, filter *types.InvitationFilter) int
 		Organizations func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.OrganizationOrder) int
 		User          func(childComplexity int) int
 	}
@@ -2172,7 +2171,6 @@ type VendorServiceResolver interface {
 }
 type ViewerResolver interface {
 	Organizations(ctx context.Context, obj *types.Viewer, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.OrganizationOrder) (*types.OrganizationConnection, error)
-	Invitations(ctx context.Context, obj *types.Viewer, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.InvitationOrder, filter *types.InvitationFilter) (*types.InvitationConnection, error)
 }
 
 type executableSchema struct {
@@ -9422,18 +9420,6 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
 
 		return e.complexity.Viewer.ID(childComplexity), true
 
-	case "Viewer.invitations":
-		if e.complexity.Viewer.Invitations == nil {
-			break
-		}
-
-		args, err := ec.field_Viewer_invitations_args(ctx, rawArgs)
-		if err != nil {
-			return 0, false
-		}
-
-		return e.complexity.Viewer.Invitations(childComplexity, args["first"].(*int), args["after"].(*page.CursorKey), args["last"].(*int), args["before"].(*page.CursorKey), args["orderBy"].(*types.InvitationOrder), args["filter"].(*types.InvitationFilter)), true
-
 	case "Viewer.organizations":
 		if e.complexity.Viewer.Organizations == nil {
 			break
@@ -12142,15 +12128,6 @@ type Viewer {
     before: CursorKey
     orderBy: OrganizationOrder
   ): OrganizationConnection! @goField(forceResolver: true)
-
-  invitations(
-    first: Int
-    after: CursorKey
-    last: Int
-    before: CursorKey
-    orderBy: InvitationOrder
-    filter: InvitationFilter
-  ): InvitationConnection! @goField(forceResolver: true)
 }
 
 # Connection Types
@@ -22840,119 +22817,6 @@ func (ec *executionContext) field_Vendor_services_argsOrderBy(
 	return zeroVal, nil
 }
 
-func (ec *executionContext) field_Viewer_invitations_args(ctx context.Context, rawArgs map[string]any) (map[string]any, error) {
-	var err error
-	args := map[string]any{}
-	arg0, err := ec.field_Viewer_invitations_argsFirst(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["first"] = arg0
-	arg1, err := ec.field_Viewer_invitations_argsAfter(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["after"] = arg1
-	arg2, err := ec.field_Viewer_invitations_argsLast(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["last"] = arg2
-	arg3, err := ec.field_Viewer_invitations_argsBefore(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["before"] = arg3
-	arg4, err := ec.field_Viewer_invitations_argsOrderBy(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["orderBy"] = arg4
-	arg5, err := ec.field_Viewer_invitations_argsFilter(ctx, rawArgs)
-	if err != nil {
-		return nil, err
-	}
-	args["filter"] = arg5
-	return args, nil
-}
-func (ec *executionContext) field_Viewer_invitations_argsFirst(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*int, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("first"))
-	if tmp, ok := rawArgs["first"]; ok {
-		return ec.unmarshalOInt2ᚖint(ctx, tmp)
-	}
-
-	var zeroVal *int
-	return zeroVal, nil
-}
-
-func (ec *executionContext) field_Viewer_invitations_argsAfter(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*page.CursorKey, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("after"))
-	if tmp, ok := rawArgs["after"]; ok {
-		return ec.unmarshalOCursorKey2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋpageᚐCursorKey(ctx, tmp)
-	}
-
-	var zeroVal *page.CursorKey
-	return zeroVal, nil
-}
-
-func (ec *executionContext) field_Viewer_invitations_argsLast(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*int, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("last"))
-	if tmp, ok := rawArgs["last"]; ok {
-		return ec.unmarshalOInt2ᚖint(ctx, tmp)
-	}
-
-	var zeroVal *int
-	return zeroVal, nil
-}
-
-func (ec *executionContext) field_Viewer_invitations_argsBefore(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*page.CursorKey, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("before"))
-	if tmp, ok := rawArgs["before"]; ok {
-		return ec.unmarshalOCursorKey2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋpageᚐCursorKey(ctx, tmp)
-	}
-
-	var zeroVal *page.CursorKey
-	return zeroVal, nil
-}
-
-func (ec *executionContext) field_Viewer_invitations_argsOrderBy(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*types.InvitationOrder, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("orderBy"))
-	if tmp, ok := rawArgs["orderBy"]; ok {
-		return ec.unmarshalOInvitationOrder2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐInvitationOrder(ctx, tmp)
-	}
-
-	var zeroVal *types.InvitationOrder
-	return zeroVal, nil
-}
-
-func (ec *executionContext) field_Viewer_invitations_argsFilter(
-	ctx context.Context,
-	rawArgs map[string]any,
-) (*types.InvitationFilter, error) {
-	ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("filter"))
-	if tmp, ok := rawArgs["filter"]; ok {
-		return ec.unmarshalOInvitationFilter2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐInvitationFilter(ctx, tmp)
-	}
-
-	var zeroVal *types.InvitationFilter
-	return zeroVal, nil
-}
-
 func (ec *executionContext) field_Viewer_organizations_args(ctx context.Context, rawArgs map[string]any) (map[string]any, error) {
 	var err error
 	args := map[string]any{}
@@ -54498,8 +54362,6 @@ func (ec *executionContext) fieldContext_Query_viewer(_ context.Context, field g
 				return ec.fieldContext_Viewer_user(ctx, field)
 			case "organizations":
 				return ec.fieldContext_Viewer_organizations(ctx, field)
-			case "invitations":
-				return ec.fieldContext_Viewer_invitations(ctx, field)
 			}
 			return nil, fmt.Errorf("no field named %q was found under type Viewer", field.Name)
 		},
@@ -71204,69 +71066,6 @@ func (ec *executionContext) fieldContext_Viewer_organizations(ctx context.Contex
 	return fc, nil
 }
 
-func (ec *executionContext) _Viewer_invitations(ctx context.Context, field graphql.CollectedField, obj *types.Viewer) (ret graphql.Marshaler) {
-	fc, err := ec.fieldContext_Viewer_invitations(ctx, field)
-	if err != nil {
-		return graphql.Null
-	}
-	ctx = graphql.WithFieldContext(ctx, fc)
-	defer func() {
-		if r := recover(); r != nil {
-			ec.Error(ctx, ec.Recover(ctx, r))
-			ret = graphql.Null
-		}
-	}()
-	resTmp, err := ec.ResolverMiddleware(ctx, func(rctx context.Context) (any, error) {
-		ctx = rctx // use context from middleware stack in children
-		return ec.resolvers.Viewer().Invitations(rctx, obj, fc.Args["first"].(*int), fc.Args["after"].(*page.CursorKey), fc.Args["last"].(*int), fc.Args["before"].(*page.CursorKey), fc.Args["orderBy"].(*types.InvitationOrder), fc.Args["filter"].(*types.InvitationFilter))
-	})
-	if err != nil {
-		ec.Error(ctx, err)
-		return graphql.Null
-	}
-	if resTmp == nil {
-		if !graphql.HasFieldError(ctx, fc) {
-			ec.Errorf(ctx, "must not be null")
-		}
-		return graphql.Null
-	}
-	res := resTmp.(*types.InvitationConnection)
-	fc.Result = res
-	return ec.marshalNInvitationConnection2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐInvitationConnection(ctx, field.Selections, res)
-}
-
-func (ec *executionContext) fieldContext_Viewer_invitations(ctx context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
-	fc = &graphql.FieldContext{
-		Object:     "Viewer",
-		Field:      field,
-		IsMethod:   true,
-		IsResolver: true,
-		Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
-			switch field.Name {
-			case "totalCount":
-				return ec.fieldContext_InvitationConnection_totalCount(ctx, field)
-			case "edges":
-				return ec.fieldContext_InvitationConnection_edges(ctx, field)
-			case "pageInfo":
-				return ec.fieldContext_InvitationConnection_pageInfo(ctx, field)
-			}
-			return nil, fmt.Errorf("no field named %q was found under type InvitationConnection", field.Name)
-		},
-	}
-	defer func() {
-		if r := recover(); r != nil {
-			err = ec.Recover(ctx, r)
-			ec.Error(ctx, err)
-		}
-	}()
-	ctx = graphql.WithFieldContext(ctx, fc)
-	if fc.Args, err = ec.field_Viewer_invitations_args(ctx, field.ArgumentMap(ec.Variables)); err != nil {
-		ec.Error(ctx, err)
-		return fc, err
-	}
-	return fc, nil
-}
-
 func (ec *executionContext) ___Directive_name(ctx context.Context, field graphql.CollectedField, obj *introspection.Directive) (ret graphql.Marshaler) {
 	fc, err := ec.fieldContext___Directive_name(ctx, field)
 	if err != nil {
@@ -98483,42 +98282,6 @@ func (ec *executionContext) _Viewer(ctx context.Context, sel ast.SelectionSet, o
 				continue
 			}
 
-			out.Concurrently(i, func(ctx context.Context) graphql.Marshaler { return innerFunc(ctx, out) })
-		case "invitations":
-			field := field
-
-			innerFunc := func(ctx context.Context, fs *graphql.FieldSet) (res graphql.Marshaler) {
-				defer func() {
-					if r := recover(); r != nil {
-						ec.Error(ctx, ec.Recover(ctx, r))
-					}
-				}()
-				res = ec._Viewer_invitations(ctx, field, obj)
-				if res == graphql.Null {
-					atomic.AddUint32(&fs.Invalids, 1)
-				}
-				return res
-			}
-
-			if field.Deferrable != nil {
-				dfs, ok := deferred[field.Deferrable.Label]
-				di := 0
-				if ok {
-					dfs.AddField(field)
-					di = len(dfs.Values) - 1
-				} else {
-					dfs = graphql.NewFieldSet([]graphql.CollectedField{field})
-					deferred[field.Deferrable.Label] = dfs
-				}
-				dfs.Concurrently(di, func(ctx context.Context) graphql.Marshaler {
-					return innerFunc(ctx, dfs)
-				})
-
-				// don't run the out.Concurrently() call below
-				out.Values[i] = graphql.Null
-				continue
-			}
-
 			out.Concurrently(i, func(ctx context.Context) graphql.Marshaler { return innerFunc(ctx, out) })
 		default:
 			panic("unknown field " + strconv.Quote(field.Name))

pkg/server/api/console/v1/types/types.go

@@ -2480,5 +2480,4 @@ type Viewer struct {
 	ID            gid.GID                 `json:"id"`
 	User          *User                   `json:"user"`
 	Organizations *OrganizationConnection `json:"organizations"`
-	Invitations   *InvitationConnection   `json:"invitations"`
 }

pkg/server/api/console/v1/v1_resolver.go

@@ -901,7 +901,9 @@ func (r *frameworkConnectionResolver) TotalCount(ctx context.Context, obj *types
 
 // Organization is the resolver for the organization field.
 func (r *invitationResolver) Organization(ctx context.Context, obj *types.Invitation) (*types.Organization, error) {
-	organization, err := r.authzSvc.GetOrganizationByInvitationID(ctx, obj.ID)
+	authz := r.AuthzService(ctx, obj.ID.TenantID())
+
+	organization, err := authz.GetOrganizationByInvitationID(ctx, obj.ID)
 	if err != nil {
 		panic(fmt.Errorf("cannot load organization: %w", err))
 	}
@@ -918,28 +920,12 @@ func (r *invitationConnectionResolver) TotalCount(ctx context.Context, obj *type
 			invitationFilter = coredata.NewInvitationFilter(obj.Filter.Statuses)
 		}
 
-		authzSvc := r.AuthzService(ctx, obj.ParentID.TenantID())
-		count, err := authzSvc.CountOrganizationInvitations(ctx, obj.ParentID, invitationFilter)
+		authz := r.AuthzService(ctx, obj.ParentID.TenantID())
+		count, err := authz.CountOrganizationInvitations(ctx, obj.ParentID, invitationFilter)
 		if err != nil {
 			panic(fmt.Errorf("cannot count organization invitations: %w", err))
 		}
 		return count, nil
-	case *viewerResolver:
-		user := UserFromContext(ctx)
-		if user == nil {
-			panic(fmt.Errorf("no authenticated user"))
-		}
-
-		invitationFilter := coredata.NewInvitationFilter(nil)
-		if obj.Filter != nil {
-			invitationFilter = coredata.NewInvitationFilter(obj.Filter.Statuses)
-		}
-
-		count, err := r.authzSvc.CountUserInvitations(ctx, user.EmailAddress, invitationFilter)
-		if err != nil {
-			panic(fmt.Errorf("cannot count user invitations: %w", err))
-		}
-		return count, nil
 	}
 
 	panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver))
@@ -1090,7 +1076,9 @@ func (r *membershipResolver) AuthMethod(ctx context.Context, obj *types.Membersh
 		return coredata.UserAuthMethodPassword, nil
 	}
 
-	authMethod, err := r.authSvc.GetUserAuthMethod(ctx, coredata.NewScope(obj.UserID.TenantID()), obj.UserID, obj.OrganizationID, session)
+	auth := r.AuthService(ctx, obj.UserID.TenantID())
+
+	authMethod, err := auth.GetUserAuthMethod(ctx, obj.UserID, obj.OrganizationID, session)
 	if err != nil {
 		return "", fmt.Errorf("cannot get user auth method: %w", err)
 	}
@@ -1101,22 +1089,26 @@ func (r *membershipResolver) AuthMethod(ctx context.Context, obj *types.Membersh
 func (r *membershipConnectionResolver) TotalCount(ctx context.Context, obj *types.MembershipConnection) (int, error) {
 	switch obj.Resolver.(type) {
 	case *organizationResolver:
-		authzSvc := r.AuthzService(ctx, obj.ParentID.TenantID())
-		count, err := authzSvc.CountOrganizationMemberships(ctx, obj.ParentID)
+		authz := r.AuthzService(ctx, obj.ParentID.TenantID())
+		count, err := authz.CountOrganizationMemberships(ctx, obj.ParentID)
 		if err != nil {
 			panic(fmt.Errorf("cannot count organization memberships: %w", err))
 		}
+
 		return count, nil
-	default:
-		panic(fmt.Errorf("unknown resolver type for membership connection"))
 	}
+
+	panic(fmt.Errorf("unknown resolver type for membership connection"))
 }
 
 // CreateOrganization is the resolver for the createOrganization field.
 func (r *mutationResolver) CreateOrganization(ctx context.Context, input types.CreateOrganizationInput) (*types.CreateOrganizationPayload, error) {
 	currentUser := UserFromContext(ctx)
 
-	prb := r.proboSvc.WithTenant(gid.NewTenantID())
+	tenantID := gid.NewTenantID()
+
+	prb := r.proboSvc.WithTenant(tenantID)
+	authz := r.authzSvc.WithTenant(tenantID)
 
 	organization, err := prb.Organizations.Create(
 		ctx,
@@ -1128,21 +1120,16 @@ func (r *mutationResolver) CreateOrganization(ctx context.Context, input types.C
 		return nil, fmt.Errorf("cannot create organization: %w", err)
 	}
 
-	err = r.authzSvc.AddUserToOrganization(
+	authz.AddUserToOrganization(
 		ctx,
 		currentUser.ID,
 		organization.ID,
-		string(authz.RoleMember),
+		"MEMBER",
 	)
 	if err != nil {
 		return nil, fmt.Errorf("cannot add user to organization: %w", err)
 	}
 
-	tenantIDs, _ := ctx.Value(userTenantContextKey).(*[]gid.TenantID)
-	*tenantIDs = append(*tenantIDs, organization.ID.TenantID())
-
-	prb = r.ProboService(ctx, organization.ID.TenantID())
-
 	_, err = prb.Peoples.Create(
 		ctx,
 		probo.CreatePeopleRequest{
@@ -1158,6 +1145,10 @@ func (r *mutationResolver) CreateOrganization(ctx context.Context, input types.C
 		return nil, fmt.Errorf("cannot create people: %w", err)
 	}
 
+	// Append tenant to allowed one
+	tenantIDs, _ := ctx.Value(userTenantContextKey).(*[]gid.TenantID)
+	*tenantIDs = append(*tenantIDs, organization.ID.TenantID())
+
 	return &types.CreateOrganizationPayload{
 		OrganizationEdge: types.NewOrganizationEdge(organization, coredata.OrganizationOrderFieldCreatedAt),
 	}, nil
@@ -1917,11 +1908,7 @@ func (r *mutationResolver) GenerateFrameworkStateOfApplicability(ctx context.Con
 // ExportFramework is the resolver for the exportFramework field.
 func (r *mutationResolver) ExportFramework(ctx context.Context, input types.ExportFrameworkInput) (*types.ExportFrameworkPayload, error) {
 	prb := r.ProboService(ctx, input.FrameworkID.TenantID())
-
 	user := UserFromContext(ctx)
-	if user == nil {
-		panic(fmt.Errorf("user not found"))
-	}
 
 	err, exportJobID := prb.Frameworks.RequestExport(
 		ctx,
@@ -2773,11 +2760,7 @@ func (r *mutationResolver) BulkExportDocuments(ctx context.Context, input types.
 	}
 
 	prb := r.ProboService(ctx, input.DocumentIds[0].TenantID())
-
 	user := UserFromContext(ctx)
-	if user == nil {
-		panic(fmt.Errorf("user not found"))
-	}
 
 	options := probo.BulkExportOptions{
 		WithWatermark:  input.WithWatermark,
@@ -3557,15 +3540,11 @@ func (r *mutationResolver) DeleteCustomDomain(ctx context.Context, input types.D
 
 // InitiateDomainVerification is the resolver for the initiateDomainVerification field.
 func (r *mutationResolver) InitiateDomainVerification(ctx context.Context, input types.InitiateDomainVerificationInput) (*types.InitiateDomainVerificationPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	organizationID := input.OrganizationID
 	tenantID := organizationID.TenantID()
 
-	config, err := r.authSvc.InitiateDomainVerification(ctx, tenantID, organizationID, input.EmailDomain)
+	authSvc := r.AuthService(ctx, tenantID)
+	config, err := authSvc.InitiateDomainVerification(ctx, organizationID, input.EmailDomain)
 	if err != nil {
 		return nil, fmt.Errorf("cannot initiate domain verification: %w", err)
 	}
@@ -3584,15 +3563,11 @@ func (r *mutationResolver) InitiateDomainVerification(ctx context.Context, input
 
 // VerifyDomain is the resolver for the verifyDomain field.
 func (r *mutationResolver) VerifyDomain(ctx context.Context, input types.VerifyDomainInput) (*types.VerifyDomainPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	configID := input.ID
 	tenantID := configID.TenantID()
 
-	config, verified, err := r.authSvc.VerifyDomain(ctx, tenantID, configID)
+	authSvc := r.AuthService(ctx, tenantID)
+	config, verified, err := authSvc.VerifyDomain(ctx, configID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot verify domain: %w", err)
 	}
@@ -3609,11 +3584,6 @@ func (r *mutationResolver) VerifyDomain(ctx context.Context, input types.VerifyD
 
 // CreateSAMLConfiguration is the resolver for the createSAMLConfiguration field.
 func (r *mutationResolver) CreateSAMLConfiguration(ctx context.Context, input types.CreateSAMLConfigurationInput) (*types.CreateSAMLConfigurationPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	organizationID := input.OrganizationID
 	tenantID := organizationID.TenantID()
 
@@ -3670,7 +3640,8 @@ func (r *mutationResolver) CreateSAMLConfiguration(ctx context.Context, input ty
 		autoSignupEnabled = *input.AutoSignupEnabled
 	}
 
-	config, err := r.authSvc.WithTenant(tenantID).CreateSAMLConfiguration(ctx, auth.CreateSAMLConfigurationRequest{
+	authSvc := r.AuthService(ctx, tenantID)
+	config, err := authSvc.CreateSAMLConfiguration(ctx, auth.CreateSAMLConfigurationRequest{
 		OrganizationID:     organizationID,
 		EmailDomain:        input.EmailDomain,
 		EnforcementPolicy:  input.EnforcementPolicy,
@@ -3699,15 +3670,11 @@ func (r *mutationResolver) CreateSAMLConfiguration(ctx context.Context, input ty
 
 // UpdateSAMLConfiguration is the resolver for the updateSAMLConfiguration field.
 func (r *mutationResolver) UpdateSAMLConfiguration(ctx context.Context, input types.UpdateSAMLConfigurationInput) (*types.UpdateSAMLConfigurationPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	configID := input.ID
 	tenantID := configID.TenantID()
 
-	updatedConfig, err := r.authSvc.WithTenant(tenantID).UpdateSAMLConfiguration(ctx, auth.UpdateSAMLConfigurationRequest{
+	authSvc := r.AuthService(ctx, tenantID)
+	updatedConfig, err := authSvc.UpdateSAMLConfiguration(ctx, auth.UpdateSAMLConfigurationRequest{
 		ID:                 configID,
 		Enabled:            input.Enabled,
 		EnforcementPolicy:  input.EnforcementPolicy,
@@ -3736,15 +3703,11 @@ func (r *mutationResolver) UpdateSAMLConfiguration(ctx context.Context, input ty
 
 // DeleteSAMLConfiguration is the resolver for the deleteSAMLConfiguration field.
 func (r *mutationResolver) DeleteSAMLConfiguration(ctx context.Context, input types.DeleteSAMLConfigurationInput) (*types.DeleteSAMLConfigurationPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	configID := input.ID
 	tenantID := configID.TenantID()
 
-	err := r.authSvc.WithTenant(tenantID).DeleteSAMLConfiguration(ctx, configID)
+	authSvc := r.AuthService(ctx, tenantID)
+	err := authSvc.DeleteSAMLConfiguration(ctx, configID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot delete SAML configuration: %w", err)
 	}
@@ -3756,15 +3719,11 @@ func (r *mutationResolver) DeleteSAMLConfiguration(ctx context.Context, input ty
 
 // EnableSaml is the resolver for the enableSAML field.
 func (r *mutationResolver) EnableSaml(ctx context.Context, input types.EnableSAMLInput) (*types.EnableSAMLPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	configID := input.ID
 	tenantID := configID.TenantID()
 
-	enabledConfig, err := r.authSvc.WithTenant(tenantID).EnableSAMLConfiguration(ctx, configID)
+	authSvc := r.AuthService(ctx, tenantID)
+	enabledConfig, err := authSvc.EnableSAMLConfiguration(ctx, configID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot enable SAML: %w", err)
 	}
@@ -3780,15 +3739,11 @@ func (r *mutationResolver) EnableSaml(ctx context.Context, input types.EnableSAM
 
 // DisableSaml is the resolver for the disableSAML field.
 func (r *mutationResolver) DisableSaml(ctx context.Context, input types.DisableSAMLInput) (*types.DisableSAMLPayload, error) {
-	user := UserFromContext(ctx)
-	if user == nil {
-		return nil, fmt.Errorf("user not authenticated")
-	}
-
 	configID := input.ID
 	tenantID := configID.TenantID()
 
-	disabledConfig, err := r.authSvc.WithTenant(tenantID).DisableSAMLConfiguration(ctx, configID)
+	authSvc := r.AuthService(ctx, tenantID)
+	disabledConfig, err := authSvc.DisableSAMLConfiguration(ctx, configID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot disable SAML: %w", err)
 	}
@@ -4551,7 +4506,8 @@ func (r *organizationResolver) CustomDomain(ctx context.Context, obj *types.Orga
 func (r *organizationResolver) SamlConfigurations(ctx context.Context, obj *types.Organization) ([]*types.SAMLConfiguration, error) {
 	tenantID := obj.ID.TenantID()
 
-	configs, err := r.authSvc.WithTenant(tenantID).GetSAMLConfigurationsByOrganizationID(ctx, obj.ID)
+	authSvc := r.AuthService(ctx, tenantID)
+	configs, err := authSvc.GetSAMLConfigurationsByOrganizationID(ctx, obj.ID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot load SAML configurations: %w", err)
 	}
@@ -5018,7 +4974,8 @@ func (r *sAMLConfigurationResolver) Organization(ctx context.Context, obj *types
 	tenantID := obj.ID.TenantID()
 	prb := r.ProboService(ctx, tenantID)
 
-	config, err := r.authSvc.WithTenant(tenantID).GetSAMLConfigurationByID(ctx, obj.ID)
+	authSvc := r.AuthService(ctx, tenantID)
+	config, err := authSvc.GetSAMLConfigurationByID(ctx, obj.ID)
 	if err != nil {
 		return nil, fmt.Errorf("cannot load SAML configuration: %w", err)
 	}
@@ -5848,35 +5805,6 @@ func (r *viewerResolver) Organizations(ctx context.Context, obj *types.Viewer, f
 	return types.NewOrganizationConnection(page), nil
 }
 
-// Invitations is the resolver for the invitations field.
-func (r *viewerResolver) Invitations(ctx context.Context, obj *types.Viewer, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.InvitationOrder, filter *types.InvitationFilter) (*types.InvitationConnection, error) {
-	user := UserFromContext(ctx)
-
-	pageOrderBy := page.OrderBy[coredata.InvitationOrderField]{
-		Field:     coredata.InvitationOrderFieldCreatedAt,
-		Direction: page.OrderDirectionDesc,
-	}
-	if orderBy != nil {
-		pageOrderBy = page.OrderBy[coredata.InvitationOrderField]{
-			Field:     orderBy.Field,
-			Direction: orderBy.Direction,
-		}
-	}
-	cursor := types.NewCursor(first, after, last, before, pageOrderBy)
-
-	invitationFilter := coredata.NewInvitationFilter(nil)
-	if filter != nil {
-		invitationFilter = coredata.NewInvitationFilter(filter.Statuses)
-	}
-
-	invitations, err := r.authzSvc.GetUserInvitations(ctx, user.EmailAddress, cursor, invitationFilter)
-	if err != nil {
-		panic(fmt.Errorf("cannot list invitations for user: %w", err))
-	}
-
-	return types.NewInvitationConnection(invitations, r, gid.GID{}, filter), nil
-}
-
 // Asset returns schema.AssetResolver implementation.
 func (r *Resolver) Asset() schema.AssetResolver { return &assetResolver{r} }
 

pkg/server/auth/accept_invitation_handler.go

@@ -36,13 +36,13 @@ type (
 	}
 )
 
-func AcceptInvitationHandler(authSvc *authsvc.Service, authzSvc *authz.Service, authCfg RoutesConfig) http.HandlerFunc {
+func AcceptInvitationHandler(authSvc *authsvc.Service, authzSvc *authz.Service, cookieName string, cookieSecret string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 
 		sessionAuthCfg := session.AuthConfig{
-			CookieName:   authCfg.CookieName,
-			CookieSecret: authCfg.CookieSecret,
+			CookieName:   cookieName,
+			CookieSecret: cookieSecret,
 		}
 
 		errorHandler := session.ErrorHandler{

pkg/server/auth/auth.go

@@ -23,20 +23,18 @@ import (
 	"github.com/getprobo/probo/pkg/filemanager"
 	"github.com/go-chi/chi/v5"
 	"go.gearno.de/kit/log"
-	"go.gearno.de/kit/pg"
 )
 
 type Config struct {
-	Auth             *authsvc.Service
-	Authz            *authz.Service
-	SAML             *authsvc.SAMLService
-	CookieName       string
-	CookieDomain     string
-	SessionDuration  time.Duration
-	CookieSecret     string
-	FileManager      *filemanager.Service
-	PGClient         *pg.Client
-	Logger           *log.Logger
+	Auth            *authsvc.Service
+	Authz           *authz.Service
+	SAML            *authsvc.SAMLService
+	CookieName      string
+	CookieDomain    string
+	SessionDuration time.Duration
+	CookieSecret    string
+	FileManager     *filemanager.Service
+	Logger          *log.Logger
 }
 
 type Server struct {
@@ -46,21 +44,21 @@ type Server struct {
 func NewServer(cfg Config) (*Server, error) {
 	router := chi.NewRouter()
 
-	MountRoutes(
-		router,
-		cfg.Auth,
-		cfg.Authz,
-		cfg.SAML,
-		RoutesConfig{
-			CookieName:      cfg.CookieName,
-			CookieDomain:    cfg.CookieDomain,
-			SessionDuration: cfg.SessionDuration,
-			CookieSecret:    cfg.CookieSecret,
-			FileManager:     cfg.FileManager,
-			PGClient:        cfg.PGClient,
-		},
-		cfg.Logger,
-	)
+	router.Post("/register", SignUpHandler(cfg.Auth, cfg.CookieName, cfg.CookieSecret))
+	router.Post("/login", SignInHandler(cfg.Auth, cfg.CookieName, cfg.CookieSecret))
+	router.Delete("/logout", SignOutHandler(cfg.Auth, cfg.CookieName, cfg.CookieSecret))
+	router.Post("/signup-from-invitation", SignupFromInvitationHandler(cfg.Auth, cfg.CookieName, cfg.CookieSecret))
+	router.Post("/forget-password", ForgetPasswordHandler(cfg.Auth))
+	router.Post("/reset-password", ResetPasswordHandler(cfg.Auth))
+	router.Post("/check-sso", SAMLCheckSSOHandler(cfg.Auth, cfg.Logger))
+	router.Get("/organizations", RequireAuth(cfg.Auth, cfg.Authz, cfg.CookieName, cfg.CookieSecret, ListOrganizationsHandler(cfg.Auth, cfg.Authz)))
+	router.Get("/organizations/{organizationID}/logo", RequireAuth(cfg.Auth, cfg.Authz, cfg.CookieName, cfg.CookieSecret, OrganizationLogoHandler(cfg.Auth, cfg.FileManager)))
+	router.Get("/invitations", RequireAuth(cfg.Auth, cfg.Authz, cfg.CookieName, cfg.CookieSecret, ListInvitationsHandler(cfg.Authz)))
+	router.Post("/invitations/accept", AcceptInvitationHandler(cfg.Auth, cfg.Authz, cfg.CookieName, cfg.CookieSecret))
+
+	router.Get("/saml/login/{samlConfigID}", SAMLLoginHandler(cfg.SAML, cfg.Auth, cfg.Logger))
+	router.Post("/saml/consume", SAMLACSHandler(cfg.SAML, cfg.Auth, cfg.Authz, cfg.CookieName, cfg.CookieSecret, cfg.SessionDuration, cfg.Logger))
+	router.Get("/saml/metadata", SAMLMetadataHandler(cfg.SAML))
 
 	return &Server{
 		router: router,

pkg/server/auth/auth_middleware.go

@@ -0,0 +1,93 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package auth
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	authsvc "github.com/getprobo/probo/pkg/auth"
+	"github.com/getprobo/probo/pkg/authz"
+	"github.com/getprobo/probo/pkg/coredata"
+	"github.com/getprobo/probo/pkg/server/session"
+	"go.gearno.de/kit/httpserver"
+)
+
+type ctxKey struct{ name string }
+
+var (
+	sessionContextKey = &ctxKey{name: "session"}
+	userContextKey    = &ctxKey{name: "user"}
+)
+
+func RequireAuth(
+	authSvc *authsvc.Service,
+	authzSvc *authz.Service,
+	cookieName string,
+	cookieSecret string,
+	next http.HandlerFunc,
+) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+
+		sessionAuthCfg := session.AuthConfig{
+			CookieName:   cookieName,
+			CookieSecret: cookieSecret,
+		}
+
+		errorHandler := session.ErrorHandler{
+			OnCookieError: func(err error) {
+				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
+			},
+			OnParseError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+				session.ClearCookie(w, authCfg)
+				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
+			},
+			OnSessionError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+				session.ClearCookie(w, authCfg)
+				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("session expired"))
+			},
+			OnUserError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
+				session.ClearCookie(w, authCfg)
+				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("user not found"))
+			},
+			OnTenantError: func(err error) {
+				panic(fmt.Errorf("cannot list tenants for user: %w", err))
+			},
+		}
+
+		authResult := session.TryAuth(ctx, w, r, authSvc, authzSvc, sessionAuthCfg, errorHandler)
+		if authResult == nil {
+			httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("authentication required"))
+			return
+		}
+
+		ctx = context.WithValue(ctx, sessionContextKey, authResult.Session)
+		ctx = context.WithValue(ctx, userContextKey, authResult.User)
+
+		next(w, r.WithContext(ctx))
+	}
+}
+
+func SessionFromContext(ctx context.Context) *coredata.Session {
+	session, _ := ctx.Value(sessionContextKey).(*coredata.Session)
+	return session
+}
+
+func UserFromContext(ctx context.Context) *coredata.User {
+	user, _ := ctx.Value(userContextKey).(*coredata.User)
+	return user
+}

pkg/server/auth/forget_password_handler.go

@@ -33,7 +33,7 @@ type (
 	}
 )
 
-func ForgetPasswordHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func ForgetPasswordHandler(authSvc *authsvc.Service) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		var req ForgetPasswordRequest
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {

pkg/server/auth/list_invitations_handler.go

@@ -15,18 +15,12 @@
 package auth
 
 import (
-	"context"
 	"fmt"
 	"net/http"
 
-	authsvc "github.com/getprobo/probo/pkg/auth"
 	"github.com/getprobo/probo/pkg/authz"
-	"github.com/getprobo/probo/pkg/coredata"
 	"github.com/getprobo/probo/pkg/gid"
-	"github.com/getprobo/probo/pkg/page"
-	"github.com/getprobo/probo/pkg/server/session"
 	"go.gearno.de/kit/httpserver"
-	"go.gearno.de/kit/pg"
 )
 
 type (
@@ -35,167 +29,56 @@ type (
 	}
 
 	InvitationResponse struct {
-		ID           gid.GID              `json:"id"`
-		Email        string               `json:"email"`
-		FullName     string               `json:"fullName"`
-		Role         string               `json:"role"`
-		ExpiresAt    string               `json:"expiresAt"`
-		AcceptedAt   *string              `json:"acceptedAt,omitempty"`
-		CreatedAt    string               `json:"createdAt"`
-		Organization OrganizationSummary  `json:"organization"`
+		ID           gid.GID                     `json:"id"`
+		Email        string                      `json:"email"`
+		FullName     string                      `json:"fullName"`
+		Role         string                      `json:"role"`
+		ExpiresAt    string                      `json:"expiresAt"`
+		AcceptedAt   *string                     `json:"acceptedAt,omitempty"`
+		CreatedAt    string                      `json:"createdAt"`
+		Organization OrganizationResponseSummary `json:"organization"`
 	}
 
-	OrganizationSummary struct {
+	OrganizationResponseSummary struct {
 		ID   gid.GID `json:"id"`
 		Name string  `json:"name"`
 	}
 )
 
-// loadOrganizationByID loads an organization by ID without tenant scope
-func loadOrganizationByID(
-	ctx context.Context,
-	conn pg.Conn,
-	orgID gid.GID,
-) (*coredata.Organization, error) {
-	query := `
-SELECT
-	id,
-	tenant_id,
-	name,
-	logo_file_id,
-	horizontal_logo_file_id,
-	description,
-	website_url,
-	email,
-	headquarter_address,
-	custom_domain_id,
-	created_at,
-	updated_at
-FROM
-	authz_organizations
-WHERE
-	id = $1
-`
-
-	row := conn.QueryRow(ctx, query, orgID)
-
-	var org coredata.Organization
-	err := row.Scan(
-		&org.ID,
-		&org.TenantID,
-		&org.Name,
-		&org.LogoFileID,
-		&org.HorizontalLogoFileID,
-		&org.Description,
-		&org.WebsiteURL,
-		&org.Email,
-		&org.HeadquarterAddress,
-		&org.CustomDomainID,
-		&org.CreatedAt,
-		&org.UpdatedAt,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("cannot load organization: %w", err)
-	}
-
-	return &org, nil
-}
-
-func ListInvitationsHandler(authSvc *authsvc.Service, authzSvc *authz.Service, authCfg RoutesConfig) http.HandlerFunc {
+func ListInvitationsHandler(authzSvc *authz.Service) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
+		user := UserFromContext(ctx)
 
-		sessionAuthCfg := session.AuthConfig{
-			CookieName:   authCfg.CookieName,
-			CookieSecret: authCfg.CookieSecret,
-		}
-
-		errorHandler := session.ErrorHandler{
-			OnCookieError: func(err error) {
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
-			},
-			OnParseError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
-			},
-			OnSessionError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("session expired"))
-			},
-			OnUserError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("user not found"))
-			},
-			OnTenantError: func(err error) {
-				panic(fmt.Errorf("cannot list tenants for user: %w", err))
-			},
-		}
-
-		authResult := session.TryAuth(ctx, w, r, authSvc, authzSvc, sessionAuthCfg, errorHandler)
-		if authResult == nil {
-			httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("authentication required"))
-			return
-		}
-
-		// Get pending invitations for the user
-		cursor := page.NewCursor(
-			1000,
-			nil,
-			page.Head,
-			page.OrderBy[coredata.InvitationOrderField]{
-				Field:     coredata.InvitationOrderFieldCreatedAt,
-				Direction: page.OrderDirectionDesc,
-			},
-		)
-
-		invitationFilter := coredata.NewInvitationFilter([]coredata.InvitationStatus{coredata.InvitationStatusPending})
-
-		invitationsPage, err := authzSvc.GetUserInvitations(ctx, authResult.User.EmailAddress, cursor, invitationFilter)
+		invitations, err := authzSvc.GetUserPendingInvitations(ctx, user.EmailAddress)
 		if err != nil {
 			panic(fmt.Errorf("cannot list invitations for user: %w", err))
 		}
 
-		// Build response
 		response := ListInvitationsResponse{
-			Invitations: make([]InvitationResponse, 0, len(invitationsPage.Data)),
+			Invitations: make([]InvitationResponse, 0, len(invitations)),
 		}
 
-		// Load organization data for each invitation
-		err = authCfg.PGClient.WithConn(ctx, func(conn pg.Conn) error {
-			for _, invitation := range invitationsPage.Data {
-				invitationResp := InvitationResponse{
-					ID:        invitation.ID,
-					Email:     invitation.Email,
-					FullName:  invitation.FullName,
-					Role:      invitation.Role,
-					ExpiresAt: invitation.ExpiresAt.Format("2006-01-02T15:04:05Z07:00"),
-					CreatedAt: invitation.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
-				}
-
-				if invitation.AcceptedAt != nil {
-					acceptedAtStr := invitation.AcceptedAt.Format("2006-01-02T15:04:05Z07:00")
-					invitationResp.AcceptedAt = &acceptedAtStr
-				}
-
-				// Load organization details
-				org, err := loadOrganizationByID(ctx, conn, invitation.OrganizationID)
-				if err != nil {
-					// Log error but continue - organization might have been deleted
-					return nil
-				}
-
-				invitationResp.Organization = OrganizationSummary{
-					ID:   org.ID,
-					Name: org.Name,
-				}
-
-				response.Invitations = append(response.Invitations, invitationResp)
+		for _, invitation := range invitations {
+			invitationResp := InvitationResponse{
+				ID:        invitation.ID,
+				Email:     invitation.Email,
+				FullName:  invitation.FullName,
+				Role:      invitation.Role,
+				ExpiresAt: invitation.ExpiresAt.Format("2006-01-02T15:04:05Z07:00"),
+				CreatedAt: invitation.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
+				Organization: OrganizationResponseSummary{
+					ID:   invitation.Organization.ID,
+					Name: invitation.Organization.Name,
+				},
 			}
 
-			return nil
-		})
-		if err != nil {
-			panic(fmt.Errorf("cannot load organization details: %w", err))
+			if invitation.AcceptedAt != nil {
+				acceptedAtStr := invitation.AcceptedAt.Format("2006-01-02T15:04:05Z07:00")
+				invitationResp.AcceptedAt = &acceptedAtStr
+			}
+
+			response.Invitations = append(response.Invitations, invitationResp)
 		}
 
 		httpserver.RenderJSON(w, http.StatusOK, response)

pkg/server/auth/list_organizations_handler.go

@@ -15,20 +15,14 @@
 package auth
 
 import (
-	"context"
-	"errors"
 	"fmt"
 	"net/http"
-	"time"
 
 	authsvc "github.com/getprobo/probo/pkg/auth"
 	"github.com/getprobo/probo/pkg/authz"
 	"github.com/getprobo/probo/pkg/coredata"
-	"github.com/getprobo/probo/pkg/filemanager"
 	"github.com/getprobo/probo/pkg/gid"
-	"github.com/getprobo/probo/pkg/server/session"
 	"go.gearno.de/kit/httpserver"
-	"go.gearno.de/kit/pg"
 )
 
 type (
@@ -54,145 +48,84 @@ const (
 	AuthStatusExpired         AuthenticationStatus = "expired"
 )
 
-// generateLogoURL generates a presigned URL for an organization's logo
-func generateLogoURL(
-	ctx context.Context,
-	fileManager *filemanager.Service,
-	conn pg.Conn,
-	logoFileID *gid.GID,
-) (*string, error) {
-	if logoFileID == nil {
-		return nil, nil
+func buildOrganizationResponse(
+	org *coredata.Organization,
+	accessResult authsvc.AccessResult,
+	sessionData coredata.SessionData,
+) OrganizationResponse {
+	// Generate logo URL path if organization has a logo
+	var logoURL *string
+	if org.LogoFileID != nil {
+		url := fmt.Sprintf("/auth/organizations/%s/logo", org.ID)
+		logoURL = &url
 	}
 
-	var file coredata.File
-	// Load file without scope since we're in auth context (cross-tenant)
-	q := `SELECT bucket_name, file_key, file_name, mime_type, file_size FROM files WHERE id = $1`
-	err := conn.QueryRow(ctx, q, logoFileID).Scan(
-		&file.BucketName,
-		&file.FileKey,
-		&file.FileName,
-		&file.MimeType,
-		&file.FileSize,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("cannot load file: %w", err)
+	orgResponse := OrganizationResponse{
+		ID:      org.ID,
+		Name:    org.Name,
+		LogoURL: logoURL,
 	}
 
-	presignedURL, err := fileManager.GenerateFileUrl(ctx, &file, 1*time.Hour)
-	if err != nil {
-		return nil, fmt.Errorf("cannot generate file URL: %w", err)
+	// User does not have required authentication
+	if !accessResult.Allowed {
+		orgResponse.AuthStatus = AuthStatusUnauthenticated
+
+		switch accessResult.MissingAuth {
+		case authsvc.AuthMethodSAML, authsvc.AuthMethodAny:
+			orgResponse.AuthenticationMethod = "saml"
+			if accessResult.SAMLConfig != nil {
+				orgResponse.LoginURL = fmt.Sprintf("/auth/saml/login/%s", accessResult.SAMLConfig.ID)
+			}
+		case authsvc.AuthMethodPassword:
+			orgResponse.AuthenticationMethod = "password"
+			orgResponse.LoginURL = "/authentication/login?method=password"
+		}
+		return orgResponse
 	}
 
-	return &presignedURL, nil
+	// User has required authentication
+	orgResponse.AuthStatus = AuthStatusAuthenticated
+
+	if sessionData.PasswordAuthenticated {
+		orgResponse.AuthenticationMethod = "password"
+		orgResponse.LoginURL = "/authentication/login?method=password"
+	} else if samlInfo, ok := sessionData.SAMLAuthenticatedOrgs[org.ID.String()]; ok {
+		orgResponse.AuthenticationMethod = "saml"
+		orgResponse.LoginURL = fmt.Sprintf("/auth/saml/login/%s", samlInfo.SAMLConfigID)
+	} else {
+		orgResponse.AuthenticationMethod = "any"
+		orgResponse.LoginURL = "/authentication/login?method=password"
+	}
+
+	return orgResponse
 }
 
-func ListOrganizationsHandler(authSvc *authsvc.Service, authzSvc *authz.Service, authCfg RoutesConfig) http.HandlerFunc {
+func ListOrganizationsHandler(authSvc *authsvc.Service, authzSvc *authz.Service) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
+		user := UserFromContext(ctx)
+		sess := SessionFromContext(ctx)
 
-		sessionAuthCfg := session.AuthConfig{
-			CookieName:   authCfg.CookieName,
-			CookieSecret: authCfg.CookieSecret,
-		}
-
-		errorHandler := session.ErrorHandler{
-			OnCookieError: func(err error) {
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
-			},
-			OnParseError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("invalid session"))
-			},
-			OnSessionError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("session expired"))
-			},
-			OnUserError: func(w http.ResponseWriter, authCfg session.AuthConfig) {
-				session.ClearCookie(w, authCfg)
-				httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("user not found"))
-			},
-			OnTenantError: func(err error) {
-				panic(fmt.Errorf("cannot list tenants for user: %w", err))
-			},
-		}
-
-		authResult := session.TryAuth(ctx, w, r, authSvc, authzSvc, sessionAuthCfg, errorHandler)
-		if authResult == nil {
-			httpserver.RenderError(w, http.StatusUnauthorized, fmt.Errorf("authentication required"))
-			return
-		}
-
-		// Get all organizations for the user (without filtering by authentication state)
-		organizations, err := authzSvc.GetAllUserOrganizations(ctx, authResult.User.ID)
+		organizations, err := authzSvc.GetAllUserOrganizations(ctx, user.ID)
 		if err != nil {
 			panic(fmt.Errorf("cannot list organizations for user: %w", err))
 		}
 
-		// Build response with authentication requirements for each organization
+		orgIDs := make([]gid.GID, len(organizations))
+		for i, org := range organizations {
+			orgIDs[i] = org.ID
+		}
+		accessResults, err := authSvc.CheckOrganizationAccess(ctx, user, orgIDs, sess)
+		if err != nil {
+			panic(fmt.Errorf("cannot check organization access: %w", err))
+		}
+
 		response := ListOrganizationsResponse{
 			Organizations: make([]OrganizationResponse, 0, len(organizations)),
 		}
-
 		for _, org := range organizations {
-			orgResponse := OrganizationResponse{
-				ID:   org.ID,
-				Name: org.Name,
-			}
-
-			// Generate logo URL if available
-			if authCfg.FileManager != nil && authCfg.PGClient != nil {
-				err := authCfg.PGClient.WithConn(ctx, func(conn pg.Conn) error {
-					logoURL, err := generateLogoURL(ctx, authCfg.FileManager, conn, org.LogoFileID)
-					if err != nil {
-						// Log error but don't fail the request
-						return nil
-					}
-					orgResponse.LogoURL = logoURL
-					return nil
-				})
-				if err != nil {
-					// Log error but continue
-				}
-			}
-
-			// Check authentication requirements for this organization
-			err := authSvc.CheckOrganizationAccess(ctx, authResult.User, org.ID, authResult.Session)
-			if err != nil {
-				// User needs additional authentication
-				var errSAMLRequired authsvc.ErrSAMLAuthRequired
-				if errors.As(err, &errSAMLRequired) {
-					orgResponse.AuthenticationMethod = "saml"
-					orgResponse.AuthStatus = AuthStatusUnauthenticated
-					orgResponse.LoginURL = fmt.Sprintf("/auth/saml/login/%s", errSAMLRequired.ConfigID)
-				} else {
-					orgResponse.AuthenticationMethod = "password"
-					orgResponse.AuthStatus = AuthStatusUnauthenticated
-					orgResponse.LoginURL = "/authentication/login?method=password"
-				}
-			} else {
-				// User has proper authentication
-				orgResponse.AuthStatus = AuthStatusAuthenticated
-
-				// Determine which auth method they used
-				if authResult.Session.Data.PasswordAuthenticated {
-					orgResponse.AuthenticationMethod = "password"
-					orgResponse.LoginURL = "/authentication/login?method=password"
-				} else if len(authResult.Session.Data.SAMLAuthenticatedOrgs) > 0 {
-					// Find SAML config for this org
-					orgResponse.AuthenticationMethod = "saml"
-					// Try to find the SAML config ID for login URL
-					if samlInfo, ok := authResult.Session.Data.SAMLAuthenticatedOrgs[org.ID.String()]; ok {
-						orgResponse.LoginURL = fmt.Sprintf("/auth/saml/login/%s", samlInfo.SAMLConfigID)
-					} else {
-						orgResponse.LoginURL = "/authentication/login?method=password"
-					}
-				} else {
-					orgResponse.AuthenticationMethod = "any"
-					orgResponse.LoginURL = "/authentication/login?method=password"
-				}
-			}
-
+			accessResult := accessResults[org.ID]
+			orgResponse := buildOrganizationResponse(org, accessResult, sess.Data)
 			response.Organizations = append(response.Organizations, orgResponse)
 		}
 

pkg/server/auth/organization_logo_handler.go

@@ -0,0 +1,58 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package auth
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"time"
+
+	authsvc "github.com/getprobo/probo/pkg/auth"
+	"github.com/getprobo/probo/pkg/coredata"
+	"github.com/getprobo/probo/pkg/gid"
+	"github.com/go-chi/chi/v5"
+)
+
+func OrganizationLogoHandler(authSvc *authsvc.Service, fileManager interface {
+	GenerateFileUrl(ctx context.Context, file *coredata.File, duration time.Duration) (string, error)
+}) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+		user := UserFromContext(ctx)
+		session := SessionFromContext(ctx)
+
+		organizationIDStr := chi.URLParam(r, "organizationID")
+		organizationID, err := gid.ParseGID(organizationIDStr)
+		if err != nil {
+			http.Error(w, "Invalid organization ID", http.StatusBadRequest)
+			return
+		}
+
+		logoFile, err := authSvc.GetOrganizationLogoFile(ctx, user, organizationID, session)
+		if err != nil {
+			panic(fmt.Errorf("cannot get organization logo: %w", err))
+		}
+
+		presignedURL, err := fileManager.GenerateFileUrl(ctx, logoFile, 1*time.Hour)
+		if err != nil {
+			panic(fmt.Errorf("cannot generate presigned URL: %w", err))
+		}
+
+		w.Header().Set("Cache-Control", "public, max-age=3600")
+
+		http.Redirect(w, r, presignedURL, http.StatusFound)
+	}
+}

pkg/server/auth/reset_password_handler.go

@@ -36,7 +36,7 @@ type (
 	}
 )
 
-func ResetPasswordHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func ResetPasswordHandler(authSvc *authsvc.Service) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		var req ResetPasswordRequest
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {

pkg/server/auth/router.go

@@ -1,60 +0,0 @@
-// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
-//
-// Permission to use, copy, modify, and/or distribute this software for any
-// purpose with or without fee is hereby granted, provided that the above
-// copyright notice and this permission notice appear in all copies.
-//
-// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
-// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-// PERFORMANCE OF THIS SOFTWARE.
-
-package auth
-
-import (
-	"time"
-
-	authsvc "github.com/getprobo/probo/pkg/auth"
-	"github.com/getprobo/probo/pkg/authz"
-	"github.com/getprobo/probo/pkg/filemanager"
-	"github.com/go-chi/chi/v5"
-	"go.gearno.de/kit/log"
-	"go.gearno.de/kit/pg"
-)
-
-type RoutesConfig struct {
-	CookieName      string
-	CookieDomain    string
-	SessionDuration time.Duration
-	CookieSecret    string
-	FileManager     *filemanager.Service
-	PGClient        *pg.Client
-}
-
-func MountRoutes(
-	r chi.Router,
-	authSvc *authsvc.Service,
-	authzSvc *authz.Service,
-	samlSvc *authsvc.SAMLService,
-	authCfg RoutesConfig,
-	logger *log.Logger,
-) {
-	r.Post("/register", SignUpHandler(authSvc, authCfg))
-	r.Post("/login", SignInHandler(authSvc, authCfg))
-	r.Delete("/logout", SignOutHandler(authSvc, authCfg))
-	r.Post("/signup-from-invitation", SignupFromInvitationHandler(authSvc, authCfg))
-	r.Post("/forget-password", ForgetPasswordHandler(authSvc, authCfg))
-	r.Post("/reset-password", ResetPasswordHandler(authSvc, authCfg))
-	r.Post("/check-sso", SAMLCheckSSOHandler(authSvc, logger))
-	r.Get("/organizations", ListOrganizationsHandler(authSvc, authzSvc, authCfg))
-	r.Get("/invitations", ListInvitationsHandler(authSvc, authzSvc, authCfg))
-	r.Post("/invitations/accept", AcceptInvitationHandler(authSvc, authzSvc, authCfg))
-
-	// SAML routes
-	r.Get("/saml/login/{samlConfigID}", SAMLLoginHandler(samlSvc, authSvc, logger))
-	r.Post("/saml/consume", SAMLACSHandler(samlSvc, authSvc, authzSvc, authCfg, logger))
-	r.Get("/saml/metadata", SAMLMetadataHandler(samlSvc))
-}

pkg/server/auth/saml_acs_handler.go

@@ -15,6 +15,7 @@
 package auth
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
 	"time"
@@ -27,11 +28,14 @@ import (
 	"go.gearno.de/kit/log"
 )
 
-func getSessionIDFromCookie(r *http.Request, authCfg RoutesConfig) (gid.GID, error) {
-	cookieValue, err := securecookie.Get(r, securecookie.DefaultConfig(
-		authCfg.CookieName,
-		authCfg.CookieSecret,
-	))
+func getSessionIDFromCookie(r *http.Request, cookieName string, cookieSecret string) (gid.GID, error) {
+	cookieValue, err := securecookie.Get(
+		r,
+		securecookie.DefaultConfig(
+			cookieName,
+			cookieSecret,
+		),
+	)
 	if err != nil {
 		return gid.GID{}, err
 	}
@@ -39,7 +43,7 @@ func getSessionIDFromCookie(r *http.Request, authCfg RoutesConfig) (gid.GID, err
 	return gid.ParseGID(cookieValue)
 }
 
-func SAMLACSHandler(samlSvc *authsvc.SAMLService, authSvc *authsvc.Service, authzSvc *authz.Service, authCfg RoutesConfig, logger *log.Logger) http.HandlerFunc {
+func SAMLACSHandler(samlSvc *authsvc.SAMLService, authSvc *authsvc.Service, authzSvc *authz.Service, cookieName string, cookieSecret string, sessionDuration time.Duration, logger *log.Logger) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 
@@ -68,10 +72,32 @@ func SAMLACSHandler(samlSvc *authsvc.SAMLService, authSvc *authsvc.Service, auth
 			return
 		}
 
-		user, err := authSvc.CreateOrGetSAMLUser(ctx, userInfo.Email, userInfo.FullName, userInfo.SAMLSubject)
+		var existingSession *coredata.Session
+		if existingSessionID, err := getSessionIDFromCookie(r, cookieName, cookieSecret); err == nil {
+			if session, err := authSvc.GetSession(ctx, existingSessionID); err == nil {
+				existingSession = session
+			}
+		}
+
+		session, user, err := authSvc.ProvisionSAMLUser(
+			ctx,
+			userInfo.SAMLConfigID,
+			userInfo.OrganizationID,
+			userInfo.Email,
+			userInfo.FullName,
+			userInfo.SAMLSubject,
+			existingSession,
+			sessionDuration,
+		)
 		if err != nil {
-			logger.ErrorCtx(ctx, "cannot create or get SAML user", log.Error(err))
-			http.Error(w, "cannot create user", http.StatusInternalServerError)
+			var autoSignupDisabledErr *authsvc.ErrSAMLAutoSignupDisabled
+			if errors.As(err, &autoSignupDisabledErr) {
+				logger.WarnCtx(ctx, "SAML auto-signup is disabled")
+				http.Error(w, "User does not exist and auto-signup is disabled for this organization", http.StatusForbidden)
+				return
+			}
+			logger.ErrorCtx(ctx, "cannot provision SAML user", log.Error(err))
+			http.Error(w, "cannot provision user", http.StatusInternalServerError)
 			return
 		}
 
@@ -83,43 +109,11 @@ func SAMLACSHandler(samlSvc *authsvc.SAMLService, authSvc *authsvc.Service, auth
 			return
 		}
 
-		var session *coredata.Session
-		if existingSessionID, err := getSessionIDFromCookie(r, authCfg); err == nil {
-			if existingSession, err := authSvc.GetSession(ctx, existingSessionID); err == nil && existingSession.UserID == user.ID {
-				session = existingSession
-			}
-		}
-
-		if session == nil {
-			session, err = authSvc.CreateSessionForUser(ctx, user.ID, authCfg.SessionDuration)
-			if err != nil {
-				logger.ErrorCtx(ctx, "cannot create session", log.Error(err), log.String("user_id", user.ID.String()))
-				http.Error(w, "cannot create session", http.StatusInternalServerError)
-				return
-			}
-		}
-
-		if session.Data.SAMLAuthenticatedOrgs == nil {
-			session.Data.SAMLAuthenticatedOrgs = make(map[string]coredata.SAMLAuthInfo)
-		}
-		session.Data.SAMLAuthenticatedOrgs[userInfo.OrganizationID.String()] = coredata.SAMLAuthInfo{
-			AuthenticatedAt: time.Now(),
-			SAMLConfigID:    userInfo.SAMLConfigID,
-			SAMLSubject:     userInfo.SAMLSubject,
-		}
-
-		err = authSvc.UpdateSessionData(ctx, session.ID, session.Data)
-		if err != nil {
-			logger.ErrorCtx(ctx, "cannot update session data", log.Error(err), log.String("session_id", session.ID.String()))
-			http.Error(w, "cannot update session", http.StatusInternalServerError)
-			return
-		}
-
 		securecookie.Set(
 			w,
 			securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
+				cookieName,
+				cookieSecret,
 			),
 			session.ID.String(),
 		)

pkg/server/auth/sign_in_handler.go

@@ -47,7 +47,7 @@ type (
 	}
 )
 
-func SignInHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func SignInHandler(authSvc *authsvc.Service, cookieName string, cookieSecret string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 
 		var req SignInRequest
@@ -57,13 +57,13 @@ func SignInHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerF
 		}
 
 		var existingSession *coredata.Session
-		if existingSessionID, err := getSessionIDFromCookie(r, authCfg); err == nil {
+		if existingSessionID, err := getSessionIDFromCookie(r, cookieName, cookieSecret); err == nil {
 			if session, err := authSvc.GetSession(r.Context(), existingSessionID); err == nil {
 				existingSession = session
 			}
 		}
 
-		session, user, err := authSvc.SignInWithExistingSession(r.Context(), req.Email, req.Password, existingSession)
+		session, user, err := authSvc.SignIn(r.Context(), req.Email, req.Password, existingSession)
 		if err != nil {
 			var ErrInvalidCredentials *authsvc.ErrInvalidCredentials
 			if errors.As(err, &ErrInvalidCredentials) {
@@ -77,8 +77,8 @@ func SignInHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerF
 		securecookie.Set(
 			w,
 			securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
+				cookieName,
+				cookieSecret,
 			),
 			session.ID.String(),
 		)

pkg/server/auth/sign_out_handler.go

@@ -24,12 +24,12 @@ import (
 	"go.gearno.de/kit/httpserver"
 )
 
-func SignOutHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func SignOutHandler(authSvc *authsvc.Service, cookieName string, cookieSecret string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 
 		sessionID, err := securecookie.Get(r, securecookie.DefaultConfig(
-			authCfg.CookieName,
-			authCfg.CookieSecret,
+			cookieName,
+			cookieSecret,
 		))
 		if err != nil {
 			httpserver.RenderError(w, http.StatusBadRequest, err)
@@ -48,8 +48,8 @@ func SignOutHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.Handler
 		}
 
 		securecookie.Clear(w, securecookie.DefaultConfig(
-			authCfg.CookieName,
-			authCfg.CookieSecret,
+			cookieName,
+			cookieSecret,
 		))
 
 		w.Header().Set("Clear-Site-Data", "*")

pkg/server/auth/sign_up_handler.go

@@ -37,7 +37,7 @@ type (
 	}
 )
 
-func SignUpHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func SignUpHandler(authSvc *authsvc.Service, cookieName string, cookieSecret string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		var req SignUpRequest
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -70,8 +70,8 @@ func SignUpHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerF
 		securecookie.Set(
 			w,
 			securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
+				cookieName,
+				cookieSecret,
 			),
 			session.ID.String(),
 		)

pkg/server/auth/signup_from_invitation_handler.go

@@ -35,7 +35,7 @@ type (
 	}
 )
 
-func SignupFromInvitationHandler(authSvc *authsvc.Service, authCfg RoutesConfig) http.HandlerFunc {
+func SignupFromInvitationHandler(authSvc *authsvc.Service, cookieName string, cookieSecret string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		var req SignupFromInvitationRequest
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -52,8 +52,8 @@ func SignupFromInvitationHandler(authSvc *authsvc.Service, authCfg RoutesConfig)
 		securecookie.Set(
 			w,
 			securecookie.DefaultConfig(
-				authCfg.CookieName,
-				authCfg.CookieSecret,
+				cookieName,
+				cookieSecret,
 			),
 			session.ID.String(),
 		)

pkg/server/server.go

@@ -109,7 +109,6 @@ func NewServer(cfg Config) (*Server, error) {
 		SessionDuration: cfg.ConsoleAuth.SessionDuration,
 		CookieSecret:    cfg.ConsoleAuth.CookieSecret,
 		FileManager:     cfg.FileManager,
-		PGClient:        cfg.PGClient,
 		Logger:          cfg.Logger.Named("auth"),
 	})
 	if err != nil {

pkg/server/session/session.go

@@ -103,15 +103,30 @@ func TryAuth(
 	allowedTenantIDs := make([]gid.TenantID, 0, len(organizations))
 	authErrors := make(map[gid.TenantID]error)
 
+	// Extract organization IDs for batch check
+	orgIDs := make([]gid.GID, len(organizations))
+	for i, org := range organizations {
+		orgIDs[i] = org.ID
+	}
+
+	// Batch check access to all organizations in a single query
+	accessResults, err := authSvc.CheckOrganizationAccess(ctx, user, orgIDs, session)
+	if err != nil {
+		if errorHandler.OnTenantError != nil {
+			errorHandler.OnTenantError(err)
+		}
+		return nil
+	}
+
+	// Process results
 	for _, org := range organizations {
-		// Check if user has the required authentication for this organization
-		err := authSvc.CheckOrganizationAccess(ctx, user, org.ID, session)
-		if err == nil {
+		result := accessResults[org.ID]
+		if result.Allowed {
 			// User has proper authentication for this org
 			allowedTenantIDs = append(allowedTenantIDs, org.ID.TenantID())
 		} else {
 			// Store the authentication error for later use
-			authErrors[org.ID.TenantID()] = err
+			authErrors[org.ID.TenantID()] = result.ToError(authSvc.BaseURL())
 		}
 	}