[IAST] New sampling mechanism

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

@@ -15,6 +15,7 @@
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Text.RegularExpressions;
+using System.Web;
 using Datadog.Trace.AppSec;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DuckTyping;
@@ -57,6 +58,7 @@ internal static partial class IastModule
     private const string OperationNameSessionTimeout = "session_timeout";
     private const string OperationNameEmailHtmlInjection = "email_html_injection";
     private const string ReferrerHeaderName = "Referrer";
+    private const int VulnerabilityStatsRouteLimit = 4096;
     internal static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(IastModule));
     private static readonly IastSettings IastSettings = Iast.Instance.Settings;
     private static readonly Lazy<EvidenceRedactor?> EvidenceRedactorLazy = new Lazy<EvidenceRedactor?>(() => CreateRedactor(IastSettings));
@@ -65,6 +67,7 @@ internal static partial class IastModule
     private static readonly SourceType[] _dbSources = [SourceType.SqlRowValue];
     private static readonly HashSet<int> LoggedAspectExceptionMessages = [];
     private static bool _showTimeoutExceptionError = true;
+    private static ConcurrentDictionary<string, VulnerabilityStats> _vulnerabilityStats = new ConcurrentDictionary<string, VulnerabilityStats>();
 
     internal static void LogTimeoutError(RegexMatchTimeoutException err)
     {
@@ -572,8 +575,9 @@ private static IastModuleResponse AddWebVulnerability(string? evidenceValue, Int
 
         var currentSpan = (tracer.ActiveScope as Scope)?.Span;
         var traceContext = currentSpan?.Context?.TraceContext;
+        var rootSpan = traceContext?.RootSpan;
 
-        if (traceContext?.IastRequestContext?.AddVulnerabilitiesAllowed() != true)
+        if (traceContext?.IastRequestContext?.AddVulnerabilityTypeAllowed(rootSpan, vulnerabilityType, GetRouteVulnerabilityStats) != true)
         {
             // we are inside a request but we don't accept more vulnerabilities or IastRequestContext is null, which means that iast is
             // not activated for this particular request
@@ -600,12 +604,42 @@ private static IastModuleResponse AddWebVulnerability(string? evidenceValue, Int
 
     public static bool AddRequestVulnerabilitiesAllowed()
     {
+        // Check for global budget only
         var currentSpan = (Tracer.Instance.ActiveScope as Scope)?.Span;
         var traceContext = currentSpan?.Context?.TraceContext;
         var isRequest = traceContext?.RootSpan?.Type == SpanTypes.Web;
         return isRequest && traceContext?.IastRequestContext?.AddVulnerabilitiesAllowed() == true;
     }
 
+    private static VulnerabilityStats GetRouteVulnerabilityStats(Span? span)
+    {
+        if (_vulnerabilityStats.Count >= VulnerabilityStatsRouteLimit)
+        {
+            _vulnerabilityStats.Clear(); // Poor man's LRU cache
+        }
+
+        if (span?.Type == SpanTypes.Web)
+        {
+            var route = span.GetTag(Tags.HttpRoute) ?? string.Empty;
+            var method = span.GetTag(Tags.HttpMethod) ?? string.Empty;
+            var key = $"{method}#{route}";
+            if (key.Length > 1)
+            {
+                return _vulnerabilityStats.GetOrAdd(key, (k) => new(k));
+            }
+        }
+
+        return new VulnerabilityStats(string.Empty);
+    }
+
+    internal static void UpdateRouteVulnerabilityStats(ref VulnerabilityStats stats)
+    {
+        if (stats.Route is not null)
+        {
+            _vulnerabilityStats[stats.Route] = stats;
+        }
+    }
+
     private static IastModuleResponse GetScope(
         string evidenceValue,
         IntegrationId integrationId,
@@ -629,7 +663,8 @@ private static IastModuleResponse GetScope(
         var scope = tracer.ActiveScope as Scope;
         var currentSpan = scope?.Span;
         var traceContext = currentSpan?.Context?.TraceContext;
-        var isRequest = traceContext?.RootSpan?.Type == SpanTypes.Web;
+        var rootSpan = traceContext?.RootSpan;
+        var isRequest = rootSpan?.Type == SpanTypes.Web;
 
         // We do not have, for now, tainted objects in console apps, so further checking is not neccessary.
         if (!isRequest && taintValidator != null)
@@ -660,14 +695,21 @@ private static IastModuleResponse GetScope(
             var unsafeRanges = Ranges.GetUnsafeRanges(ranges, exclusionSecureMarks, safeSources);
             if (unsafeRanges is null || unsafeRanges.Length == 0)
             {
-                OnSupressedVulnerabilityTelemetry(VulnerabilityTypeUtils.GetTag(vulnerabilityType));
+                OnSupressedVulnerabilityTelemetry(VulnerabilityTypeUtils.FromName(vulnerabilityType));
                 return IastModuleResponse.Empty;
             }
 
             // Contains at least one range that is not safe (when analyzing a vulnerability that can have secure marks)
             ranges = unsafeRanges;
         }
 
+        if (isRequest && traceContext?.IastRequestContext?.AddVulnerabilityTypeAllowed(rootSpan, vulnerabilityType, GetRouteVulnerabilityStats) != true)
+        {
+            // we are inside a request but we don't accept more vulnerabilities or IastRequestContext is null, which means that iast is
+            // not activated for this particular request
+            return IastModuleResponse.Empty;
+        }
+
         var location = addLocation ? GetLocation(externalStack, currentSpan) : null;
         if (addLocation && location is null)
         {

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

@@ -7,18 +7,16 @@
 
 using System;
 using System.Collections.Generic;
-using System.Collections.Specialized;
 using System.Globalization;
+using System.Reflection;
 using System.Threading;
 using System.Web;
 #if !NETFRAMEWORK
 using Microsoft.AspNetCore.Http;
 #endif
 using Datadog.Trace.AppSec;
-using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK;
 using Datadog.Trace.Iast.Telemetry;
 using Datadog.Trace.Logging;
-using Datadog.Trace.Telemetry.Metrics;
 using Datadog.Trace.Util;
 using static Datadog.Trace.Telemetry.Metrics.MetricTags;
 
@@ -35,6 +33,10 @@ internal class IastRequestContext
     private ExecutedTelemetryHelper? _executedTelemetryHelper = ExecutedTelemetryHelper.Enabled() ? new ExecutedTelemetryHelper() : null;
     private int _lastVulnerabilityStackId = 0;
 
+    private bool _routeVulnerabilityStatsDirty = false;
+    private VulnerabilityStats _routeVulnerabilityStats;
+    private VulnerabilityStats _requestVulnerabilityStats;
+
     internal static void AddIastDisabledFlagToSpan(Span span)
     {
         span.Tags.SetTag(Tags.IastEnabled, "0");
@@ -46,6 +48,24 @@ internal void AddIastVulnerabilitiesToSpan(Span span)
         {
             span.Tags.SetTag(Tags.IastEnabled, "1");
 
+            if (_routeVulnerabilityStatsDirty)
+            {
+                if (AddVulnerabilitiesAllowed())
+                {
+                    // Global budget not depleted. Reset route stats so vulns can be detected again.
+                    Log.Debug("Clearing Vulnerability Stats for Route {Route} (Budget left)", _routeVulnerabilityStats.Route);
+                    _routeVulnerabilityStats = new VulnerabilityStats(_requestVulnerabilityStats.Route);
+                }
+                else
+                {
+                    // Global budget depleted. Update route stats so vulns new can be detected.
+                    Log.Debug("Updating Vulnerability Stats for Route {Route} (Budget depleted)", _routeVulnerabilityStats.Route);
+                    _routeVulnerabilityStats.TransferNewVulns(ref _requestVulnerabilityStats);
+                }
+
+                IastModule.UpdateRouteVulnerabilityStats(ref _routeVulnerabilityStats);
+            }
+
             if (_vulnerabilityBatch != null)
             {
                 if (Iast.Instance.IsMetaStructSupported())
@@ -79,9 +99,53 @@ internal void AddIastVulnerabilitiesToSpan(Span span)
 
     internal bool AddVulnerabilitiesAllowed()
     {
+        // Check global budget
         return ((_vulnerabilityBatch?.Vulnerabilities.Count ?? 0) < Iast.Instance.Settings.VulnerabilitiesPerRequest);
     }
 
+    internal bool AddVulnerabilityTypeAllowed(Span? span, string vulnerabilityType, Func<Span?, VulnerabilityStats> getForCurrentRoute)
+    {
+        // Check global budget
+        if (AddVulnerabilitiesAllowed())
+        {
+            if (_requestVulnerabilityStats.Route is null)
+            {
+                // Init the route vulnerability stats
+                _routeVulnerabilityStats = getForCurrentRoute(span);
+                _requestVulnerabilityStats = new(_routeVulnerabilityStats.Route);
+            }
+
+            if (_requestVulnerabilityStats.Route is { Length: > 0 })
+            {
+                // Check route budget
+                var index = (int)VulnerabilityTypeUtils.FromName(vulnerabilityType);
+                _requestVulnerabilityStats[index]++;
+
+                string debugTxt = string.Empty;
+                if (Log.IsEnabled(Vendors.Serilog.Events.LogEventLevel.Debug))
+                {
+                    var currentVulns = (_vulnerabilityBatch?.Vulnerabilities.Count ?? 0) + 1;
+                    debugTxt = $"Vulnerability {vulnerabilityType} detected for Route {_requestVulnerabilityStats.Route}. Current count: {_requestVulnerabilityStats[index]}  Route count: {_routeVulnerabilityStats[index]}  Budget: {currentVulns} out of {Iast.Instance.Settings.VulnerabilitiesPerRequest}";
+                }
+
+                _routeVulnerabilityStatsDirty = true;
+                if (_requestVulnerabilityStats[index] > _routeVulnerabilityStats[index] || _routeVulnerabilityStats[index] == 0)
+                {
+                    Log.Debug("Vulnerability Sampler ACCEPTED: {Txt}", debugTxt);
+                    return true;
+                }
+
+                Log.Debug("Vulnerability Sampler SKIPPED: {Txt}", debugTxt);
+            }
+            else
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     internal void AddVulnerability(Vulnerability vulnerability)
     {
         lock (_vulnerabilityLock)

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

@@ -0,0 +1,73 @@
+// <copyright file="VulnerabilityStats.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.InteropServices;
+using System.Text;
+using System.Threading.Tasks;
+
+#if !NETCOREAPP
+using Datadog.Trace.VendoredMicrosoftCode.System.Runtime.CompilerServices.Unsafe;
+#else
+using System.Runtime.CompilerServices;
+#endif
+
+namespace Datadog.Trace.Iast;
+
+[StructLayout(LayoutKind.Sequential, Pack = 1)]
+internal unsafe struct VulnerabilityStats
+{
+    private const int ElementSize = 2; // sizeof(ushort)
+    private const int Size = ElementSize * Length;
+    private fixed byte data[Size];
+
+    public const int Length = VulnerabilityTypeUtils.Count;
+
+    public VulnerabilityStats(string? route)
+    {
+        Route = route;
+    }
+
+    public string? Route { get; }
+
+    public ushort this[int index]
+    {
+        get
+        {
+            if (index < 0 || index >= Length)
+            {
+                throw new IndexOutOfRangeException();
+            }
+
+            return Unsafe.ReadUnaligned<ushort>(ref data[index * ElementSize]);
+        }
+
+        set
+        {
+            if (index < 0 || index >= Length)
+            {
+                throw new IndexOutOfRangeException();
+            }
+
+            Unsafe.WriteUnaligned<ushort>(ref data[index * ElementSize], value);
+        }
+    }
+
+    public void TransferNewVulns(ref VulnerabilityStats stats)
+    {
+        for (int x = 0; x < Length; x++)
+        {
+            var value = stats[x];
+            if (value > this[x])
+            {
+                this[x] = value;
+            }
+        }
+    }
+}

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

@@ -4,6 +4,8 @@
 // </copyright>
 
 #nullable enable
+using Datadog.Trace.SourceGenerators;
+
 namespace Datadog.Trace.Iast;
 
 /// <summary>
@@ -90,5 +92,5 @@ internal enum VulnerabilityType
     SessionTimeout = 25,
 
     /// <summary> EMAIL_HTML_INJECTION vulnerability </summary>
-    EmailHtmlInjection = 26
+    EmailHtmlInjection = 26,
 }

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

@@ -5,6 +5,7 @@
 
 #nullable enable
 using System.ComponentModel;
+using Datadog.Trace.Telemetry.Metrics;
 using static Datadog.Trace.Telemetry.Metrics.MetricTags;
 
 namespace Datadog.Trace.Iast;
@@ -14,6 +15,9 @@ namespace Datadog.Trace.Iast;
 /// </summary>
 internal static class VulnerabilityTypeUtils
 {
+    /// <summary> The number of vulnerability types </summary>
+    public const int Count = 27;
+
     /// <summary> Undefined vulnerability </summary>
     public const string None = "NONE";
 
@@ -130,7 +134,7 @@ public static string GetName(VulnerabilityType vulnerability)
         };
     }
 
-    public static IastVulnerabilityType GetTag(string vulnerability)
+    public static IastVulnerabilityType FromName(string vulnerability)
     {
         return vulnerability switch
         {