Unify rack middleware integrations

lib/appsignal/integrations/grape.rb

@@ -4,52 +4,42 @@ module Appsignal
   # @todo Move to sub-namespace
   # @api private
   module Grape
-    class Middleware < ::Grape::Middleware::Base
-      def call(env)
-        if Appsignal.active?
-          call_with_appsignal_monitoring(env)
-        else
-          app.call(env)
+    class GrapeInstrumentation < ::Appsignal::Rack::GenericInstrumentation
+      def set_transaction_attributes_from_request(transaction, request)
+        env = request.env
+
+        request_method = request.request_method.to_s.upcase
+        path = request.path # Path without namespaces
+        endpoint = env["api.endpoint"]
+
+        # Do not set error if "grape.skip_appsignal_error" is set to `true`.
+        transaction.set_error(error) unless env["grape.skip_appsignal_error"]
+
+        if endpoint&.options
+          options = endpoint.options
+          request_method = options[:method].first.to_s.upcase
+          klass = options[:for]
+          namespace = endpoint.namespace
+          namespace = "" if namespace == "/"
+
+          path = options[:path].first.to_s
+          path = "/#{path}" if path[0] != "/"
+          path = "#{namespace}#{path}"
+
+          transaction.set_action_if_nil("#{request_method}::#{klass}##{path}")
         end
+
+        transaction.set_http_or_background_queue_start
+        transaction.set_metadata("path", path)
+        transaction.set_metadata("method", request_method)
       end
+    end
 
-      def call_with_appsignal_monitoring(env)
-        request     = ::Rack::Request.new(env)
-        transaction = Appsignal::Transaction.create(
-          SecureRandom.uuid,
-          Appsignal::Transaction::HTTP_REQUEST,
-          request
-        )
-        begin
-          app.call(env)
-        rescue Exception => error # rubocop:disable Lint/RescueException
-          # Do not set error if "grape.skip_appsignal_error" is set to `true`.
-          transaction.set_error(error) unless env["grape.skip_appsignal_error"]
-          raise error
-        ensure
-          request_method = request.request_method.to_s.upcase
-          path = request.path # Path without namespaces
-          endpoint = env["api.endpoint"]
-
-          if endpoint&.options
-            options = endpoint.options
-            request_method = options[:method].first.to_s.upcase
-            klass = options[:for]
-            namespace = endpoint.namespace
-            namespace = "" if namespace == "/"
-
-            path = options[:path].first.to_s
-            path = "/#{path}" if path[0] != "/"
-            path = "#{namespace}#{path}"
-
-            transaction.set_action_if_nil("#{request_method}::#{klass}##{path}")
-          end
-
-          transaction.set_http_or_background_queue_start
-          transaction.set_metadata("path", path)
-          transaction.set_metadata("method", request_method)
-          Appsignal::Transaction.complete_current!
-        end
+    # Grape middleware has a slightly different API than Rack middleware,
+    # but nothing forbids us from embedding actual Rack middleware inside of it
+    class Middleware < ::Grape::Middleware::Base
+      def call(env)
+        GrapeInstrumentation.new(app).call(env)
       end
     end
   end

lib/appsignal/rack/generic_instrumentation.rb

@@ -1,41 +1,104 @@
 # frozen_string_literal: true
 
 require "rack"
+require "delegate"
 
 module Appsignal
   # @api private
   module Rack
     class GenericInstrumentation
+      # This wrapper is necessary because the request may get initialized in an outer middleware, before
+      # the ActionDispatch request gets created. In that case, the Rack request will take precedence
+      # and the Rails filtered params are not going to be honored, potentially revealing user data.
+      # This wrapper will check whether the actual Rack env contains an ActionDispatch request
+      # and delegate the "filtered_params" call to it when the transaction gets completed. This will
+      # ensure that if there was a Rails request somewhere in the handling flow upstream from this
+      # middleware, its filtered params will end up used for the Appsignal transaction. So barring
+      # a few edge cases even when GenericInstrumentation gets mounted from a Rails app' `config.ru`
+      # there will still be no involuntary disclosure of data that is supposed to get filtered
+      class FilteredParamsWrapper < SimpleDelegator
+        def filtered_params
+          actual_request = __getobj__
+          if actual_request.respond_to?(:filtered_params)
+            actual_request.filtered_params
+          elsif has_rails_filtered_params?(actual_request)
+            # This will delegate to the request itself in case of Rails
+            actual_request.env["action_dispatch.request"].filtered_params
+          else
+            actual_request.params
+          end
+        end
+
+        def has_rails_filtered_params?(req)
+          # `env` is available on both ActionDispatch::Request and Rack::Request
+          req.respond_to?(:env) && req.env["action_dispatch.request"] && req.env["action_dispatch.request"].respond_to?(:filtered_params)
+        end
+      end
+
       def initialize(app, options = {})
-        Appsignal.internal_logger.debug "Initializing Appsignal::Rack::GenericInstrumentation"
+        Appsignal.internal_logger.debug "Initializing #{self.class}"
         @app = app
         @options = options
+        @instrument_span_name = "process_action.generic"
       end
 
       def call(env)
-        if Appsignal.active?
-          call_with_appsignal_monitoring(env)
+        if Appsignal.active? && !env["appsignal.transaction"]
+          call_with_new_appsignal_transaction(env)
+        elsif Appsignal.active?
+          call_with_existing_appsignal_transaction(env)
         else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+          call_without_appsignal_transaction(env)
         end
       end
 
-      def call_with_appsignal_monitoring(env)
-        request = ::Rack::Request.new(env)
-        transaction = Appsignal::Transaction.create(
+      def call_without_appsignal_transaction(env)
+        nil_transaction = Appsignal::Transaction::NilTransaction.new
+        status, headers, obody = @app.call(env)
+        [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+      end
+
+      def parse_or_reuse_request(env)
+        request_class = @options.fetch(:request_class, ::Rack::Request)
+        request_class.new(env)
+      end
+
+      def create_transaction_from_request(request)
+        # A middleware may support more options besides the ones supported by the Transaction
+        default_options = {:force => false, :params_method => :filtered_params}
+        options_for_transaction = default_options.merge!(@options.slice(:force, :params_method))
+
+        Appsignal::Transaction.create(
           SecureRandom.uuid,
           Appsignal::Transaction::HTTP_REQUEST,
-          request
+          FilteredParamsWrapper.new(request),
+          options_for_transaction
         )
+      end
+
+
+      def call_with_existing_appsignal_transaction(env)
+        request = parse_or_reuse_request(env)
+        transaction = env["appsignal.transaction"]
+        Appsignal.instrument(@instrument_span_name) { @app.call(env) }
+        # Rescuing the exception from `call` will be handled by the middleware which
+        # added the transaction to the Rack env
+      ensure
+        set_transaction_attributes_from_request(transaction, request)
+      end
+
+      def call_with_new_appsignal_transaction(env)
+        request = parse_or_reuse_request(env)
+        transaction = create_transaction_from_request(request)
+        env["appsignal.transaction"] = transaction
+
         # We need to complete the transaction if there is an exception inside the `call`
         # of the app. If there isn't one and the app returns us a Rack response triplet, we let
         # the BodyWrapper complete the transaction when #close gets called on it
         # (guaranteed by the webserver)
         complete_transaction_without_body = false
         begin
-          Appsignal.instrument("process_action.generic") do
+          Appsignal.instrument(@instrument_span_name) do
             status, headers, obody = @app.call(env)
             [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
           end
@@ -44,17 +107,20 @@ def call_with_appsignal_monitoring(env)
           complete_transaction_without_body = true
           raise error
         ensure
-          default_action = env["appsignal.route"] || env["appsignal.action"] || "unknown"
-          transaction.set_action_if_nil(default_action)
-          transaction.set_metadata("path", request.path)
-          transaction.set_metadata("method", request.request_method)
-          transaction.set_http_or_background_queue_start
-
+          set_transaction_attributes_from_request(transaction, request)
           # Transaction gets completed when the body gets read out, except in cases when
           # the app failed before returning us the Rack response triplet.
           Appsignal::Transaction.complete_current! if complete_transaction_without_body
         end
       end
+
+      def set_transaction_attributes_from_request(transaction, request)
+        default_action = request.env["appsignal.route"] || request.env["appsignal.action"] || "unknown"
+        transaction.set_action_if_nil(action_name_for_transaction)
+        transaction.set_metadata("path", request.path)
+        transaction.set_metadata("method", request.request_method)
+        transaction.set_http_or_background_queue_start
+      end
     end
   end
 end

lib/appsignal/rack/rails_instrumentation.rb

@@ -6,59 +6,27 @@ module Appsignal
   # @api private
   module Rack
     class RailsInstrumentation
-      def initialize(app, options = {})
-        Appsignal.internal_logger.debug "Initializing Appsignal::Rack::RailsInstrumentation"
-        @app = app
-        @options = options
+      def initialize(app, options={})
+        super
+        @options[:request_class] = ActionDispatch::Request
+        @instrument_span_name = "process_action.rails"
       end
 
-      def call(env)
-        if Appsignal.active?
-          call_with_appsignal_monitoring(env)
-        else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+      def set_transaction_attributes_from_request(transaction, request)
+        controller = request.env["action_controller.instance"]
+        if controller
+          transaction.set_action_if_nil("#{controller.class}##{controller.action_name}")
         end
+        super
       end
 
-      def call_with_appsignal_monitoring(env)
-        request = ActionDispatch::Request.new(env)
-        transaction = Appsignal::Transaction.create(
-          request_id(env),
+      def create_transaction_from_request(request)
+        Appsignal::Transaction.create(
+          request_id(request.env),
           Appsignal::Transaction::HTTP_REQUEST,
           request,
           :params_method => :filtered_parameters
         )
-        # We need to complete the transaction if there is an exception exception inside the `call`
-        # of the app. If there isn't one and the app returns us a Rack response triplet, we let
-        # the BodyWrapper complete the transaction when #close gets called on it
-        # (guaranteed by the webserver)
-        complete_transaction_without_body = false
-        begin
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
-        rescue Exception => error # rubocop:disable Lint/RescueException
-          transaction.set_error(error)
-          complete_transaction_without_body = true
-          raise error
-        ensure
-          controller = env["action_controller.instance"]
-          if controller
-            transaction.set_action_if_nil("#{controller.class}##{controller.action_name}")
-          end
-          transaction.set_http_or_background_queue_start
-          transaction.set_metadata("path", request.path)
-          begin
-            transaction.set_metadata("method", request.request_method)
-          rescue => error
-            Appsignal.internal_logger.error("Unable to report HTTP request method: '#{error}'")
-          end
-
-          # Transaction gets completed when the body gets read out, except in cases when
-          # the app failed before returning us the Rack response triplet.
-          Appsignal::Transaction.complete_current! if complete_transaction_without_body
-        end
       end
 
       def request_id(env)