Add better error handling for policyDefinitions dict

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

@@ -1,18 +1,12 @@
 #nullable enable
 
-using System.Collections.Immutable;
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
-using Bit.Core.AdminConsole.Models.Data.Organizations.Policies;
 using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
 using Bit.Core.AdminConsole.Repositories;
-using Bit.Core.Entities;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
-using Bit.Core.Models.Data.Organizations.OrganizationUsers;
-using Bit.Core.Repositories;
 using Bit.Core.Services;
-using Bit.Core.Settings;
 
 namespace Bit.Core.AdminConsole.Services.Implementations;
 
@@ -33,9 +27,16 @@ public class PolicyServicevNext : IPolicyServicevNext
         _eventService = eventService;
         _policyRepository = policyRepository;
 
-        _policyDefinitions = ImmutableDictionary.CreateRange(
+        var policyDefinitionsDict = new Dictionary<PolicyType, IPolicyDefinition>();
-            policyDefinitions.Select(def => KeyValuePair.Create(def.Type, def))
+        foreach (var policyDefinition in policyDefinitions)
-        );
+        {
+            if (!policyDefinitionsDict.TryAdd(policyDefinition.Type, policyDefinition))
+            {
+                throw new Exception($"Duplicate PolicyDefinition for {policyDefinition.Type} policy.");
+            }
+        }
+
+        _policyDefinitions = policyDefinitionsDict;
     }
 
     public async Task SaveAsync(Policy policy, IUserService userService, IOrganizationService organizationService,
@@ -52,7 +53,7 @@ public class PolicyServicevNext : IPolicyServicevNext
             throw new BadRequestException("This organization cannot use policies.");
         }
 
-        var policyDefinition = _policyDefinitions[policy.Type];
+        var policyDefinition = GetDefinition(policy.Type);
         var allSavedPolicies = await _policyRepository.GetManyByOrganizationIdAsync(org.Id);
         var currentPolicy = allSavedPolicies.SingleOrDefault(p => p.Id == policy.Id);
 
@@ -108,4 +109,14 @@ public class PolicyServicevNext : IPolicyServicevNext
         await _policyRepository.UpsertAsync(policy);
         await _eventService.LogPolicyEventAsync(policy, EventType.Policy_Updated);
     }
+
+    private IPolicyDefinition GetDefinition(PolicyType type)
+    {
+        if (!_policyDefinitions.TryGetValue(type, out var result))
+        {
+            throw new Exception($"No PolicyDefinition found for {type} policy.");
+        }
+
+        return result;
+    }
 }

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

@@ -18,9 +18,21 @@ namespace Bit.Core.Test.AdminConsole.Services;
 
 public class PolicyServicevNextTests
 {
+    [Fact]
+    public void Constructor_DuplicatePolicyDefinitions_Throws()
+    {
+        var exception = Assert.Throws<Exception>(() =>
+            new PolicyServicevNext(
+                Substitute.For<IApplicationCacheService>(),
+                Substitute.For<IEventService>(),
+                Substitute.For<IPolicyRepository>(),
+                [new FakeSingleOrgPolicyDefinition(), new FakeSingleOrgPolicyDefinition()]
+            ));
+        Assert.Contains("Duplicate PolicyDefinition for SingleOrg policy", exception.Message);
+    }
+
     [Theory, BitAutoData]
-    public async Task SaveAsync_OrganizationDoesNotExist_ThrowsBadRequest(
+    public async Task SaveAsync_OrganizationDoesNotExist_ThrowsBadRequest(Policy policy)
-        Policy policy)
     {
         var sutProvider = SutProviderFactory();
         sutProvider.GetDependency<IApplicationCacheService>()
@@ -45,8 +57,7 @@ public class PolicyServicevNextTests
     }
 
     [Theory, BitAutoData]
-    public async Task SaveAsync_OrganizationCannotUsePolicies_ThrowsBadRequest(
+    public async Task SaveAsync_OrganizationCannotUsePolicies_ThrowsBadRequest(Policy policy)
-        Policy policy)
     {
         var sutProvider = SutProviderFactory();
         sutProvider.GetDependency<IApplicationCacheService>()
@@ -74,6 +85,35 @@ public class PolicyServicevNextTests
             .LogPolicyEventAsync(default, default, default);
     }
 
+    [Theory, BitAutoData]
+    public async Task SaveAsync_PolicyDefinitionNotFound_Throws([Policy(PolicyType.SingleOrg)]Policy policy)
+    {
+        var sutProvider = SutProviderFactory();
+        sutProvider.GetDependency<IApplicationCacheService>()
+            .GetOrganizationAbilityAsync(policy.OrganizationId)
+            .Returns(new OrganizationAbility
+            {
+                Id = policy.OrganizationId,
+                UsePolicies = true
+            });
+
+        var exception = await Assert.ThrowsAsync<Exception>(
+            () => sutProvider.Sut.SaveAsync(policy,
+                Substitute.For<IUserService>(),
+                Substitute.For<IOrganizationService>(),
+                Guid.NewGuid()));
+
+        Assert.Contains("No PolicyDefinition found for SingleOrg policy", exception.Message, StringComparison.OrdinalIgnoreCase);
+
+        await sutProvider.GetDependency<IPolicyRepository>()
+            .DidNotReceiveWithAnyArgs()
+            .UpsertAsync(default);
+
+        await sutProvider.GetDependency<IEventService>()
+            .DidNotReceiveWithAnyArgs()
+            .LogPolicyEventAsync(default, default, default);
+    }
+
     [Theory, BitAutoData]
     public async Task SaveAsync_ThrowsOnValidationError([AdminConsoleFixtures.Policy(PolicyType.SingleOrg)] Policy policy)
     {