wip: testing for ldap exceptions

SpecterOps · Aug 29, 2024 · cce2c5e · cce2c5e
1 parent 1ebeebe
commit cce2c5e
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 6 deletions.
diff --git a/src/CommonLib/LdapConnectionPool.cs b/src/CommonLib/LdapConnectionPool.cs
@@ -90,7 +90,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
                 }
                 try {
                     _log.LogTrace("Sending ldap request - {Info}", queryParameters.GetQueryInfo());
-                    response = (SearchResponse)connectionWrapper.Connection.SendRequest(searchRequest);
+                    response = (SearchResponse)connectionWrapper.SendRequest(searchRequest);
 
                     if (response != null) {
                         querySuccess = true;
@@ -114,6 +114,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
                      * since non-paged queries do not require same server connections
                      */
                     queryRetryCount++;
+                    _log.LogDebug("Query - Attempting to get new connection after ServerDown for query {Info}", queryParameters.GetQueryInfo());
                     ReleaseConnection(connectionWrapper, true);
 
                     for (var retryCount = 0; retryCount < MaxRetries; retryCount++) {
@@ -145,6 +146,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
                      */
                     busyRetryCount++;
                     var backoffDelay = GetNextBackoff(busyRetryCount);
+                    _log.LogDebug("Query - Executing {Time} second backoff for query {Info}", backoffDelay.TotalSeconds, queryParameters.GetQueryInfo());
                     await Task.Delay(backoffDelay, cancellationToken);
                 } catch (LdapException le) {
                     /*
@@ -228,7 +230,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
                 SearchResponse response = null;
                 try {
                     _log.LogTrace("Sending paged ldap request - {Info}", queryParameters.GetQueryInfo());
-                    response = (SearchResponse)connectionWrapper.Connection.SendRequest(searchRequest);
+                    response = (SearchResponse)connectionWrapper.SendRequest(searchRequest);
                     if (response != null) {
                         pageResponse = (PageResultResponseControl)response.Controls
                             .Where(x => x is PageResultResponseControl).DefaultIfEmpty(null).FirstOrDefault();
@@ -420,7 +422,7 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
                     _log.LogTrace("RangedRetrieval entered semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
                 }
                 try {
-                    response = (SearchResponse)connectionWrapper.Connection.SendRequest(searchRequest);
+                    response = (SearchResponse)connectionWrapper.SendRequest(searchRequest);
                 } catch (LdapException le) when (le.ErrorCode == (int)ResultCode.Busy && busyRetryCount < MaxRetries) {
                     busyRetryCount++;
                     var backoffDelay = GetNextBackoff(busyRetryCount);
@@ -587,7 +589,7 @@ private bool CallDsGetDcName(string domainName, out NetAPIStructs.DomainControll
             return true;
         }
 
-        public async Task<(bool Success, LdapConnectionWrapper ConnectionWrapper, string Message)> GetConnectionAsync() {
+        public virtual async Task<(bool Success, LdapConnectionWrapper ConnectionWrapper, string Message)> GetConnectionAsync() {
             if (!_connections.TryTake(out var connectionWrapper)) {
                 var (success, connection, message) = await CreateNewConnection();
                 if (!success) {
@@ -605,11 +607,10 @@ private bool CallDsGetDcName(string domainName, out NetAPIStructs.DomainControll
             return CreateNewConnectionForServer(server, globalCatalog);
         }
 
-        public async Task<(bool Success, LdapConnectionWrapper ConnectionWrapper, string Message)> GetGlobalCatalogConnectionAsync() {
+        public virtual async Task<(bool Success, LdapConnectionWrapper ConnectionWrapper, string Message)> GetGlobalCatalogConnectionAsync() {
             if (!_globalCatalogConnection.TryTake(out var connectionWrapper)) {
                 var (success, connection, message) = await CreateNewConnection(true);
                 if (!success) {
-                    //If we didn't get a connection, immediately release the semaphore so we don't have hanging ones
                     return (false, null, message);
                 }
 

diff --git a/src/CommonLib/LdapConnectionWrapper.cs b/src/CommonLib/LdapConnectionWrapper.cs
@@ -32,6 +32,10 @@ public string GetServer() {
             return _server;
         }
 
+        public virtual DirectoryResponse SendRequest(SearchRequest searchRequest) {
+            return Connection.SendRequest(searchRequest);
+        }
+
         public bool GetSearchBase(NamingContext context, out string searchBase) {
             searchBase = GetSavedContext(context);
             if (searchBase != null) {

diff --git a/test/unit/LdapConnectionTests.cs b/test/unit/LdapConnectionTests.cs
@@ -0,0 +1,11 @@
+using System.Threading.Tasks;
+using Xunit;
+
+namespace CommonLibTest;
+
+public class LdapConnectionTests {
+    [Fact]
+    public async Task LdapConnection_TestExceptionHandling_ServerDown() {
+
+    }
+}