From 74d48340ceb2c2facfd29684b2baa8d5daf333d3 Mon Sep 17 00:00:00 2001 From: Daniel Romano Date: Thu, 13 Aug 2026 13:22:38 +0200 Subject: [PATCH] [IAST] Fix entry loss in DefaultTaintedMap under concurrency Put published the chain head with a read-modify-write, and RemoveDeadKeys collected dead keys and removed them unconditionally in a second pass. Both drop entries when a request is served by more than one thread. Publish the head with a compare-and-swap, and make purging use compare-and-remove / compare-and-update so it never deletes an entry inserted concurrently. Co-Authored-By: Claude Opus 5 --- .../Datadog.Trace/Iast/DefaultTaintedMap.cs | 115 +++++--- .../DefaultTaintedMapConcurrencyTests.cs | 262 ++++++++++++++++++ 2 files changed, 337 insertions(+), 40 deletions(-) create mode 100644 tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapConcurrencyTests.cs diff --git a/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs b/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs index 98565d1761e3..36102c930c3c 100644 --- a/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs +++ b/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs @@ -26,6 +26,10 @@ internal sealed class DefaultTaintedMap : ITaintedMap private const int PurgeMask = PurgeCount - 1; // Map containing the tainted objects private ConcurrentDictionary _map; + // Same instance as _map. ConcurrentDictionary implements ICollection>.Remove as an + // atomic compare-and-remove, which is the only portable way of getting that behaviour on net461 + // and netstandard2.0 (TryRemove(KeyValuePair<,>) only exists from netcoreapp2.0 onwards). + private ICollection> _mapAsCollection; // Bitmask for fast modulo with table length. private int _lengthMask; // Flag to ensure we do not run multiple purges concurrently. @@ -39,6 +43,7 @@ internal sealed class DefaultTaintedMap : ITaintedMap public DefaultTaintedMap() { _map = new ConcurrentDictionary(); + _mapAsCollection = _map; _lengthMask = DefaultCapacity - 1; _flatModeThreshold = DefaultFlatModeThresold; } @@ -100,27 +105,45 @@ public void Put(ITaintedObject entry) var index = Index(entry.PositiveHashCode); - if (!IsFlat) + if (IsFlat) + { + // If we flipped to flat mode: + // - Always override elements ignoring chaining. + // - Stop updating the estimated size. + _map[index] = entry; + } + else { // By default, add the new entry to the head of the chain. // We do not control duplicate entries. - _map.TryGetValue(index, out var existingValue); - entry.Next = existingValue; - - // If there are two callers calling Put on the same map and the objects have the same index and we are not flat, - // then one of the ITaintedObjects could potentially be lost because of racing conditions. - // We assume that the corresponding lock mechanism benefits would not compensate the performance loss. + // The head is published with a compare-and-swap and retried on conflict: a plain + // read-modify-write here silently dropped entries whenever two callers hit the same + // index concurrently, which is the common case on a request served by several threads + // (async continuations), not a rare one. + while (true) + { + if (_map.TryGetValue(index, out var existingValue)) + { + entry.Next = existingValue; + if (_map.TryUpdate(index, entry, existingValue)) + { + break; + } + } + else + { + entry.Next = null; + if (_map.TryAdd(index, entry)) + { + break; + } + } + } // We only count the entries if we are not in flat mode Interlocked.Increment(ref _entriesCount); } - // If we flipped to flat mode: - // - Always override elements ignoring chaining. - // - Stop updating the estimated size. - - _map[index] = entry; - if ((entry.PositiveHashCode & PurgeMask) == 0) { Purge(); @@ -171,49 +194,61 @@ internal void Purge() private int RemoveDeadKeys() { var removed = 0; - List deadKeys = new(); - ITaintedObject? previous; foreach (var key in _map.Keys.ToArray()) { - var current = _map[key]; - previous = null; + if (!_map.TryGetValue(key, out var current)) + { + // Removed by a concurrent Clear(). + continue; + } + + ITaintedObject? previous = null; while (current is not null) { - if (!current.IsAlive) + var next = current.Next; + + if (current.IsAlive) { - if (previous is null) - { - // We can delete the map key - if (current.Next is null) - { - deadKeys.Add(key); - } - else - { - _map[key] = current.Next; - } - } - else + previous = current; + } + else if (previous is not null) + { + // Unlinking a node in the middle of the chain only ever moves `Next` forward, so a + // concurrent Get() walking this chain can neither loop nor go backwards. Dead nodes it + // still reaches have a null Value, so they simply never match. + previous.Next = next; + removed++; + } + else if (next is null) + { + // The whole chain is dead, so the key can go. Compare-and-remove: a plain TryRemove + // would also delete an entry that a concurrent Put() published on this key in the + // meantime. + if (_mapAsCollection.Remove(new KeyValuePair(key, current))) { - previous.Next = current.Next; + removed++; } - current = current.Next; + // Either we removed the dead chain, or a concurrent Put() replaced the head with a + // live entry that already chains the rest. Nothing left to do on this key. + break; + } + else if (_map.TryUpdate(key, next, current)) + { + // Dropped the dead head, `previous` stays null because `next` is the new head. removed++; } else { - previous = current; - current = current.Next; + // A concurrent Put() replaced the head; its chain still contains the dead nodes, + // so the next purge picks them up. + break; } - } - } - foreach (var key in deadKeys) - { - _map.TryRemove(key, out _); + current = next; + } } return removed; diff --git a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapConcurrencyTests.cs b/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapConcurrencyTests.cs new file mode 100644 index 000000000000..734de0056e1f --- /dev/null +++ b/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/Tainted/DefaultTaintedMapConcurrencyTests.cs @@ -0,0 +1,262 @@ +// +// 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. +// + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Datadog.Trace.Iast; +using FluentAssertions; +using Xunit; + +namespace Datadog.Trace.Security.Unit.Tests.Iast.Tainted; + +public class DefaultTaintedMapConcurrencyTests +{ + // A bucket index whose hash does not trigger an implicit Purge() from Put() + // ((hash & PurgeMask) != 0), so that Put/Put races can be tested in isolation. + private const int NonPurgingHash = 1; + + [Fact] + public void GivenATaintedObjectMap_WhenConcurrentPutsCollideOnTheSameBucket_NoEntryIsLost() + { + const int threads = 8; + const int perThread = 400; + + var map = new DefaultTaintedMap(); + var entries = new TestTaintedObject[threads][]; + + for (var t = 0; t < threads; t++) + { + entries[t] = new TestTaintedObject[perThread]; + for (var i = 0; i < perThread; i++) + { + entries[t][i] = new TestTaintedObject(new HashedValue(NonPurgingHash)); + } + } + + RunConcurrently(threads, t => + { + foreach (var entry in entries[t]) + { + map.Put(entry); + } + }); + + var lost = 0; + foreach (var perThreadEntries in entries) + { + foreach (var entry in perThreadEntries) + { + if (map.Get(entry.Key) is null) + { + lost++; + } + } + } + + lost.Should().Be(0, "concurrent Put on the same bucket must not drop entries from the chain"); + map.GetEstimatedSize().Should().Be(threads * perThread); + } + + [Fact] + public void GivenATaintedObjectMap_WhenPurgingWhileInsertingOnTheSameBuckets_LiveEntriesAreNotLost() + { + // Every dead entry is alone in its bucket, so RemoveDeadKeys() collects all of these keys + // as "dead keys" and only removes them in a second pass. A live entry inserted on one of + // those buckets in between is silently deleted by that second pass. + const int buckets = 2048; + const int iterations = 10; + + for (var iteration = 0; iteration < iterations; iteration++) + { + var map = new DefaultTaintedMap(); + var live = new List(buckets); + + for (var i = 0; i < buckets; i++) + { + // Odd hashes only, so that no Put triggers an implicit Purge(). + var hash = (i * 2) + 1; + + // Insert alive and then invalidate, mimicking a collected WeakReference target. + var dead = new TestTaintedObject(new HashedValue(hash)); + map.Put(dead); + dead.Invalidate(); + + live.Add(new TestTaintedObject(new HashedValue(hash))); + } + + RunConcurrently(2, t => + { + if (t == 0) + { + map.Purge(); + } + else + { + foreach (var entry in live) + { + map.Put(entry); + } + } + }); + + var lost = 0; + foreach (var entry in live) + { + if (map.Get(entry.Key) is null) + { + lost++; + } + } + + lost.Should().Be(0, $"Purge() must not remove entries inserted concurrently (iteration {iteration})"); + } + } + + [Fact] + public void GivenATaintedObjectMap_WhenGettingWhilePurgingAndInserting_LiveEntriesStayReachable() + { + const int buckets = 512; + const int entriesPerBucket = 4; + const int readerIterations = 100; + + var map = new DefaultTaintedMap(); + var live = new List(); + + // Chains that mix live and dead entries, so RemoveDeadKeys() has to splice them. + for (var i = 0; i < buckets; i++) + { + var hash = (i * 2) + 1; + for (var e = 0; e < entriesPerBucket; e++) + { + var entry = new TestTaintedObject(new HashedValue(hash)); + map.Put(entry); + + if (e % 2 == 0) + { + entry.Invalidate(); + } + else + { + live.Add(entry); + } + } + } + + RunConcurrently(4, t => + { + switch (t) + { + case 0: + for (var i = 0; i < 20; i++) + { + map.Purge(); + } + + break; + + case 1: + for (var i = 0; i < buckets; i++) + { + var entry = new TestTaintedObject(new HashedValue((i * 2) + 1)); + map.Put(entry); + entry.Invalidate(); + } + + break; + + default: + for (var i = 0; i < readerIterations; i++) + { + foreach (var entry in live) + { + map.Get(entry.Key); + } + } + + break; + } + }); + + var lost = 0; + foreach (var entry in live) + { + if (map.Get(entry.Key) is null) + { + lost++; + } + } + + lost.Should().Be(0, "purging dead entries must not drop the live ones sharing their chain"); + } + + private static void RunConcurrently(int threads, Action body) + { + using var start = new ManualResetEventSlim(false); + var tasks = new Task[threads]; + + for (var t = 0; t < threads; t++) + { + var index = t; + tasks[t] = Task.Factory.StartNew( + () => + { + start.Wait(); + body(index); + }, + TaskCreationOptions.LongRunning); + } + + start.Set(); + Task.WaitAll(tasks); + } + + /// + /// Value with a controlled hash code, so that entries can be forced into a chosen bucket. + /// Equality stays reference-based, which is what DefaultTaintedMap.Get relies on. + /// + private class HashedValue + { + private readonly int _hash; + + public HashedValue(int hash) + { + _hash = hash; + } + + public override bool Equals(object obj) => ReferenceEquals(this, obj); + + public override int GetHashCode() => _hash; + } + + /// + /// ITaintedObject with deterministic liveness, so purging does not depend on the GC. + /// + private class TestTaintedObject : ITaintedObject + { + private readonly HashedValue _value; + private bool _isAlive = true; + + public TestTaintedObject(HashedValue value) + { + _value = value; + PositiveHashCode = IastUtils.IdentityHashCode(value) & DefaultTaintedMap.PositiveMask; + } + + public object Value => _isAlive ? _value : null; + + public bool IsAlive => _isAlive; + + public int PositiveHashCode { get; } + + public ITaintedObject Next { get; set; } + + /// Gets the value, regardless of liveness, for lookups from the test. + public object Key => _value; + + public void Invalidate() => _isAlive = false; + } +}