From dd62d84affecf7a4d381bc1a038c4b36d4291e31 Mon Sep 17 00:00:00 2001
From: Alex Nemeth <anemeth@specterops.io>
Date: Thu, 27 Jun 2024 09:26:52 -0700
Subject: [PATCH] RangedRetrieval polish continued

---
 src/CommonLib/LdapUtilsNew.cs | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/CommonLib/LdapUtilsNew.cs b/src/CommonLib/LdapUtilsNew.cs
index 011c23c4..c2fe54e8 100644
--- a/src/CommonLib/LdapUtilsNew.cs
+++ b/src/CommonLib/LdapUtilsNew.cs
@@ -160,7 +160,7 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(
                 await HandleBusyServer(busyRetryCount++, cancellationToken);
             }
             catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.ServerDown && queryRetryCount < MaxRetries) {
-                var newConnection = await HandleServerDown(connectionWrapper, queryParameters, cancellationToken);
+                var newConnection = await HandleServerDown(domain, connectionWrapper, cancellationToken);
                 if (newConnection.Success)
                 {
                     connectionWrapper = newConnection.Wrapper;
@@ -199,6 +199,26 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(
         return (complete, attr, values);
     }
 
+    private async Task<(bool Success, ConnectionWrapper Wrapper, Result<string> Error)> HandleServerDown(
+        string domain,
+        ConnectionWrapper connectionWrapper,
+        CancellationToken cancellationToken) {
+        _connectionPool.ReleaseConnection(connectionWrapper, true);
+
+        for (var retryCount = 0; retryCount < MaxRetries; retryCount++) {
+            await Task.Delay(GetNextBackoff(retryCount), cancellationToken);
+            var (success, newConnectionWrapper, message) = await _connectionPool.GetLdapConnection(domain, false);
+            
+            if (success) {
+                _log.LogDebug("RangedRetrieval - Recovered from ServerDown successfully, connection made to {NewServer}", newConnectionWrapper.GetServer());
+                return (true, newConnectionWrapper, null);
+            }
+        }
+
+        _log.LogError("RangedRetrieval - Failed to get a new connection after ServerDown");
+        return (false, null, Result<string>.Fail("RangedRetrieval - Failed to get a new connection after ServerDown."));
+    }
+
     public async IAsyncEnumerable<LdapResult<ISearchResultEntry>> Query(LdapQueryParameters queryParameters,
         [EnumeratorCancellation] CancellationToken cancellationToken = new()) {
         var setupResult = await SetupLdapQuery(queryParameters);
@@ -339,10 +359,6 @@ public async IAsyncEnumerable<LdapResult<ISearchResultEntry>> PagedQuery(
             var queryResult = await ExecutePagedQuery(searchRequest, connectionWrapper, queryParameters, serverName,
                 cancellationToken);
             if (!queryResult.Success) {
-                if (queryResult.Error != null) {
-                    yield return queryResult.Error;
-                }
-
                 yield break;
             }