[ASM] Added 1 hour timer to vulnerability deduplication (#4411)

* Add 1hr timer to deduplication * Update tracer/src/Datadog.Trace/Iast/HashBasedDeduplication.cs Co-authored-by: Daniel Romano <108014683+daniel-romano-DD@users.noreply.github.com> * Remove empty space --------- Co-authored-by: Daniel Romano <108014683+daniel-romano-DD@users.noreply.github.com>
DataDog · Jul 19, 2023 · 55d50ce · 55d50ce
1 parent 77a8cb8
commit 55d50ce
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 8 deletions.
diff --git a/tracer/src/Datadog.Trace/Iast/HashBasedDeduplication.cs b/tracer/src/Datadog.Trace/Iast/HashBasedDeduplication.cs
@@ -3,38 +3,53 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
+using System;
 using System.Collections.Generic;
 
+#nullable enable
+
 namespace Datadog.Trace.Iast;
 
 internal class HashBasedDeduplication
 {
     public const int MaximumSize = 1000;
-    private HashSet<int> vulnerabilityHashes = new();
+    public const int MinutesToClearCache = 60;
+    private HashSet<int> _vulnerabilityHashes = new();
+    private DateTime _cacheClearedTime;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="HashBasedDeduplication"/> class.
     /// For testing only.
     /// Note that this API does NOT replace the global HashBasedDeduplication instance.
     /// </summary>
-    internal HashBasedDeduplication()
+    internal HashBasedDeduplication(DateTime? currentTime = null)
     {
+        _cacheClearedTime = currentTime ?? DateTime.Now;
     }
 
     public static HashBasedDeduplication Instance { get; } = new();
 
-    public bool Add(Vulnerability vulnerability)
+    public bool Add(Vulnerability vulnerability, DateTime? addTime = null)
     {
         var hashCode = vulnerability.GetHashCode();
+        var currentTime = addTime ?? DateTime.Now;
 
         bool newVulnerability;
-        lock (vulnerabilityHashes)
+        lock (_vulnerabilityHashes)
         {
-            newVulnerability = vulnerabilityHashes.Add(hashCode);
-            if (newVulnerability && vulnerabilityHashes.Count > MaximumSize)
+            if ((currentTime - _cacheClearedTime).TotalMinutes >= MinutesToClearCache)
+            {
+                _vulnerabilityHashes.Clear();
+                _cacheClearedTime = currentTime;
+            }
+
+            newVulnerability = _vulnerabilityHashes.Add(hashCode);
+
+            if (newVulnerability && _vulnerabilityHashes.Count > MaximumSize)
             {
-                vulnerabilityHashes.Clear();
-                vulnerabilityHashes.Add(hashCode);
+                _vulnerabilityHashes.Clear();
+                _cacheClearedTime = currentTime;
+                _vulnerabilityHashes.Add(hashCode);
             }
         }
 

diff --git a/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/HashBasedDeduplicationTests.cs b/tracer/test/Datadog.Trace.Security.Unit.Tests/IAST/HashBasedDeduplicationTests.cs
@@ -4,6 +4,7 @@
 // </copyright>
 
 using Datadog.Trace.Iast;
+using FluentAssertions;
 using Xunit;
 
 namespace Datadog.Trace.Security.Unit.Tests.Iast;
@@ -183,4 +184,22 @@ public void GivenManyVulnerabilities_WhenAddedToDeduplication_CacheIsCleared()
 
         Assert.True(instance.Add(vulnerability1));
     }
+
+    [Theory]
+    [InlineData(10, false, 55, true)]
+    [InlineData(60, true, 55, false)]
+    [InlineData(6, false, 5, false)]
+    [InlineData(61, true, 65, true)]
+    public void GivenHashBasedDeduplication_WhenTestDeduplicationTimeout_ResultIsOk(int minutesAfter1, bool expectedResult1, int minutesAfter2, bool expectedResult2)
+    {
+        var date = new System.DateTime(2001, 1, 1, 1, 0, 0);
+        var instance = new HashBasedDeduplication(date);
+        Assert.True(instance.Add(new Vulnerability(VulnerabilityTypeName.NoSameSiteCookie, null, new Evidence("value")), date));
+        date = date.AddMinutes(minutesAfter1);
+        var result = instance.Add(new Vulnerability(VulnerabilityTypeName.NoSameSiteCookie, null, new Evidence("value")), date);
+        result.Should().Be(expectedResult1);
+        date = date.AddMinutes(minutesAfter2);
+        result = instance.Add(new Vulnerability(VulnerabilityTypeName.NoSameSiteCookie, null, new Evidence("value")), date);
+        result.Should().Be(expectedResult2);
+    }
 }