Skip to content

Commit

Permalink
Disable Semaphore and fix some flaky tests (#157)
Browse files Browse the repository at this point in the history 
* fix: temporarily disable semaphore until we can iron out deadlocks
chore: some misc test fixes to fix flakiness
chore: fix an async function that wasn't async

* chore: bump version

* chore: improve retry logging at debug level
  • Loading branch information
@rvazarkar
rvazarkar authored Aug 29, 2024
1 parent 33f990a commit 3f8f147
Show file tree
Hide file tree
Showing 9 changed files with 120 additions and 110 deletions.
4 changes: 2 additions & 2 deletions src/CommonLib/ConnectionPoolManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,13 +99,13 @@ private bool GetPool(string identifier, out LdapConnectionPool pool) {
return await pool.GetConnectionAsync();
}

public async Task<(bool Success, LdapConnectionWrapper connectionWrapper, string Message)> GetLdapConnectionForServer(
public (bool Success, LdapConnectionWrapper connectionWrapper, string Message) GetLdapConnectionForServer(
string identifier, string server, bool globalCatalog) {
if (!GetPool(identifier, out var pool)) {
return (false, default, $"Unable to resolve a pool for {identifier}");
}

return await pool.GetConnectionForSpecificServerAsync(server, globalCatalog);
return pool.GetConnectionForSpecificServerAsync(server, globalCatalog);
}

private string ResolveIdentifier(string identifier) {
Expand Down
56 changes: 43 additions & 13 deletions src/CommonLib/LdapConnectionPool.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,13 @@ internal class LdapConnectionPool : IDisposable{
public LdapConnectionPool(string identifier, string poolIdentifier, LdapConfig config, PortScanner scanner = null, NativeMethods nativeMethods = null, ILogger log = null) {
_connections = new ConcurrentBag<LdapConnectionWrapper>();
_globalCatalogConnection = new ConcurrentBag<LdapConnectionWrapper>();
if (config.MaxConcurrentQueries > 0) {
_semaphore = new SemaphoreSlim(config.MaxConcurrentQueries, config.MaxConcurrentQueries);
} else {
//If MaxConcurrentQueries is 0, we'll just disable the semaphore entirely
_semaphore = null;
}
//TODO: Re-enable this once we track down the semaphore deadlock
// if (config.MaxConcurrentQueries > 0) {
// _semaphore = new SemaphoreSlim(config.MaxConcurrentQueries, config.MaxConcurrentQueries);
// } else {
// //If MaxConcurrentQueries is 0, we'll just disable the semaphore entirely
// _semaphore = null;
// }

_identifier = identifier;
_poolIdentifier = poolIdentifier;
Expand Down Expand Up @@ -83,7 +84,9 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
while (!cancellationToken.IsCancellationRequested) {
//Grab our semaphore here to take one of our query slots
if (_semaphore != null){
_log.LogTrace("Query entering semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
await _semaphore.WaitAsync(cancellationToken);
_log.LogTrace("Query entered semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
try {
_log.LogTrace("Sending ldap request - {Info}", queryParameters.GetQueryInfo());
Expand Down Expand Up @@ -111,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 recover from ServerDown for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), queryRetryCount);
ReleaseConnection(connectionWrapper, true);

for (var retryCount = 0; retryCount < MaxRetries; retryCount++) {
Expand Down Expand Up @@ -141,6 +145,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
* The expectation is that given enough time, the server should stop being busy and service our query appropriately
*/
busyRetryCount++;
_log.LogDebug("Query - Executing busy backoff for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), busyRetryCount);
var backoffDelay = GetNextBackoff(busyRetryCount);
await Task.Delay(backoffDelay, cancellationToken);
} catch (LdapException le) {
Expand All @@ -159,7 +164,11 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> Query(LdapQueryParam
queryParameters);
} finally {
// Always release our semaphore to prevent deadlocks
_semaphore?.Release();
if (_semaphore != null) {
_log.LogTrace("Query releasing semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
_semaphore.Release();
_log.LogTrace("Query released semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
}

//If we have a tempResult set it means we hit an error we couldn't recover from, so yield that result and then break out of the function
Expand Down Expand Up @@ -214,7 +223,9 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery

while (!cancellationToken.IsCancellationRequested) {
if (_semaphore != null){
_log.LogTrace("PagedQuery entering semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
await _semaphore.WaitAsync(cancellationToken);
_log.LogTrace("PagedQuery entered semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
SearchResponse response = null;
try {
Expand Down Expand Up @@ -249,13 +260,15 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
ReleaseConnection(connectionWrapper, true);
yield break;
}


_log.LogDebug("PagedQuery - Attempting to recover from ServerDown for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), queryRetryCount);

ReleaseConnection(connectionWrapper, true);
for (var retryCount = 0; retryCount < MaxRetries; retryCount++) {
var backoffDelay = GetNextBackoff(retryCount);
await Task.Delay(backoffDelay, cancellationToken);
var (success, ldapConnectionWrapperNew, _) =
await GetConnectionForSpecificServerAsync(serverName, queryParameters.GlobalCatalog);
GetConnectionForSpecificServerAsync(serverName, queryParameters.GlobalCatalog);

if (success) {
_log.LogDebug("PagedQuery - Recovered from ServerDown successfully");
Expand All @@ -277,6 +290,7 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
* The expectation is that given enough time, the server should stop being busy and service our query appropriately
*/
busyRetryCount++;
_log.LogDebug("PagedQuery - Executing busy backoff for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), busyRetryCount);
var backoffDelay = GetNextBackoff(busyRetryCount);
await Task.Delay(backoffDelay, cancellationToken);
} catch (LdapException le) {
Expand All @@ -288,7 +302,11 @@ public async IAsyncEnumerable<LdapResult<IDirectoryObject>> PagedQuery(LdapQuery
LdapResult<IDirectoryObject>.Fail($"PagedQuery - Caught unrecoverable exception: {e.Message}",
queryParameters);
} finally {
_semaphore?.Release();
if (_semaphore != null) {
_log.LogTrace("PagedQuery releasing semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
_semaphore.Release();
_log.LogTrace("PagedQuery released semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
}

if (tempResult != null) {
Expand Down Expand Up @@ -402,17 +420,21 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
while (!cancellationToken.IsCancellationRequested) {
SearchResponse response = null;
if (_semaphore != null){
_log.LogTrace("RangedRetrieval entering semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
await _semaphore.WaitAsync(cancellationToken);
_log.LogTrace("RangedRetrieval entered semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
try {
response = (SearchResponse)connectionWrapper.Connection.SendRequest(searchRequest);
} catch (LdapException le) when (le.ErrorCode == (int)ResultCode.Busy && busyRetryCount < MaxRetries) {
busyRetryCount++;
_log.LogDebug("RangedRetrieval - Executing busy backoff for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), busyRetryCount);
var backoffDelay = GetNextBackoff(busyRetryCount);
await Task.Delay(backoffDelay, cancellationToken);
} catch (LdapException le) when (le.ErrorCode == (int)LdapErrorCodes.ServerDown &&
queryRetryCount < MaxRetries) {
queryRetryCount++;
_log.LogDebug("RangedRetrieval - Attempting to recover from ServerDown for query {Info} (Attempt {Count})", queryParameters.GetQueryInfo(), queryRetryCount);
ReleaseConnection(connectionWrapper, true);
for (var retryCount = 0; retryCount < MaxRetries; retryCount++) {
var backoffDelay = GetNextBackoff(retryCount);
Expand Down Expand Up @@ -446,7 +468,11 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
tempResult =
LdapResult<string>.Fail($"Caught unrecoverable exception: {e.Message}", queryParameters);
} finally {
_semaphore?.Release();
if (_semaphore != null) {
_log.LogTrace("RangedRetrieval releasing semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
_semaphore.Release();
_log.LogTrace("RangedRetrieval released semaphore with {Count} remaining for query {Info}", _semaphore.CurrentCount, queryParameters.GetQueryInfo());
}
}

//If we have a tempResult set it means we hit an error we couldn't recover from, so yield that result and then break out of the function
Expand All @@ -471,13 +497,17 @@ public async IAsyncEnumerable<Result<string>> RangedRetrieval(string distinguish
step = entry.Attributes[currentRange].Count;
}

//Release our connection before we iterate
if (complete) {
ReleaseConnection(connectionWrapper);
}

foreach (string dn in entry.Attributes[currentRange].GetValues(typeof(string))) {
yield return Result<string>.Ok(dn);
index++;
}

if (complete) {
ReleaseConnection(connectionWrapper);
yield break;
}

Expand Down Expand Up @@ -577,7 +607,7 @@ private bool CallDsGetDcName(string domainName, out NetAPIStructs.DomainControll
return (true, connectionWrapper, null);
}

public async Task<(bool Success, LdapConnectionWrapper connectionWrapper, string Message)>
public (bool Success, LdapConnectionWrapper connectionWrapper, string Message)
GetConnectionForSpecificServerAsync(string server, bool globalCatalog) {
return CreateNewConnectionForServer(server, globalCatalog);
}
Expand Down
13 changes: 10 additions & 3 deletions src/CommonLib/LdapQueryParameters.cs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
using System;
using System.DirectoryServices.Protocols;
using System.Threading;
using SharpHoundCommonLib.Enums;

namespace SharpHoundCommonLib {
public class LdapQueryParameters
{
public class LdapQueryParameters {
private static int _queryIDIndex;
private string _searchBase;
private string _relativeSearchBase;
public string LDAPFilter { get; set; }
Expand All @@ -14,6 +15,12 @@ public class LdapQueryParameters
public bool GlobalCatalog { get; set; }
public bool IncludeSecurityDescriptor { get; set; } = false;
public bool IncludeDeleted { get; set; } = false;
private int QueryID { get; }

public LdapQueryParameters() {
QueryID = _queryIDIndex;
Interlocked.Increment(ref _queryIDIndex);
}

public string SearchBase {
get => _searchBase;
Expand All @@ -35,7 +42,7 @@ public string RelativeSearchBase {

public string GetQueryInfo()
{
return $"Query Information - Filter: {LDAPFilter}, Domain: {DomainName}, GlobalCatalog: {GlobalCatalog}, ADSPath: {SearchBase}";
return $"Query Information - Filter: {LDAPFilter}, Domain: {DomainName}, GlobalCatalog: {GlobalCatalog}, ADSPath: {SearchBase}, ID: {QueryID}";
}
}
}
2 changes: 1 addition & 1 deletion src/CommonLib/Processors/ComputerAvailability.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ public class ComputerAvailability
private readonly bool _skipPasswordCheck;
private readonly bool _skipPortScan;

public ComputerAvailability(int timeout = 500, int computerExpiryDays = 60, bool skipPortScan = false,
public ComputerAvailability(int timeout = 10000, int computerExpiryDays = 60, bool skipPortScan = false,
bool skipPasswordCheck = false, ILogger log = null)
{
_scanner = new PortScanner();
Expand Down
7 changes: 4 additions & 3 deletions src/CommonLib/Processors/ComputerSessionProcessor.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@ public class ComputerSessionProcessor {
private readonly string _localAdminUsername;
private readonly string _localAdminPassword;

public ComputerSessionProcessor(ILdapUtils utils, string currentUserName = null,
NativeMethods nativeMethods = null, ILogger log = null, bool doLocalAdminSessionEnum = false,
public ComputerSessionProcessor(ILdapUtils utils,
NativeMethods nativeMethods = null, ILogger log = null, string currentUserName = null,
bool doLocalAdminSessionEnum = false,
string localAdminUsername = null, string localAdminPassword = null) {
_utils = utils;
_nativeMethods = nativeMethods ?? new NativeMethods();
Expand Down Expand Up @@ -92,7 +93,7 @@ await SendComputerStatus(new CSVComputerStatus {
return ret;
}

_log.LogTrace("NetSessionEnum succeeded on {ComputerName}", computerName);
_log.LogDebug("NetSessionEnum succeeded on {ComputerName}", computerName);
await SendComputerStatus(new CSVComputerStatus {
Status = CSVComputerStatus.StatusSuccess,
Task = "NetSessionEnum",
Expand Down
2 changes: 1 addition & 1 deletion src/CommonLib/SharpHoundCommonLib.csproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
<PackageDescription>Common library for C# BloodHound enumeration tasks</PackageDescription>
<PackageLicenseExpression>GPL-3.0-only</PackageLicenseExpression>
<RepositoryUrl>https://github.com/BloodHoundAD/SharpHoundCommon</RepositoryUrl>
<Version>4.0.5</Version>
<Version>4.0.6</Version>
<AssemblyName>SharpHoundCommonLib</AssemblyName>
<RootNamespace>SharpHoundCommonLib</RootNamespace>
</PropertyGroup>
Expand Down
Loading

0 comments on commit 3f8f147

Please sign in to comment.