From 96cce998209a96185afab7390610847052e547bc Mon Sep 17 00:00:00 2001
From: rvazarkar <rvazarkar@specterops.io>
Date: Tue, 10 Oct 2023 12:31:38 -0400
Subject: [PATCH] feat: add logic if we hit a serverdown to create a new ldap
 connection and retry the search

---
 src/CommonLib/Enums/LdapErrorCodes.cs |  1 +
 src/CommonLib/LDAPUtils.cs            | 53 ++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/src/CommonLib/Enums/LdapErrorCodes.cs b/src/CommonLib/Enums/LdapErrorCodes.cs
index 942b4546..31e61890 100644
--- a/src/CommonLib/Enums/LdapErrorCodes.cs
+++ b/src/CommonLib/Enums/LdapErrorCodes.cs
@@ -4,5 +4,6 @@ public enum LdapErrorCodes : int
     {
         Success = 0,
         Busy = 51,
+        ServerDown = 81
     }
 }
\ No newline at end of file
diff --git a/src/CommonLib/LDAPUtils.cs b/src/CommonLib/LDAPUtils.cs
index 59b89947..8945f989 100644
--- a/src/CommonLib/LDAPUtils.cs
+++ b/src/CommonLib/LDAPUtils.cs
@@ -837,6 +837,22 @@ public IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, SearchScope
                     if (response != null)
                         pageResponse = (PageResultResponseControl) response.Controls
                             .Where(x => x is PageResultResponseControl).DefaultIfEmpty(null).FirstOrDefault();
+                }catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.ServerDown &&
+                                                retryCount < MaxRetries)
+                {
+                    retryCount++;
+                    Thread.Sleep(backoffDelay);
+                    backoffDelay = TimeSpan.FromSeconds(Math.Min(
+                        backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
+                    conn = CreateNewConnection(domainName, globalCatalog, skipCache);
+                    if (conn == null)
+                    {
+                        _log.LogError("Unable to create replacement ldap connection for ServerDown exception. Breaking loop");
+                        yield break;
+                    }
+                    
+                    _log.LogInformation("Created new LDAP connection after receiving ServerDown from server");
+                    continue;
                 }catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.Busy && retryCount < MaxRetries) {
                     retryCount++;
                     Thread.Sleep(backoffDelay);
@@ -887,6 +903,22 @@ public IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, SearchScope
             }
         }
 
+        private LdapConnection CreateNewConnection(string domainName = null, bool globalCatalog = false, bool skipCache = false)
+        {
+            var task = globalCatalog
+                ? Task.Run(() => CreateGlobalCatalogConnection(domainName, _ldapConfig.AuthType))
+                : Task.Run(() => CreateLDAPConnection(domainName, skipCache, _ldapConfig.AuthType));
+
+            try
+            {
+                return task.ConfigureAwait(false).GetAwaiter().GetResult();
+            }
+            catch
+            {
+                return null;
+            }
+        }
+
         /// <summary>
         ///     Performs an LDAP query using the parameters specified by the user.
         /// </summary>
@@ -937,9 +969,9 @@ public virtual IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, Sear
                 try
                 {
                     _log.LogTrace("Sending LDAP request for {Filter}", ldapFilter);
-                    response = (SearchResponse) conn.SendRequest(request);
+                    response = (SearchResponse)conn.SendRequest(request);
                     if (response != null)
-                        pageResponse = (PageResultResponseControl) response.Controls
+                        pageResponse = (PageResultResponseControl)response.Controls
                             .Where(x => x is PageResultResponseControl).DefaultIfEmpty(null).FirstOrDefault();
                 }
                 catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.Busy && retryCount < MaxRetries)
@@ -950,6 +982,23 @@ public virtual IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, Sear
                         backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
                     continue;
                 }
+                catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.ServerDown &&
+                                               retryCount < MaxRetries)
+                {
+                    retryCount++;
+                    Thread.Sleep(backoffDelay);
+                    backoffDelay = TimeSpan.FromSeconds(Math.Min(
+                        backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
+                    conn = CreateNewConnection(domainName, globalCatalog, skipCache);
+                    if (conn == null)
+                    {
+                        _log.LogError("Unable to create replacement ldap connection for ServerDown exception. Breaking loop");
+                        yield break;
+                    }
+                    
+                    _log.LogInformation("Created new LDAP connection after receiving ServerDown from server");
+                    continue;
+                }
                 catch (LdapException le)
                 {
                     if (le.ErrorCode != 82)