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

require 'contrast/agent/assess/policy/propagation_node'
require 'contrast/components/interface'

# In order to instrument some difficult methods like String#gsub, as it
# returns an enumerator, we need to instrument methods on Fiber.
# Specifically, we instrument 'rb_fiber_yield' and 'rb_fiber_new' in
# order to track what happens within Enumerator#next.
module Contrast
  module Extension
    module Assess
      # This Class provides us with a way to invoke Regexp propagation for those
      # methods which are too complex to fit into one of the standard
      # Contrast::Agent::Assess::Policy::Propagator molds without cluttering up the
      # Fiber Class or exposing our methods there.
      class FiberPropagator
        include Contrast::Components::Interface

        access_component :analysis, :logging, :scope

        # we use funchook to patch rb_fiber_new the initialize method is not exposed by Ruby core
        FIBER_NEW_NODE_HASH = {
            'class_name' => 'Fiber',
            'instance_method' => true,
            'method_visibility' => 'public',
            'method_name' => 'c_new',
            'action' => 'CUSTOM',
            'source' => 'O',
            'target' => 'R',
            'patch_class' => 'Contrast::Extension::Assess::FiberPropagator',
            'patch_method' => 'track_rb_fiber_new'
        }.cs__freeze
        FIBER_NEW_NODE = Contrast::Agent::Assess::Policy::PropagationNode.new(FIBER_NEW_NODE_HASH)
        private_constant :FIBER_NEW_NODE_HASH
        private_constant :FIBER_NEW_NODE

        FIBER_YIELD_NODE_HASH = {
            'class_name' => 'Fiber',
            'instance_method' => true,
            'method_visibility' => 'public',
            'method_name' => 'c_yield',
            'action' => 'CUSTOM',
            'source' => 'O',
            'target' => 'R',
            'patch_class' => 'Contrast::Extension::Assess::FiberPropagator',
            'patch_method' => 'track_rb_fiber_yield'
        }.cs__freeze
        FIBER_YIELD_NODE = Contrast::Agent::Assess::Policy::PropagationNode.new(FIBER_YIELD_NODE_HASH)
        private_constant :FIBER_YIELD_NODE_HASH
        private_constant :FIBER_YIELD_NODE

        class << self
          def track_rb_fiber_yield fiber, _method, results
            return unless ASSESS.enabled?

            # results will be nil if StopIteration was raised,
            # otherwise an Array of the yielded arguments
            return unless results.cs__is_a?(Array)

            with_contrast_scope do
              results.each do |result|
                next unless (result_properties = Contrast::Agent::Assess::Tracker.properties!(result))

                result_properties.splat_from(fiber, result)
                result_properties.build_event(FIBER_YIELD_NODE, result, fiber, result, [])
              end
            end
          rescue Exception => e # rubocop:disable Lint/RescueException
            logger.error('Unable to propagate during Fiber#yield', e)
          end

          def track_rb_fiber_new fiber, _enum, _enum_method, underlying, _underlying_method
            return unless ASSESS.enabled?
            return unless underlying.is_a?(String) && !underlying.empty?

            with_contrast_scope do
              properties = Contrast::Agent::Assess::Tracker.properties!(fiber)
              return unless properties

              properties.splat_from(underlying, fiber)
              properties.build_event(FIBER_NEW_NODE, fiber, underlying, fiber, [])
            end
          rescue Exception => e # rubocop:disable Lint/RescueException
            logger.error('Unable to propagate during Fiber.new', e)
          end

          def instrument_fiber_track
            @_instrument_fiber_variables ||= begin
              require 'cs__assess_fiber_track/cs__assess_fiber_track' if Funchook.available?
              true
            end
          rescue StandardError, LoadError => e
            logger.error('Error loading fiber track patch', e)
            false
          end
        end
      end
    end
  end
end