[BED-5972] Add cancellation tokens to timeout logic (#213)

* WIP new Timeout replacement

* Trying to mitigate task failures from timeout expiring before task makes it into the thread pool

* feat: Timeout with cancellation token

* Adding Cancellation Tokens where appropriate

* Put timeout on LookupPrincipalBySid; correct tests; mark WindowsOnlyFacts

* Add new ExecuteWithTimeout tests, fix DCLdapProcessor_CheckScan_Timeout

* chore: add tests for early exit of ExecuteWithTimeout tasks; formatting

* Add test failure message to 'Timeout_Cancel' tests for clarity of intent

* Add comments to describe the double-ExecuteWithTimeout decision

* Add LongRunning hint to synchronous ExecuteWithTimeout

* Add LongRunning hint to async ExecuteWithTimeout as well

* Reverting last commit: Too many potential pitfalls

* Decided to add cancellation token to Task.Factory.StartNew after all

* I think technically this will more comprehensively capture completed work tasks

Author: definitelynotagoblin

Dockerfile

@@ -1,7 +1,7 @@
 FROM mono:6.12.0
 
 # Install .NET SDK
-ENV DOTNET_VERSION=7.0
+ENV DOTNET_VERSION=8.0
 
 RUN curl -sSL https://dot.net/v1/dotnet-install.sh  \
     | bash -s -- -Channel $DOTNET_VERSION -InstallDir /usr/share/dotnet \

src/CommonLib/Helpers.cs

@@ -11,6 +11,10 @@
 using System.Security;
 using SharpHoundCommonLib.Processors;
 using Microsoft.Win32;
+using System.Threading.Tasks;
+using System.Threading;
+using SharpHoundRPC.NetAPINative;
+using SharpHoundRPC.Shared;
 
 namespace SharpHoundCommonLib {
     public static class Helpers {
@@ -135,7 +139,8 @@ public static string DistinguishedNameToDomain(string distinguishedName) {
             int idx;
             if (distinguishedName.ToUpper().Contains("DELETED OBJECTS")) {
                 idx = distinguishedName.IndexOf("DC=", 3, StringComparison.Ordinal);
-            } else {
+            }
+            else {
                 idx = distinguishedName.IndexOf("DC=",
                     StringComparison.CurrentCultureIgnoreCase);
             }
@@ -193,7 +198,8 @@ public static long ConvertFileTimeToUnixEpoch(string ldapTime) {
 
             try {
                 toReturn = (long)Math.Floor(DateTime.FromFileTimeUtc(time).Subtract(EpochDiff).TotalSeconds);
-            } catch {
+            }
+            catch {
                 toReturn = -1;
             }
 
@@ -209,7 +215,8 @@ public static long ConvertTimestampToUnixEpoch(string ldapTime) {
             try {
                 var dt = DateTime.ParseExact(ldapTime, "yyyyMMddHHmmss.0K", CultureInfo.CurrentCulture).ToUniversalTime();
                 return (long)dt.Subtract(EpochDiff).TotalSeconds;
-            } catch {
+            }
+            catch {
                 return 0;
             }
         }
@@ -263,19 +270,23 @@ public static RegistryResult GetRegistryKeyData(string target, string subkey, st
                 data.Value = value;
 
                 data.Collected = true;
-            } catch (IOException e) {
+            }
+            catch (IOException e) {
                 log.LogDebug(e, "Error getting data from registry for {Target}: {RegSubKey}:{RegValue}",
                     target, subkey, subvalue);
                 data.FailureReason = "Target machine was not found or not connectable";
-            } catch (SecurityException e) {
+            }
+            catch (SecurityException e) {
                 log.LogDebug(e, "Error getting data from registry for {Target}: {RegSubKey}:{RegValue}",
                     target, subkey, subvalue);
                 data.FailureReason = "User does not have the proper permissions to perform this operation";
-            } catch (UnauthorizedAccessException e) {
+            }
+            catch (UnauthorizedAccessException e) {
                 log.LogDebug(e, "Error getting data from registry for {Target}: {RegSubKey}:{RegValue}",
                     target, subkey, subvalue);
                 data.FailureReason = "User does not have the necessary registry rights";
-            } catch (Exception e) {
+            }
+            catch (Exception e) {
                 log.LogDebug(e, "Error getting data from registry for {Target}: {RegSubKey}:{RegValue}",
                     target, subkey, subvalue);
                 data.FailureReason = e.Message;
@@ -300,7 +311,7 @@ public static IRegistryKey OpenRemoteRegistry(string target) {
             CommonOids.ClientAuthentication,
             CommonOids.AnyPurpose
         };
-        
+
         public static string DumpDirectoryObject(this IDirectoryObject directoryObject) {
             var builder = new StringBuilder();
             builder.AppendLine("PropertyName : PropertyValue");
@@ -310,6 +321,109 @@ public static string DumpDirectoryObject(this IDirectoryObject directoryObject)
 
             return builder.ToString();
         }
+
+        /// <summary>
+        /// Returns a Fail result if a task runs longer than its budgeted time.
+        /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="timeout"></param>
+        /// <param name="func"></param>
+        /// <returns></returns>
+        public static async Task<Result<T>> ExecuteWithTimeout<T>(TimeSpan timeout, Func<CancellationToken, T> func) {
+            var cts = new CancellationTokenSource();
+            var task = Task.Factory.StartNew(() => func(cts.Token), cts.Token, TaskCreationOptions.LongRunning, TaskScheduler.Default);
+            await Task.WhenAny(task, Task.Delay(timeout, cts.Token));
+            cts.Cancel();
+
+            if (task.IsCompleted) {
+                try {
+                    return Result<T>.Ok(await task);
+                }
+                catch (OperationCanceledException) { }
+            }
+
+            return Result<T>.Fail("Timeout");
+        }
+
+        // These two ExecuteWithTimeout functions should perform equivalently -
+        // they both create a new task from a function arg
+        // But where the one below can invoke an async function directly to spawn the Task
+        // The one above spawns a Task from a synchronous function.
+        // The caller shouldn't have to worry about which they're using however,
+        // the compiler should figure it out intrinsically
+
+        /// <summary>
+        /// Returns a Fail result if a task runs longer than its budgeted time.
+        /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="timeout"></param>
+        /// <param name="func"></param>
+        /// <returns></returns>
+        public static async Task<Result<T>> ExecuteWithTimeout<T>(TimeSpan timeout, Func<CancellationToken, Task<T>> func) {
+            var cts = new CancellationTokenSource();
+            var task = func.Invoke(cts.Token);
+            await Task.WhenAny(task, Task.Delay(timeout, cts.Token));
+            cts.Cancel();
+
+            if (task.IsCompleted) {
+                try {
+                    return Result<T>.Ok(await task);
+                }
+                catch (OperationCanceledException) { }
+            }
+
+            return Result<T>.Fail("Timeout");
+        }
+
+        /// <summary>
+        /// Returns a Fail result if a task runs longer than its budgeted time.
+        /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="timeout"></param>
+        /// <param name="func"></param>
+        /// <returns></returns>
+        public static async Task<NetAPIResult<T>> ExecuteNetAPIWithTimeout<T>(TimeSpan timeout, Func<CancellationToken, NetAPIResult<T>> func) {
+            var result = await ExecuteWithTimeout(timeout, func);
+            if (result.IsSuccess)
+                return result.Value;
+            else
+                return NetAPIResult<T>.Fail(result.Error);
+        }
+
+        /// <summary>
+        /// Returns a Fail result if a task runs longer than its budgeted time.
+        /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="timeout"></param>
+        /// <param name="func"></param>
+        /// <returns></returns>
+        public static async Task<SharpHoundRPC.Result<T>> ExecuteRPCWithTimeout<T>(TimeSpan timeout, Func<CancellationToken, SharpHoundRPC.Result<T>> func) {
+            var result = await ExecuteWithTimeout(timeout, func);
+            if (result.IsSuccess)
+                return result.Value;
+            else
+                return SharpHoundRPC.Result<T>.Fail(result.Error);
+        }
+
+        /// <summary>
+        /// Returns a Fail result if a task runs longer than its budgeted time.
+        /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="timeout"></param>
+        /// <param name="func"></param>
+        /// <returns></returns>
+        public static async Task<SharpHoundRPC.Result<T>> ExecuteRPCWithTimeout<T>(TimeSpan timeout, Func<CancellationToken, Task<SharpHoundRPC.Result<T>>> func) {
+            var result = await ExecuteWithTimeout(timeout, func);
+            if (result.IsSuccess)
+                return result.Value;
+            else
+                return SharpHoundRPC.Result<T>.Fail(result.Error);
+        }
     }
 
     public class ParsedGPLink {

src/CommonLib/Ntlm/NtlmAuthenticationHandler.cs

@@ -1,13 +1,13 @@
 using Microsoft.Extensions.Logging;
 using SharpHoundCommonLib.Processors;
 using SharpHoundCommonLib.ThirdParty.PSOpenAD;
-using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace SharpHoundCommonLib.Ntlm;
 
 interface INtlmAuthenticationHandler {
-    Task<object> PerformNtlmAuthenticationAsync(INtlmTransport transport);
+    Task<object> PerformNtlmAuthenticationAsync(INtlmTransport transport, CancellationToken cancellationToken = default);
 }
 
 /// <summary>
@@ -28,7 +28,7 @@ public NtlmAuthenticationHandler(string targetService, ILogger logger = null) {
         };
     }
 
-    public virtual async Task<object> PerformNtlmAuthenticationAsync(INtlmTransport transport) {
+    public virtual async Task<object> PerformNtlmAuthenticationAsync(INtlmTransport transport, CancellationToken cancellationToken = default) {
         using var context = new SspiContext(
             null,
             null,
@@ -39,12 +39,16 @@ public virtual async Task<object> PerformNtlmAuthenticationAsync(INtlmTransport
             Options.Signing
         );
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         // NEGOTIATE
         var negotiateMsgBytes = context.Step();
 
         // CHALLENGE
         var challengeMessageBytes = await transport.NegotiateAsync(negotiateMsgBytes);
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         // AUTHENTICATE
         var authenticateMsgBytes = context.Step(challengeMessageBytes);
 

src/CommonLib/Processors/ComputerSessionProcessor.cs

@@ -8,7 +8,6 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Win32;
 using SharpHoundCommonLib.OutputTypes;
-using SharpHoundRPC;
 using SharpHoundRPC.NetAPINative;
 
 namespace SharpHoundCommonLib.Processors {
@@ -58,7 +57,7 @@ public async Task<SessionAPIResult> ReadUserSessions(string computerName, string
 
             _log.LogDebug("Running NetSessionEnum for {ObjectName}", computerName);
 
-            var result = await Task.Run(() => {
+            var result = await Helpers.ExecuteNetAPIWithTimeout(timeout, (timeoutToken) => {
                 NetAPIResult<IEnumerable<NetSessionEnumResults>> result;
                 if (_doLocalAdminSessionEnum) {
                     // If we are authenticating using a local admin, we need to impersonate for this
@@ -67,19 +66,22 @@ public async Task<SessionAPIResult> ReadUserSessions(string computerName, string
                         result = _nativeMethods.NetSessionEnum(computerName);
                     }
 
+                    timeoutToken.ThrowIfCancellationRequested();
+
                     if (result.IsFailed) {
                         // Fall back to default User
                         _log.LogDebug(
                             "NetSessionEnum failed on {ComputerName} with local admin credentials: {Status}. Fallback to default user.",
                             computerName, result.Status);
                         result = _nativeMethods.NetSessionEnum(computerName);
                     }
-                } else {
+                }
+                else {
                     result = _nativeMethods.NetSessionEnum(computerName);
                 }
 
                 return result;
-            }).TimeoutAfter(timeout);
+            });
 
             if (result.IsFailed) {
                 await SendComputerStatus(new CSVComputerStatus {
@@ -121,7 +123,7 @@ await SendComputerStatus(new CSVComputerStatus {
                     username.Equals("anonymous logon", StringComparison.CurrentCultureIgnoreCase)) {
                     continue;
                 }
-                
+
                 //Filter out domains that are "."
                 if (computerDomain.Equals(".")) {
                     continue;
@@ -147,7 +149,8 @@ await SendComputerStatus(new CSVComputerStatus {
                 if (matchSuccess) {
                     results.AddRange(
                         sids.Select(s => new Session { ComputerSID = resolvedComputerSID, UserSID = s }));
-                } else {
+                }
+                else {
                     var res = await _utils.ResolveAccountName(username, computerDomain);
                     if (res.Success)
                         results.Add(new Session {
@@ -180,7 +183,7 @@ public async Task<SessionAPIResult> ReadUserSessionsPrivileged(string computerNa
 
             _log.LogDebug("Running NetWkstaUserEnum for {ObjectName}", computerName);
 
-            var result = await Task.Run(() => {
+            var result = await Helpers.ExecuteNetAPIWithTimeout(timeout, (timeoutToken) => {
                 NetAPIResult<IEnumerable<NetWkstaUserEnumResults>>
                     result;
                 if (_doLocalAdminSessionEnum) {
@@ -190,19 +193,22 @@ public async Task<SessionAPIResult> ReadUserSessionsPrivileged(string computerNa
                         result = _nativeMethods.NetWkstaUserEnum(computerName);
                     }
 
+                    timeoutToken.ThrowIfCancellationRequested();
+
                     if (result.IsFailed) {
                         // Fall back to default User
                         _log.LogDebug(
                             "NetWkstaUserEnum failed on {ComputerName} with local admin credentials: {Status}. Fallback to default user.",
                             computerName, result.Status);
                         result = _nativeMethods.NetWkstaUserEnum(computerName);
                     }
-                } else {
+                }
+                else {
                     result = _nativeMethods.NetWkstaUserEnum(computerName);
                 }
 
                 return result;
-            }).TimeoutAfter(timeout);
+            });
 
             if (result.IsFailed) {
                 await SendComputerStatus(new CSVComputerStatus {
@@ -244,7 +250,7 @@ await SendComputerStatus(new CSVComputerStatus {
                 if (string.IsNullOrWhiteSpace(username) || username.EndsWith("$", StringComparison.Ordinal)) {
                     continue;
                 }
-                
+
                 //Filter out domains that are "."
                 if (domain.Equals(".")) {
                     continue;
@@ -312,7 +318,8 @@ await SendComputerStatus(new CSVComputerStatus {
                 ret.Results = results.ToArray();
 
                 return ret;
-            } catch (Exception e) {
+            }
+            catch (Exception e) {
                 _log.LogTrace("Registry session enum failed on {ComputerName}: {Status}", computerName, e.Message);
                 await SendComputerStatus(new CSVComputerStatus {
                     Status = e.Message,
@@ -322,7 +329,8 @@ await SendComputerStatus(new CSVComputerStatus {
                 ret.Collected = false;
                 ret.FailureReason = e.Message;
                 return ret;
-            } finally {
+            }
+            finally {
                 key?.Dispose();
             }
         }