[APPSEC-56187] Rework WAF span metrics collection

lib/datadog/appsec.rb

@@ -40,6 +40,11 @@ def reconfigure_lock(&block)
         appsec_component.reconfigure_lock(&block)
       end
 
+      def api_security_enabled?
+        Datadog.configuration.appsec.api_security.enabled &&
+          Datadog.configuration.appsec.api_security.sample_rate.sample?
+      end
+
       private
 
       def components

lib/datadog/appsec/context.rb

@@ -6,11 +6,9 @@ module AppSec
     # interface sufficient for instrumentation to perform threat detection.
     class Context
       ActiveContextError = Class.new(StandardError)
+      WAFMetrics = Struct.new(:timeouts, :duration_ns, :duration_ext_ns, keyword_init: true)
 
-      attr_reader :trace, :span
-
-      # NOTE: This is an intermediate state and will be changed
-      attr_reader :waf_runner
+      attr_reader :trace, :span, :events, :waf_metrics
 
       class << self
         def activate(context)
@@ -34,18 +32,33 @@ def active
       def initialize(trace, span, security_engine)
         @trace = trace
         @span = span
+        @events = []
         @security_engine = security_engine
-        @waf_runner = security_engine.new_context
+        @waf_runner = security_engine.new_runner
+        @waf_metrics = WAFMetrics.new(timeouts: 0, duration_ns: 0, duration_ext_ns: 0)
+        @mutex = Mutex.new
       end
 
       def run_waf(persistent_data, ephemeral_data, timeout = WAF::LibDDWAF::DDWAF_RUN_TIMEOUT)
-        @waf_runner.run(persistent_data, ephemeral_data, timeout)
+        result = @waf_runner.run(persistent_data, ephemeral_data, timeout)
+
+        @mutex.synchronize do
+          @waf_metrics.timeouts += 1 if result.timeout?
+          @waf_metrics.duration_ns += result.duration_ns
+          @waf_metrics.duration_ext_ns += result.duration_ext_ns
+        end
+
+        result
       end
 
       def run_rasp(_type, persistent_data, ephemeral_data, timeout = WAF::LibDDWAF::DDWAF_RUN_TIMEOUT)
         @waf_runner.run(persistent_data, ephemeral_data, timeout)
       end
 
+      def extract_schema
+        @waf_runner.run({ 'waf.context.processor' => { 'extract-schema' => true } }, {})
+      end
+
       def finalize
         @waf_runner.finalize
       end

lib/datadog/appsec/contrib/active_record/instrumentation.rb

@@ -25,7 +25,7 @@ def detect_sql_injection(sql, adapter_name)
             waf_timeout = Datadog.configuration.appsec.waf_timeout
             result = context.run_rasp(Ext::RASP_SQLI, {}, ephemeral_data, waf_timeout)
 
-            if result.status == :match
+            if result.match?
               Datadog::AppSec::Event.tag_and_keep!(context, result)
 
               event = {
@@ -35,7 +35,7 @@ def detect_sql_injection(sql, adapter_name)
                 sql: sql,
                 actions: result.actions
               }
-              context.waf_runner.events << event
+              context.events << event
             end
           end
 

lib/datadog/appsec/contrib/graphql/gateway/watcher.rb

@@ -38,7 +38,7 @@ def watch_multiplex(gateway = Instrumentation.gateway)
                       }
 
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
 
                     block = GraphQL::Reactive::Multiplex.publish(engine, gateway_multiplex)

lib/datadog/appsec/contrib/graphql/reactive/multiplex.rb

@@ -32,7 +32,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/contrib/rack/gateway/watcher.rb

@@ -30,7 +30,7 @@ def watch_request(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Rack::Reactive::Request.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -43,7 +43,7 @@ def watch_request(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 
@@ -61,7 +61,7 @@ def watch_response(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Rack::Reactive::Response.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -74,7 +74,7 @@ def watch_response(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 
@@ -92,7 +92,7 @@ def watch_request_body(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Rack::Reactive::RequestBody.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -105,7 +105,7 @@ def watch_request_body(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 

lib/datadog/appsec/contrib/rack/reactive/request.rb

@@ -55,7 +55,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/contrib/rack/reactive/request_body.rb

@@ -33,7 +33,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/contrib/rack/reactive/response.rb

@@ -39,7 +39,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/contrib/rack/request_middleware.rb

@@ -92,20 +92,20 @@ def call(env)
 
             http_response = AppSec::Response.negotiate(env, block_actions).to_rack if block_actions
 
-            if (result = ctx.waf_runner.extract_schema)
-              ctx.waf_runner.events << {
+            if AppSec.api_security_enabled?
+              ctx.events << {
                 trace: ctx.trace,
                 span: ctx.span,
-                waf_result: result,
+                waf_result: ctx.extract_schema,
               }
             end
 
-            ctx.waf_runner.events.each do |e|
+            ctx.events.each do |e|
               e[:response] ||= gateway_response
               e[:request]  ||= gateway_request
             end
 
-            AppSec::Event.record(ctx.span, *ctx.waf_runner.events)
+            AppSec::Event.record(ctx.span, *ctx.events)
 
             http_response
           ensure
@@ -200,15 +200,13 @@ def add_request_tags(context, env)
 
           def add_waf_runtime_tags(context)
             span = context.span
-            context = context.waf_runner
-
-            return unless span && context
+            return unless span
 
-            span.set_tag('_dd.appsec.waf.timeouts', context.timeouts)
+            span.set_tag('_dd.appsec.waf.timeouts', context.waf_metrics.timeouts)
 
             # these tags expect time in us
-            span.set_tag('_dd.appsec.waf.duration', context.time_ns / 1000.0)
-            span.set_tag('_dd.appsec.waf.duration_ext', context.time_ext_ns / 1000.0)
+            span.set_tag('_dd.appsec.waf.duration', context.waf_metrics.duration_ns / 1000.0)
+            span.set_tag('_dd.appsec.waf.duration_ext', context.waf_metrics.duration_ext_ns / 1000.0)
           end
 
           def to_rack_header(header)

lib/datadog/appsec/contrib/rails/gateway/watcher.rb

@@ -26,7 +26,7 @@ def watch_request_action(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Rails::Reactive::Action.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -39,7 +39,7 @@ def watch_request_action(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 

lib/datadog/appsec/contrib/rails/reactive/action.rb

@@ -39,7 +39,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/contrib/sinatra/gateway/watcher.rb

@@ -28,7 +28,7 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Rack::Reactive::RequestBody.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -41,7 +41,7 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 
@@ -59,7 +59,7 @@ def watch_request_routed(gateway = Instrumentation.gateway)
                   engine = AppSec::Reactive::Engine.new
 
                   Sinatra::Reactive::Routed.subscribe(engine, context) do |result|
-                    if result.status == :match
+                    if result.match?
                       # TODO: should this hash be an Event instance instead?
                       event = {
                         waf_result: result,
@@ -72,7 +72,7 @@ def watch_request_routed(gateway = Instrumentation.gateway)
                       # We want to keep the trace in case of security event
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
-                      context.waf_runner.events << event
+                      context.events << event
                     end
                   end
 

lib/datadog/appsec/contrib/sinatra/reactive/routed.rb

@@ -34,7 +34,7 @@ def self.subscribe(engine, context)
                 waf_timeout = Datadog.configuration.appsec.waf_timeout
                 result = context.run_waf(persistent_data, {}, waf_timeout)
 
-                next if result.status != :match
+                next unless result.match?
 
                 yield result
                 throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/monitor/gateway/watcher.rb

@@ -24,7 +24,7 @@ def watch_user_id(gateway = Instrumentation.gateway)
                 engine = AppSec::Reactive::Engine.new
 
                 Monitor::Reactive::SetUser.subscribe(engine, context) do |result|
-                  if result.status == :match
+                  if result.match?
                     # TODO: should this hash be an Event instance instead?
                     event = {
                       waf_result: result,
@@ -37,7 +37,7 @@ def watch_user_id(gateway = Instrumentation.gateway)
                     # We want to keep the trace in case of security event
                     context.trace.keep! if context.trace
                     Datadog::AppSec::Event.tag_and_keep!(context, result)
-                    context.waf_runner.events << event
+                    context.events << event
                   end
                 end
 

lib/datadog/appsec/monitor/reactive/set_user.rb

@@ -32,7 +32,7 @@ def self.subscribe(engine, context)
               waf_timeout = Datadog.configuration.appsec.waf_timeout
               result = context.run_waf(persistent_data, {}, waf_timeout)
 
-              next if result.status != :match
+              next unless result.match?
 
               yield result
               throw(:block, true) unless result.actions.empty?

lib/datadog/appsec/processor.rb

@@ -1,10 +1,11 @@
 # frozen_string_literal: true
 
-require_relative 'processor/context'
+require_relative 'security_engine/runner'
 
 module Datadog
   module AppSec
     # Processor integrates libddwaf into datadog/appsec
+    # NOTE: This class will be moved under AppSec::SecurityEngine namespace
     class Processor
       attr_reader :diagnostics, :addresses
 
@@ -29,8 +30,8 @@ def finalize
         @handle.finalize
       end
 
-      def new_context
-        Context.new(@handle, telemetry: @telemetry)
+      def new_runner
+        SecurityEngine::Runner.new(@handle, telemetry: @telemetry)
       end
 
       private