fix: Make ACLProcessor guid cache singleton and thread-safe - BED-9236

src/CommonLib/Processors/ACLProcessor.cs

@@ -13,23 +13,99 @@
 using SharpHoundCommonLib.LDAPQueries;
 using SharpHoundCommonLib.OutputTypes;
 using System.Linq;
+using System.Threading;
 
 namespace SharpHoundCommonLib.Processors {
+    /// <summary>
+    ///     Owns state shared by processor instances and gives that state an explicit lifetime.
+    /// </summary>
+    public sealed class ACLProcessorContext : IDisposable {
+        private readonly ACLProcessor.GuidCache _aclGuidCache = new();
+        private int _disposed;
+
+        /// <summary>
+        ///     Creates an <see cref="ACLProcessor"/> that shares its GUID cache with other
+        ///     ACL processors created by this context.
+        /// </summary>
+        public ACLProcessor CreateACLProcessor(ILdapUtils utils, ILogger log = null) {
+            if (Volatile.Read(ref _disposed) != 0) {
+                throw new ObjectDisposedException(nameof(ACLProcessorContext));
+            }
+
+            return new ACLProcessor(utils, _aclGuidCache, log);
+        }
+
+        /// <summary>
+        ///     Clears the shared processor state. Processors created by this context must not
+        ///     be used after the context is disposed.
+        /// </summary>
+        public void Dispose() {
+            if (Interlocked.Exchange(ref _disposed, 1) != 0) {
+                return;
+            }
+
+            _aclGuidCache.Dispose();
+        }
+    }
+
     public class ACLProcessor {
         private static readonly Dictionary<Label, string> BaseGuids;
-        private readonly ConcurrentDictionary<string, string> _guidMap = new();
         private readonly ILogger _log;
         private readonly ILdapUtils _utils;
-        private readonly ConcurrentHashSet _builtDomainCaches = new(StringComparer.OrdinalIgnoreCase);
         private readonly ConcurrentDictionary<string, string[]> _exchangeTrusteeSidCache = new(StringComparer.OrdinalIgnoreCase);
-        private readonly object _lock = new();
         // These Exchange principals commonly carry product-added deny ACEs that we intentionally suppress.
         private static readonly HashSet<string> ExchangeTrusteeNames = new(StringComparer.OrdinalIgnoreCase) {
             "Exchange Windows Permissions",
             "Exchange Trusted Subsystem",
             "Exchange Servers",
             "Organization Management"
         };
+        private readonly GuidCache _guidCache;
+
+        internal sealed class GuidCache : IDisposable {
+            private readonly ConcurrentDictionary<string, string> _guidMap = new();
+            private readonly ConcurrentDictionary<string, Lazy<Task>> _buildTasks =
+                new(StringComparer.OrdinalIgnoreCase);
+            private int _disposed;
+
+            public Lazy<Task> GetOrAddBuildTask(string domain, Func<Lazy<Task>> buildTaskFactory) {
+                ThrowIfDisposed();
+                return _buildTasks.GetOrAdd(domain, _ => buildTaskFactory());
+            }
+
+            public bool RemoveBuildTask(string domain, Lazy<Task> buildTask) {
+                // Remove only this instance so a delayed fault cannot remove a newer retry task.
+                // _buildTasks.TryRemove does not guarantee that the value is the same as the one being removed, so we need to cast to ICollection and use Remove instead.
+                // This lets us conditionally remove the task only if it is the same instance as the one we expect.
+                return ((ICollection<KeyValuePair<string, Lazy<Task>>>)_buildTasks)
+                    .Remove(new KeyValuePair<string, Lazy<Task>>(domain, buildTask));
+            }
+
+            public void AddGuid(string guid, string name) {
+                ThrowIfDisposed();
+                _guidMap.TryAdd(guid, name);
+            }
+
+            public bool TryGetGuid(string guid, out string name) {
+                ThrowIfDisposed();
+                return _guidMap.TryGetValue(guid, out name);
+            }
+
+            public void Dispose() {
+                if (Interlocked.Exchange(ref _disposed, 1) != 0) {
+                    return;
+                }
+
+                _buildTasks.Clear();
+                _guidMap.Clear();
+            }
+
+            private void ThrowIfDisposed() {
+                if (Volatile.Read(ref _disposed) != 0) {
+                    throw new ObjectDisposedException(nameof(ACLProcessorContext));
+                }
+            }
+        }
 
         static ACLProcessor() {
             //Create a dictionary with the base GUIDs of each object type
@@ -51,9 +127,12 @@ static ACLProcessor() {
             };
         }
 
-        public ACLProcessor(ILdapUtils utils, ILogger log = null)
-        {
+        public ACLProcessor(ILdapUtils utils, ILogger log = null) : this(utils, new GuidCache(), log) {
+        }
+
+        internal ACLProcessor(ILdapUtils utils, GuidCache guidCache, ILogger log = null) {
             _utils = utils;
+            _guidCache = guidCache;
             _log = log ?? Logging.LogProvider.CreateLogger("ACLProc");
         }
 
@@ -120,14 +199,20 @@ public override string ToString() {
         ///     LAPS
         /// </summary>
         private async Task BuildGuidCache(string domain) {
-            lock (_lock) {
-                if (_builtDomainCaches.Contains(domain)) {
-                    return;
-                }
+            var buildTask = _guidCache.GetOrAddBuildTask(domain,
+                // The ExecutionAndPublication mode ensures that only one thread can execute the factory method at a time, and all other threads will wait for the result of that execution. This prevents multiple threads from building the cache simultaneously for the same domain.
+                () => new Lazy<Task>(() => BuildGuidCacheCore(domain), LazyThreadSafetyMode.ExecutionAndPublication));
 
-                _builtDomainCaches.Add(domain);
+            try {
+                await buildTask.Value;
             }
+            catch {
+                _guidCache.RemoveBuildTask(domain, buildTask);
+                throw;
+            }
+        }
 
+        private async Task BuildGuidCacheCore(string domain) {
             _log.LogInformation("Building GUID Cache for {Domain}", domain);
             await foreach (var result in _utils.PagedQuery(new LdapQueryParameters {
                 DomainName = domain,
@@ -155,7 +240,7 @@ private async Task BuildGuidCache(string domain) {
 
                     if (name is LDAPProperties.LAPSPlaintextPassword or LDAPProperties.LAPSEncryptedPassword or LDAPProperties.LegacyLAPSPassword) {
                         _log.LogInformation("Found GUID for ACL Right {Name}: {Guid} in domain {Domain}", name, guid, domain);
-                        _guidMap.TryAdd(guid, name);
+                        _guidCache.AddGuid(guid, name);
                     }
                 } else {
                     _log.LogDebug("Error while building GUID cache for {Domain}: {Message}", domain, result.Error);
@@ -770,7 +855,7 @@ await CountCustomDenyAce(ace, customDenyAceAccumulator, objectDomain, objectType
                                     IsPermissionForOwnerRightsSid = isPermissionForOwnerRightsSid,
                                     IsInheritedPermissionForOwnerRightsSid = isInheritedPermissionForOwnerRightsSid,
                                 };
-                            else if (_guidMap.TryGetValue(aceType, out var lapsAttribute)) {
+                            else if (_guidCache.TryGetGuid(aceType, out var lapsAttribute)) {
                                 // Compare the retrieved attribute name against LDAPProperties values
                                 if (lapsAttribute == LDAPProperties.LegacyLAPSPassword ||
                                     lapsAttribute == LDAPProperties.LAPSPlaintextPassword ||

test/unit/ACLProcessorTest.cs

@@ -58,6 +58,95 @@ public void SanityCheck() {
             Assert.True(true);
         }
 
+        [Fact]
+        public async Task ProcessorContext_ACLProcessors_QueryOncePerDomain() {
+            var mockLdapUtils = new Mock<ILdapUtils>();
+            mockLdapUtils
+                .Setup(x => x.PagedQuery(It.IsAny<LdapQueryParameters>(), It.IsAny<CancellationToken>()))
+                .Returns(Array.Empty<LdapResult<IDirectoryObject>>().ToAsyncEnumerable);
+            var domain = $"{Guid.NewGuid():N}.TEST";
+            using var context = new ACLProcessorContext();
+            var processors = Enumerable.Range(0, 50)
+                .Select(_ => context.CreateACLProcessor(mockLdapUtils.Object))
+                .ToArray();
+
+            await Task.WhenAll(processors.Select(processor =>
+                processor.ProcessACL(null, domain, Label.Computer, false).ToArrayAsync()));
+
+            mockLdapUtils.Verify(
+                x => x.PagedQuery(It.Is<LdapQueryParameters>(parameters => parameters.DomainName == domain),
+                    It.IsAny<CancellationToken>()),
+                Times.Once);
+        }
+
+        [Fact]
+        public async Task ProcessorContext_ACLProcessors_RetriesGuidCacheBuildAfterFailure() {
+            var mockLdapUtils = new Mock<ILdapUtils>();
+            var queryAttempts = 0;
+            mockLdapUtils
+                .Setup(x => x.PagedQuery(It.IsAny<LdapQueryParameters>(), It.IsAny<CancellationToken>()))
+                .Returns(() => {
+                    if (Interlocked.Increment(ref queryAttempts) == 1) {
+                        throw new InvalidOperationException("Expected test failure");
+                    }
+
+                    return Array.Empty<LdapResult<IDirectoryObject>>().ToAsyncEnumerable();
+                });
+            var domain = $"{Guid.NewGuid():N}.TEST";
+            using var context = new ACLProcessorContext();
+            var processor = context.CreateACLProcessor(mockLdapUtils.Object);
+
+            await Assert.ThrowsAsync<InvalidOperationException>(() =>
+                processor.ProcessACL(null, domain, Label.Computer, false).ToArrayAsync());
+
+            await processor.ProcessACL(null, domain, Label.Computer, false).ToArrayAsync();
+
+            mockLdapUtils.Verify(
+                x => x.PagedQuery(It.Is<LdapQueryParameters>(parameters => parameters.DomainName == domain),
+                    It.IsAny<CancellationToken>()),
+                Times.Exactly(2));
+        }
+
+        [Fact]
+        public async Task ProcessorContext_ACLProcessors_DoNotShareCacheAcrossContexts() {
+            var mockLdapUtils = new Mock<ILdapUtils>();
+            mockLdapUtils
+                .Setup(x => x.PagedQuery(It.IsAny<LdapQueryParameters>(), It.IsAny<CancellationToken>()))
+                .Returns(Array.Empty<LdapResult<IDirectoryObject>>().ToAsyncEnumerable);
+            var domain = $"{Guid.NewGuid():N}.TEST";
+            using var firstContext = new ACLProcessorContext();
+            using var secondContext = new ACLProcessorContext();
+
+            await Task.WhenAll(
+                firstContext.CreateACLProcessor(mockLdapUtils.Object)
+                    .ProcessACL(null, domain, Label.Computer, false).ToArrayAsync(),
+                secondContext.CreateACLProcessor(mockLdapUtils.Object)
+                    .ProcessACL(null, domain, Label.Computer, false).ToArrayAsync());
+
+            mockLdapUtils.Verify(
+                x => x.PagedQuery(It.Is<LdapQueryParameters>(parameters => parameters.DomainName == domain),
+                    It.IsAny<CancellationToken>()),
+                Times.Exactly(2));
+        }
+
+        [Fact]
+        public void ProcessorContext_CreateACLProcessor_AfterDispose_Throws() {
+            var context = new ACLProcessorContext();
+            context.Dispose();
+
+            Assert.Throws<ObjectDisposedException>(() => context.CreateACLProcessor(new MockLdapUtils()));
+        }
+
+        [Fact]
+        public async Task ProcessorContext_ACLProcessor_AfterDispose_Throws() {
+            var context = new ACLProcessorContext();
+            var processor = context.CreateACLProcessor(new MockLdapUtils());
+            context.Dispose();
+
+            await Assert.ThrowsAsync<ObjectDisposedException>(() =>
+                processor.ProcessACL(null, "TEST.LOCAL", Label.Computer, false).ToArrayAsync());
+        }
+
         [Fact]
         public void ACLProcessor_IsACLProtected_NullNTSD_ReturnsFalse() {
             var processor = new ACLProcessor(new MockLdapUtils());