[IAST] Small string cache bugfix

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

@@ -9,19 +9,51 @@ namespace Datadog.Trace.Iast
 {
     internal class TaintedObjects
     {
+        private static readonly bool _largeNumericCache = false;
         private readonly ITaintedMap _map;
 
+        static TaintedObjects()
+        {
+            // From Net 8.0 onwards first 300 digit strings are cached instead of first 10
+            _largeNumericCache = object.ReferenceEquals(299.ToString(), 299.ToString());
+        }
+
         public TaintedObjects()
         {
             _map = new DefaultTaintedMap();
         }
 
-        public void TaintInputString(string stringToTaint, Source source)
+        public bool TaintInputString(string stringToTaint, Source source)
         {
-            if (!string.IsNullOrEmpty(stringToTaint))
+            if (!IsFiltered(stringToTaint))
             {
                 _map.Put(new TaintedObject(stringToTaint, IastUtils.GetRangesForString(stringToTaint, source)));
+                return true;
+            }
+
+            bool IsFiltered(string arg)
+            {
+                // Try to bail out ASAP
+                if (string.IsNullOrEmpty(arg)) { return true; }
+
+                // 0 - 9 are cached only
+                if (!_largeNumericCache)
+                {
+                    return arg.Length == 1 && char.IsDigit(arg[0]);
+                }
+
+                return arg.Length switch
+                {
+                    > 3 => false, // Only 0 - 299 are cached
+                    _ when !char.IsDigit(arg[0]) => false, // Not a number
+                    > 1 when !char.IsDigit(arg[1]) => false, // Not a number
+                    > 2 when !char.IsDigit(arg[2]) => false, // Not a number
+                    > 2 when arg[0] - '0' >= 3 => false, // Bigger than 299
+                    _ => true
+                };
             }
+
+            return false;
         }
 
         public void Taint(object objectToTaint, Range[] ranges)

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

@@ -76,4 +76,63 @@ public void GivenATaintedString_WhenGetSameValueDifferentReference_NullIsReturne
         Assert.Equal(tainted1, taintedObjects.Get(tainted1).Value);
         Assert.Null(taintedObjects.Get(tainted2));
     }
+
+    [Fact]
+    public void GivenInputStrings_WhenTainted_DigitsAreFiltered()
+    {
+        TaintedObjects taintedObjects = new();
+        Source source = new Source(SourceType.RequestBody, "name", "value");
+
+        void TaintInput(bool mustBeTainted, string s)
+        {
+            if (mustBeTainted)
+            {
+                Assert.True(taintedObjects.TaintInputString(s, source), s + " should be tainted");
+            }
+            else
+            {
+                Assert.False(taintedObjects.TaintInputString(s, source), s + " should NOT be tainted");
+            }
+        }
+
+        TaintInput(false, string.Empty);
+
+        TaintInput(false, "0");
+        TaintInput(false, "1");
+        TaintInput(false, "2");
+        TaintInput(false, "3");
+        TaintInput(false, "4");
+        TaintInput(false, "5");
+        TaintInput(false, "6");
+        TaintInput(false, "7");
+        TaintInput(false, "8");
+        TaintInput(false, "9");
+
+#if NET8_0_OR_GREATER
+        bool largeCache = true;
+#else
+        bool largeCache = false;
+
+#endif
+        TaintInput(!largeCache, "10");
+        TaintInput(!largeCache, "100");
+        TaintInput(!largeCache, "200");
+        TaintInput(!largeCache, "299");
+
+        TaintInput(true, "300");
+        TaintInput(true, "500");
+        TaintInput(true, "1000");
+
+        TaintInput(true, "N");
+        TaintInput(true, "No");
+        TaintInput(true, "5N");
+        TaintInput(true, "5ot");
+        TaintInput(true, "50t");
+        TaintInput(true, "Not");
+        TaintInput(true, "Not Numeric");
+
+        TaintInput(true, "-1");
+        TaintInput(true, "-29");
+        TaintInput(true, "-35");
+    }
 }