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

require 'contrast/agent/patching/policy/after_load_patch'
require 'contrast/components/logger'
require 'contrast/framework/manager'
require 'contrast/extension/extension'

module Contrast
  module Agent
    module Patching
      module Policy
        # Some modules diverge from our generic instrumentation and require custom instrumentation
        # after they've been loaded
        module AfterLoadPatcher
          include Contrast::Components::Logger::InstanceMethods

          # After initialization run a catchup check to instrument any already loaded modules we care about
          def catchup_after_load_patches
            apply_require_patches!
            apply_direct_patches!
            apply_load_patches!
          end

          private

          # These patches need to be applied directly, not from our policy, so
          # do so and do so only once. This should be the new standard so that
          # there are no require time side effects of loading our core
          # extensions.
          def apply_direct_patches!
            @_apply_direct_patches ||= begin
              paths = %w[
                array
                basic_object
                module
                fiber_track
                hash
                kernel
                marshal_module
                regexp
                string
                string_interpolation26
              ].cs__freeze
              paths.each do |p|
                path_part = "cs__assess_#{ p }"
                Contrast::Extension::Assess::InstrumentHelper.instrument "#{ path_part }/#{ path_part }"
              end
              true
            end
          end

          def apply_load_patches!
            after_load_patches.each do |after_load_patch|
              next unless after_load_patch.target_defined?
              next if ::Contrast::AGENT.skip_instrumentation?(after_load_patch.module_name)

              logger.trace('Catching up on already loaded afterload patch - applying instrumentation',
                           module: after_load_patch.module_name)
              after_load_patch.instrument!
            rescue NameError => e
              logger.error('Method undefined in afterload patch', e, module: after_load_patch.module_name,
                                                                     method: after_load_patch.method_to_instrument)
            rescue StandardError => e
              logger.error('Afterload patch failed to apply', e, module: after_load_patch.module_name,
                                                                 method: after_load_patch.method_to_instrument)
            end
            after_load_patches.delete_if(&:applied?)
          end

          # These patches need to be applied directly, not from our policy, and
          # are applied as a result of requiring the file as they alias methods
          # directly, allowing us to control things like scope and exception
          # handling
          def apply_require_patches!
            @_apply_require_patches ||= begin
              require 'contrast/extension/thread'
              require 'contrast/extension/kernel'
              true
            rescue LoadError, StandardError => e
              logger.error('failed instrumenting apply_require_patches!', e)
              false
            end
          end

          def after_load_patches
            @_after_load_patches ||= Contrast::Agent.framework_manager.find_after_load_patches
          end

          # Use for any checks after we've initialized
          def load_patches_for_module module_name
            return if after_load_patches.empty?

            patch = after_load_patches.find { |after_load_patch| after_load_patch.applies?(module_name) }
            return unless patch

            logger.trace('Detected loading of afterload patch - applying instrumentation', module: module_name)
            patch.instrument!
            after_load_patches.delete_if(&:applied?)
          end
        end
      end
    end
  end
end