add more tests

DataDog · Dec 26, 2024 · cab8f96 · cab8f96
1 parent 67f3178
commit cab8f96
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 10 deletions.
diff --git a/tracer/src/Datadog.Trace/AppSec/Coordinator/SecurityCoordinator.cs b/tracer/src/Datadog.Trace/AppSec/Coordinator/SecurityCoordinator.cs
@@ -101,10 +101,6 @@ internal readonly partial struct SecurityCoordinator
             if (additiveContext?.ShouldRunWith(_security, userId, userLogin, userSessionId, fromSdk) is { Count: > 0 } userAddresses)
             {
                 addresses = userAddresses.ToDictionary(k => k.Key, object (v) => v.Value);
-                if (addresses.IsEmpty())
-                {
-                    return null;
-                }
 
                 if (otherTags is not null)
                 {
@@ -123,6 +119,11 @@ internal readonly partial struct SecurityCoordinator
                 result = additiveContext.Run(addresses, _security.Settings.WafTimeoutMicroSeconds);
                 additiveContext.CommitUserRuns(userAddresses, fromSdk);
                 RecordTelemetry(result);
+
+                if (_localRootSpan.Context.TraceContext is not null)
+                {
+                    _localRootSpan.Context.TraceContext.WafExecuted = true;
+                }
             }
         }
         catch (Exception ex) when (ex is not BlockException)
@@ -139,11 +140,6 @@ internal readonly partial struct SecurityCoordinator
             }
         }
 
-        if (_localRootSpan.Context.TraceContext is not null)
-        {
-            _localRootSpan.Context.TraceContext.WafExecuted = true;
-        }
-
         return result;
     }
 

diff --git a/tracer/src/Datadog.Trace/AppSec/Waf/Context.cs b/tracer/src/Datadog.Trace/AppSec/Waf/Context.cs
@@ -123,7 +123,7 @@ public void CommitUserRuns(IDictionary<string, string> addresses, bool fromSdk)
                 ephemeralParameters);
         }
 
-        // not restart cause it's the total runtime over runs, and we run several * during request
+        // not restart because it's the total runtime over runs, and we run several * during request
         _stopwatch.Start();
         WafReturnCode code;
         lock (_stopwatch)

diff --git a/tracer/test/Datadog.Trace.Security.Unit.Tests/ContextUserEventTests.cs b/tracer/test/Datadog.Trace.Security.Unit.Tests/ContextUserEventTests.cs
@@ -17,6 +17,68 @@ namespace Datadog.Trace.Security.Unit.Tests;
 
 public class ContextUserEventTests
 {
+    [Fact]
+    public void NewValuesTests()
+    {
+        var iWaf = new Mock<IWaf>().Object;
+        var encoder = new Mock<Encoder>().Object;
+        var context = Context.GetContext(IntPtr.Zero, iWaf, new Mock<IWafLibraryInvoker>().Object, encoder);
+        var security = new Mock<IDatadogSecurity>();
+        security.Setup(s => s.AddressEnabled(AddressesConstants.UserId)).Returns(true);
+        var userId = "toto";
+        var addresses = context!.ShouldRunWith(security.Object, userId: userId);
+        addresses.Should().HaveCount(1);
+        addresses.Should().Contain(new KeyValuePair<string, string>(AddressesConstants.UserId, userId));
+        context.CommitUserRuns(addresses, false);
+        userId = "tata";
+        // should run with a different value
+        addresses = context!.ShouldRunWith(security.Object, userId: userId);
+        addresses.Should().Contain(new KeyValuePair<string, string>(AddressesConstants.UserId, userId));
+        addresses.Should().HaveCount(1);
+        context.CommitUserRuns(addresses, false);
+
+        // should not run with same value
+        addresses = context!.ShouldRunWith(security.Object, userId: userId);
+        addresses.Should().HaveCount(0);
+    }
+
+    [Fact]
+    public void AddressDisabledNo()
+    {
+        var iWaf = new Mock<IWaf>().Object;
+        var encoder = new Mock<Encoder>().Object;
+        var context = Context.GetContext(IntPtr.Zero, iWaf, new Mock<IWafLibraryInvoker>().Object, encoder);
+        var security = new Mock<IDatadogSecurity>();
+        security.Setup(s => s.AddressEnabled(AddressesConstants.UserId)).Returns(false);
+        security.Setup(s => s.AddressEnabled(AddressesConstants.UserSessionId)).Returns(true);
+        var userId = "toto";
+        var userSessionId = "123";
+        var addresses = context!.ShouldRunWith(security.Object, userId: userId);
+        // waf shouldn't run with a disabled address
+        addresses.Should().HaveCount(0);
+        context.CommitUserRuns(addresses, false);
+
+        // should run with a different value
+        addresses = context!.ShouldRunWith(security.Object, userId: userId, userSessionId: userSessionId);
+        addresses.Should().Contain(new KeyValuePair<string, string>(AddressesConstants.UserSessionId, userSessionId));
+        addresses.Should().HaveCount(1);
+        context.CommitUserRuns(addresses, false);
+    }
+
+    [Fact]
+    public void NullValueNo()
+    {
+        var iWaf = new Mock<IWaf>().Object;
+        var encoder = new Mock<Encoder>().Object;
+        var context = Context.GetContext(IntPtr.Zero, iWaf, new Mock<IWafLibraryInvoker>().Object, encoder);
+        var security = new Mock<IDatadogSecurity>();
+        security.Setup(s => s.AddressEnabled(AddressesConstants.UserId)).Returns(true);
+        var addresses = context!.ShouldRunWith(security.Object, userId: null);
+        // waf shouldn't run with a disabled address
+        addresses.Should().HaveCount(0);
+        context.CommitUserRuns(addresses, false);
+    }
+
     [Fact]
     public void SdkOverrideTest()
     {
@@ -29,5 +91,8 @@ public void SdkOverrideTest()
         var addresses = context!.ShouldRunWith(security.Object, userId: userId);
         addresses.Should().HaveCount(1);
         addresses.Should().Contain(new KeyValuePair<string, string>(AddressesConstants.UserId, userId));
+        context.CommitUserRuns(addresses, true);
+        addresses = context!.ShouldRunWith(security.Object, userId: "other");
+        addresses.Should().HaveCount(0);
     }
 }