From cce2c5e946a0299bad4a7a826214a6f235b226ad Mon Sep 17 00:00:00 2001
From: rvazarkar <rvazarkar@specterops.io>
Date: Thu, 29 Aug 2024 15:35:33 -0400
Subject: [PATCH] wip: testing for ldap exceptions

---
 src/CommonLib/LdapConnectionPool.cs    | 13 +++++++------
 src/CommonLib/LdapConnectionWrapper.cs |  4 ++++
 test/unit/LdapConnectionTests.cs       | 11 +++++++++++
 3 files changed, 22 insertions(+), 6 deletions(-)
 create mode 100644 test/unit/LdapConnectionTests.cs

diff --git a/src/CommonLib/LdapConnectionPool.cs b/src/CommonLib/LdapConnectionPool.cs
index 490806cb..6396f36b 100644
--- 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
index 17314518..8d94d274 100644
--- 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
new file mode 100644
index 00000000..0747a523
--- /dev/null
+++ 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() {
+        
+    }
+}
\ No newline at end of file