# Copyright (c) 2020 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

cs__scoped_require 'ipaddr'
cs__scoped_require 'json'
cs__scoped_require 'rack'

cs__scoped_require 'contrast/security_exception'
cs__scoped_require 'contrast/utils/object_share'
cs__scoped_require 'contrast/components/interface'
cs__scoped_require 'contrast/utils/heap_dump_util'
cs__scoped_require 'contrast/agent/request_handler'
cs__scoped_require 'contrast/agent/static_analysis'

cs__scoped_require 'contrast/utils/timer'
cs__scoped_require 'contrast/utils/freeze_util'

module Contrast
  module Agent
    # This class allows the Agent to plug into the Rack middleware stack. When the
    # application is first started, we initialize ourselves as a rack middleware
    # inside of #initialize. Afterwards, we process each http request and response
    # as it goes through the middleware stack inside of #call.
    class Middleware
      include Contrast::Components::Interface
      access_component :agent, :config, :logging, :scope, :settings

      attr_reader :app

      # Allows the Agent to function as a middleware. We perform all our one-time whole-app routines in here
      # since we're only going to be initialized a single time. Our initialization order is:
      # - capture the application
      # - setup the Agent
      # - startup the Agent
      #
      # @param app [Rack::Application] the application to be instrumented
      # @param _legacy_param [nil] was a flag we no longer need, but Sinatra may call it
      def initialize app, _legacy_param = nil
        @app = app  # THIS MUST BE FIRST AND ALWAYS SET!
        setup_agent # THIS MUST BE SECOND AND ALWAYS CALLED!
        unless AGENT.enabled?
          logger.error(
              'The Agent was unable to initialize before the application middleware was initialized. Disabling permanently.')
          AGENT.disable! # ensure the agent is disabled (probably redundant)
          return
        end
        agent_startup_routine
      end

      # Startup the Agent as part of the initialization process:
      # - start the service sending thread, responsible for sending and
      #   processing messages
      # - start the heartbeat thread, which triggers service startup
      # - start instrumenting libraries and do a 'catchup' patch for everything
      #   we didn't see get loaded
      # - enable TracePoint, which handles all class loads and required
      #   instrumentation going forward
      def agent_startup_routine
        logger.debug_with_time('middleware: starting service') do
          Contrast::Agent.thread_watcher.ensure_running?
        end

        logger.debug_with_time('middleware: instrument shared libraries and patch') do
          Contrast::Agent::Patching::Policy::Patcher.patch
        end

        AGENT.enable_tracepoint
        Contrast::Agent::AtExitHook.exit_hook
      end

      # This is where we're hooked into the middleware stack. If the agent is enabled, we're ready
      # to do some processing on a per request basis. If not, we just pass the request along to the
      # next middleware in the stack.
      #
      # @param env [Hash] the various variables stored by this and other Middlewares to know the state
      #   and values of this Request
      # @return [Array,Rack::Response] the Response of this and subsequent Middlewares to be passed back
      #   to the user up the Rack framework.
      def call env
        Contrast::Utils::HeapDumpUtil.run

        if AGENT.enabled?
          Contrast::Agent::StaticAnalysis.catchup
          call_with_agent(env)
        else
          app.call(env)
        end
      end

      private

      def setup_agent
        SETTINGS.reset_state

        if CONFIG.invalid?
          AGENT.disable!
          logger.error('!!! CONFIG FILE IS INVALID - DISABLING CONTRAST AGENT !!!')
        elsif CONFIG.disabled?
          AGENT.disable!
          logger.warn('Contrast disabled by configuration. Continuing without instrumentation.')
        else
          AGENT.enable!
        end
      end

      # This is where we process each request we intercept as a middleware. We make the request context
      # available globally so that it can be accessed from anywhere. A RequestHandler object is made
      # for each request, which handles prefilter and postfilter operations.
      def call_with_agent env
        Contrast::Agent.thread_watcher.ensure_running?
        return unless AGENT.enabled?

        framework_request = Contrast::Agent.framework_manager.retrieve_request(env)
        context = Contrast::Agent::RequestContext.new(framework_request)
        response = nil

        # make the context available for the lifecycle of this request
        Contrast::Agent::REQUEST_TRACKER.lifespan(context) do
          request_handler = Contrast::Agent::RequestHandler.new(context)

          logger.debug_with_time('HTTP request cycle') do
            # prefilter sequence
            with_contrast_scope do
              context.service_extract_request
              request_handler.ruleset.prefilter
            end

            response = application_code(env) # pass request down the Rack chain with original env

            # postfilter sequence
            with_contrast_scope do
              context.extract_after(response) # update context with final response information

              if Contrast::Agent.framework_manager.streaming?(env)
                context.reset_activity
                request_handler.stream_safe_postfilter
              else
                request_handler.ruleset.postfilter
                # return response stored in the context in case any postfilter rules updated the response data
                response = context.response&.rack_response || response
                request_handler.send_activity_messages
              end
            end
          end
        end

        response
      rescue StandardError => e
        handle_exception(e)
      end

      def application_code env
        logger.trace_with_time('application') do
          app.call(env)
        end
      rescue Contrast::SecurityException => e
        logger.trace('Security Exception raised during application lifecycle to prevent an attack', e)
        raise e
      end

      SECURITY_EXCEPTION_MARKER = 'Contrast::SecurityException'
      # We're only going to suppress SecurityExceptions indicating a blocked attack.
      # And, only if the config.agent.ruby.exceptions.capture? is set
      def handle_exception exception
        if exception.is_a?(Contrast::SecurityException) ||
              exception.message&.include?(SECURITY_EXCEPTION_MARKER)

          exception_control = AGENT.exception_control
          raise exception unless exception_control[:enable]

          [exception_control[:status], {}, [exception_control[:message]]]
        else
          logger.debug('Re-throwing original error', exception)
          raise exception
        end
      end
    end
  end
end