Rack middleware support

lib/newrelic_security/agent/configuration/manager.rb

@@ -37,6 +37,7 @@ def initialize
           @cache[:'security.detection.deserialization.enabled'] = ::NewRelic::Agent.config[:'security.detection.deserialization.enabled'].nil? ? true : ::NewRelic::Agent.config[:'security.detection.deserialization.enabled']
           @cache[:'security.scan_controllers.iast_scan_request_rate_limit'] = ::NewRelic::Agent.config[:'security.scan_controllers.iast_scan_request_rate_limit'].to_i
           @cache[:framework] = detect_framework
+          @cache[:app_class] = detect_app_class
           @cache[:'security.application_info.port'] = ::NewRelic::Agent.config[:'security.application_info.port'].to_i
           @cache[:listen_port] = nil
           @cache[:process_start_time] = current_time_millis # TODO: Ruby doesn't provide process start time in pure ruby implementation using agent loading time for now.
@@ -155,6 +156,16 @@ def detect_framework
           return :sinatra if defined?(::Sinatra)
           return :roda if defined?(::Roda)
           return :grape if defined?(::Grape)
+          return :rack if defined?(::Rack) && defined?(Rack::Builder)
+        end
+
+        def detect_app_class
+          target_class = nil
+          ObjectSpace.each_object(::Rack::Builder) do |z| target_class = z.instance_variable_get(:@run).target end
+          target_class
+        rescue StandardError => exception
+          NewRelic::Security::Agent.logger.error "Exception in detect_app_class : #{exception.inspect} #{exception.backtrace}"
+          nil
         end
 
         def generate_uuid

lib/newrelic_security/agent/control/collector.rb

@@ -45,7 +45,7 @@ def collect(case_type, args, event_category = nil, **keyword_args)
           # In rails 5 method name keeps chaning for same api call (ex: _app_views_sqli_sqlinjectionattackcase_html_erb__1999281606898621405_2624809100).
           # Hence, considering only frame absolute_path & lineno for apiId calculation.
           user_frame_index = get_user_frame_index(stk)
-          event.apiId = "#{case_type}-#{calculate_api_id(stk[0..user_frame_index].map { |frame| "#{frame.absolute_path}:#{frame.lineno}" }, event.httpRequest[:method], route)}"
+          event.apiId = "#{case_type}-#{calculate_api_id(stk[0..user_frame_index].map { |frame| "#{frame.absolute_path}:#{frame.lineno}" }, event.httpRequest[:method], route.gsub(/\d+$/, EMPTY_STRING))}"
           stk.delete_if { |frame| frame.path.match?(/newrelic_security/) || frame.path.match?(/new_relic/) }
           user_frame_index = get_user_frame_index(stk)
           return if case_type != REFLECTED_XSS && user_frame_index == -1 # TODO: Add log message here: "Filtered because User Stk frame NOT FOUND   \r\n"

lib/newrelic_security/agent/utils/agent_utils.rb

@@ -132,6 +132,8 @@ def get_app_routes(framework, router = nil)
           router.owner.superclass.public_instance_methods(false).each do |m|
             NewRelic::Security::Agent.agent.route_map << "*@/#{router.owner}/#{m}"
           end
+        when :rack
+          # TODO: API enpointes(routes) extraction for rack
         else
           NewRelic::Security::Agent.logger.error "Unable to get app routes as Framework not detected"
         end

lib/newrelic_security/instrumentation-security/rack/chain.rb

@@ -0,0 +1,24 @@
+module NewRelic::Security
+  module Instrumentation
+    module Rack
+      module Builder
+        module Chain
+          def self.instrument!
+            ::Rack::Builder.class_eval do 
+              include NewRelic::Security::Instrumentation::Rack::Builder
+
+              alias_method :call_without_security, :call
+
+              def call(env)
+                retval = nil
+                event = call_on_enter(env) { retval = call_without_security(env) }
+                call_on_exit(event, retval) { return retval }
+              end
+            end
+          end
+        end
+      end
+    end
+
+  end
+end

lib/newrelic_security/instrumentation-security/rack/instrumentation.rb

@@ -0,0 +1,44 @@
+require_relative 'prepend'
+require_relative 'chain'
+
+module NewRelic::Security
+  module Instrumentation
+    module Rack::Builder
+
+      def call_on_enter(env)
+        event = nil
+        NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
+        return unless NewRelic::Security::Agent.config[:enabled]
+        NewRelic::Security::Agent.config.update_port = NewRelic::Security::Agent::Utils.app_port(env) unless NewRelic::Security::Agent.config[:listen_port]
+        NewRelic::Security::Agent::Utils.get_app_routes(:rack) if NewRelic::Security::Agent.agent.route_map.empty?
+        NewRelic::Security::Agent::Control::HTTPContext.set_context(env)
+        ctxt = NewRelic::Security::Agent::Control::HTTPContext.get_context
+        ctxt.route = "#{env[REQUEST_METHOD]}@#{env[PATH_INFO]}" if ctxt
+        NewRelic::Security::Agent::Utils.parse_fuzz_header(NewRelic::Security::Agent::Control::HTTPContext.get_context)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+        return event
+      end
+
+      def call_on_exit(event, retval)
+        NewRelic::Security::Agent.logger.debug "OnExit :  #{self.class}.#{__method__}"
+        # NewRelic::Security::Agent.logger.debug "\n\nHTTP Context : #{::NewRelic::Agent::Tracer.current_transaction.instance_variable_get(:@security_context_data).inspect}\n\n"
+        NewRelic::Security::Agent::Control::ReflectedXSS.check_xss(NewRelic::Security::Agent::Control::HTTPContext.get_context, retval) if NewRelic::Security::Agent.config[:'security.detection.rxss.enabled']
+        NewRelic::Security::Agent::Utils.delete_created_files(NewRelic::Security::Agent::Control::HTTPContext.get_context)
+        NewRelic::Security::Agent.agent.error_reporting&.report_unhandled_or_5xx_exceptions(NewRelic::Security::Agent::Control::HTTPContext.get_current_transaction, NewRelic::Security::Agent::Control::HTTPContext.get_context, retval[0])
+        NewRelic::Security::Agent::Control::HTTPContext.reset_context
+        NewRelic::Security::Agent.logger.debug "Exit event : #{event}"
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+      end
+
+    end
+
+  end
+end
+
+NewRelic::Security::Instrumentation::InstrumentationLoader.install_instrumentation(:rack, NewRelic::Security::Agent.config[:app_class].class, ::NewRelic::Security::Instrumentation::Rack::Builder) if NewRelic::Security::Agent.config[:framework] == :rack

lib/newrelic_security/instrumentation-security/rack/prepend.rb

@@ -0,0 +1,18 @@
+module NewRelic::Security
+  module Instrumentation
+    module Rack
+      module Builder
+        module Prepend
+          include NewRelic::Security::Instrumentation::Rack::Builder
+
+          def call(env, &block)
+            retval = nil
+            event = call_on_enter(env) { retval = super }
+            call_on_exit(event, retval) { return retval }
+          end
+
+        end
+      end
+    end
+  end
+end