[IAST] Fix flaky Interpolated Strings tests

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

@@ -17,4 +17,6 @@ internal interface ITaintedObject
     public ITaintedObject? Next { get; set; }
 
     public int PositiveHashCode { get; }
+
+    public void Invalidate();
 }

tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs

@@ -8,21 +8,28 @@
 #nullable enable
 
 using System;
-using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
-using System.Text;
-using Datadog.Trace.VendoredMicrosoftCode.System.Runtime.InteropServices;
+using System.Collections.Generic;
 
 namespace Datadog.Trace.Iast.Propagation;
 
 internal static class DefaultInterpolatedStringHandlerModuleImpl
 {
-    public static unsafe void Append(IntPtr target, string? value)
+    private const int MaxStackSize = 4;
+
+    [ThreadStatic]
+    private static Stack<object>? _taintedRefStructs;
+
+    private static Stack<object> TaintedRefStructs
+    {
+        get => _taintedRefStructs ??= new Stack<object>(MaxStackSize);
+    }
+
+    public static void Append(IntPtr target, string? value)
     {
         FullTaintIfAnyTainted(target, value);
     }
 
-    public static unsafe void FullTaintIfAnyTainted(IntPtr target, string? input)
+    public static void FullTaintIfAnyTainted(IntPtr target, string? input)
     {
         try
         {
@@ -51,7 +58,16 @@ public static unsafe void FullTaintIfAnyTainted(IntPtr target, string? input)
             var rangesResult = new[] { new Range(0, 0, tainted!.Ranges[0].Source, tainted.Ranges[0].SecureMarks) };
             if (!targetIsTainted)
             {
-                taintedObjects.Taint(target, rangesResult);
+                // Safe guard to avoid memory leak
+                if (TaintedRefStructs.Count >= MaxStackSize)
+                {
+                    TaintedRefStructs.Clear();
+                }
+
+                object targetObj = target;
+                TaintedRefStructs.Push(targetObj);
+
+                taintedObjects.Taint(targetObj, rangesResult);
             }
             else
             {
@@ -91,6 +107,11 @@ public static unsafe void FullTaintIfAnyTainted(IntPtr target, string? input)
 
             var range = new Range(0, result.Length, taintedSelf.Ranges[0].Source, taintedSelf.Ranges[0].SecureMarks);
             taintedObjects.Taint(result, [range]);
+            taintedSelf.Invalidate();
+            if (TaintedRefStructs.Count > 0)
+            {
+                TaintedRefStructs.Pop();
+            }
         }
         catch (Exception err)
         {

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

@@ -29,4 +29,9 @@ public TaintedObject(object value, Range[] ranges)
     public Range[] Ranges { get; set; }
 
     public ITaintedObject? Next { get; set; }
+
+    public void Invalidate()
+    {
+        _weak.Target = null;
+    }
 }

tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapTests.cs

@@ -235,7 +235,7 @@ public void GivenATaintedObjectMap_WhenPutObjectsThatGetDisposed_ObjectsArePurge
 
         foreach (var item in disposedObjects)
         {
-            (map.Get(item) as TaintedForTest).SetAlive(false);
+            (map.Get(item) as TaintedForTest).Invalidate();
         }
 
         map.Purge();
@@ -263,7 +263,7 @@ public void GivenATaintedObjectMap_WhenASingleObjectDisposed_IsPurged()
         var tainted = new TaintedForTest(testString, null);
         map.Put(tainted);
         map.Get(testString).Should().NotBeNull();
-        (map.Get(testString) as TaintedForTest).SetAlive(false);
+        (map.Get(testString) as TaintedForTest).Invalidate();
         map.Purge();
         map.Get(testString).Should().BeNull();
     }
@@ -279,7 +279,7 @@ public void GivenATaintedObjectMap_WhenHasCausesPurge_IsPurged()
             testString.Hash = 1;
             var tainted = new TaintedForTest(testString, null);
             map.Put(tainted);
-            tainted.SetAlive(false);
+            tainted.Invalidate();
             map.Get(testString).Should().NotBeNull();
         }
 
@@ -313,7 +313,7 @@ public void GivenATaintedObjectMap_WhenDisposedInSameHashPosition_IsPurged(int t
             addedObjects.Add(testString);
         }
 
-        (map.Get(addedObjects[disposedIndex]) as TaintedForTest).SetAlive(false);
+        (map.Get(addedObjects[disposedIndex]) as TaintedForTest).Invalidate();
         map.Purge();
 
         for (int i = 0; i < totalObjects; i++)

tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/TaintedForTest.cs

@@ -15,7 +15,7 @@ public class TaintedForTest : ITaintedObject
 {
     private bool _alive = true;
     private Range[] _ranges;
-    private object _value;
+    private object? _value;
 
     internal TaintedForTest(object value, Range[] ranges)
     {
@@ -28,7 +28,7 @@ internal TaintedForTest(object value, Range[] ranges)
 
     ITaintedObject? ITaintedObject.Next { get; set; }
 
-    public object Value => _value;
+    public object? Value => _value;
 
     public int PositiveHashCode { get; set; }
 
@@ -37,6 +37,11 @@ public void SetAlive(bool isAlive)
         _alive = isAlive;
     }
 
+    public void Invalidate()
+    {
+        _alive = false;
+    }
+
     internal Range[]? GetRanges()
     {
         return _ranges;

tracer/test/snapshots/Iast.SqliInterpolatedString.AspNetCore5.IastEnabled.verified.txt

@@ -38,15 +38,11 @@
             "value": "INSERT INTO Orders (CustomerId, EmployeeId, OrderDate, RequiredDate, ShipVia, Freight, ShipName, ShipAddress, ShipCity, ShipPostalCode, ShipCountry) VALUES ('VINET','5','2021-01-01','2021-01-01',"
           },
           {
-            "value": "'3','32.38','Vins et alcools Chevalier','John',",
+            "value": "'3','32','Vins et alcools Chevalier','John',",
             "source": 0
           },
           {
-            "value": "'Reims','51100','France')",
-            "source": 0
-          },
-          {
-            "value": ";\nSELECT OrderID FROM Orders ORDER BY OrderID DESC LIMIT 1;"
+            "value": "'Reims','51100','France');\nSELECT OrderID FROM Orders ORDER BY OrderID DESC LIMIT 1;"
           }
         ]
       }

tracer/test/test-applications/integrations/Samples.InstrumentedTests/Propagation/String/DefaultInterpolatedStringHandlerTests.cs

@@ -183,7 +183,22 @@ public void GivenAnExplicitInterpolatedString_WhenAddingTaintedValueMultipleValu
         AssertTainted(test.ToStringAndClear());
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
+    public void GivenAnExplicitInterpolatedString_WhenAddingTaintedValuesAndGCCollects_GetString_Vulnerable()
+    {
+        var test = new DefaultInterpolatedStringHandler();
+        test.AppendLiteral(UntaintedValue);
+        test.AppendFormatted(new ReadOnlySpan<char>([' ', 'w', 'o', 'r', 'l', 'd', ' ']));
+        test.AppendFormatted(TaintedValue);
+        test.AppendFormatted(42);
+
+        GC.Collect();
+
+        AssertTainted(test.ToStringAndClear());
+    }
+
+
+    [Fact]
     public void GivenAnImplicitInterpolatedString_WhenAddingTaintedValue_GetString_Vulnerable()
     {
         var number = 5;
@@ -192,7 +207,7 @@ public void GivenAnImplicitInterpolatedString_WhenAddingTaintedValue_GetString_V
         AssertTainted(str);
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
     public void GivenAnImplicitInterpolatedString_WhenAddingUntaintedValue_GetString_NonVulnerable()
     {
         var number = 5;
@@ -201,7 +216,7 @@ public void GivenAnImplicitInterpolatedString_WhenAddingUntaintedValue_GetString
         AssertNotTainted(str);
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
     public void GivenAnImplicitInterpolatedString_WhenAddingTaintedValueAsObject_GetString_Vulnerable()
     {
         var number = 5;
@@ -210,7 +225,7 @@ public void GivenAnImplicitInterpolatedString_WhenAddingTaintedValueAsObject_Get
         AssertTainted(str);
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
     public void GivenAnImplicitInterpolatedString_WhenAddingMultipleValuesWithTaintedValues_GetString_Vulnerable()
     {
         var order = new
@@ -240,7 +255,7 @@ public void GivenAnImplicitInterpolatedString_WhenAddingMultipleValuesWithTainte
         AssertTainted(sql);
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
     public void GivenImplicitInterpolatedString_WhenAddingTaintedValuesNested_GetString_Vulnerable()
     {
         const int number = 42;
@@ -259,7 +274,7 @@ public void GivenImplicitInterpolatedString_WhenAddingTaintedValuesNested_GetStr
         AssertTainted(finalString);
     }
 
-    [Fact (Skip = "Flaky test")]
+    [Fact]
     public void GivenImplicitInterpolatedString_WhenAddingTaintedValuesComplex_GetString_Vulnerable()
     {
         var interpolatedString = $"""

tracer/test/test-applications/security/Samples.Security.AspNetCore5/Controllers/IastController.cs

@@ -1209,7 +1209,7 @@ public IActionResult InterpolatedSqlString(string name)
                 OrderDate = new DateTime(2021, 1, 1),
                 RequiredDate = new DateTime(2021, 1, 1),
                 ShipVia = 3,
-                Freight = 32.38M,
+                Freight = 32,
                 ShipName = "Vins et alcools Chevalier",
                 ShipAddress = name,
                 ShipCity = "Reims",