[IAST] Add Stack trace to vuln location

tracer/src/Datadog.Trace/AppSec/AppSecRequestContext.cs

@@ -17,6 +17,7 @@ internal class AppSecRequestContext
 {
     private const string StackKey = "_dd.stack";
     private const string ExploitStackKey = "exploit";
+    private const string VulnerabilityStackKey = "vulnerability";
     private const string AppsecKey = "appsec";
     private readonly object _sync = new();
     private readonly List<object> _wafSecurityEvents = new();
@@ -68,21 +69,31 @@ internal void AddWafSecurityEvents(IReadOnlyCollection<object> events)
     }
 
     internal void AddRaspStackTrace(Dictionary<string, object> stackTrace, int maxStackTraces)
+    {
+        AddStackTrace(ExploitStackKey, stackTrace, maxStackTraces);
+    }
+
+    internal void AddVulnerabilityStackTrace(Dictionary<string, object> stackTrace, int maxStackTraces)
+    {
+        AddStackTrace(VulnerabilityStackKey, stackTrace, maxStackTraces);
+    }
+
+    internal void AddStackTrace(string stackCategory, Dictionary<string, object> stackTrace, int maxStackTraces)
     {
         lock (_sync)
         {
             _raspStackTraces ??= new();
 
-            if (!_raspStackTraces.ContainsKey(ExploitStackKey))
+            if (!_raspStackTraces.ContainsKey(stackCategory))
             {
-                _raspStackTraces.Add(ExploitStackKey, new());
+                _raspStackTraces.Add(stackCategory, new());
             }
-            else if (maxStackTraces > 0 && _raspStackTraces[ExploitStackKey].Count >= maxStackTraces)
+            else if (maxStackTraces > 0 && _raspStackTraces[stackCategory].Count >= maxStackTraces)
             {
                 return;
             }
 
-            _raspStackTraces[ExploitStackKey].Add(stackTrace);
+            _raspStackTraces[stackCategory].Add(stackTrace);
         }
     }
 }

tracer/src/Datadog.Trace/AppSec/Rasp/RaspModule.cs

@@ -177,7 +177,7 @@ private static void SendStack(Span rootSpan, string id)
 
         if (stack is not null)
         {
-            rootSpan.Context.TraceContext.AddStackTraceElement(stack, Security.Instance.Settings.MaxStackTraces);
+            rootSpan.Context.TraceContext.AddRaspStackTraceElement(stack, Security.Instance.Settings.MaxStackTraces);
         }
     }
 

tracer/src/Datadog.Trace/Iast/IastModule.cs

@@ -25,6 +25,7 @@
 using Datadog.Trace.Iast.Telemetry;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Sampling;
+using Datadog.Trace.VendoredMicrosoftCode.System;
 using static Datadog.Trace.Configuration.ConfigurationKeys;
 using static Datadog.Trace.Telemetry.Metrics.MetricTags;
 
@@ -422,15 +423,6 @@ public static void OnHardcodedSecret(Vulnerability vulnerability)
         }
     }
 
-    public static void OnHardcodedSecret(List<Vulnerability> vulnerabilities)
-    {
-        if (Iast.Instance.Settings.Enabled)
-        {
-            // We provide a hash value for the vulnerability instead of calculating one, following the agreed conventions
-            AddVulnerabilityAsSingleSpan(Tracer.Instance, IntegrationId.HardcodedSecret, OperationNameHardcodedSecret, vulnerabilities).SingleSpan?.Dispose();
-        }
-    }
-
     public static IastModuleResponse OnInsecureAuthProtocol(string authHeader, IntegrationId integrationId)
     {
         OnExecutedSinkTelemetry(IastInstrumentedSinks.InsecureAuthProtocol);
@@ -667,7 +659,7 @@ private static IastModuleResponse GetScope(string evidenceValue, IntegrationId i
             return IastModuleResponse.Empty;
         }
 
-        var location = addLocation ? GetLocation(externalStack, currentSpan?.SpanId) : null;
+        var location = addLocation ? GetLocation(externalStack, currentSpan) : null;
         if (addLocation && location is null)
         {
             return IastModuleResponse.Empty;
@@ -686,6 +678,8 @@ private static IastModuleResponse GetScope(string evidenceValue, IntegrationId i
                 traceContext?.Tags.SetTag(Tags.Propagated.AppSec, "1");
 
                 traceContext?.IastRequestContext?.AddVulnerability(vulnerability);
+                vulnerability.Location?.ReportStack(currentSpan);
+
                 return IastModuleResponse.Vulnerable;
             }
             else
@@ -697,50 +691,50 @@ private static IastModuleResponse GetScope(string evidenceValue, IntegrationId i
         return IastModuleResponse.Empty;
     }
 
-    private static Location? GetLocation(StackTrace? externalStack = null, ulong? currentSpanId = null)
+    private static Location? GetLocation(StackTrace? stack = null, Span? currentSpan = null)
     {
-        var frameInfo = StackWalker.GetFrame(externalStack);
-        if (!frameInfo.IsValid)
+        var stackFrame = StackWalker.GetFrame(ref stack);
+        if (stackFrame is null)
         {
             return null;
         }
 
-        return new Location(frameInfo.StackFrame, currentSpanId);
-    }
-
-    private static IastModuleResponse AddVulnerabilityAsSingleSpan(Tracer tracer, IntegrationId integrationId, string operationName, List<Vulnerability> vulnerabilities)
-    {
-        // we either are not in a request or the distributed tracer returned a scope that cannot be casted to Scope and we cannot access the root span.
-        var batch = GetVulnerabilityBatch();
-        foreach (var vulnerability in vulnerabilities)
+        string? stackId = null;
+        if (stack != null && Security.Instance.Settings.StackTraceEnabled)
         {
-            batch.Add(vulnerability);
+            if (currentSpan is null)
+            {
+                stackId = "1";
+            }
+            else
+            {
+                stackId = currentSpan.Context.TraceContext.GetNextVulnerabilityStackTraceId();
+            }
         }
 
-        return AddVulnerabilityAsSingleSpan(tracer, integrationId, operationName, batch.ToJson());
+        return new Location(stackFrame, stack, stackId, currentSpan?.SpanId);
     }
 
     private static IastModuleResponse AddVulnerabilityAsSingleSpan(Tracer tracer, IntegrationId integrationId, string operationName, Vulnerability vulnerability)
     {
         // we either are not in a request or the distributed tracer returned a scope that cannot be casted to Scope and we cannot access the root span.
         var batch = GetVulnerabilityBatch();
         batch.Add(vulnerability);
-        return AddVulnerabilityAsSingleSpan(tracer, integrationId, operationName, batch.ToJson());
-    }
 
-    private static IastModuleResponse AddVulnerabilityAsSingleSpan(Tracer tracer, IntegrationId integrationId, string operationName, string vulnsJson)
-    {
         var tags = new IastTags()
         {
-            IastJson = vulnsJson,
+            IastJson = batch.ToJson(),
             IastEnabled = "1"
         };
 
         var scope = tracer.StartActiveInternal(operationName, tags: tags);
-        scope.Span.Type = SpanTypes.IastVulnerability;
+        var span = scope.Span;
+        var traceContext = span.Context.TraceContext;
+        span.Type = SpanTypes.IastVulnerability;
         tracer.TracerManager.Telemetry.IntegrationGeneratedSpan(integrationId);
-        scope.Span.Context.TraceContext?.SetSamplingPriority(SamplingPriorityValues.UserKeep, SamplingMechanism.Asm);
-        scope.Span.Context.TraceContext?.Tags.SetTag(Tags.Propagated.AppSec, "1");
+        traceContext?.SetSamplingPriority(SamplingPriorityValues.UserKeep, SamplingMechanism.Asm);
+        traceContext?.Tags.SetTag(Tags.Propagated.AppSec, "1");
+        vulnerability.Location?.ReportStack(span);
         return new IastModuleResponse(scope);
     }
 

tracer/src/Datadog.Trace/Iast/IastRequestContext.cs

@@ -8,6 +8,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Specialized;
+using System.Threading;
 using System.Web;
 #if !NETFRAMEWORK
 using Microsoft.AspNetCore.Http;
@@ -31,6 +32,7 @@ internal class IastRequestContext
     private bool _routedParametersAdded = false;
     private bool _querySourcesAdded = false;
     private ExecutedTelemetryHelper? _executedTelemetryHelper = ExecutedTelemetryHelper.Enabled() ? new ExecutedTelemetryHelper() : null;
+    private int _lastVulnerabilityStackId = 0;
 
     internal static void AddIastDisabledFlagToSpan(Span span)
     {
@@ -379,4 +381,9 @@ internal void OnExecutedPropagationTelemetry()
     {
         _executedTelemetryHelper?.AddExecutedPropagation();
     }
+
+    internal string GetNextVulnerabilityStackId()
+    {
+        return Interlocked.Increment(ref _lastVulnerabilityStackId).ToString();
+    }
 }

tracer/src/Datadog.Trace/Iast/Location.cs

@@ -6,12 +6,18 @@
 #nullable enable
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
+using System.Linq;
+using Datadog.Trace.AppSec;
+using Datadog.Trace.AppSec.Rasp;
 
 namespace Datadog.Trace.Iast;
 
 internal readonly struct Location
 {
+    internal readonly StackTrace? _stack = null;
+
     public Location(string method)
     {
         var index = method.LastIndexOf("::", StringComparison.Ordinal);
@@ -29,7 +35,7 @@ public Location(string method)
         }
     }
 
-    public Location(StackFrame? stackFrame, ulong? spanId)
+    public Location(StackFrame? stackFrame, StackTrace? stack, string? stackId, ulong? spanId)
     {
         var method = stackFrame?.GetMethod();
         Path = method?.DeclaringType?.FullName;
@@ -38,9 +44,12 @@ public Location(StackFrame? stackFrame, ulong? spanId)
         Line = line > 0 ? line : null;
 
         SpanId = spanId == 0 ? null : spanId;
+
+        _stack = stack;
+        StackId = stackId;
     }
 
-    public Location(string? typeName, string? methodName, int? line, ulong? spanId)
+    internal Location(string? typeName, string? methodName, int? line, ulong? spanId) // For testing purposes only
     {
         this.Path = typeName;
         this.Method = methodName;
@@ -57,9 +66,25 @@ public Location(string? typeName, string? methodName, int? line, ulong? spanId)
 
     public int? Line { get; }
 
+    public string? StackId { get; }
+
     public override int GetHashCode()
     {
         // We do not calculate the hash including the spanId nor the line
         return IastUtils.GetHashCode(Path, Method);
     }
+
+    internal void ReportStack(Span? span)
+    {
+        if (span is not null && StackId is not null && _stack is not null && _stack.FrameCount > 0)
+        {
+#pragma warning disable CS8620 // Argument cannot be used for parameter due to differences in the nullability of reference types.
+            var stack = StackReporter.GetStack(Security.Instance.Settings.MaxStackTraceDepth, StackId, _stack.GetFrames());
+#pragma warning restore CS8620 // Argument cannot be used for parameter due to differences in the nullability of reference types.
+            if (stack is not null)
+            {
+                span.Context.TraceContext?.AddVulnerabilityStackTraceElement(stack, Security.Instance.Settings.MaxStackTraces);
+            }
+        }
+    }
 }

tracer/src/Datadog.Trace/Iast/StackWalker.cs

@@ -16,7 +16,13 @@ namespace Datadog.Trace.Iast;
 internal static class StackWalker
 {
     private const int DefaultSkipFrames = 2;
-    private static readonly string[] ExcludeSpanGenerationTypes = { "Datadog.Trace.Debugger.Helpers.StringExtensions", "Microsoft.AspNetCore.Razor.Language.StreamSourceDocument", "System.Security.IdentityHelper" };
+    private static readonly string[] ExcludeSpanGenerationTypes =
+    {
+        "Datadog.Trace.Debugger.Helpers.StringExtensions",
+        "Microsoft.AspNetCore.Razor.Language.StreamSourceDocument",
+        "System.Security.IdentityHelper"
+    };
+
     private static readonly string[] AssemblyNamesToSkip =
     {
         "Datadog.Trace",
@@ -41,9 +47,10 @@ internal static class StackWalker
 
     private static readonly ConcurrentDictionary<string, bool> ExcludedAssemblyCache = new ConcurrentDictionary<string, bool>();
 
-    public static StackFrameInfo GetFrame(StackTrace? externalStack = null)
+    public static StackFrame? GetFrame(ref StackTrace? externalStack)
     {
         var stackTrace = externalStack ?? new StackTrace(DefaultSkipFrames, true);
+        externalStack = stackTrace;
 
         foreach (var frame in stackTrace.GetFrames())
         {
@@ -53,26 +60,21 @@ public static StackFrameInfo GetFrame(StackTrace? externalStack = null)
             {
                 if (excludeType == declaringType?.FullName)
                 {
-                    return new StackFrameInfo(null, false);
+                    return null;
                 }
             }
 
-            if (ExcludeSpanGenerationTypes.Contains(declaringType?.FullName))
-            {
-                return new StackFrameInfo(null, false);
-            }
-
             var assembly = declaringType?.Assembly.GetName().Name;
-            if (assembly != null && !AssemblyExcluded(assembly))
+            if (assembly != null && !MustSkipAssembly(assembly))
             {
-                return new StackFrameInfo(frame, true);
+                return frame;
             }
         }
 
-        return new StackFrameInfo(null, true);
+        return null;
     }
 
-    public static bool AssemblyExcluded(string assembly)
+    public static bool MustSkipAssembly(string assembly)
     {
         if (ExcludedAssemblyCache.TryGetValue(assembly, out bool excluded))
         {
@@ -83,34 +85,34 @@ public static bool AssemblyExcluded(string assembly)
         ExcludedAssemblyCache[assembly] = excluded;
 
         return excluded;
-    }
 
-    // For performance reasons, we are not supporting wildcards fully. We just need to use '.' at the end for now. We can use regular expressions
-    // if in the future we need a more sophisticated wildcard support
-    private static bool IsExcluded(string assembly)
-    {
-        foreach (var assemblyToSkip in AssemblyNamesToSkip)
+        // For performance reasons, we are not supporting wildcards fully. We just need to use '.' at the end for now. We can use regular expressions
+        // if in the future we need a more sophisticated wildcard support
+        static bool IsExcluded(string assembly)
         {
+            foreach (var assemblyToSkip in AssemblyNamesToSkip)
+            {
 #if NETCOREAPP3_1_OR_GREATER
             if (assemblyToSkip.EndsWith('.'))
 #else
-            if (assemblyToSkip.EndsWith("."))
+                if (assemblyToSkip.EndsWith("."))
 #endif
-            {
-                if (assembly.StartsWith(assemblyToSkip, StringComparison.OrdinalIgnoreCase))
                 {
-                    return true;
+                    if (assembly.StartsWith(assemblyToSkip, StringComparison.OrdinalIgnoreCase))
+                    {
+                        return true;
+                    }
                 }
-            }
-            else
-            {
-                if (assembly.Equals(assemblyToSkip, StringComparison.OrdinalIgnoreCase))
+                else
                 {
-                    return true;
+                    if (assembly.Equals(assemblyToSkip, StringComparison.OrdinalIgnoreCase))
+                    {
+                        return true;
+                    }
                 }
             }
-        }
 
-        return false;
+            return false;
+        }
     }
 }