From dcd2b3f468b2a7dbfae0c7da1597874fba0cba0f Mon Sep 17 00:00:00 2001 From: anemeth Date: Fri, 28 Aug 2026 13:00:49 -0700 Subject: [PATCH] chore: Apply context pattern to other processors using shared state --- .../Processors/GPOLocalGroupProcessor.cs | 145 ++++++++++++++---- src/CommonLib/Processors/PortScanner.cs | 91 +++++++++-- test/unit/GPOLocalGroupProcessorTest.cs | 105 ++++++++++++- test/unit/PortScannerTest.cs | 43 +++++- 4 files changed, 336 insertions(+), 48 deletions(-) diff --git a/src/CommonLib/Processors/GPOLocalGroupProcessor.cs b/src/CommonLib/Processors/GPOLocalGroupProcessor.cs index 28a6996f6..b34e96cf3 100644 --- a/src/CommonLib/Processors/GPOLocalGroupProcessor.cs +++ b/src/CommonLib/Processors/GPOLocalGroupProcessor.cs @@ -5,6 +5,7 @@ using System.IO; using System.Linq; using System.Text.RegularExpressions; +using System.Threading; using System.Threading.Tasks; using System.Xml.XPath; using Microsoft.Extensions.Logging; @@ -13,6 +14,38 @@ using SharpHoundCommonLib.OutputTypes; namespace SharpHoundCommonLib.Processors { + /// + /// Owns state shared by GPOLocalGroupProcessor instances and gives that state an explicit lifetime. + /// + public sealed class GPOLocalGroupProcessorContext : IDisposable { + private readonly GPOLocalGroupProcessor.ActionCache _actionCache = new(); + private int _disposed; + + /// + /// Creates a that shares its GPO action cache with other + /// processors created by this context. + /// + public GPOLocalGroupProcessor CreateGPOLocalGroupProcessor(ILdapUtils utils, ILogger log = null) { + if (Volatile.Read(ref _disposed) != 0) { + throw new ObjectDisposedException(nameof(GPOLocalGroupProcessorContext)); + } + + return new GPOLocalGroupProcessor(utils, _actionCache, log); + } + + /// + /// Clears the shared processor state. Processors created by this context must not + /// be used after the context is disposed. + /// + public void Dispose() { + if (Interlocked.Exchange(ref _disposed, 1) != 0) { + return; + } + + _actionCache.Dispose(); + } + } + public class GPOLocalGroupProcessor { private static readonly Regex KeyRegex = new(@"(.+?)\s*=(.*)", RegexOptions.Compiled); @@ -30,8 +63,6 @@ public class GPOLocalGroupProcessor { private static readonly Regex ExtractRid = new(@"S-1-5-32-([0-9]{3})", RegexOptions.Compiled | RegexOptions.IgnoreCase); - private static readonly ConcurrentDictionary> GpoActionCache = new(); - private static readonly Dictionary ValidGroupNames = new(StringComparer.OrdinalIgnoreCase) { { "Administrators", LocalGroupRids.Administrators }, @@ -43,9 +74,46 @@ public class GPOLocalGroupProcessor { private readonly ILogger _log; private readonly ILdapUtils _utils; + private readonly ActionCache _actionCache; + + internal sealed class ActionCache : IDisposable { + private readonly ConcurrentDictionary>>> _buildTasks = + new(StringComparer.OrdinalIgnoreCase); + private int _disposed; + + public Lazy>> GetOrAddBuildTask(string distinguishedName, + Func>>> buildTaskFactory) { + ThrowIfDisposed(); + return _buildTasks.GetOrAdd(distinguishedName, _ => buildTaskFactory()); + } + + public void RemoveBuildTask(string distinguishedName, Lazy>> buildTask) { + ThrowIfDisposed(); + ((ICollection>>>>)_buildTasks).Remove( + new KeyValuePair>>>(distinguishedName, buildTask)); + } + + public void Dispose() { + if (Interlocked.Exchange(ref _disposed, 1) != 0) { + return; + } + + _buildTasks.Clear(); + } - public GPOLocalGroupProcessor(ILdapUtils utils, ILogger log = null) { + private void ThrowIfDisposed() { + if (Volatile.Read(ref _disposed) != 0) { + throw new ObjectDisposedException(nameof(GPOLocalGroupProcessorContext)); + } + } + } + + public GPOLocalGroupProcessor(ILdapUtils utils, ILogger log = null) : this(utils, new ActionCache(), log) { + } + + internal GPOLocalGroupProcessor(ILdapUtils utils, ActionCache actionCache, ILogger log = null) { _utils = utils; + _actionCache = actionCache; _log = log ?? Logging.LogProvider.CreateLogger("GPOLocalGroupProc"); } @@ -124,36 +192,22 @@ public async Task ReadGPOLocalGroups(string gpLink, string foreach (var rid in Enum.GetValues(typeof(LocalGroupRids))) data[(LocalGroupRids)rid] = new GroupResults(); foreach (var linkDn in orderedLinks) { - if (!GpoActionCache.TryGetValue(linkDn.ToLower(), out var actions)) { - actions = new List(); - - var gpoDomain = Helpers.DistinguishedNameToDomain(linkDn); - var result = await _utils.Query(new LdapQueryParameters() { - LDAPFilter = new LdapFilter().AddAllObjects().GetFilter(), - SearchScope = SearchScope.Base, - Attributes = [LDAPProperties.GPCFileSYSPath, LDAPProperties.Flags], - SearchBase = linkDn, - DomainName = gpoDomain - }).DefaultIfEmpty(LdapResult.Fail()).FirstOrDefaultAsync(); - - if (!result.IsSuccess) { - continue; - } - - if (!result.Value.TryGetProperty(LDAPProperties.GPCFileSYSPath, out var filePath) || - // Filter out GPOs that are disabled or the computer configuration is disabled - (result.Value.TryGetProperty(LDAPProperties.Flags, out var flags) && flags is "2" or "3")) { - GpoActionCache.TryAdd(linkDn, actions); - continue; - } - - //Add the actions for each file. The GPO template file actions will override the XML file actions - await foreach (var item in ProcessGPOXmlFile(filePath, gpoDomain)) actions.Add(item); - await foreach (var item in ProcessGPOTemplateFile(filePath, gpoDomain)) actions.Add(item); + var buildTask = _actionCache.GetOrAddBuildTask(linkDn, + () => new Lazy>>(() => BuildGPOActionCache(linkDn), + LazyThreadSafetyMode.ExecutionAndPublication)); + List actions; + try { + actions = await buildTask.Value; + } catch { + _actionCache.RemoveBuildTask(linkDn, buildTask); + throw; } - //Cache the actions for this GPO for later - GpoActionCache.TryAdd(linkDn.ToLower(), actions); + // Query failures are not cached so a later attempt can retry the GPO. + if (actions == null) { + _actionCache.RemoveBuildTask(linkDn, buildTask); + continue; + } //If there are no actions, then we can move on from this GPO if (actions.Count == 0) @@ -248,6 +302,33 @@ public async Task ReadGPOLocalGroups(string gpLink, string return ret; } + private async Task> BuildGPOActionCache(string linkDn) { + var actions = new List(); + var gpoDomain = Helpers.DistinguishedNameToDomain(linkDn); + var result = await _utils.Query(new LdapQueryParameters() { + LDAPFilter = new LdapFilter().AddAllObjects().GetFilter(), + SearchScope = SearchScope.Base, + Attributes = [LDAPProperties.GPCFileSYSPath, LDAPProperties.Flags], + SearchBase = linkDn, + DomainName = gpoDomain + }).DefaultIfEmpty(LdapResult.Fail()).FirstOrDefaultAsync(); + + if (!result.IsSuccess) { + return null; + } + + if (!result.Value.TryGetProperty(LDAPProperties.GPCFileSYSPath, out var filePath) || + // Filter out GPOs that are disabled or the computer configuration is disabled + (result.Value.TryGetProperty(LDAPProperties.Flags, out var flags) && flags is "2" or "3")) { + return actions; + } + + //Add the actions for each file. The GPO template file actions will override the XML file actions + await foreach (var item in ProcessGPOXmlFile(filePath, gpoDomain)) actions.Add(item); + await foreach (var item in ProcessGPOTemplateFile(filePath, gpoDomain)) actions.Add(item); + return actions; + } + /// /// Parses a GPO GptTmpl.inf file and pulls group membership changes out /// @@ -576,4 +657,4 @@ internal enum LocalGroupRids { PSRemote = 580 } } -} \ No newline at end of file +} diff --git a/src/CommonLib/Processors/PortScanner.cs b/src/CommonLib/Processors/PortScanner.cs index 6ff9cbd2f..6615d069d 100644 --- a/src/CommonLib/Processors/PortScanner.cs +++ b/src/CommonLib/Processors/PortScanner.cs @@ -1,23 +1,90 @@ using System; using System.Collections.Concurrent; using System.Net.Sockets; +using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.Logging; using SharpHoundRPC.PortScanner; namespace SharpHoundCommonLib.Processors { + /// + /// Owns state shared by PortScanner instances and gives that state an explicit lifetime. + /// + public sealed class PortScannerContext : IDisposable { + private readonly PortScanner.ScanCache _scanCache = new(); + private int _disposed; + + /// + /// Creates a that shares its scan cache with other scanners + /// created by this context. + /// + public PortScanner CreatePortScanner(ILogger log = null, int maxTimeout = 10000) { + if (Volatile.Read(ref _disposed) != 0) { + throw new ObjectDisposedException(nameof(PortScannerContext)); + } + + return new PortScanner(_scanCache, log, maxTimeout); + } + + /// + /// Clears the shared scanner state. Scanners created by this context must not + /// be used after the context is disposed. + /// + public void Dispose() { + if (Interlocked.Exchange(ref _disposed, 1) != 0) { + return; + } + + _scanCache.Dispose(); + } + } + public class PortScanner : IPortScanner { - private static readonly ConcurrentDictionary PortScanCache = new(); private readonly ILogger _log; private readonly AdaptiveTimeout _adaptiveTimeout; + private readonly ScanCache _scanCache; - public PortScanner() : this(null) { + internal sealed class ScanCache : IDisposable { + private readonly ConcurrentDictionary _portScanCache = new(); + private int _disposed; + + public bool TryGet(PingCacheKey key, out bool status) { + ThrowIfDisposed(); + return _portScanCache.TryGetValue(key, out status); + } + + public void Add(PingCacheKey key, bool status) { + ThrowIfDisposed(); + _portScanCache.TryAdd(key, status); + } + + public void Dispose() { + if (Interlocked.Exchange(ref _disposed, 1) != 0) { + return; + } + + _portScanCache.Clear(); + } + + private void ThrowIfDisposed() { + if (Volatile.Read(ref _disposed) != 0) { + throw new ObjectDisposedException(nameof(PortScannerContext)); + } + } + } + + public PortScanner() : this((ILogger)null) { } - public PortScanner(ILogger log = null, int maxTimeout = 10000) { + public PortScanner(ILogger log = null, int maxTimeout = 10000) : this( + new ScanCache(), log, maxTimeout) { + } + + internal PortScanner(ScanCache scanCache, ILogger log = null, int maxTimeout = 10000) { + _scanCache = scanCache; _log = log ?? Logging.LogProvider.CreateLogger("PortScanner"); - _adaptiveTimeout = new AdaptiveTimeout(maxTimeout: TimeSpan.FromMilliseconds(maxTimeout), _log); + _adaptiveTimeout = new AdaptiveTimeout(TimeSpan.FromMilliseconds(maxTimeout), _log); } /// @@ -35,7 +102,7 @@ public virtual async Task CheckPort(string hostname, int port = 445, HostName = hostname }; - if (PortScanCache.TryGetValue(key, out var status)) { + if (_scanCache.TryGet(key, out var status)) { _log.LogTrace("Port scan cache hit for {HostName}:{Port}: {Status}", hostname, port, status); return status; } @@ -48,12 +115,12 @@ public virtual async Task CheckPort(string hostname, int port = 445, if (throwError) { throw new TimeoutException(ca.Error); } - PortScanCache.TryAdd(key, false); + _scanCache.Add(key, false); return false; } _log.LogTrace("CheckPort Succeeded for {HostName}:{Port}", hostname, port); - PortScanCache.TryAdd(key, true); + _scanCache.Add(key, true); return true; } catch (Exception e) { @@ -63,16 +130,12 @@ public virtual async Task CheckPort(string hostname, int port = 445, throw; } - PortScanCache.TryAdd(key, false); + _scanCache.Add(key, false); return false; } } - public static void ClearCache() { - PortScanCache.Clear(); - } - - private class PingCacheKey { + internal class PingCacheKey { internal string HostName { get; set; } internal int Port { get; set; } @@ -94,4 +157,4 @@ public override int GetHashCode() { } } } -} \ No newline at end of file +} diff --git a/test/unit/GPOLocalGroupProcessorTest.cs b/test/unit/GPOLocalGroupProcessorTest.cs index 07107a12e..3eee539d9 100644 --- a/test/unit/GPOLocalGroupProcessorTest.cs +++ b/test/unit/GPOLocalGroupProcessorTest.cs @@ -93,6 +93,86 @@ public GPOLocalGroupProcessorTest(ITestOutputHelper testOutputHelper) { _testOutputHelper = testOutputHelper; } + [Fact] + public async Task GPOLocalGroupProcessorContext_Processors_QueryGPOOnce() { + var (mockLdapUtils, gpLink, linkDn) = CreateContextTestData(); + using var context = new GPOLocalGroupProcessorContext(); + var processors = Enumerable.Range(0, 50) + .Select(_ => context.CreateGPOLocalGroupProcessor(mockLdapUtils.Object)) + .ToArray(); + + await Task.WhenAll(processors.Select(processor => + processor.ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL"))); + + mockLdapUtils.Verify(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddAllObjects().GetFilter() && + parameters.SearchBase == linkDn), + It.IsAny()), Times.Once); + } + + [Fact] + public async Task GPOLocalGroupProcessorContext_Processors_DoNotShareCacheAcrossContexts() { + var (mockLdapUtils, gpLink, linkDn) = CreateContextTestData(); + using var firstContext = new GPOLocalGroupProcessorContext(); + using var secondContext = new GPOLocalGroupProcessorContext(); + + await Task.WhenAll( + firstContext.CreateGPOLocalGroupProcessor(mockLdapUtils.Object) + .ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL"), + secondContext.CreateGPOLocalGroupProcessor(mockLdapUtils.Object) + .ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL")); + + mockLdapUtils.Verify(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddAllObjects().GetFilter() && + parameters.SearchBase == linkDn), + It.IsAny()), Times.Exactly(2)); + } + + [Fact] + public async Task GPOLocalGroupProcessorContext_QueryFailure_IsRetried() { + var (mockLdapUtils, gpLink, linkDn) = CreateContextTestData(); + mockLdapUtils.SetupSequence(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddAllObjects().GetFilter()), + It.IsAny())) + .Returns(new[] { LdapResult.Fail() }.ToAsyncEnumerable) + .Returns(new[] { LdapResult.Ok(new Mock().Object) } + .ToAsyncEnumerable); + using var context = new GPOLocalGroupProcessorContext(); + var processor = context.CreateGPOLocalGroupProcessor(mockLdapUtils.Object); + + await processor.ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL"); + await processor.ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL"); + + mockLdapUtils.Verify(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddAllObjects().GetFilter() && + parameters.SearchBase == linkDn), + It.IsAny()), Times.Exactly(2)); + } + + [Fact] + public void GPOLocalGroupProcessorContext_CreateProcessor_AfterDispose_Throws() { + var context = new GPOLocalGroupProcessorContext(); + context.Dispose(); + + Assert.Throws(() => + context.CreateGPOLocalGroupProcessor(new MockLdapUtils())); + } + + [Fact] + public async Task GPOLocalGroupProcessorContext_Processor_AfterDispose_Throws() { + var (mockLdapUtils, gpLink, _) = CreateContextTestData(); + var context = new GPOLocalGroupProcessorContext(); + var processor = context.CreateGPOLocalGroupProcessor(mockLdapUtils.Object); + context.Dispose(); + + await Assert.ThrowsAsync(() => + processor.ReadGPOLocalGroups(gpLink, "DC=TEST,DC=LOCAL")); + } + [Fact] public async Task GPOLocalGroupProcessor_ReadGPOLocalGroups_Null_GPLink() { var mockLDAPUtils = new Mock(); @@ -351,6 +431,29 @@ public async Task GPOLocalGroupProcess_ProcessGPOXMLFile_NoFile() { Assert.Empty(actual); } + private static (Mock LdapUtils, string GPLink, string LinkDn) CreateContextTestData() { + var mockLdapUtils = new Mock(); + var computerEntry = new Mock(); + var computerSid = $"S-1-5-21-{Random.Shared.Next()}-{Random.Shared.Next()}-{Random.Shared.Next()}-1000"; + computerEntry.Setup(x => x.TryGetSecurityIdentifier(out computerSid)).Returns(true); + var computerResults = new[] { LdapResult.Ok(computerEntry.Object) }; + var gpoResults = new[] { LdapResult.Ok(new Mock().Object) }; + + mockLdapUtils.Setup(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddComputersNoMSAs().GetFilter()), + It.IsAny())) + .Returns(computerResults.ToAsyncEnumerable); + mockLdapUtils.Setup(x => x.Query( + It.Is(parameters => + parameters.LDAPFilter == new LdapFilter().AddAllObjects().GetFilter()), + It.IsAny())) + .Returns(gpoResults.ToAsyncEnumerable); + + var linkDn = $"CN={Guid.NewGuid():N},CN=Policies,CN=System,DC=TEST,DC=LOCAL"; + return (mockLdapUtils, $"[LDAP://{linkDn};0]", linkDn); + } + [Fact] public async Task GPOLocalGroupProcess_ProcessGPOXMLFile_Disabled() { var mockLDAPUtils = new Mock(); @@ -467,4 +570,4 @@ public void GPOLocalGroupProcess_GroupAction() { str); } } -} \ No newline at end of file +} diff --git a/test/unit/PortScannerTest.cs b/test/unit/PortScannerTest.cs index 2e37a7e74..99eb11940 100644 --- a/test/unit/PortScannerTest.cs +++ b/test/unit/PortScannerTest.cs @@ -1,5 +1,6 @@ using System; using System.Diagnostics.CodeAnalysis; +using System.Threading.Tasks; using SharpHoundCommonLib.Processors; using Xunit; @@ -7,6 +8,46 @@ namespace CommonLibTest; [SuppressMessage("Interoperability", "CA1416:Validate platform compatibility")] public class PortScannerTest { + [Fact] + public async Task PortScannerContext_ScannersShareCache() { + using var context = new PortScannerContext(); + var firstScanner = context.CreatePortScanner(); + var secondScanner = context.CreatePortScanner(); + + // An invalid port is cached as false. A cache miss with throwError enabled would throw, + // so returning false demonstrates that the second scanner used the first scanner's result. + Assert.False(await firstScanner.CheckPort("localhost", -1)); + Assert.False(await secondScanner.CheckPort("localhost", -1, throwError: true)); + } + + [Fact] + public async Task PortScannerContext_ScannersDoNotShareCacheAcrossContexts() { + using var firstContext = new PortScannerContext(); + using var secondContext = new PortScannerContext(); + + Assert.False(await firstContext.CreatePortScanner().CheckPort("localhost", -1)); + // The second context has no cached result and therefore attempts the invalid scan. + await Assert.ThrowsAnyAsync(() => + secondContext.CreatePortScanner().CheckPort("localhost", -1, throwError: true)); + } + + [Fact] + public void PortScannerContext_CreateScanner_AfterDispose_Throws() { + var context = new PortScannerContext(); + context.Dispose(); + + Assert.Throws(() => context.CreatePortScanner()); + } + + [Fact] + public async Task PortScannerContext_Scanner_AfterDispose_Throws() { + var context = new PortScannerContext(); + var scanner = context.CreatePortScanner(); + context.Dispose(); + + await Assert.ThrowsAsync(() => scanner.CheckPort("localhost")); + } + //// Throws "no such host is known" exception // [Fact] // public void PortScanner_CheckPort_TimeoutException() { @@ -16,4 +57,4 @@ public class PortScannerTest { // var ex = Assert.ThrowsAsync(() => scanner.CheckPort(hostname, port, 1, true)); // Assert.Equal("Timed Out", ex.Result.Message); // } -} \ No newline at end of file +}