[PM-37259] - Update Sso Request Validator

src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyRequirements/RequireSsoPolicyRequirement.cs

@@ -28,13 +28,16 @@ public class RequireSsoPolicyRequirement : IPolicyRequirement
     /// that has the Require SSO policy enabled.
     /// </remarks>
     public bool SsoRequired { get; init; }
-}
 
+    /// <summary>
+    /// Organizations that require SSO login
+    /// </summary>
+    public ICollection<Guid> OrganizationIds { get; init; } = Array.Empty<Guid>();
+}
 
 public class RequireSsoPolicyRequirementFactory(GlobalSettings globalSettings, IFeatureService featureService)
     : BasePolicyRequirementFactory<RequireSsoPolicyRequirement>
 {
-
     public override PolicyType PolicyType => PolicyType.RequireSso;
 
     protected override IEnumerable<OrganizationUserType> ExemptRoles =>
@@ -48,14 +51,18 @@ public override RequireSsoPolicyRequirement Create(IEnumerable<PolicyDetails> po
 
         var acceptedFeatureFlagEnabled = featureService.IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState);
 
+        var ssoRequiredDetails = policyDetails.Where(p => !acceptedFeatureFlagEnabled
+            ? p.OrganizationUserStatus is OrganizationUserStatusType.Confirmed
+            : p.OrganizationUserStatus is OrganizationUserStatusType.Accepted
+                or OrganizationUserStatusType.Confirmed).ToArray();
+
         var result = new RequireSsoPolicyRequirement
         {
             CanUsePasskeyLogin = !policyDetails.Any(p =>
-                p.OrganizationUserStatus is OrganizationUserStatusType.Accepted or OrganizationUserStatusType.Confirmed),
-
-            SsoRequired = policyDetails.Any(p => !acceptedFeatureFlagEnabled
-                ? p.OrganizationUserStatus is OrganizationUserStatusType.Confirmed
-                : p.OrganizationUserStatus is OrganizationUserStatusType.Accepted or OrganizationUserStatusType.Confirmed)
+                p.OrganizationUserStatus is OrganizationUserStatusType.Accepted
+                    or OrganizationUserStatusType.Confirmed),
+            SsoRequired = ssoRequiredDetails.Length > 0,
+            OrganizationIds = [.. ssoRequiredDetails.Select(x => x.OrganizationId)]
         };
 
         return result;

src/Core/Auth/UserFeatures/UserServiceCollectionExtensions.cs

@@ -1,5 +1,4 @@
-using Bit.Core.Auth.Sso;
-using Bit.Core.Auth.UserFeatures.EmergencyAccess.Commands;
+using Bit.Core.Auth.UserFeatures.EmergencyAccess.Commands;
 using Bit.Core.Auth.UserFeatures.EmergencyAccess.Interfaces;
 using Bit.Core.Auth.UserFeatures.Registration;
 using Bit.Core.Auth.UserFeatures.Registration.Implementations;
@@ -34,7 +33,6 @@ public static void AddUserServices(this IServiceCollection services, IGlobalSett
         services.AddWebAuthnLoginCommands();
         services.AddTdeOffboardingPasswordCommands();
         services.AddTwoFactorCommandsQueries();
-        services.AddSsoQueries();
         services.AddUserApiKeyCommands();
     }
 
@@ -92,8 +90,4 @@ private static void AddTwoFactorCommandsQueries(this IServiceCollection services
         services.AddScoped<IResetUserTwoFactorCommand, ResetUserTwoFactorCommand>();
     }
 
-    private static void AddSsoQueries(this IServiceCollection services)
-    {
-        services.AddScoped<IUserSsoOrganizationIdentifierQuery, UserSsoOrganizationIdentifierQuery>();
-    }
 }

src/Identity/IdentityServer/RequestValidators/SsoRequestValidator.cs

@@ -1,7 +1,7 @@
 using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
-using Bit.Core.Auth.Sso;
 using Bit.Core.Entities;
 using Bit.Core.Models.Api;
+using Bit.Core.Repositories;
 using Bit.Identity.IdentityServer.RequestValidationConstants;
 using Duende.IdentityModel;
 using Duende.IdentityServer.Validation;
@@ -12,7 +12,7 @@ namespace Bit.Identity.IdentityServer.RequestValidators;
 /// Validates whether a user is required to authenticate via SSO based on organization policies.
 /// </summary>
 public class SsoRequestValidator(
-    IUserSsoOrganizationIdentifierQuery _userSsoOrganizationIdentifierQuery,
+    IOrganizationRepository _organizationRepository,
     IPolicyRequirementQuery _policyRequirementQuery) : ISsoRequestValidator
 {
     /// <summary>
@@ -26,7 +26,21 @@ public class SsoRequestValidator(
     /// <returns>true if the user can proceed with authentication; false if SSO is required and the user must be redirected to SSO flow.</returns>
     public async Task<bool> ValidateAsync(User user, ValidatedTokenRequest request, CustomValidatorRequestContext context)
     {
-        context.SsoRequired = await RequireSsoAuthenticationAsync(user, request.GrantType);
+        // Check if the user is required to authenticate via SSO. If the user requires SSO, but they are
+        // logging in using an API Key (client_credentials) then they are allowed to bypass the SSO requirement.
+        // If the GrantType is authorization_code or client_credentials we know the user is trying to log in
+        // using the SSO flow so they are allowed to continue.
+        if (request.GrantType is OidcConstants.GrantTypes.AuthorizationCode or OidcConstants.GrantTypes.ClientCredentials)
+        {
+            // SSO is not required for users already using SSO to authenticate which uses the authorization_code grant type,
+            // or logging-in via API key which is the client_credentials grant type.
+            // Allow user to continue request validation
+            context.SsoRequired = false;
+            return true;
+        }
+
+        var requiredSsoRequirement = await _policyRequirementQuery.GetAsyncVNext<RequireSsoPolicyRequirement>(user.Id);
+        context.SsoRequired = requiredSsoRequirement.SsoRequired;
 
         if (!context.SsoRequired)
         {
@@ -41,48 +55,29 @@ public async Task<bool> ValidateAsync(User user, ValidatedTokenRequest request,
         // evaluated, and recovery will have been performed if requested.
         // We will send a descriptive message in these cases so clients can give the appropriate feedback and redirect
         // to /login.
-        if (context.TwoFactorRequired && context.TwoFactorRecoveryRequested)
+        if (context is { TwoFactorRequired: true, TwoFactorRecoveryRequested: true })
         {
-            await SetContextCustomResponseSsoErrorAsync(context, SsoConstants.RequestErrors.SsoTwoFactorRecoveryDescription);
+            await SetContextCustomResponseSsoErrorAsync(context, requiredSsoRequirement, SsoConstants.RequestErrors.SsoTwoFactorRecoveryDescription);
             return false;
         }
 
-        await SetContextCustomResponseSsoErrorAsync(context, SsoConstants.RequestErrors.SsoRequiredDescription);
+        await SetContextCustomResponseSsoErrorAsync(context, requiredSsoRequirement, SsoConstants.RequestErrors.SsoRequiredDescription);
         return false;
     }
 
-    /// <summary>
-    /// Check if the user is required to authenticate via SSO. If the user requires SSO, but they are
-    /// logging in using an API Key (client_credentials) then they are allowed to bypass the SSO requirement.
-    /// If the GrantType is authorization_code or client_credentials we know the user is trying to log in
-    /// using the SSO flow so they are allowed to continue.
-    /// </summary>
-    /// <param name="user">user trying to log in</param>
-    /// <param name="grantType">magic string identifying the grant type requested</param>
-    /// <returns>true if sso required; false if not required or already in process</returns>
-    private async Task<bool> RequireSsoAuthenticationAsync(User user, string grantType)
-    {
-        if (grantType == OidcConstants.GrantTypes.AuthorizationCode ||
-            grantType == OidcConstants.GrantTypes.ClientCredentials)
-        {
-            // SSO is not required for users already using SSO to authenticate which uses the authorization_code grant type,
-            // or logging-in via API key which is the client_credentials grant type.
-            // Allow user to continue request validation
-            return false;
-        }
-
-        // Check if user belongs to any organization with an active SSO policy
-        return (await _policyRequirementQuery.GetAsyncVNext<RequireSsoPolicyRequirement>(user.Id)).SsoRequired;
-    }
-
     /// <summary>
     /// Sets the customResponse in the context with the error result for the SSO validation failure.
     /// </summary>
     /// <param name="context">The validator context to update with error details.</param>
+    /// <param name="requireSsoPolicyRequirement">Require Sso policy requirement for user.</param>
     /// <param name="errorMessage">The error message to return to the client.</param>
-    private async Task SetContextCustomResponseSsoErrorAsync(CustomValidatorRequestContext context, string errorMessage)
+    private async Task SetContextCustomResponseSsoErrorAsync(CustomValidatorRequestContext context, RequireSsoPolicyRequirement requireSsoPolicyRequirement, string errorMessage)
     {
-        var ssoOrganizationIdentifier = await _userSsoOrganizationIdentifierQuery.GetSsoOrganizationIdentifierAsync(context.User.Id);
+        var organization = requireSsoPolicyRequirement.OrganizationIds.Count == 1
+            ? await _organizationRepository.GetByIdAsync(requireSsoPolicyRequirement.OrganizationIds.First())
+            : null;
+
+        var ssoOrganizationIdentifier = organization?.Identifier;
 
         context.ValidationErrorResult = new ValidationResult
         {

test/Core.Test/AdminConsole/OrganizationFeatures/Policies/PolicyRequirements/RequireSsoPolicyRequirementFactoryTests.cs

@@ -149,4 +149,138 @@ public void SsoRequired_PoliciesInAcceptedStateDisabled_Confirmed_ReturnsTrue(
 
         Assert.True(actual.SsoRequired);
     }
+
+    [Theory, BitAutoData]
+    public void OrganizationIds_WithNoPolicies_ReturnsEmpty(
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        var actual = sutProvider.Sut.Create([]);
+
+        Assert.Empty(actual.OrganizationIds);
+    }
+
+    [Theory]
+    [BitAutoData(OrganizationUserStatusType.Accepted)]
+    [BitAutoData(OrganizationUserStatusType.Confirmed)]
+    public void OrganizationIds_PoliciesInAcceptedStateEnabled_SinglePolicy_ReturnsThatOrgId(
+        OrganizationUserStatusType userStatus,
+        Guid organizationId,
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState)
+            .Returns(true);
+
+        var actual = sutProvider.Sut.Create(
+        [
+            new PolicyDetails
+            {
+                OrganizationId = organizationId,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = userStatus
+            }
+        ]);
+
+        Assert.Equal(new[] { organizationId }, actual.OrganizationIds);
+    }
+
+    [Theory, BitAutoData]
+    public void OrganizationIds_PoliciesInAcceptedStateEnabled_MultiplePolicies_ReturnsAllOrgIds(
+        Guid organizationIdA,
+        Guid organizationIdB,
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState)
+            .Returns(true);
+
+        var actual = sutProvider.Sut.Create(
+        [
+            new PolicyDetails
+            {
+                OrganizationId = organizationIdA,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = OrganizationUserStatusType.Confirmed
+            },
+            new PolicyDetails
+            {
+                OrganizationId = organizationIdB,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = OrganizationUserStatusType.Accepted
+            }
+        ]);
+
+        Assert.Equal(2, actual.OrganizationIds.Count);
+        Assert.Contains(organizationIdA, actual.OrganizationIds);
+        Assert.Contains(organizationIdB, actual.OrganizationIds);
+    }
+
+    [Theory, BitAutoData]
+    public void OrganizationIds_PoliciesInAcceptedStateDisabled_Accepted_ReturnsEmpty(
+        Guid organizationId,
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState)
+            .Returns(false);
+
+        var actual = sutProvider.Sut.Create(
+        [
+            new PolicyDetails
+            {
+                OrganizationId = organizationId,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = OrganizationUserStatusType.Accepted
+            }
+        ]);
+
+        Assert.Empty(actual.OrganizationIds);
+    }
+
+    [Theory, BitAutoData]
+    public void OrganizationIds_PoliciesInAcceptedStateDisabled_Confirmed_ReturnsThatOrgId(
+        Guid organizationId,
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState)
+            .Returns(false);
+
+        var actual = sutProvider.Sut.Create(
+        [
+            new PolicyDetails
+            {
+                OrganizationId = organizationId,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = OrganizationUserStatusType.Confirmed
+            }
+        ]);
+
+        Assert.Equal(new[] { organizationId }, actual.OrganizationIds);
+    }
+
+    [Theory]
+    [BitAutoData(OrganizationUserStatusType.Revoked)]
+    [BitAutoData(OrganizationUserStatusType.Invited)]
+    public void OrganizationIds_WithExemptStatus_ReturnsEmpty(
+        OrganizationUserStatusType userStatus,
+        Guid organizationId,
+        SutProvider<RequireSsoPolicyRequirementFactory> sutProvider)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.PoliciesInAcceptedState)
+            .Returns(true);
+
+        var actual = sutProvider.Sut.Create(
+        [
+            new PolicyDetails
+            {
+                OrganizationId = organizationId,
+                PolicyType = PolicyType.RequireSso,
+                OrganizationUserStatus = userStatus
+            }
+        ]);
+
+        Assert.Empty(actual.OrganizationIds);
+    }
 }