From 43d174bc529b87d9427208c759d80f79b50c9dc8 Mon Sep 17 00:00:00 2001 From: Daniel Romano <108014683+daniel-romano-DD@users.noreply.github.com> Date: Fri, 31 Jan 2025 10:07:22 +0100 Subject: [PATCH] [IAST] Fix flaky Interpolated Strings tests (#6555) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary of changes Fix test flakiness by keeping alive the tainted boxed object during the lifespan of the ref struct ## Reason for change If a GC collect happened between the tainting of the `DefaultInterpolatedStringHandler` and the call to `ToString` the boxed object was collected and the weak reference was lost, causing the loss of the tainted status. Thanks to @e-n-0 for finding what was happening IDK who said flakiness was bad ## Implementation details Added a thread local queue to hold the boxed object of the tainted ref structs ## Test coverage Added a unit test ## Other details --------- Co-authored-by: Flavien Darche --- .../src/Datadog.Trace/Iast/ITaintedObject.cs | 2 ++ ...aultInterpolatedStringHandlerModuleImpl.cs | 35 +++++++++++++++---- .../src/Datadog.Trace/Iast/TaintedObject.cs | 5 +++ .../IAST/Tainted/DefaultTaintedMapTests.cs | 8 ++--- .../IAST/Tainted/TaintedForTest.cs | 9 +++-- ...tring.AspNetCore5.IastEnabled.verified.txt | 8 ++--- .../DefaultInterpolatedStringHandlerTests.cs | 27 ++++++++++---- .../Controllers/IastController.cs | 2 +- 8 files changed, 70 insertions(+), 26 deletions(-) diff --git a/tracer/src/Datadog.Trace/Iast/ITaintedObject.cs b/tracer/src/Datadog.Trace/Iast/ITaintedObject.cs index e11080dd6fb4..687def539f4a 100644 --- a/tracer/src/Datadog.Trace/Iast/ITaintedObject.cs +++ b/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(); } diff --git a/tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs b/tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs index e73a341fa630..6d665b47e784 100644 --- a/tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs +++ b/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? _taintedRefStructs; + + private static Stack TaintedRefStructs + { + get => _taintedRefStructs ??= new Stack(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) { diff --git a/tracer/src/Datadog.Trace/Iast/TaintedObject.cs b/tracer/src/Datadog.Trace/Iast/TaintedObject.cs index 72618f8c8f9b..a783eb41d918 100644 --- a/tracer/src/Datadog.Trace/Iast/TaintedObject.cs +++ b/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; + } } diff --git a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapTests.cs b/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapTests.cs index 081fda1bab48..13ad0543e4f8 100644 --- a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapTests.cs +++ b/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++) diff --git a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/TaintedForTest.cs b/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/TaintedForTest.cs index b16f39180cf3..8e510a4f5e22 100644 --- a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/TaintedForTest.cs +++ b/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; diff --git a/tracer/test/snapshots/Iast.SqliInterpolatedString.AspNetCore5.IastEnabled.verified.txt b/tracer/test/snapshots/Iast.SqliInterpolatedString.AspNetCore5.IastEnabled.verified.txt index a3d022e4a938..74f895952514 100644 --- a/tracer/test/snapshots/Iast.SqliInterpolatedString.AspNetCore5.IastEnabled.verified.txt +++ b/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;" } ] } diff --git a/tracer/test/test-applications/integrations/Samples.InstrumentedTests/Propagation/String/DefaultInterpolatedStringHandlerTests.cs b/tracer/test/test-applications/integrations/Samples.InstrumentedTests/Propagation/String/DefaultInterpolatedStringHandlerTests.cs index 197314fc807e..39db739db8f8 100644 --- a/tracer/test/test-applications/integrations/Samples.InstrumentedTests/Propagation/String/DefaultInterpolatedStringHandlerTests.cs +++ b/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([' ', '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 = $""" diff --git a/tracer/test/test-applications/security/Samples.Security.AspNetCore5/Controllers/IastController.cs b/tracer/test/test-applications/security/Samples.Security.AspNetCore5/Controllers/IastController.cs index 7ca498b17979..22e843d71358 100644 --- a/tracer/test/test-applications/security/Samples.Security.AspNetCore5/Controllers/IastController.cs +++ b/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",