From 9844ba8512c11e7811aec3efa6ae0af445fed2ce Mon Sep 17 00:00:00 2001
From: Thomas Rittson <trittson@bitwarden.com>
Date: Mon, 7 Oct 2024 08:23:20 +1000
Subject: [PATCH] Add better error handling for policyDefinitions dict

---
 .../Implementations/PolicyServicevNext.cs     | 31 ++++++++----
 .../Services/PolicyServicevNextTests.cs       | 48 +++++++++++++++++--
 2 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/src/Core/AdminConsole/Services/Implementations/PolicyServicevNext.cs b/src/Core/AdminConsole/Services/Implementations/PolicyServicevNext.cs
index b70e55875f..63c1cdafe5 100644
--- a/src/Core/AdminConsole/Services/Implementations/PolicyServicevNext.cs
+++ b/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(
-            policyDefinitions.Select(def => KeyValuePair.Create(def.Type, def))
-        );
+        var policyDefinitionsDict = new Dictionary<PolicyType, IPolicyDefinition>();
+        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;
+    }
 }
diff --git a/test/Core.Test/AdminConsole/Services/PolicyServicevNextTests.cs b/test/Core.Test/AdminConsole/Services/PolicyServicevNextTests.cs
index 8401f59e83..8faa4e34ef 100644
--- a/test/Core.Test/AdminConsole/Services/PolicyServicevNextTests.cs
+++ b/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(
-        Policy policy)
+    public async Task SaveAsync_OrganizationDoesNotExist_ThrowsBadRequest(Policy policy)
     {
         var sutProvider = SutProviderFactory();
         sutProvider.GetDependency<IApplicationCacheService>()
@@ -45,8 +57,7 @@ public class PolicyServicevNextTests
     }
 
     [Theory, BitAutoData]
-    public async Task SaveAsync_OrganizationCannotUsePolicies_ThrowsBadRequest(
-        Policy policy)
+    public async Task SaveAsync_OrganizationCannotUsePolicies_ThrowsBadRequest(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)
     {