[ASM] Context disposed related exceptions when calling the WAF

tracer/src/Datadog.Trace/AppSec/Coordinator/HttpTransportBase.cs

@@ -43,6 +43,8 @@ internal void DisposeAdditiveContext()
 
     internal bool IsAdditiveContextDisposed() => _isAdditiveContextDisposed;
 
+    protected void SetAdditiveContextDisposed(bool value) => _isAdditiveContextDisposed = value;
+
     internal abstract void SetAdditiveContext(IContext additiveContext);
 
     internal abstract IHeadersCollection GetRequestHeaders();

tracer/src/Datadog.Trace/AppSec/Coordinator/SecurityCoordinator.Core.cs

@@ -6,14 +6,15 @@
 #nullable enable
 #pragma warning disable CS0282
 #if !NETFRAMEWORK
+using System;
 using System.Collections;
 using System.Collections.Generic;
 using System.Linq;
-using System.Runtime.CompilerServices;
 using Datadog.Trace.AppSec.Waf;
 using Datadog.Trace.Headers;
 using Datadog.Trace.Util.Http;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.Extensions.Primitives;
 
@@ -176,13 +177,32 @@ internal override bool ReportedExternalWafsRequestHeaders
 
         internal override void MarkBlocked() => Context.Items[BlockingAction.BlockDefaultActionName] = true;
 
-        internal override IContext GetAdditiveContext() => Context.Features.Get<IContext>();
+        internal override IContext? GetAdditiveContext() => IsAdditiveContextDisposed() ? null : GetContextFeatures()?.Get<IContext>();
 
         internal override void SetAdditiveContext(IContext additiveContext) => Context.Features.Set(additiveContext);
 
         internal override IHeadersCollection GetRequestHeaders() => new HeadersCollectionAdapter(Context.Request.Headers);
 
         internal override IHeadersCollection GetResponseHeaders() => new HeadersCollectionAdapter(Context.Response.Headers);
+
+        // In some edge situations we can get an ObjectDisposedException when accessing the context features or other
+        // properties such as Context.Items or Context.Response.Headers that ultimatelly rely on features
+        // This means that the context has been uninitiallized and we should not try to access it anymore
+        // Unfortunatelly, there is no way to know that but catching the exception or using reflection
+        private IFeatureCollection? GetContextFeatures()
+        {
+            try
+            {
+                return Context.Features;
+            }
+            catch (ObjectDisposedException)
+            {
+                Log.Debug("ObjectDisposedException while trying to access a Context.");
+                SetAdditiveContextDisposed(true);
+                return null;
+            }
+        }
     }
 }
 #endif
+

tracer/src/Datadog.Trace/AppSec/Coordinator/SecurityCoordinator.cs

@@ -10,6 +10,7 @@
 using Datadog.Trace.AppSec.Waf;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Util;
+
 #if !NETFRAMEWORK
 using Microsoft.AspNetCore.Http;
 #else
@@ -47,22 +48,13 @@ internal readonly partial struct SecurityCoordinator
     {
         SecurityReporter.LogAddressIfDebugEnabled(args);
         IResult? result = null;
+
         try
         {
-            var additiveContext = _httpTransport.GetAdditiveContext();
+            var additiveContext = GetOrCreateAdditiveContext();
 
-            if (additiveContext == null)
-            {
-                additiveContext = _security.CreateAdditiveContext();
-                // prevent very cases where waf has been disposed between here and has been passed as argument until the 2nd line of constructor..
-                if (additiveContext != null)
-                {
-                    _httpTransport.SetAdditiveContext(additiveContext);
-                }
-            }
-            else if (_httpTransport.IsAdditiveContextDisposed())
+            if (additiveContext is null)
             {
-                Log.Warning("Waf could not run as waf additive context is disposed");
                 return null;
             }
 
@@ -173,4 +165,27 @@ private static Dictionary<string, object> ExtractHeaders(ICollection<string> key
 
         return headersDic;
     }
+
+    private IContext? GetOrCreateAdditiveContext()
+    {
+        var additiveContext = _httpTransport.GetAdditiveContext();
+
+        if (_httpTransport.IsAdditiveContextDisposed())
+        {
+            Log.Warning("Waf could not run as waf additive context is disposed");
+            return null;
+        }
+
+        if (additiveContext == null)
+        {
+            additiveContext = _security.CreateAdditiveContext();
+
+            if (additiveContext is not null)
+            {
+                _httpTransport.SetAdditiveContext(additiveContext);
+            }
+        }
+
+        return additiveContext;
+    }
 }

tracer/test/Datadog.Trace.Security.Unit.Tests/RASP/RaspWafTests.cs

@@ -18,7 +18,6 @@
 using Datadog.Trace.Vendors.Serilog.Events;
 using FluentAssertions;
 using Xunit;
-
 using Action = Datadog.Trace.AppSec.Rcm.Models.Asm.Action;
 
 namespace Datadog.Trace.Security.Unit.Tests;

tracer/test/Datadog.Trace.Security.Unit.Tests/SecurityCoordinatorTests.cs

@@ -3,9 +3,16 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
+using System;
+using Datadog.Trace.AppSec;
 using Datadog.Trace.AppSec.Coordinator;
 using FluentAssertions;
+#if NETCOREAPP
+using Microsoft.AspNetCore.Http;
+#endif
+using Moq;
 using Xunit;
+using static Datadog.Trace.AppSec.Coordinator.SecurityCoordinator;
 
 namespace Datadog.Trace.Security.Unit.Tests
 {
@@ -19,5 +26,31 @@ public void DefaultBehavior()
             var secCoord = SecurityCoordinator.TryGet(target, span);
             secCoord.Should().BeNull();
         }
+
+#if NETCOREAPP
+        [Fact]
+        public void GivenHttpTransportInstanceWithDisposedContext_WhenGetContextUninitialized_ThenResultIsTrue()
+        {
+            var contextMoq = new Mock<HttpContext>();
+            contextMoq.Setup(x => x.Features).Throws(new ObjectDisposedException("Test exception"));
+            var context = contextMoq.Object;
+            HttpTransport transport = new(context);
+            transport.GetAdditiveContext().Should().BeNull();
+            transport.IsAdditiveContextDisposed().Should().BeTrue();
+        }
+
+        [Fact]
+        public void GivenSecurityCoordinatorInstanceWithDisposedContext_WheRunWaf_ThenResultIsNull()
+        {
+            var contextMoq = new Mock<HttpContext>();
+            contextMoq.Setup(x => x.Features).Throws(new ObjectDisposedException("Test exception"));
+            var context = contextMoq.Object;
+            CoreHttpContextStore.Instance.Set(context);
+            var span = new Span(new SpanContext(1, 1), new DateTimeOffset());
+            var securityCoordinator = SecurityCoordinator.TryGet(AppSec.Security.Instance, span);
+            var result = securityCoordinator.Value.RunWaf(new(), runWithEphemeral: true, isRasp: true);
+            result.Should().BeNull();
+        }
+#endif
     }
 }

tracer/test/test-applications/security/Samples.Security.AspNetCore5/Controllers/IastController.cs

@@ -2,37 +2,34 @@
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Data;
-using System.Data.Common;
 using System.Data.SqlClient;
 using System.Data.SQLite;
 using System.Diagnostics;
 using System.DirectoryServices;
-using System.Drawing.Drawing2D;
 using System.IO;
 using System.Linq;
 using System.Net;
 using System.Net.Http;
 using System.Net.Mail;
-using System.Reflection;
-using System.Runtime.Versioning;
 using System.Security.Cryptography;
-using System.Text;
 #if NETCOREAPP3_0_OR_GREATER
 using System.Text.Json;
 #endif
-using System.Threading.Tasks;
+using System.Threading;
 using System.Xml;
-using System.Xml.Linq;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.AspNetCore.Mvc.Filters;
 using Microsoft.Data.Sqlite;
-using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Primitives;
 using MongoDB.Bson;
 using MongoDB.Driver;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
+#if NET5_0_OR_GREATER
+using System.Runtime.Versioning;
+#endif
+
 #if NETCOREAPP3_0_OR_GREATER
 using MySql.Data.MySqlClient;
 using Npgsql;
@@ -584,6 +581,73 @@ public IActionResult GetDirectoryContent(string directory)
             }
         }
 
+        // This method actually performs some file operations after the request has been normally closed.
+        [HttpGet("GetFileContentThread")]
+        [Route("GetFileContentThread")]
+        public IActionResult GetFileContentThread(string file, int numThreads = 100, int delayPerThread = 50)
+        {
+            for (int i = 0; i < numThreads; i++)
+            {
+                var thread = new Thread(() => { GetFileAux(file, i * delayPerThread); });
+                thread.Start();
+            }
+
+            return Content("Ok");
+        }
+
+        private void GetFileAux(string file, int delay)
+        {
+            try
+            {
+                if (delay > 0)
+                {
+                    Thread.Sleep(delay);
+                }
+                GetFileContent(file);
+            }
+            catch (Exception ex)
+            {
+                if (!ex.Message.Contains("BlockException"))
+                {
+                    throw;
+                }
+            }
+        }
+
+#if NET5_0_OR_GREATER
+        // This method tests some edge conditions that can happen
+        [HttpGet("GetFileContentEdgeConditions")]
+        [Route("GetFileContentEdgeConditions")]
+        public IActionResult GetFileContentEdgeConditions(string file, bool uninitializeContext = true, bool setStatusCode = true, bool setContent = true, bool abortContext = true)
+        {
+            if (setStatusCode)
+            {
+                Response.StatusCode = 200;
+            }
+
+            if (setContent)
+            {
+                Response.ContentType = "text/plain";
+                Response.WriteAsync("This is a dummy content.").Wait();
+            }
+
+            if (abortContext)
+            {
+                HttpContext.Abort();
+            }
+
+            if (uninitializeContext)
+            {
+                (HttpContext as DefaultHttpContext)?.Uninitialize();
+            }
+
+            // call RASP and IAST
+            GetFileAux(file, 0);
+
+            return Content("Ok");
+        }
+#endif
+
         [HttpGet("GetFileContent")]
         [Route("GetFileContent")]
         public IActionResult GetFileContent(string file)
@@ -1133,7 +1197,7 @@ public IActionResult ReflectedXssEscaped(string param)
             return View("Xss");
         }
 
-        #if NET6_0_OR_GREATER
+#if NET6_0_OR_GREATER
         [HttpGet("InterpolatedSqlString")]
         [Route("InterpolatedSqlString")]
         public IActionResult InterpolatedSqlString(string name)
@@ -1173,7 +1237,7 @@ public IActionResult InterpolatedSqlString(string name)
 
             return Content("Yey");
         }
-        #endif
+#endif
 
         [HttpGet("TestJsonTagSizeExceeded")]
         [Route("TestJsonTagSizeExceeded")]

tracer/test/test-applications/security/Samples.Security.AspNetCore5/Views/Home/Index.cshtml

@@ -88,6 +88,10 @@
 
     <div><a href="/Iast/GetFileContent?file=/nonexisting.txt">GET /Iast/GetFileContent?file=/nonexisting.txt</a></div>
 
+    <div><a href="/Iast/GetFileContentEdgeConditions?file=/nonexisting.txt&uninitializeContext=true&setStatusCode=true&setContent=true&abortContext=true">GET /Iast/GetFileContentEdgeConditions?file=/nonexisting.txt&uninitializeContext=true&setStatusCode=true&setContent=true&abortContext=true</a></div>
+
+    <div><a href="/Iast/GetFileContentThread?file=/nonexisting.txt&numThreads=100&delayPerThread=50">GET /Iast/GetFileContentThread?file=/nonexisting.txt&numThreads=100&delayPerThread=50</a></div>
+
     <div><a href="javascript:post('/Iast/ExecuteQueryFromBodyText','&quot;SELECT Surname from Persons where name=\'Vicent\'&quot;')">POST /Iast/ExecuteQueryFromBodyText</a></div>
 
     <div><a href="javascript:post('/Iast/ExecuteQueryFromBodyQueryData', '{%22InnerQuery%22: {%22Arguments%22: [%22SELECT Surname from Persons where name=\'Vicent\'%22]}, %22Query%22: %22SELECT Surname from Persons where name=\'Vicent\'%22,%22QueryNotUsed%22: %22SELECT Surname from Persons where name=\'Vicent\'%22,%22IntField%22: 1,%22Arguments%22: [%22SELECT Surname from Persons where name=\'Vicent\'%22, %22SELECT Surname from Persons where name=\'Mark\'%22],%22StringMap%22: {%22query1%22: %22SELECT Surname from Persons where name=\'Vicent\'%22,%22query2%22: %22SELECT Surname from Persons where name=\'Vicent\'%22},%22StringArrayArguments%22: [%22SELECT Surname from Persons where name=\'Vicent\'%22, %22SELECT Surname from Persons where name=\'Mark\'%22]}')">POST /Iast/ExecuteQueryFromBodyQueryData</a></div>

tracer/test/test-applications/security/aspnet/Samples.Security.AspNetMvc5/Controllers/IastController.cs

@@ -15,6 +15,7 @@
 using System.Web.Script.Serialization;
 using System.Xml;
 using System.Net.Mail;
+using System.Threading;
 
 namespace Samples.Security.AspNetCore5.Controllers
 {
@@ -307,6 +308,63 @@ public ActionResult GetDirectoryContent(string directory)
             }
         }
 
+        // This method actually performs some file operations after the request has been normally closed.
+        [Route("GetFileContentThread")]
+        public ActionResult GetFileContentThread(string file, int numThreads = 100, int delayPerThread = 50)
+        {
+            for (int i = 0; i < numThreads; i++)
+            {
+                var thread = new Thread(() => { GetFileAux(file, i * delayPerThread); });
+                thread.Start();
+            }
+
+            return Content("Ok");
+        }
+
+        private void GetFileAux(string file, int i)
+        {
+            try
+            {
+                Thread.Sleep(i);
+                GetFileContent(file);
+            }
+            catch (Exception ex)
+            {
+                if (!ex.Message.Contains("BlockException"))
+                {
+                    throw;
+                }
+            }
+        }
+
+        // This method tests some edge conditions that can happen
+        [Route("GetFileContentEdgeConditions")]
+        public ActionResult GetFileContentEdgeConditions(string file, bool endRequest = true, bool setStatusCode = true, bool setContent = true)
+        {
+            if (setStatusCode)
+            {
+                Response.StatusCode = 200;
+            }
+
+            if (setContent)
+            {
+                Response.ContentType = "text/plain";
+                Response.Write("This is a dummy content.");
+            }
+
+            if (endRequest)
+            {
+                Response.End();
+                Response.Close();
+                HttpContext.ApplicationInstance.CompleteRequest();
+            }
+
+            // call RASP and IAST
+            GetFileAux(file, 0);
+
+            return Content("Ok");
+        }
+
         [Route("GetFileContent")]
         public ActionResult GetFileContent(string file)
         {