From 8a9ef4e7709b0f413c6827c117d0317d56ab1d7b Mon Sep 17 00:00:00 2001 From: anemeth Date: Mon, 29 Jul 2024 15:59:41 -0500 Subject: [PATCH 1/3] WIP: utils query tests, especially around exception handling --- src/CommonLib/ConnectionPoolManager.cs | 12 ++- src/CommonLib/LdapUtils.cs | 9 +- test/unit/LdapUtilsQueryTest.cs | 135 +++++++++++++++++++++++++ 3 files changed, 154 insertions(+), 2 deletions(-) create mode 100644 test/unit/LdapUtilsQueryTest.cs diff --git a/src/CommonLib/ConnectionPoolManager.cs b/src/CommonLib/ConnectionPoolManager.cs index f7f015f8..e649a653 100644 --- a/src/CommonLib/ConnectionPoolManager.cs +++ b/src/CommonLib/ConnectionPoolManager.cs @@ -7,7 +7,17 @@ using SharpHoundCommonLib.Processors; namespace SharpHoundCommonLib { - public class ConnectionPoolManager : IDisposable{ + public interface ILdapConnectionProvider { + Task<(bool Success, string Message)> TestDomainConnection(string identifier, bool globalCatalog); + Task<(bool Success, LdapConnectionWrapper ConnectionWrapper, string Message)> GetLdapConnection( + string identifier, bool globalCatalog); + Task<(bool Success, LdapConnectionWrapper connectionWrapper, string Message)> GetLdapConnectionForServer( + string identifier, string server, bool globalCatalog); + void ReleaseConnection(LdapConnectionWrapper connectionWrapper, bool connectionFaulted = false); + void Dispose(); + } + + public class ConnectionPoolManager : ILdapConnectionProvider, IDisposable{ private readonly ConcurrentDictionary _pools = new(); private readonly LdapConfig _ldapConfig; private readonly string[] _translateNames = { "Administrator", "admin" }; diff --git a/src/CommonLib/LdapUtils.cs b/src/CommonLib/LdapUtils.cs index bc087b69..5aedc4aa 100644 --- a/src/CommonLib/LdapUtils.cs +++ b/src/CommonLib/LdapUtils.cs @@ -53,7 +53,7 @@ private readonly ConcurrentDictionary private readonly string[] _translateNames = { "Administrator", "admin" }; private LdapConfig _ldapConfig = new(); - private ConnectionPoolManager _connectionPool; + private ILdapConnectionProvider _connectionPool; private static readonly TimeSpan MinBackoffDelay = TimeSpan.FromSeconds(2); private static readonly TimeSpan MaxBackoffDelay = TimeSpan.FromSeconds(20); @@ -82,6 +82,13 @@ public LdapUtils() { _connectionPool = new ConnectionPoolManager(_ldapConfig, _log); } + public LdapUtils(ILdapConnectionProvider ldapConnectionProvider) { + _nativeMethods = new NativeMethods(); + _portScanner = new PortScanner(); + _log = Logging.LogProvider.CreateLogger("LDAPUtils"); + _connectionPool = ldapConnectionProvider; + } + public LdapUtils(NativeMethods nativeMethods = null, PortScanner scanner = null, ILogger log = null) { _nativeMethods = nativeMethods ?? new NativeMethods(); _portScanner = scanner ?? new PortScanner(); diff --git a/test/unit/LdapUtilsQueryTest.cs b/test/unit/LdapUtilsQueryTest.cs new file mode 100644 index 00000000..c7a2b54a --- /dev/null +++ b/test/unit/LdapUtilsQueryTest.cs @@ -0,0 +1,135 @@ +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Xunit; +using Moq; +using System.DirectoryServices.Protocols; +using SharpHoundCommonLib; + +public class RangedRetrievalTests +{ + private Mock _mockConnectionPool; + private LdapUtils _utils; + + public RangedRetrievalTests() + { + _mockConnectionPool = new Mock(); + _utils = new LdapUtils(); + } + + // [Fact] + // public async Task RangedRetrieval_SuccessfulRetrieval_ReturnsExpectedResults() + // { + // // Arrange + // var distinguishedName = "CN=TestUser,DC=example,DC=com"; + // var attributeName = "member"; + // var domain = "example.com"; + + // var connectionWrapper = new Mock(); + // var connection = new Mock(); + // connectionWrapper.SetupGet(x => x.Connection).Returns(connection.Object); + + // _mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) + // .ReturnsAsync((true, connectionWrapper.Object, null)); + + // var searchResponse = new Mock(); + // var entry = new SearchResultEntry + // { + // Attributes = + // { + // new DirectoryAttribute("member;range=0-*", "CN=Member1,DC=example,DC=com", "CN=Member2,DC=example,DC=com") + // } + // }; + // searchResponse.Entries.Add(entry); + + // connection.Setup(x => x.SendRequest(It.IsAny())) + // .Returns(searchResponse); + + // // Act + // var results = new List>(); + // await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName)) + // { + // results.Add(result); + // } + + // // Assert + // Assert.Equal(2, results.Count); + // Assert.True(results[0].IsSuccess); + // Assert.Equal("CN=Member1,DC=example,DC=com", results[0].Value); + // Assert.True(results[1].IsSuccess); + // Assert.Equal("CN=Member2,DC=example,DC=com", results[1].Value); + // } + + [Fact] + public async Task RangedRetrieval_ConnectionFailure_ReturnsFailResult() + { + // Arrange + var distinguishedName = "CN=TestUser,DC=example,DC=com"; + var attributeName = "member"; + + // Act + var results = new List>(); + await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName)) + { + results.Add(result); + } + + // Assert + Assert.Single(results); + Assert.False(results[0].IsSuccess); + Assert.Equal("All attempted connections failed", results[0].Error); + } + + // [Fact] + // public async Task RangedRetrieval_ServerDown_RetriesAndRecovers() + // { + // // Arrange + // var distinguishedName = "CN=TestUser,DC=example,DC=com"; + // var attributeName = "member"; + // var domain = "example.com"; + + // var connectionWrapper = new Mock(); + // var connection = new Mock(); + + // // TODO : setup + + // // Act + // var results = new List>(); + // await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName)) + // { + // results.Add(result); + // } + + // // TODO Assert + // } + + [Fact] + public async Task RangedRetrieval_CancellationRequested_StopsRetrieval() + { + // Arrange + var distinguishedName = "CN=TestUser,DC=example,DC=com"; + var attributeName = "member"; + var domain = "example.com"; + + var connectionWrapper = new Mock(null, null, false, string.Empty); + var connection = new Mock(); + + _mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) + .ReturnsAsync((true, connectionWrapper.Object, null)); + + _utils = new LdapUtils(_mockConnectionPool.Object); + + var cts = new CancellationTokenSource(); + cts.Cancel(); + + // Act + var results = new List>(); + await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName, cts.Token)) + { + results.Add(result); + } + + // Assert + Assert.False(results[0].IsSuccess); + } +} \ No newline at end of file From f4cca011454ae6076f3e1365a8c3bfbfe81a588e Mon Sep 17 00:00:00 2001 From: anemeth Date: Mon, 29 Jul 2024 22:54:55 -0500 Subject: [PATCH 2/3] Add Query error tests --- test/unit/CommonLibTest.csproj | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/unit/CommonLibTest.csproj b/test/unit/CommonLibTest.csproj index 717d5c93..cd4ca2f3 100644 --- a/test/unit/CommonLibTest.csproj +++ b/test/unit/CommonLibTest.csproj @@ -1,7 +1,7 @@ - net7.0 + net8.0 false true ..\..\docfx\coverage\ From d244a6e1281fb6a9a0fd4e099689d471994b49ef Mon Sep 17 00:00:00 2001 From: anemeth Date: Mon, 29 Jul 2024 22:57:31 -0500 Subject: [PATCH 3/3] Add Query error tests --- test/unit/CommonLibTest.csproj | 2 +- test/unit/LdapUtilsQueryTest.cs | 79 ++++++++++++++++++++++++++------- 2 files changed, 63 insertions(+), 18 deletions(-) diff --git a/test/unit/CommonLibTest.csproj b/test/unit/CommonLibTest.csproj index cd4ca2f3..717d5c93 100644 --- a/test/unit/CommonLibTest.csproj +++ b/test/unit/CommonLibTest.csproj @@ -1,7 +1,7 @@ - net8.0 + net7.0 false true ..\..\docfx\coverage\ diff --git a/test/unit/LdapUtilsQueryTest.cs b/test/unit/LdapUtilsQueryTest.cs index c7a2b54a..b5dd30cb 100644 --- a/test/unit/LdapUtilsQueryTest.cs +++ b/test/unit/LdapUtilsQueryTest.cs @@ -6,17 +6,8 @@ using System.DirectoryServices.Protocols; using SharpHoundCommonLib; -public class RangedRetrievalTests +public class LdapUtilsQueryTest { - private Mock _mockConnectionPool; - private LdapUtils _utils; - - public RangedRetrievalTests() - { - _mockConnectionPool = new Mock(); - _utils = new LdapUtils(); - } - // [Fact] // public async Task RangedRetrieval_SuccessfulRetrieval_ReturnsExpectedResults() // { @@ -28,10 +19,13 @@ public RangedRetrievalTests() // var connectionWrapper = new Mock(); // var connection = new Mock(); // connectionWrapper.SetupGet(x => x.Connection).Returns(connection.Object); - - // _mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) + + // var mockConnectionPool = new Mock(); + // mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) // .ReturnsAsync((true, connectionWrapper.Object, null)); + // var utils = new LdapUtils(mockConnectionPool.Object); + // var searchResponse = new Mock(); // var entry = new SearchResultEntry // { @@ -47,7 +41,7 @@ public RangedRetrievalTests() // // Act // var results = new List>(); - // await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName)) + // await foreach (var result in utils.RangedRetrieval(distinguishedName, attributeName)) // { // results.Add(result); // } @@ -67,9 +61,11 @@ public async Task RangedRetrieval_ConnectionFailure_ReturnsFailResult() var distinguishedName = "CN=TestUser,DC=example,DC=com"; var attributeName = "member"; + var utils = new LdapUtils(); + // Act var results = new List>(); - await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName)) + await foreach (var result in utils.RangedRetrieval(distinguishedName, attributeName)) { results.Add(result); } @@ -113,18 +109,19 @@ public async Task RangedRetrieval_CancellationRequested_StopsRetrieval() var connectionWrapper = new Mock(null, null, false, string.Empty); var connection = new Mock(); + var mockConnectionPool = new Mock(); - _mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) + mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) .ReturnsAsync((true, connectionWrapper.Object, null)); - _utils = new LdapUtils(_mockConnectionPool.Object); + var utils = new LdapUtils(mockConnectionPool.Object); var cts = new CancellationTokenSource(); cts.Cancel(); // Act var results = new List>(); - await foreach (var result in _utils.RangedRetrieval(distinguishedName, attributeName, cts.Token)) + await foreach (var result in utils.RangedRetrieval(distinguishedName, attributeName, cts.Token)) { results.Add(result); } @@ -132,4 +129,52 @@ public async Task RangedRetrieval_CancellationRequested_StopsRetrieval() // Assert Assert.False(results[0].IsSuccess); } + + [Fact] + public async Task Query_ConnectionFailure_ReturnsEmptyResult() + { + // Arrange + var queryParams = new LdapQueryParameters(); + var utils = new LdapUtils(); + + // Act + var results = new List>(); + await foreach (var result in utils.Query(queryParams)) + { + results.Add(result); + } + + // Assert + Assert.Empty(results); + } + + [Fact] + public async Task Query_CancellationRequested_StopsRetrieval() + { + // Arrange + var queryParams = new LdapQueryParameters(); + var domain = "example.com"; + + var connectionWrapper = new Mock(null, null, false, string.Empty); + var connection = new Mock(); + var mockConnectionPool = new Mock(); + + mockConnectionPool.Setup(x => x.GetLdapConnection(domain, false)) + .ReturnsAsync((true, connectionWrapper.Object, null)); + + var utils = new LdapUtils(mockConnectionPool.Object); + + var cts = new CancellationTokenSource(); + cts.Cancel(); + + // Act + var results = new List>(); + await foreach (var result in utils.Query(queryParams, cts.Token)) + { + results.Add(result); + } + + // Assert + Assert.Empty(results); + } } \ No newline at end of file