Skip to content

Commit

Permalink
fix: fix ldap backoff retry logic to actually create new connections,…
Browse files Browse the repository at this point in the history 
… and clean up edge cases
  • Loading branch information
@rvazarkar
rvazarkar committed Feb 2, 2024
1 parent ab2e21a commit de0ffda
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 119 deletions.
3 changes: 2 additions & 1 deletion src/CommonLib/Enums/LdapErrorCodes.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ public enum LdapErrorCodes : int
{
Success = 0,
Busy = 51,
ServerDown = 81
ServerDown = 81,
LocalError = 82
}
}
181 changes: 63 additions & 118 deletions src/CommonLib/LDAPUtils.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
using System.DirectoryServices;
using System.DirectoryServices.ActiveDirectory;
using System.DirectoryServices.Protocols;
using System.IO.Compression;
using System.Linq;
using System.Net;
using System.Net.Sockets;
Expand All @@ -18,9 +19,7 @@
using SharpHoundCommonLib.LDAPQueries;
using SharpHoundCommonLib.OutputTypes;
using SharpHoundCommonLib.Processors;
using SharpHoundRPC;
using SharpHoundRPC.NetAPINative;
using SharpHoundRPC.Wrappers;
using Domain = System.DirectoryServices.ActiveDirectory.Domain;
using SearchScope = System.DirectoryServices.Protocols.SearchScope;
using SecurityMasks = System.DirectoryServices.Protocols.SecurityMasks;
Expand Down Expand Up @@ -53,8 +52,8 @@ private static readonly ConcurrentDictionary<string, ResolvedWellKnownPrincipal>
private readonly ConcurrentDictionary<string, Domain> _domainCache = new();
private readonly ConcurrentDictionary<string, string> _domainControllerCache = new();
private static readonly TimeSpan MinBackoffDelay = TimeSpan.FromSeconds(2);
private static readonly TimeSpan MaxBackoffDelay = TimeSpan.FromSeconds(10);
private static readonly TimeSpan BackoffDelayMultiplier = TimeSpan.FromSeconds(2);
private static readonly TimeSpan MaxBackoffDelay = TimeSpan.FromSeconds(20);
private static readonly int BackoffDelayMultiplier = 2;
private const int MaxRetries = 3;

private readonly ConcurrentDictionary<string, LdapConnection> _globalCatalogConnections = new();
Expand All @@ -66,6 +65,7 @@ private static readonly ConcurrentDictionary<string, ResolvedWellKnownPrincipal>
private readonly ConcurrentDictionary<string, string> _netbiosCache = new();
private readonly PortScanner _portScanner;
private LDAPConfig _ldapConfig = new();
private readonly SemaphoreSlim _semaphoreSlim = new(1, 1);

/// <summary>
/// Creates a new instance of LDAP Utils with defaults
Expand Down Expand Up @@ -517,8 +517,7 @@ public IEnumerable<string> DoRangedRetrieval(string distinguishedName, string at
//Allow three retries with a backoff on each one if we get a "Server is Busy" error
retryCount++;
Thread.Sleep(backoffDelay);
backoffDelay = TimeSpan.FromSeconds(Math.Min(
backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
backoffDelay = GetNextBackoff(retryCount);
continue;
}
catch (Exception e)
Expand Down Expand Up @@ -874,38 +873,65 @@ public IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, SearchScope
}catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.ServerDown &&
retryCount < MaxRetries)
{
var isSemaphoreHeld = _semaphoreSlim.CurrentCount == 0;
retryCount++;
Thread.Sleep(backoffDelay);
backoffDelay = TimeSpan.FromSeconds(Math.Min(
backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
conn = CreateNewConnection(domainName, globalCatalog, skipCache);
if (conn == null)

_semaphoreSlim.Wait(cancellationToken);
try
{
_log.LogError("Unable to create replacement ldap connection for ServerDown exception. Breaking loop");
yield break;
if (!isSemaphoreHeld)
{
Thread.Sleep(backoffDelay);
backoffDelay = GetNextBackoff(retryCount);
conn = CreateNewConnection(domainName, globalCatalog, true);
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");
}
else
{
backoffDelay = GetNextBackoff(retryCount);
conn = CreateNewConnection(domainName, globalCatalog);
}
}
finally
{
_semaphoreSlim.Release();
}

_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);
backoffDelay = TimeSpan.FromSeconds(Math.Min(
backoffDelay.TotalSeconds * BackoffDelayMultiplier.TotalSeconds, MaxBackoffDelay.TotalSeconds));
backoffDelay = GetNextBackoff(retryCount);
continue;
}
catch (LdapException le)
{
if (le.ErrorCode != 82)
if (le.ErrorCode != (int)LdapErrorCodes.LocalError)
{
if (throwException)
{
throw new LDAPQueryException(
$"LDAP Exception in Loop: {le.ErrorCode}. {le.ServerErrorMessage}. {le.Message}. Filter: {ldapFilter}. Domain: {domainName}.",
$"LDAP Exception in Loop: {le.ErrorCode}. {le.ServerErrorMessage}. {le.Message}. Filter: {ldapFilter}. Domain: {domainName}",
le);
else
_log.LogWarning(le,
"LDAP Exception in Loop: {ErrorCode}. {ServerErrorMessage}. {Message}. Filter: {Filter}. Domain: {Domain}",
le.ErrorCode, le.ServerErrorMessage, le.Message, ldapFilter, domainName);
}

_log.LogWarning(le,
"LDAP Exception in Loop: {ErrorCode}. {ServerErrorMessage}. {Message}. Filter: {Filter}. Domain: {Domain}",
le.ErrorCode, le.ServerErrorMessage, le.Message, ldapFilter, domainName);
}

if (le.ErrorCode == (int)LdapErrorCodes.ServerDown)
{
throw new LDAPQueryException(
$"LDAP Exception in Loop: {le.ErrorCode}. {le.ServerErrorMessage}. {le.Message}. Filter: {ldapFilter}. Domain: {domainName}",
le);
}

yield break;
}
catch (Exception e)
Expand Down Expand Up @@ -979,99 +1005,15 @@ public virtual IEnumerable<ISearchResultEntry> QueryLDAP(string ldapFilter, Sear
string[] props, string domainName = null, bool includeAcl = false, bool showDeleted = false,
string adsPath = null, bool globalCatalog = false, bool skipCache = false, bool throwException = false)
{
var queryParams = SetupLDAPQueryFilter(
ldapFilter, scope, props, includeAcl, domainName, includeAcl, adsPath, globalCatalog, skipCache);

if (queryParams.Exception != null)
{
if (throwException) throw queryParams.Exception;

_log.LogWarning(queryParams.Exception, "Failed to setup LDAP Query Filter");
yield break;
}
var conn = queryParams.Connection;
var request = queryParams.SearchRequest;
var pageControl = queryParams.PageControl;

PageResultResponseControl pageResponse = null;

var backoffDelay = MinBackoffDelay;
var retryCount = 0;

while (true)
{
SearchResponse response;

try
{
_log.LogTrace("Sending LDAP request for {Filter}", ldapFilter);
response = (SearchResponse)conn.SendRequest(request);
if (response != null)
pageResponse = (PageResultResponseControl)response.Controls
.Where(x => x is PageResultResponseControl).DefaultIfEmpty(null).FirstOrDefault();
}
catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.Busy && retryCount < MaxRetries)
{
retryCount++;
Thread.Sleep(backoffDelay);
backoffDelay = TimeSpan.FromSeconds(Math.Min(
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)
if (throwException)
throw new LDAPQueryException(
$"LDAP Exception in Loop: {le.ErrorCode}. {le.ServerErrorMessage}. {le.Message}. Filter: {ldapFilter}. Domain: {domainName}",
le);
else
_log.LogWarning(le,
"LDAP Exception in Loop: {ErrorCode}. {ServerErrorMessage}. {Message}. Filter: {Filter}. Domain: {Domain}",
le.ErrorCode, le.ServerErrorMessage, le.Message, ldapFilter, domainName);
yield break;
}
catch (Exception e)
{
if (throwException)
throw new LDAPQueryException(
$"Exception in LDAP loop for {ldapFilter} and {domainName ?? "Default Domain"}", e);

_log.LogWarning(e, "Exception in LDAP loop for {Filter} and {Domain}", ldapFilter,
domainName ?? "Default Domain");
yield break;
}

if (response == null || pageResponse == null) continue;

if (response.Entries == null)
yield break;

foreach (SearchResultEntry entry in response.Entries)
yield return new SearchResultEntryWrapper(entry, this);

if (pageResponse.Cookie.Length == 0 || response.Entries.Count == 0)
yield break;
return QueryLDAP(ldapFilter, scope, props, new CancellationToken(), domainName, includeAcl, showDeleted,
adsPath, globalCatalog, skipCache, throwException);
}

pageControl.Cookie = pageResponse.Cookie;
}
private static TimeSpan GetNextBackoff(int retryCount)
{
return TimeSpan.FromSeconds(Math.Min(
BackoffDelayMultiplier * (retryCount + 1) * retryCount,
MaxBackoffDelay.TotalSeconds));
}

/// <summary>
Expand Down Expand Up @@ -1560,8 +1502,11 @@ private async Task<LdapConnection> CreateLDAPConnection(string domainName = null

connection.AuthType = authType;

if (!skipCache)
_ldapConnections.TryAdd(targetServer, connection);
_ldapConnections.AddOrUpdate(targetServer, connection, (s, ldapConnection) =>
{
ldapConnection.Dispose();
return connection;
});

return connection;
}
Expand Down

0 comments on commit de0ffda

Please sign in to comment.