Create Gateway::Argument

- Replace contrib Rack::Request and Rack::Response with Gateway::Request and Gateway::Response

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

@@ -0,0 +1,104 @@
+# typed: false
+# frozen_string_literal: true
+
+require_relative '../../../instrumentation/gateway/argument'
+require_relative '../../../../tracing/client_ip'
+require_relative '../../../../tracing/contrib/rack/header_collection'
+
+
+module Datadog
+  module AppSec
+    module Contrib
+      module Rack
+        module Gateway
+          class Request < Instrumentation::Gateway::Argument
+            attr_reader :env
+
+            def initialize(env)
+              @env = env
+            end
+
+            def request
+              @request ||= ::Rack::Request.new(env)
+            end
+
+            def query
+              # Downstream libddwaf expects keys and values to be extractable
+              # separately so we can't use [[k, v], ...]. We also want to allow
+              # duplicate keys, so we use [{k, v}, ...] instead.
+              request.query_string.split('&').map do |e|
+                k, v = e.split('=').map { |s| CGI.unescape(s) }
+
+                { k => v }
+              end
+            end
+
+            # Rack < 2.0 does not have :each_header
+            # TODO: We need access to Rack here. We must make sure we are able to load AppSec without Rack,
+            # TODO: while still ensure correctness in ths code path.
+            if defined?(::Rack) && ::Rack::Request.instance_methods.include?(:each_header)
+              def headers
+                request.each_header.each_with_object({}) do |(k, v), h|
+                  h[k.gsub(/^HTTP_/, '').downcase.tr('_', '-')] = v if k =~ /^HTTP_/
+                end
+              end
+            else
+              def headers
+                request.env.each_with_object({}) do |(k, v), h|
+                  h[k.gsub(/^HTTP_/, '').downcase.tr('_', '-')] = v if k =~ /^HTTP_/
+                end
+              end
+            end
+
+            def body
+              request.body.read.tap { request.body.rewind }
+            end
+
+            def url
+              request.url
+            end
+
+            def cookies
+              request.cookies
+            end
+
+            def host
+              request.host
+            end
+
+            def user_agent
+              request.user_agent
+            end
+
+            def remote_addr
+              env['REMOTE_ADDR']
+            end
+
+            def form_hash
+              # force form data processing
+              request.POST if request.form_data?
+
+              # usually Hash<String,String> but can be a more complex
+              # Hash<String,String||Array||Hash> when e.g coming from JSON
+              env['rack.request.form_hash']
+            end
+
+            def client_ip
+              remote_ip = remote_addr
+              headers = Datadog::Tracing::Contrib::Rack::Header::RequestHeaderCollection.new(env)
+
+              result = Datadog::Tracing::ClientIp.raw_ip_from_request(headers, remote_ip)
+
+              if result.raw_ip
+                ip = Datadog::Tracing::ClientIp.strip_decorations(result.raw_ip)
+                return unless Datadog::Tracing::ClientIp.valid_ip?(ip)
+
+                ip
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,32 @@
+# typed: false
+# frozen_string_literal: true
+
+require_relative '../../../instrumentation/gateway/argument'
+
+module Datadog
+  module AppSec
+    module Contrib
+      module Rack
+        module Gateway
+          class Response < Instrumentation::Gateway::Argument
+            attr_reader :body, :status, :headers
+
+            def initialize(body, status, headers)
+              @body = body
+              @status = status
+              @headers = headers.each_with_object({}) { |(k, v), h| h[k.downcase] = v }
+            end
+
+            def response
+              @response ||= ::Rack::Response.new(body, status, headers)
+            end
+
+            def cookies
+              response.cookies
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -22,10 +22,10 @@ def watch
               end
 
               def watch_request(gateway = Instrumentation.gateway)
-                gateway.watch('rack.request', :appsec) do |stack, request|
+                gateway.watch('rack.request', :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('rack.request') do |op|
                     trace = active_trace
@@ -38,7 +38,7 @@ def watch_request(gateway = Instrumentation.gateway)
                           waf_result: result,
                           trace: trace,
                           span: span,
-                          request: request,
+                          request: gateway_request,
                           actions: result.actions
                         }
 
@@ -48,12 +48,12 @@ def watch_request(gateway = Instrumentation.gateway)
                       end
                     end
 
-                    _result, block = Rack::Reactive::Request.publish(op, request)
+                    _result, block = Rack::Reactive::Request.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 ||= []
@@ -65,10 +65,10 @@ def watch_request(gateway = Instrumentation.gateway)
               end
 
               def watch_response(gateway = Instrumentation.gateway)
-                gateway.watch('rack.response', :appsec) do |stack, response|
+                gateway.watch('rack.response', :appsec) do |stack, gateway_response|
                   block = false
                   event = nil
-                  waf_context = response.instance_eval { @waf_context }
+                  waf_context = AppSec::Processor.active_context
 
                   AppSec::Reactive::Operation.new('rack.response') do |op|
                     trace = active_trace
@@ -81,7 +81,7 @@ def watch_response(gateway = Instrumentation.gateway)
                           waf_result: result,
                           trace: trace,
                           span: span,
-                          response: response,
+                          response: gateway_response,
                           actions: result.actions
                         }
 
@@ -91,12 +91,12 @@ def watch_response(gateway = Instrumentation.gateway)
                       end
                     end
 
-                    _result, block = Rack::Reactive::Response.publish(op, response)
+                    _result, block = Rack::Reactive::Response.publish(op, gateway_response)
                   end
 
                   next [nil, [[:block, event]]] if block
 
-                  ret, res = stack.call(response)
+                  ret, res = stack.call(gateway_response.response)
 
                   if event
                     res ||= []
@@ -108,10 +108,10 @@ def watch_response(gateway = Instrumentation.gateway)
               end
 
               def watch_request_body(gateway = Instrumentation.gateway)
-                gateway.watch('rack.request.body', :appsec) do |stack, request|
+                gateway.watch('rack.request.body', :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('rack.request.body') do |op|
                     trace = active_trace
@@ -124,7 +124,7 @@ def watch_request_body(gateway = Instrumentation.gateway)
                           waf_result: result,
                           trace: trace,
                           span: span,
-                          request: request,
+                          request: gateway_request,
                           actions: result.actions
                         }
 
@@ -134,12 +134,12 @@ def watch_request_body(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 ||= []

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

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require_relative '../request'
-
 module Datadog
   module AppSec
     module Contrib
@@ -18,13 +16,13 @@ module Request
             ].freeze
             private_constant :ADDRESSES
 
-            def self.publish(op, request)
+            def self.publish(op, gateway_request)
               catch(:block) do
-                op.publish('request.query', Rack::Request.query(request))
-                op.publish('request.headers', Rack::Request.headers(request))
-                op.publish('request.uri.raw', Rack::Request.url(request))
-                op.publish('request.cookies', Rack::Request.cookies(request))
-                op.publish('request.client_ip', Rack::Request.client_ip(request))
+                op.publish('request.query', gateway_request.query)
+                op.publish('request.headers', gateway_request.headers)
+                op.publish('request.uri.raw', gateway_request.url)
+                op.publish('request.cookies', gateway_request.cookies)
+                op.publish('request.client_ip', gateway_request.client_ip)
 
                 nil
               end

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

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require_relative '../request'
-
 module Datadog
   module AppSec
     module Contrib
@@ -14,10 +12,10 @@ module RequestBody
             ].freeze
             private_constant :ADDRESSES
 
-            def self.publish(op, request)
+            def self.publish(op, gateway_request)
               catch(:block) do
                 # params have been parsed from the request body
-                op.publish('request.body', Rack::Request.form_hash(request))
+                op.publish('request.body', gateway_request.form_hash)
 
                 nil
               end

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

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require_relative '../response'
-
 module Datadog
   module AppSec
     module Contrib
@@ -14,9 +12,9 @@ module Response
             ].freeze
             private_constant :ADDRESSES
 
-            def self.publish(op, response)
+            def self.publish(op, gateway_response)
               catch(:block) do
-                op.publish('response.status', Rack::Response.status(response))
+                op.publish('response.status', gateway_response.status)
 
                 nil
               end