Skip to content

Commit

Permalink
[IAST] Fix flaky Interpolated Strings tests (#6555)
Browse files Browse the repository at this point in the history
## 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
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

---------

Co-authored-by: Flavien Darche <[email protected]>
  • Loading branch information
daniel-romano-DD and e-n-0 authored Jan 31, 2025
1 parent 642f1cb commit 43d174b
Show file tree
Hide file tree
Showing 8 changed files with 70 additions and 26 deletions.
2 changes: 2 additions & 0 deletions tracer/src/Datadog.Trace/Iast/ITaintedObject.cs
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,6 @@ internal interface ITaintedObject
public ITaintedObject? Next { get; set; }

public int PositiveHashCode { get; }

public void Invalidate();
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down Expand Up @@ -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
{
Expand Down Expand Up @@ -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)
{
Expand Down
5 changes: 5 additions & 0 deletions tracer/src/Datadog.Trace/Iast/TaintedObject.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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();
}
Expand All @@ -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();
}

Expand Down Expand Up @@ -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++)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand All @@ -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; }

Expand All @@ -37,6 +37,11 @@ public void SetAlive(bool isAlive)
_alive = isAlive;
}

public void Invalidate()
{
_alive = false;
}

internal Range[]? GetRanges()
{
return _ranges;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;"
}
]
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand All @@ -259,7 +274,7 @@ public void GivenImplicitInterpolatedString_WhenAddingTaintedValuesNested_GetStr
AssertTainted(finalString);
}

[Fact (Skip = "Flaky test")]
[Fact]
public void GivenImplicitInterpolatedString_WhenAddingTaintedValuesComplex_GetString_Vulnerable()
{
var interpolatedString = $"""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down

0 comments on commit 43d174b

Please sign in to comment.