DataDog · anna-git · Aug 31, 2022 · Aug 8, 2022 · Aug 12, 2022 · Aug 9, 2022
@@ -20,6 +20,7 @@ internal static class AddressesConstants
         public const string RequestBodyFileNames = "server.request.body.filenames";
         public const string RequestPathParams = "server.request.path_params";
         public const string RequestBodyCombinedFileSize = "server.request.body.combined_file_size";
+        public const string RequestClientIp = "http.client_ip";
         public const string ResponseStatus = "server.response.status";
 
         public const string ResponseBodyRaw = "server.response.body.raw";

@@ -0,0 +1,31 @@
+// <copyright file="BlockException.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+using System;
+
+namespace Datadog.Trace.AppSec
+{
+    internal class BlockException : Exception
+    {
+        internal BlockException()
+        {
+        }
+
+        internal BlockException(string message)
+            : base(message)
+        {
+        }
+
+        internal BlockException(string message, Exception inner)
+            : base(message, inner)
+        {
+        }
+
+        internal BlockException(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context)
+            : base(info, context)
+        {
+        }
+    }
+}
@@ -10,6 +10,7 @@
 using System.Web.Routing;
 #endif
 using Datadog.Trace.AppSec.Transports.Http;
+using Datadog.Trace.Configuration;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Util.Http;
 using Datadog.Trace.Vendors.Serilog.Events;
@@ -26,35 +27,42 @@ internal partial class InstrumentationGateway
 
         public event EventHandler<InstrumentationGatewaySecurityEventArgs> PathParamsAvailable;
 
+        public event EventHandler<InstrumentationGatewaySecurityEventArgs> StartRequest;
+
         public event EventHandler<InstrumentationGatewaySecurityEventArgs> EndRequest;
 
         public event EventHandler<InstrumentationGatewaySecurityEventArgs> BodyAvailable;
 
         public event EventHandler<InstrumentationGatewayEventArgs> LastChanceToWriteTags;
 
-        public void RaiseEndRequest(HttpContext context, HttpRequest request, Span relatedSpan)
+        public event EventHandler<InstrumentationGatewayBlockingEventArgs> BlockingOpportunity;
+
+        public void RaiseRequestEnd(HttpContext context, HttpRequest request, Span relatedSpan)
         {
             var getEventData = () =>
             {
-                var eventData = request.PrepareArgsForWaf();
+                var eventData = request.PrepareArgsForWaf(relatedSpan);
                 eventData.Add(AddressesConstants.ResponseStatus, context.Response.StatusCode.ToString());
                 return eventData;
             };
 
             RaiseEvent(context, relatedSpan, getEventData, EndRequest);
         }
 
+        public void RaiseRequestStart(HttpContext context, HttpRequest request, Span relatedSpan)
+        {
+            var getEventData = () => request.PrepareArgsForWaf(relatedSpan);
+            RaiseEvent(context, relatedSpan, getEventData, StartRequest);
+        }
+
         public void RaisePathParamsAvailable(HttpContext context, Span relatedSpan, IDictionary<string, object> pathParams, bool eraseExistingAddress = true) => RaiseEvent(context, relatedSpan, () => new Dictionary<string, object> { { AddressesConstants.RequestPathParams, pathParams } }, PathParamsAvailable, eraseExistingAddress);
 
         public void RaiseBodyAvailable(HttpContext context, Span relatedSpan, object body)
         {
             var getEventData = () =>
             {
                 var keysAndValues = BodyExtractor.Extract(body);
-                var eventData = new Dictionary<string, object>
-                {
-                    { AddressesConstants.RequestBody, keysAndValues }
-                };
+                var eventData = new Dictionary<string, object> { { AddressesConstants.RequestBody, keysAndValues } };
                 return eventData;
             };
 
@@ -95,10 +103,15 @@ private void RaiseEvent(HttpContext context, Span relatedSpan, Func<IDictionary<
                 LogAddressIfDebugEnabled(eventData);
                 eventHandler.Invoke(this, new InstrumentationGatewaySecurityEventArgs(eventData, transport, relatedSpan, eraseExistingAddress));
             }
-            catch (Exception ex)
+            catch (Exception ex) when (ex is not BlockException)
             {
                 Log.Error(ex, "DDAS-0004-00: AppSec failed to process request.");
             }
         }
+
+        internal void RaiseBlockingOpportunity(HttpContext context, Scope scope, ImmutableTracerSettings tracerSettings, Action<InstrumentationGatewayBlockingEventArgs> doBeforeActualBlocking = null)
+        {
+            BlockingOpportunity?.Invoke(this, new InstrumentationGatewayBlockingEventArgs(context, scope, tracerSettings, doBeforeActualBlocking));
+        }
     }
 }
@@ -0,0 +1,48 @@
+// <copyright file="InstrumentationGatewayBlockingEventArgs.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+using System;
+using Datadog.Trace.AppSec.Transports;
+using Datadog.Trace.AppSec.Transports.Http;
+using Datadog.Trace.Configuration;
+#if NETFRAMEWORK
+using System.Web;
+using System.Web.Routing;
+#else
+using Microsoft.AspNetCore.Http;
+#endif
+
+namespace Datadog.Trace.AppSec
+{
+    internal class InstrumentationGatewayBlockingEventArgs : EventArgs
+    {
+        internal InstrumentationGatewayBlockingEventArgs(HttpContext context, Scope scope, ImmutableTracerSettings settings, Action<InstrumentationGatewayBlockingEventArgs> doBeforeBlocking = null)
+        {
+            Context = context;
+            Scope = scope;
+            TracerSettings = settings;
+            DoBeforeBlocking = doBeforeBlocking;
+            Transport = new HttpTransport(context);
+        }
+
+        internal HttpContext Context { get; }
+
+        internal ITransport Transport { get; }
+
+        internal Scope Scope { get; }
+
+        internal ImmutableTracerSettings TracerSettings { get; }
+
+        private Action<InstrumentationGatewayBlockingEventArgs> DoBeforeBlocking { get; }
+
+        internal void InvokeDoBeforeBlocking()
+        {
+            if (Scope != null && DoBeforeBlocking != null)
+            {
+                DoBeforeBlocking(this);
+            }
+        }
+    }
+}
@@ -8,6 +8,7 @@
 using System.Runtime.CompilerServices;
 using System.Threading;
 using Datadog.Trace.AppSec.Transports;
+using Datadog.Trace.AppSec.Transports.Http;
 using Datadog.Trace.AppSec.Waf;
 using Datadog.Trace.AppSec.Waf.ReturnTypes.Managed;
 using Datadog.Trace.ClrProfiler;
@@ -17,6 +18,7 @@
 using Datadog.Trace.Sampling;
 using Datadog.Trace.Vendors.Newtonsoft.Json;
 using Datadog.Trace.Vendors.Serilog.Events;
+using Datadog.Trace.Vendors.StatsdClient;
 
 namespace Datadog.Trace.AppSec
 {
@@ -69,10 +71,7 @@ static Security()
 
             ResponseHeaders = new()
             {
-                { "content-length", string.Empty },
-                { "content-type", string.Empty },
-                { "Content-Encoding", string.Empty },
-                { "Content-Language", string.Empty },
+                { "content-length", string.Empty }, { "content-type", string.Empty }, { "Content-Encoding", string.Empty }, { "Content-Language", string.Empty },
             };
         }
 
@@ -97,9 +96,11 @@ private Security(SecuritySettings settings = null, InstrumentationGateway instru
                     _waf = waf ?? Waf.Waf.Create(_settings.ObfuscationParameterKeyRegex, _settings.ObfuscationParameterValueRegex, _settings.Rules);
                     if (_waf?.InitializedSuccessfully ?? false)
                     {
-                        _instrumentationGateway.EndRequest += RunWaf;
-                        _instrumentationGateway.PathParamsAvailable += AggregateAddressesInContext;
-                        _instrumentationGateway.BodyAvailable += AggregateAddressesInContext;
+                        _instrumentationGateway.StartRequest += RunWafAndReact;
+                        _instrumentationGateway.EndRequest += RunWafAndReactAndCleanup;
+                        _instrumentationGateway.PathParamsAvailable += RunWafAndReact;
+                        _instrumentationGateway.BodyAvailable += RunWafAndReact;
+                        _instrumentationGateway.BlockingOpportunity += MightStopRequest;
 #if NETFRAMEWORK
                         try
                         {
@@ -125,7 +126,7 @@ private Security(SecuritySettings settings = null, InstrumentationGateway instru
                         _settings.Enabled = false;
                     }
 
-                    _instrumentationGateway.EndRequest += ReportWafInitInfoOnce;
+                    _instrumentationGateway.StartRequest += ReportWafInitInfoOnce;
                     LifetimeManager.Instance.AddShutdownTask(RunShutdown);
                     _rateLimiter = new AppSecRateLimiter(_settings.TraceRateLimit);
                 }
@@ -142,10 +143,7 @@ private Security(SecuritySettings settings = null, InstrumentationGateway instru
         /// </summary>
         public static Security Instance
         {
-            get
-            {
-                return LazyInitializer.EnsureInitialized(ref _instance, ref _globalInstanceInitialized, ref _globalInstanceLock);
-            }
+            get => LazyInitializer.EnsureInitialized(ref _instance, ref _globalInstanceInitialized, ref _globalInstanceLock);
 
             set
             {
@@ -278,6 +276,11 @@ private void InstrumentationGateway_AddHeadersResponseTags(object sender, Instru
         private void Report(ITransport transport, Span span, IResult result, bool blocked)
         {
             span.SetTag(Tags.AppSecEvent, "true");
+            if (blocked)
+            {
+                span.SetTag(Tags.AppSecBlocked, "true");
+            }
+
             var resultData = result.Data;
             SetTraceSamplingPriority(span);
 
@@ -332,46 +335,61 @@ private IContext GetOrCreateContext(ITransport transport)
             return additiveContext;
         }
 
-        private void AggregateAddressesInContext(object sender, InstrumentationGatewaySecurityEventArgs e)
+        private void MightStopRequest(object sender, InstrumentationGatewayBlockingEventArgs args)
+        {
+            if (args.Transport.Blocked)
+            {
+                AddResponseHeaderTags(args.Transport, args.Scope.Span);
+                args.InvokeDoBeforeBlocking();
+                var transport = args.Transport;
+                var additiveContext = GetOrCreateContext(transport);
+                additiveContext.Dispose();
+                throw new BlockException();
+            }
+        }
+
+        private void RunWafAndReactAndCleanup(object sender, InstrumentationGatewaySecurityEventArgs e)
         {
-            var context = GetOrCreateContext(e.Transport);
-            context.AggregateAddresses(e.EventData, e.OverrideExistingAddress);
+            RunWafAndReact(sender, e);
+            e.Transport.DisposeAdditiveContext();
         }
 
-        // NOTE: This method disposes of the WAF context, so it should be run once,
-        // and only once, at the end of the request.
-        private void RunWaf(object sender, InstrumentationGatewaySecurityEventArgs e)
+        private void RunWafAndReact(object sender, InstrumentationGatewaySecurityEventArgs e)
         {
             try
             {
-                using var additiveContext = GetOrCreateContext(e.Transport);
-                additiveContext.AggregateAddresses(e.EventData, e.OverrideExistingAddress);
+                if (e.Transport.Blocked)
+                {
+                    return;
+                }
+
+                var additiveContext = GetOrCreateContext(e.Transport);
                 var span = GetLocalRootSpan(e.RelatedSpan);
 
                 AnnotateSpan(span);
 
                 // run the WAF and execute the results
-                using var wafResult = additiveContext.Run(_settings.WafTimeoutMicroSeconds);
-                if (wafResult.ReturnCode == ReturnCode.Monitor || wafResult.ReturnCode == ReturnCode.Block)
+                using var wafResult = additiveContext.Run(e.EventData, _settings.WafTimeoutMicroSeconds);
+                if (wafResult.ReturnCode is ReturnCode.Monitor or ReturnCode.Block)
                 {
-                    var block = wafResult.ReturnCode == ReturnCode.Block;
+                    var block = wafResult.ReturnCode == ReturnCode.Block || wafResult.Data.Contains("ublock");
                     if (block)
                     {
-                        // blocking has been removed, waiting a better implementation
+                        e.Transport.WriteBlockedResponse(_settings.BlockedJsonTemplate, _settings.BlockedHtmlTemplate);
                     }
 
                     Report(e.Transport, span, wafResult, block);
                 }
             }
-            catch (Exception ex)
+            catch (Exception ex) when (ex is not BlockException)
             {
                 Log.Error(ex, "Call into the security module failed");
             }
         }
 
         private void ReportWafInitInfoOnce(object sender, InstrumentationGatewaySecurityEventArgs e)
         {
-            _instrumentationGateway.EndRequest -= ReportWafInitInfoOnce;
+            _instrumentationGateway.StartRequest -= ReportWafInitInfoOnce;
             var span = e.RelatedSpan.Context.TraceContext.RootSpan ?? e.RelatedSpan;
             span.Context.TraceContext?.SetSamplingPriority(SamplingPriorityValues.UserKeep, SamplingMechanism.Asm);
             span.SetMetric(Metrics.AppSecWafInitRulesLoaded, _waf.InitializationResult.LoadedRules);
@@ -388,9 +406,11 @@ private void RunShutdown()
         {
             if (_instrumentationGateway != null)
             {
-                _instrumentationGateway.PathParamsAvailable -= AggregateAddressesInContext;
-                _instrumentationGateway.BodyAvailable -= AggregateAddressesInContext;
-                _instrumentationGateway.EndRequest -= RunWaf;
+                _instrumentationGateway.PathParamsAvailable -= RunWafAndReact;
+                _instrumentationGateway.BodyAvailable -= RunWafAndReact;
+                _instrumentationGateway.StartRequest -= RunWafAndReact;
+                _instrumentationGateway.EndRequest -= RunWafAndReactAndCleanup;
+                _instrumentationGateway.BlockingOpportunity -= MightStopRequest;
 
 #if NETFRAMEWORK
                 if (_usingIntegratedPipeline)