Merge pull request #4459 from DataDog/appsec-56614-track-triggered-instrumentation-events

[APPSEC-56614] Add Instrumentation::Gateway#pushed? method

Author: Strech

lib/datadog/appsec/instrumentation/gateway.rb

@@ -1,35 +1,29 @@
 # frozen_string_literal: true
 
+require_relative 'gateway/middleware'
+
 module Datadog
   module AppSec
     # Instrumentation for AppSec
     module Instrumentation
       # Instrumentation gateway implementation
       class Gateway
-        # Instrumentation gateway middleware
-        class Middleware
-          attr_reader :key, :block
-
-          def initialize(key, &block)
-            @key = key
-            @block = block
-          end
-
-          def call(stack, env)
-            @block.call(stack, env)
-          end
-        end
-
-        private_constant :Middleware
-
         def initialize
           @middlewares = Hash.new { |h, k| h[k] = [] }
+          @pushed_events = {}
         end
 
+        # NOTE: Be careful with pushed names because every pushed event name
+        #       is recorded in order to provide an ability to any subscriber
+        #       to check wether an arbitrary event had happened.
+        #
+        # WARNING: If we start pushing generated names we should consider
+        #          limiting the storage of pushed names.
         def push(name, env, &block)
-          block ||= -> {}
+          @pushed_events[name] = true
 
-          middlewares_for_name = middlewares[name]
+          block ||= -> {}
+          middlewares_for_name = @middlewares[name]
 
           return [block.call, nil] if middlewares_for_name.empty?
 
@@ -48,14 +42,15 @@ def push(name, env, &block)
         end
 
         def watch(name, key, &block)
-          @middlewares[name] << Middleware.new(key, &block) unless middlewares[name].any? { |m| m.key == key }
+          @middlewares[name] << Middleware.new(key, &block) unless @middlewares[name].any? { |m| m.key == key }
         end
 
-        private
-
-        attr_reader :middlewares
+        def pushed?(name)
+          @pushed_events.key?(name)
+        end
       end
 
+      # NOTE: This left as-is and will be depricated soon.
       def self.gateway
         @gateway ||= Gateway.new # TODO: not thread safe
       end

lib/datadog/appsec/instrumentation/gateway/middleware.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module Instrumentation
+      class Gateway
+        # NOTE: This class extracted as-is and will be deprecated
+        # Instrumentation gateway middleware
+        class Middleware
+          attr_reader :key, :block
+
+          def initialize(key, &block)
+            @key = key
+            @block = block
+          end
+
+          def call(stack, env)
+            @block.call(stack, env)
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/appsec/instrumentation/gateway.rbs

@@ -7,26 +7,17 @@ module Datadog
         type final_call = ^() -> final_call_result
         type stack = ::Proc
 
-        # TODO: this should be a struct
         type stack_result = ::Array[nil | middleware_result | final_call_result]
 
-        class Middleware
-          attr_reader key: ::Symbol
-          attr_reader block: ::Proc
-
-          def initialize: (::Symbol key) { (stack next, Instrumentation::Gateway::Argument env) -> stack_result } -> void
-          def call: (stack next, Gateway::Argument? env) -> stack_result
-        end
-
         @middlewares: ::Hash[::String, ::Array[Middleware]]
 
         def initialize: () -> void
+
         def push: (::String name, Gateway::Argument env) ?{ () -> final_call_result } -> stack_result
-        def watch: (::String name, ::Symbol key) { (stack next, Gateway::Argument env) -> stack_result } -> void
 
-        private
+        def watch: (::String name, ::Symbol key) { (stack next, Gateway::Argument env) -> stack_result } -> void
 
-        attr_reader middlewares: ::Hash[::String, ::Array[Middleware]]
+        def pushed?: (::String name) -> bool
       end
 
       self.@gateway: Gateway

sig/datadog/appsec/instrumentation/gateway/middleware.rbs

@@ -0,0 +1,15 @@
+module Datadog
+  module AppSec
+    module Instrumentation
+      class Gateway
+        class Middleware
+          attr_reader key: ::Symbol
+          attr_reader block: ::Proc
+
+          def initialize: (::Symbol key) { (Gateway::stack next, Gateway::Argument env) -> stack_result } -> void
+          def call: (Gateway::stack next, Gateway::Argument? env) -> stack_result
+        end
+      end
+    end
+  end
+end

spec/datadog/appsec/instrumentation/gateway_spec.rb

@@ -5,7 +5,7 @@
 
 RSpec.describe Datadog::AppSec::Instrumentation::Gateway do
   subject(:gateway) { described_class.new }
-  let(:middlewares) { gateway.send(:middlewares) }
+  let(:middlewares) { gateway.instance_variable_get(:@middlewares) }
 
   describe '#watch' do
     it 'stores middleware' do
@@ -81,4 +81,19 @@
       expect(env_2).to eq({ a: :b, c: :d })
     end
   end
+
+  describe '#pushed?' do
+    it { expect(gateway.pushed?('event.0')).to be(false) }
+
+    it 'returns true if event was pushed' do
+      expect { gateway.push('event.1', {}) }.to change { gateway.pushed?('event.1') }
+        .from(false).to(true)
+
+      expect { gateway.push('event.2', {}) }.to change { gateway.pushed?('event.2') }
+        .from(false).to(true)
+
+      expect { gateway.push('event.2', {}) }.not_to change { gateway.pushed?('event.2') }
+        .from(true)
+    end
+  end
 end