Introduce AppSec::Instrumentation::Gateway::Argument

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

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative '../../../instrumentation/gateway/argument'
+
+module Datadog
+  module AppSec
+    module Contrib
+      module Sinatra
+        module Gateway
+          # Gateway Route Params argument.
+          class RouteParams < Instrumentation::Gateway::Argument
+            attr_reader :params
+
+            def initialize(params)
+              super
+              @params = params
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -20,10 +20,10 @@ def watch
               end
 
               def watch_request_dispatch(gateway = Instrumentation.gateway)
-                gateway.watch('sinatra.request.dispatch', :appsec) do |stack, request|
+                gateway.watch('sinatra.request.dispatch', :appsec) do |stack, gateway_request|
                   block = false
                   event = nil
-                  waf_context = request.env['datadog.waf.context']
+                  waf_context = gateway_request.env['datadog.waf.context']
 
                   AppSec::Reactive::Operation.new('sinatra.request.dispatch') do |op|
                     trace = active_trace
@@ -36,7 +36,7 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
                           waf_result: result,
                           trace: trace,
                           span: span,
-                          request: request,
+                          request: gateway_request,
                           actions: result.actions
                         }
 
@@ -46,12 +46,12 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
                       end
                     end
 
-                    _result, block = Rack::Reactive::RequestBody.publish(op, request)
+                    _result, block = Rack::Reactive::RequestBody.publish(op, gateway_request)
                   end
 
                   next [nil, [[:block, event]]] if block
 
-                  ret, res = stack.call(request)
+                  ret, res = stack.call(gateway_request.request)
 
                   if event
                     res ||= []
@@ -63,10 +63,10 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
               end
 
               def watch_request_routed(gateway = Instrumentation.gateway)
-                gateway.watch('sinatra.request.routed', :appsec) do |stack, (request, route_params)|
+                gateway.watch('sinatra.request.routed', :appsec) do |stack, (gateway_request, gateway_route_params)|
                   block = false
                   event = nil
-                  waf_context = request.env['datadog.waf.context']
+                  waf_context = gateway_request.env['datadog.waf.context']
 
                   AppSec::Reactive::Operation.new('sinatra.request.routed') do |op|
                     trace = active_trace
@@ -79,7 +79,7 @@ def watch_request_routed(gateway = Instrumentation.gateway)
                           waf_result: result,
                           trace: trace,
                           span: span,
-                          request: request,
+                          request: gateway_request,
                           actions: result.actions
                         }
 
@@ -89,12 +89,12 @@ def watch_request_routed(gateway = Instrumentation.gateway)
                       end
                     end
 
-                    _result, block = Sinatra::Reactive::Routed.publish(op, [request, route_params])
+                    _result, block = Sinatra::Reactive::Routed.publish(op, [gateway_request, gateway_route_params])
                   end
 
                   next [nil, [[:block, event]]] if block
 
-                  ret, res = stack.call(request)
+                  ret, res = stack.call(gateway_request.request)
 
                   if event
                     res ||= []

lib/datadog/appsec/contrib/sinatra/patcher.rb

@@ -3,8 +3,10 @@
 require_relative '../patcher'
 require_relative '../../response'
 require_relative '../rack/request_middleware'
+require_relative '../rack/gateway/request'
 require_relative 'framework'
 require_relative 'gateway/watcher'
+require_relative 'gateway/route_params'
 require_relative '../../../tracing/contrib/sinatra/framework'
 
 module Datadog
@@ -55,7 +57,9 @@ def dispatch!
 
             # TODO: handle exceptions, except for super
 
-            request_return, request_response = Instrumentation.gateway.push('sinatra.request.dispatch', request) do
+            gateway_request = Datadog::AppSec::Contrib::Rack::Gateway::Request.new(env)
+
+            request_return, request_response = Instrumentation.gateway.push('sinatra.request.dispatch', gateway_request) do
               # handle process_route interruption
               catch(Ext::ROUTE_INTERRUPT) { super }
             end
@@ -90,7 +94,13 @@ def process_route(*)
               # At this point params has both route params and normal params.
               route_params = params.each.with_object({}) { |(k, v), h| h[k] = v unless base_params.key?(k) }
 
-              _, request_response = Instrumentation.gateway.push('sinatra.request.routed', [request, route_params])
+              gateway_request = Datadog::AppSec::Contrib::Rack::Gateway::Request.new(env)
+              gateway_route_params = Gateway::RouteParams.new(route_params)
+
+              _, request_response = Instrumentation.gateway.push(
+                'sinatra.request.routed',
+                [gateway_request, gateway_route_params]
+              )
 
               if request_response && request_response.any? { |action, _event| action == :block }
                 self.response = AppSec::Response.negotiate(env).to_sinatra_response

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

@@ -16,7 +16,7 @@ def self.publish(op, data)
               _request, route_params = data
 
               catch(:block) do
-                op.publish('sinatra.request.route_params', route_params)
+                op.publish('sinatra.request.route_params', route_params.params)
 
                 nil
               end

sig/datadog/appsec/contrib/sinatra/gateway/route_params.rbs

@@ -0,0 +1,15 @@
+module Datadog
+  module AppSec
+    module Contrib
+      module Sinatra
+        module Gateway
+          class RouteParams < Instrumentation::Gateway::Argument
+            attr_reader params: untyped
+
+            def initialize: (untyped params) -> void
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/appsec/contrib/sinatra/reactive/routed.rbs

@@ -6,7 +6,7 @@ module Datadog
           module Routed
             ADDRESSES: ::Array[::String]
 
-            def self.publish: (untyped op, untyped data) -> untyped
+            def self.publish: (untyped op, ::Array[Datadog::AppSec::Instrumentation::Gateway::Argument] data) -> untyped
 
             def self.subscribe: (untyped op, untyped waf_context) { (untyped) -> untyped } -> untyped
           end

spec/datadog/appsec/contrib/sinatra/reactive/routed_spec.rb

@@ -3,24 +3,27 @@
 require 'datadog/appsec/spec_helper'
 require 'datadog/appsec/reactive/operation'
 require 'datadog/appsec/contrib/sinatra/reactive/routed'
+require 'datadog/appsec/contrib/rack/gateway/request'
+require 'datadog/appsec/contrib/sinatra/gateway/route_params'
 require 'rack'
 
 RSpec.describe Datadog::AppSec::Contrib::Sinatra::Reactive::Routed do
   let(:operation) { Datadog::AppSec::Reactive::Operation.new('test') }
   let(:request) do
-    Rack::Request.new(
+    Datadog::AppSec::Contrib::Rack::Gateway::Request.new(
       Rack::MockRequest.env_for(
         'http://example.com:8080/?a=foo',
         { 'REMOTE_ADDR' => '10.10.10.10', 'HTTP_CONTENT_TYPE' => 'text/html' }
       )
     )
   end
+  let(:routed_params) { Datadog::AppSec::Contrib::Sinatra::Gateway::RouteParams.new({ id: '1234' }) }
 
   describe '.publish' do
     it 'propagates routed params attributes to the operation' do
       expect(operation).to receive(:publish).with('sinatra.request.route_params', { id: '1234' })
 
-      described_class.publish(operation, [request, { id: '1234' }])
+      described_class.publish(operation, [request, routed_params])
     end
   end
 
@@ -49,7 +52,7 @@
           Datadog::AppSec.settings.waf_timeout
         ).and_return(waf_result)
         described_class.subscribe(operation, waf_context)
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -64,7 +67,7 @@
           expect(result).to eq(waf_result)
           expect(block).to eq(false)
         end
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
 
@@ -77,7 +80,7 @@
           expect(result).to eq(waf_result)
           expect(block).to eq(true)
         end
-        publish_result, publish_block = described_class.publish(operation, [request, { id: '1234' }])
+        publish_result, publish_block = described_class.publish(operation, [request, routed_params])
         expect(publish_result).to eq(waf_result)
         expect(publish_block).to eq(true)
       end
@@ -90,7 +93,7 @@
         waf_result = double(:waf_result, status: :ok, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -102,7 +105,7 @@
         waf_result = double(:waf_result, status: :invalid_call, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -114,7 +117,7 @@
         waf_result = double(:waf_result, status: :invalid_rule, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -126,7 +129,7 @@
         waf_result = double(:waf_result, status: :invalid_flow, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -138,7 +141,7 @@
         waf_result = double(:waf_result, status: :no_rule, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end
@@ -150,7 +153,7 @@
         waf_result = double(:waf_result, status: :foo, timeout: false)
         expect(waf_context).to receive(:run).and_return(waf_result)
         expect { |b| described_class.subscribe(operation, waf_context, &b) }.not_to yield_control
-        result = described_class.publish(operation, [request, { id: '1234' }])
+        result = described_class.publish(operation, [request, routed_params])
         expect(result).to be_nil
       end
     end