Move extract_schema into Context

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

@@ -57,6 +57,10 @@ def run_rasp(_type, persistent_data, ephemeral_data, timeout = WAF::LibDDWAF::DD
         @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/rack/request_middleware.rb

@@ -92,11 +92,11 @@ def call(env)
 
             http_response = AppSec::Response.negotiate(env, block_actions).to_rack if block_actions
 
-            if (result = ctx.waf_runner.extract_schema)
+            if AppSec.api_security_enabled?
               ctx.waf_runner.events << {
                 trace: ctx.trace,
                 span: ctx.span,
-                waf_result: result,
+                waf_result: ctx.extract_schema,
               }
             end
 

lib/datadog/appsec/security_engine/runner.rb

@@ -17,17 +17,6 @@ def initialize(handle, telemetry:)
           @debug_tag = "libddwaf:#{WAF::VERSION::STRING} method:ddwaf_run"
         end
 
-        # TODO: Replace nil return value with SecurityEngine::Result::Ok
-        def extract_schema
-          return generic_ok_result unless extract_schema?
-
-          persistent_data = {
-            'waf.context.processor' => { 'extract-schema' => true }
-          }
-
-          run(persistent_data, {})
-        end
-
         def run(persistent_data, ephemeral_data, timeout = WAF::LibDDWAF::DDWAF_RUN_TIMEOUT)
           @mutex.lock
 
@@ -78,7 +67,7 @@ def try_run(persistent_data, ephemeral_data, timeout)
           Datadog.logger.debug { "#{@debug_tag} execution error: #{e} backtrace: #{e.backtrace&.first(3)}" }
           @telemetry.report(e, description: 'libddwaf-rb internal low-level error')
 
-          fallback_waf_error_result
+          [:err_internal, WAF::Result.new(:err_internal, [], 0, false, [], [])]
         end
 
         def report_execution(result)
@@ -93,23 +82,6 @@ def report_execution(result)
             @telemetry.error(message)
           end
         end
-
-        # NOTE: This configuration reads should be a part of the SecurityEngine
-        #       configuration instead.
-        def extract_schema?
-          Datadog.configuration.appsec.api_security.enabled &&
-            Datadog.configuration.appsec.api_security.sample_rate.sample?
-        end
-
-        def fallback_waf_error_result
-          [:err_internal, WAF::Result.new(:err_internal, [], 0, false, [], [])]
-        end
-
-        def generic_ok_result
-          Result::Ok.new(
-            events: [], actions: [], derivatives: [], timeout: false, duration_ns: 0, duration_ext_ns: 0
-          )
-        end
       end
     end
   end

sig/datadog/appsec.rbs

@@ -11,6 +11,8 @@ module Datadog
 
     def self.reconfigure_lock: () { () -> void } -> void
 
+    def self.api_security_enabled?: () -> bool
+
     private
 
     def self.components: () -> Datadog::Core::Configuration::Components

sig/datadog/appsec/contrib/rack/gateway/response.rbs

@@ -10,9 +10,9 @@ module Datadog
 
             attr_reader headers: untyped
 
-            attr_reader active_context: Datadog::AppSec::Processor::Context
+            attr_reader context: AppSec::Context
 
-            def initialize: (untyped body, untyped status, untyped headers, active_context: Datadog::AppSec::Processor::Context) -> void
+            def initialize: (untyped body, untyped status, untyped headers, context: AppSec::Context) -> void
 
             def response: () -> untyped
           end

sig/datadog/appsec/security_engine/runner.rbs

@@ -16,8 +16,6 @@ module Datadog
 
         def initialize: (WAF::Handle handle, telemetry: Core::Telemetry::Component) -> void
 
-        def extract_schema: () -> SecurityEngine::result
-
         def run: (SecurityEngine::data persistent_data, SecurityEngine::data ephemeral_data, ?::Integer timeout) -> SecurityEngine::result
 
         def finalize: () -> void
@@ -27,12 +25,6 @@ module Datadog
         def try_run: (SecurityEngine::data persistent_data, SecurityEngine::data ephemeral_data, untyped timeout) -> waf_run_result
 
         def report_execution: (WAF::Result result) -> void
-
-        def extract_schema?: () -> bool
-
-        def fallback_waf_error_result: () -> waf_run_result
-
-        def generic_ok_result: () -> SecurityEngine::result
       end
     end
   end

spec/datadog/appsec/context_spec.rb

@@ -126,6 +126,16 @@
     end
   end
 
+  describe '#extract_schema' do
+    it 'calls the waf runner with specific addresses' do
+      expect_any_instance_of(Datadog::AppSec::SecurityEngine::Runner).to receive(:run)
+        .with({ 'waf.context.processor' => { 'extract-schema' => true } }, {})
+        .and_call_original
+
+      expect(context.extract_schema).to be_instance_of(Datadog::AppSec::SecurityEngine::Result::Ok)
+    end
+  end
+
   describe '#waf_metrics' do
     context 'when multiple calls were successful' do
       let!(:run_results) do

spec/datadog/appsec/processor_spec.rb

@@ -284,9 +284,9 @@ def diagnostics
     end
   end
 
-  describe '#new_context' do
+  describe '#new_runner' do
     let(:processor) { described_class.new(ruleset: ruleset, telemetry: telemetry) }
 
-    it { expect(processor.new_context).to be_instance_of(described_class::Context) }
+    it { expect(processor.new_runner).to be_instance_of(Datadog::AppSec::SecurityEngine::Runner) }
   end
 end

spec/datadog/appsec/security_engine/runner_spec.rb

@@ -182,54 +182,4 @@
       end
     end
   end
-
-  describe '#extract_schema' do
-    context 'when api security is enabled' do
-      before { stub_const('Datadog::AppSec::WAF::LibDDWAF::DDWAF_RUN_TIMEOUT', 1_000) }
-      around do |example|
-        ClimateControl.modify(
-          'DD_EXPERIMENTAL_API_SECURITY_ENABLED' => 'true',
-          'DD_API_SECURITY_REQUEST_SAMPLE_RATE' => '1'
-        ) do
-          example.run
-        end
-      end
-
-      let(:waf_result) do
-        instance_double(
-          Datadog::AppSec::WAF::Result,
-          status: :ok,
-          events: [],
-          actions: [],
-          derivatives: [],
-          timeout: false,
-          total_runtime: 10
-        )
-      end
-
-      it 'calls the libddwaf with specific addresses' do
-        expect(waf_context).to receive(:run)
-          .with({ 'waf.context.processor' => { 'extract-schema' => true } }, {}, 1_000)
-          .and_return([waf_result.status, waf_result])
-
-        expect(runner.extract_schema).to be_instance_of(Datadog::AppSec::SecurityEngine::Result::Ok)
-      end
-    end
-
-    context 'when api security is disabled' do
-      around do |example|
-        ClimateControl.modify(
-          'DD_EXPERIMENTAL_API_SECURITY_ENABLED' => 'false',
-          'DD_API_SECURITY_REQUEST_SAMPLE_RATE' => '1'
-        ) do
-          example.run
-        end
-      end
-
-      it 'does not call the libddwaf' do
-        expect(waf_context).not_to receive(:run)
-        expect(runner.extract_schema).to be_instance_of(Datadog::AppSec::SecurityEngine::Result::Ok)
-      end
-    end
-  end
 end