Some fixes

DataDog · Nov 29, 2024 · 458dbd5 · 458dbd5
1 parent 2df7434
commit 458dbd5
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 19 deletions.
diff --git a/...r/src/Datadog.Trace/Iast/Aspects/System.Runtime/DefaultInterpolatedStringHandlerAspect.cs b/...r/src/Datadog.Trace/Iast/Aspects/System.Runtime/DefaultInterpolatedStringHandlerAspect.cs
@@ -226,6 +226,7 @@ public static string ToStringAndClear(ref DefaultInterpolatedStringHandler targe
     private static IntPtr ToPointer(ref DefaultInterpolatedStringHandler ts)
     {
         Ldarg(nameof(ts));
+        Conv_U();
         return IL.Return<IntPtr>();
     }
 }

diff --git a/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs b/tracer/src/Datadog.Trace/Iast/DefaultTaintedMap.cs
@@ -62,26 +62,68 @@ public DefaultTaintedMap()
         }
 
         _map.TryGetValue(IndexObject(objectToFind), out var entry);
-        bool isString = objectToFind is string;
 
         while (entry != null)
         {
-            if (isString)
+            if (objectToFind.Equals(entry.Value))
             {
-                if (objectToFind == entry.Value)
-                {
-                    return entry;
-                }
+                return entry;
             }
-            else
+
+            entry = entry.Next;
+        }
+
+        return null;
+    }
+
+    public ITaintedObject? Pop(object objectToFind)
+    {
+        if (objectToFind is null)
+        {
+            return null;
+        }
+
+        var index = IndexObject(objectToFind);
+
+        _map.TryGetValue(index, out var entry);
+
+        var current = entry;
+        ITaintedObject? previous = null;
+
+        while (current != null)
+        {
+            if (objectToFind.Equals(current.Value))
             {
-                if (objectToFind.Equals(entry.Value))
+                if (previous == null)
+                {
+                    // Removing the head of the chain
+                    if (current.Next == null)
+                    {
+                        // No more entries in the chain, remove the key
+                        _map.TryRemove(index, out _);
+                    }
+                    else
+                    {
+                        // Update the map with the new head of the chain
+                        _map[index] = current.Next;
+                    }
+                }
+                else
+                {
+                    // Removing an entry in the middle or end of the chain
+                    previous.Next = current.Next;
+                }
+
+                if (!IsFlat)
                 {
-                    return entry;
+                    Interlocked.Decrement(ref _entriesCount);
                 }
+
+                return current;
             }
 
-            entry = entry.Next;
+            previous = current;
+            current = current.Next;
         }
 
         return null;

diff --git a/tracer/src/Datadog.Trace/Iast/ITaintedMap.cs b/tracer/src/Datadog.Trace/Iast/ITaintedMap.cs
@@ -13,6 +13,8 @@ internal interface ITaintedMap
 
     public ITaintedObject Get(object obj);
 
+    public ITaintedObject Pop(object obj);
+
     public List<ITaintedObject> GetListValues();
 
     public int GetEstimatedSize();

diff --git a/tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs b/tracer/src/Datadog.Trace/Iast/Propagation/DefaultInterpolatedStringHandlerModuleImpl.cs
@@ -17,12 +17,12 @@ namespace Datadog.Trace.Iast.Propagation;
 
 internal static class DefaultInterpolatedStringHandlerModuleImpl
 {
-    public static unsafe void Append(IntPtr target, string? value)
+    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
         {
@@ -64,29 +64,29 @@ public static unsafe void FullTaintIfAnyTainted(IntPtr target, string? input)
         }
     }
 
-    public static object? PropagateTaint(object? input, string? result)
+    public static void PropagateTaint(IntPtr? input, string? result)
     {
         try
         {
             IastModule.OnExecutedPropagationTelemetry();
 
             if (result is null || input is null)
             {
-                return result;
+                return;
             }
 
             var iastContext = IastModule.GetIastContext();
             if (iastContext == null)
             {
-                return result;
+                return;
             }
 
             var taintedObjects = iastContext.GetTaintedObjects();
-            var taintedSelf = taintedObjects.Get(input);
+            var taintedSelf = taintedObjects.Pop(input);
 
             if (taintedSelf == null)
             {
-                return result;
+                return;
             }
 
             var range = new Range(0, result.Length, taintedSelf.Ranges[0].Source, taintedSelf.Ranges[0].SecureMarks);
@@ -96,8 +96,6 @@ public static unsafe void FullTaintIfAnyTainted(IntPtr target, string? input)
         {
             IastModule.Log.Error(err, $"{nameof(DefaultInterpolatedStringHandlerModuleImpl)}.{nameof(PropagateTaint)} exception");
         }
-
-        return result;
     }
 
     private static TaintedObject? GetTaintedWithRanges(TaintedObjects taintedObjects, object? value)

diff --git a/tracer/src/Datadog.Trace/Iast/TaintedObjects.cs b/tracer/src/Datadog.Trace/Iast/TaintedObjects.cs
@@ -73,6 +73,11 @@ public void Taint(object objectToTaint, Range[] ranges)
             return _map.Get(objectToFind) as TaintedObject;
         }
 
+        public TaintedObject? Pop(object objectToPop)
+        {
+            return _map.Pop(objectToPop) as TaintedObject;
+        }
+
         public int GetEstimatedSize()
         {
             return _map.GetEstimatedSize();