[PM-33043] Refactor PolicyService, CipherService, and TwoFactorAuthenticationValidator

src/Core/AdminConsole/Services/Implementations/PolicyService.cs

@@ -8,6 +8,7 @@
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
+using Bit.Core.Models.Data.Organizations;
 using Bit.Core.Models.Data.Organizations.OrganizationUsers;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -90,13 +91,35 @@ private async Task<IEnumerable<OrganizationUserPolicyDetails>> QueryOrganization
     {
         var organizationUserPolicyDetails = await _organizationUserRepository.GetByUserIdWithPolicyDetailsAsync(userId, policyType);
         var excludedUserTypes = GetUserTypesExcludedFromPolicy(policyType);
-        var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
-        return organizationUserPolicyDetails.Where(o =>
-            (!orgAbilities.TryGetValue(o.OrganizationId, out var orgAbility) || orgAbility.UsePolicies) &&
-            o.PolicyEnabled &&
-            !excludedUserTypes.Contains(o.OrganizationUserType) &&
-            o.OrganizationUserStatus >= minStatus &&
-            !o.IsProvider);
+
+        var filteredPolicyDetails = organizationUserPolicyDetails
+            .Where(o => !o.IsProvider)
+            .Where(o => o.OrganizationUserStatus >= minStatus)
+            .Where(o => !excludedUserTypes.Contains(o.OrganizationUserType))
+            .Where(o => o.PolicyEnabled)
+            .ToList();
+
+        var orgAbilities = await GetOrganizationAbilitiesAsync(filteredPolicyDetails);
+
+        return filteredPolicyDetails.Where(userPolicyDetails =>
+        {
+            if (orgAbilities.TryGetValue(userPolicyDetails.OrganizationId, out var orgAbility) && !orgAbility.UsePolicies)
+            {
+                return false;
+            }
+
+            return true;
+        });
+    }
+
+    private async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync(List<OrganizationUserPolicyDetails> filteredPolicyDetails)
+    {
+        var orgIds = filteredPolicyDetails
+            .Select(o => o.OrganizationId)
+            .Distinct()
+            .ToList();
+        var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync(orgIds);
+        return orgAbilities;
     }
 
     private OrganizationUserType[] GetUserTypesExcludedFromPolicy(PolicyType policyType)

src/Core/Vault/Services/Implementations/CipherService.cs

@@ -10,6 +10,7 @@
 using Bit.Core.Billing.Pricing;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
+using Bit.Core.Models.Data.Organizations;
 using Bit.Core.Platform.Push;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -1153,11 +1154,15 @@ private async Task<List<T>> FilterCiphersByDeletePermission<T>(
         Guid userId) where T : CipherDetails
     {
         var user = await _userService.GetUserByIdAsync(userId);
-        var organizationAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
 
-        var filteredCiphers = ciphers
+        var groupedCiphers = ciphers
             .Where(c => cipherIdsSet.Contains(c.Id))
             .GroupBy(c => c.OrganizationId)
+            .ToList();
+
+        var organizationAbilities = await GetOrganizationAbilitiesAsync(groupedCiphers);
+
+        var filteredCiphers = groupedCiphers
             .SelectMany(group =>
             {
                 var organizationAbility = group.Key.HasValue &&
@@ -1170,4 +1175,16 @@ private async Task<List<T>> FilterCiphersByDeletePermission<T>(
 
         return filteredCiphers;
     }
+
+    private async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync<T>(IEnumerable<IGrouping<Guid?, T>> groupedCiphers) where T : CipherDetails
+    {
+        var organizationIds = groupedCiphers
+            .Select(group => group.Key)
+            .Where(id => id.HasValue)
+            .Select(id => id!.Value)
+            .ToList();
+
+        var organizationAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync(organizationIds);
+        return organizationAbilities;
+    }
 }

src/Identity/IdentityServer/RequestValidators/TwoFactorAuthenticationValidator.cs

@@ -60,7 +60,7 @@ await _userManager.GetTwoFactorEnabledAsync(user) &&
         var orgs = (await _currentContext.OrganizationMembershipAsync(_organizationUserRepository, user.Id)).ToList();
         if (orgs.Count > 0)
         {
-            var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
+            var orgAbilities = await GetOrganizationAbilitiesAsync(orgs);
             var twoFactorOrgs = orgs.Where(o => OrgUsing2fa(orgAbilities, o.Id));
             if (twoFactorOrgs.Any())
             {
@@ -73,6 +73,13 @@ await _userManager.GetTwoFactorEnabledAsync(user) &&
         return new Tuple<bool, Organization>(individualRequired || firstEnabledOrg != null, firstEnabledOrg);
     }
 
+    private async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync(List<CurrentContextOrganization> orgs)
+    {
+        var organizationIds = orgs.Select(organization => organization.Id).ToList();
+        var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync(organizationIds);
+        return orgAbilities;
+    }
+
     public async Task<Dictionary<string, object>> BuildTwoFactorResultAsync(User user, Organization organization)
     {
         var enabledProviders = await GetEnabledTwoFactorProvidersAsync(user, organization);

test/Core.Test/AdminConsole/Services/PolicyServiceTests.cs

@@ -7,6 +7,7 @@
 using Bit.Core.AdminConsole.Services.Implementations;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
+using Bit.Core.Models.Data.Organizations;
 using Bit.Core.Models.Data.Organizations.OrganizationUsers;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -155,6 +156,38 @@ public async Task GetMasterPasswordPolicyForUserAsync_WithFeatureFlagDisabled_Ev
         await sutProvider.GetDependency<IPolicyRequirementQuery>().DidNotReceive().GetAsync<MasterPasswordPolicyRequirement>(user.Id);
     }
 
+    [Theory, BitAutoData]
+    public async Task GetPoliciesApplicableToUserAsync_OnlyFetchesAbilitiesForFilteredOrgs(
+        Guid userId, SutProvider<PolicyService> sutProvider)
+    {
+        var includedOrgId = Guid.NewGuid();
+        var excludedOrgId = Guid.NewGuid(); // filtered out because IsProvider = true
+
+        sutProvider.GetDependency<IOrganizationUserRepository>()
+            .GetByUserIdWithPolicyDetailsAsync(userId, PolicyType.DisableSend)
+            .Returns(new List<OrganizationUserPolicyDetails>
+            {
+                new() { OrganizationId = includedOrgId, PolicyType = PolicyType.DisableSend, PolicyEnabled = true, OrganizationUserType = OrganizationUserType.User, OrganizationUserStatus = OrganizationUserStatusType.Invited, IsProvider = false },
+                new() { OrganizationId = excludedOrgId, PolicyType = PolicyType.DisableSend, PolicyEnabled = true, OrganizationUserType = OrganizationUserType.User, OrganizationUserStatus = OrganizationUserStatusType.Invited, IsProvider = true }
+            });
+
+        sutProvider.GetDependency<IApplicationCacheService>()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns(new Dictionary<Guid, OrganizationAbility>
+            {
+                { includedOrgId, new OrganizationAbility { Id = includedOrgId, UsePolicies = true } }
+            });
+
+        await sutProvider.Sut.GetPoliciesApplicableToUserAsync(userId, PolicyType.DisableSend, OrganizationUserStatusType.Invited);
+
+        // Assert - only the non-provider org ID should be requested
+        await sutProvider.GetDependency<IApplicationCacheService>()
+            .Received(1)
+            .GetOrganizationAbilitiesAsync(Arg.Is<IEnumerable<Guid>>(ids =>
+                ids.Contains(includedOrgId) &&
+                !ids.Contains(excludedOrgId)));
+    }
+
     private static void SetupOrg(SutProvider<PolicyService> sutProvider, Guid organizationId, Organization organization)
     {
         sutProvider.GetDependency<IOrganizationRepository>()

test/Core.Test/Vault/Services/CipherServiceTests.cs

@@ -1076,7 +1076,7 @@ public async Task RestoreManyAsync_WithManagePermission_RestoresCiphers(
             .GetUserByIdAsync(restoringUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {
@@ -1138,7 +1138,7 @@ public async Task RestoreManyAsync_WithoutManagePermission_DoesNotRestoreCiphers
             .GetUserByIdAsync(restoringUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {
@@ -1906,7 +1906,7 @@ public async Task DeleteManyAsync_WithoutManagePermission_DoesNotDeleteCiphers(
             .GetUserByIdAsync(deletingUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {
@@ -1954,7 +1954,7 @@ public async Task DeleteManyAsync_WithManagePermission_DeletesCiphers(
             .GetUserByIdAsync(deletingUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {
@@ -1980,6 +1980,52 @@ await sutProvider.GetDependency<IPushNotificationService>()
             .PushSyncCiphersAsync(deletingUserId);
     }
 
+    [Theory]
+    [OrganizationCipherCustomize]
+    [BitAutoData]
+    public async Task DeleteManyAsync_WithOrgCipherNotFoundInCache_ThrowsNotFoundException(
+        Guid deletingUserId, List<CipherDetails> ciphers, User user, SutProvider<CipherService> sutProvider)
+    {
+        var targetOrgId = Guid.NewGuid();
+        var orgIdNotInCache = Guid.NewGuid();
+        var cipherDetailsNotInCache = new CipherDetails { Id = Guid.NewGuid(), OrganizationId = orgIdNotInCache, Manage = true };
+
+        foreach (var cipher in ciphers)
+        {
+            cipher.OrganizationId = targetOrgId;
+            cipher.Manage = true;
+        }
+
+        var cipherIds = ciphers.Concat([cipherDetailsNotInCache]).Select(c => c.Id).ToArray();
+
+        var allCiphers = ciphers.Concat([cipherDetailsNotInCache]).ToList();
+
+        sutProvider.GetDependency<ICipherRepository>()
+            .GetManyByUserIdAsync(deletingUserId)
+            .Returns(allCiphers);
+        sutProvider.GetDependency<IUserService>()
+            .GetUserByIdAsync(deletingUserId)
+            .Returns(user);
+        sutProvider.GetDependency<IApplicationCacheService>()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns(new Dictionary<Guid, OrganizationAbility>
+            {
+                { targetOrgId, new OrganizationAbility { Id = targetOrgId, LimitItemDeletion = true } }
+            });
+
+        // Assert
+        var exception = await Assert.ThrowsAsync<Exception>(() =>
+             sutProvider.Sut.DeleteManyAsync(cipherIds, deletingUserId));
+
+        Assert.Contains("Cipher does not belong to the input organization.", exception.Message);
+
+        await sutProvider.GetDependency<IApplicationCacheService>()
+             .Received(1)
+             .GetOrganizationAbilitiesAsync(Arg.Is<IEnumerable<Guid>>(ids =>
+                 ids.Contains(targetOrgId) &&
+                 ids.Contains(orgIdNotInCache)));
+    }
+
     [Theory]
     [BitAutoData]
     public async Task SoftDeleteAsync_WithPersonalCipherOwner_SoftDeletesCipher(
@@ -2253,7 +2299,7 @@ public async Task SoftDeleteManyAsync_WithoutManagePermission_DoesNotDeleteCiphe
             .GetUserByIdAsync(deletingUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {
@@ -2302,7 +2348,7 @@ public async Task SoftDeleteManyAsync_WithManagePermission_SoftDeletesCiphers(
             .GetUserByIdAsync(deletingUserId)
             .Returns(user);
         sutProvider.GetDependency<IApplicationCacheService>()
-            .GetOrganizationAbilitiesAsync()
+            .GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>
             {
                 {

test/Identity.Test/IdentityServer/TwoFactorAuthenticationValidatorTests.cs

@@ -135,7 +135,7 @@ public async void RequiresTwoFactorAsync_IndividualFalse_OrganizationRequired_Re
         _currentContext.OrganizationMembershipAsync(Arg.Any<IOrganizationUserRepository>(), Arg.Any<Guid>())
             .Returns(Task.FromResult(organizationCollection));
 
-        _applicationCacheService.GetOrganizationAbilitiesAsync()
+        _applicationCacheService.GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
             .Returns(new Dictionary<Guid, OrganizationAbility>()
             {
                 { orgUser.OrganizationId, new OrganizationAbility(organization)}
@@ -152,6 +152,58 @@ public async void RequiresTwoFactorAsync_IndividualFalse_OrganizationRequired_Re
         Assert.IsType<Organization>(result.Item2);
     }
 
+    [Theory]
+    [BitAutoData("password")]
+    [BitAutoData("authorization_code")]
+    public async void RequiresTwoFactorAsync_OrganizationRequired_OnlyFetchesAbilitiesForUserOrgs(
+        string grantType,
+        [AuthFixtures.ValidatedTokenRequest] ValidatedTokenRequest request,
+        User user,
+        OrganizationUserOrganizationDetails orgUser,
+        Organization organization,
+        OrganizationUserOrganizationDetails orgUserNotInCache,
+        Organization organizationNotInCache,
+        ICollection<CurrentContextOrganization> organizationCollection)
+    {
+        // Arrange
+        request.GrantType = grantType;
+        orgUser.UserId = user.Id;
+        organization.Id = orgUser.OrganizationId;
+        organization.Use2fa = true;
+        organization.TwoFactorProviders = GetTwoFactorOrganizationDuoProviderJson();
+        organization.Enabled = true;
+
+        orgUserNotInCache.UserId = user.Id;
+        organizationNotInCache.Id = orgUserNotInCache.OrganizationId;
+        orgUserNotInCache.Permissions = "{}";
+
+        organizationCollection.Clear();
+        orgUser.Permissions = "{}";
+        organizationCollection.Add(new CurrentContextOrganization(orgUser));
+        organizationCollection.Add(new CurrentContextOrganization(orgUserNotInCache));
+
+        _currentContext.OrganizationMembershipAsync(Arg.Any<IOrganizationUserRepository>(), Arg.Any<Guid>())
+            .Returns(Task.FromResult(organizationCollection));
+
+        _applicationCacheService.GetOrganizationAbilitiesAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns(new Dictionary<Guid, OrganizationAbility>()
+            {
+                { orgUser.OrganizationId, new OrganizationAbility(organization) }
+            });
+
+        _organizationRepository.GetManyByUserIdAsync(Arg.Any<Guid>()).Returns([organization]);
+
+        // Act
+        await _sut.RequiresTwoFactorAsync(user, request);
+
+        // Assert - verify both org IDs were requested, not just one
+        await _applicationCacheService.Received(1).GetOrganizationAbilitiesAsync(
+            Arg.Is<IEnumerable<Guid>>(ids =>
+                ids.Contains(orgUser.OrganizationId) &&
+                ids.Contains(orgUserNotInCache.OrganizationId) &&
+                ids.Count() == 2));
+    }
+
     [Theory]
     [BitAutoData]
     public async void BuildTwoFactorResultAsync_NoProviders_ReturnsNull(