From b165565c783ed0051ba2479a8eef5f5d4f957e83 Mon Sep 17 00:00:00 2001
From: Sergey Fedorov <sergey.fedorov@datadoghq.com>
Date: Fri, 25 Oct 2024 14:19:01 +0200
Subject: [PATCH] Update AppSec Processor and Context RBS

---
 lib/datadog/appsec/processor/context.rb  |  1 -
 sig/datadog/appsec/processor.rbs         | 25 +++++++++++++++--------
 sig/datadog/appsec/processor/context.rbs | 26 ++++++++++++++++++------
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/lib/datadog/appsec/processor/context.rb b/lib/datadog/appsec/processor/context.rb
index 356291e6aa7..ef3469b4e0e 100644
--- a/lib/datadog/appsec/processor/context.rb
+++ b/lib/datadog/appsec/processor/context.rb
@@ -41,7 +41,6 @@ def run(input, timeout = WAF::LibDDWAF::DDWAF_RUN_TIMEOUT)
           @time_ext_ns += (stop_ns - start_ns)
           @timeouts += 1 if res.timeout
 
-          # TODO: handle the response
           res
         ensure
           @run_mutex.unlock
diff --git a/sig/datadog/appsec/processor.rbs b/sig/datadog/appsec/processor.rbs
index 1ff5ca24a65..43a75531137 100644
--- a/sig/datadog/appsec/processor.rbs
+++ b/sig/datadog/appsec/processor.rbs
@@ -1,29 +1,38 @@
 module Datadog
   module AppSec
     class Processor
-      def self.active_context: () -> Context
+      type ruleset = ::Hash[::String, untyped]
 
-      private
+      @telemetry: Core::Telemetry::Component
+
+      @diagnostics: WAF::LibDDWAF::Object?
+
+      @addresses: ::Array[::String]
+
+      @handle: WAF::Handle
 
       attr_reader diagnostics: WAF::LibDDWAF::Object?
+
       attr_reader addresses: ::Array[::String]
 
-      @handle: WAF::Handle
-      @ruleset: ::Hash[::String, untyped]
-      @addresses: ::Array[::String]
+      def initialize: (ruleset: ruleset, telemetry: Core::Telemetry::Component) -> void
 
-      def initialize: (ruleset: ::Hash[untyped, untyped], telemetry: Core::Telemetry::Component) -> void
       def ready?: () -> bool
+
       def finalize: () -> void
 
-      attr_reader handle: untyped
+      def new_context: () -> Context
 
       private
 
       def require_libddwaf: () -> bool
+
       def libddwaf_provides_waf?: () -> bool
-      def create_waf_handle: (Core::Configuration::Settings::_AppSec settings, ::Hash[String, untyped] ruleset) -> bool
+
+      def create_waf_handle: (Core::Configuration::Settings::_AppSec settings, ruleset ruleset) -> bool
+
       def libddwaf_platform: () -> ::String
+
       def ruby_platforms: () -> ::Array[::String]
     end
   end
diff --git a/sig/datadog/appsec/processor/context.rbs b/sig/datadog/appsec/processor/context.rbs
index 95843dc391d..55c2be12c99 100644
--- a/sig/datadog/appsec/processor/context.rbs
+++ b/sig/datadog/appsec/processor/context.rbs
@@ -2,24 +2,38 @@ module Datadog
   module AppSec
     class Processor
       class Context
-        type event = untyped
-        type data = WAF::data
+        @context: WAF::Context
+
+        @telemetry: Core::Telemetry::Component
+
+        @time_ns: ::Float
+
+        @time_ext_ns: ::Float
+
+        @timeouts: ::Integer
+
+        @events: ::Array[untyped]
+
+        @run_mutex: ::Thread::Mutex
 
         attr_reader time_ns: ::Float
+
         attr_reader time_ext_ns: ::Float
+
         attr_reader timeouts: ::Integer
-        attr_reader events: ::Array[event]
 
-        @context: WAF::Context
+        attr_reader events: ::Array[untyped]
 
-        @run_mutex: ::Thread::Mutex
+        def initialize: (WAF::Handle handle, telemetry: Core::Telemetry::Component) -> void
 
-        def initialize: (Processor processor) -> void
         def run: (Hash[untyped, untyped] input, ?::Integer timeout) -> WAF::Result
+
         def extract_schema: () -> WAF::Result?
+
         def finalize: () -> void
 
         private
+
         def extract_schema?: () -> bool
       end
     end