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

cs__scoped_require 'contrast/components/interface'

cs__scoped_require 'fiber'

# 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.
# TODO: RUBY-531 move these out of Fiber class
class Fiber
  include Contrast::CoreExtensions::Assess::AssessExtension
  include Contrast::Components::Interface

  access_component :analysis, :scope

  FIBER_NEW_NODE_HASH = {
      'class_name' => 'Fiber',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'c_new', # TODO: Why do we patch new here and not initalize? Historically, new has been problematic.
      'action' => 'CUSTOM',
      'source' => 'O',
      'target' => 'R',
      'patch_class' => 'NOOP',
      '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' => 'NOOP',
      '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
      # results will be nil if StopIteration was raised,
      # otherwise an Array of the yielded arguments
      return if results.nil?

      return unless ASSESS.enabled?

      with_contrast_scope do
        results.each do |result|
          next unless Contrast::Utils::DuckUtils.quacks_like_cs_patched?(result)
          next if result.cs__frozen?

          cs__splat_tags(result, fiber)
          result.cs__properties.build_event(
              FIBER_YIELD_NODE,
              result,
              fiber,
              result,
              [])
        end
      end
    end

    def track_rb_fiber_new fiber, _enum, _enum_method, underlying, _underlying_method
      return unless Contrast::Utils::DuckUtils.quacks_like_cs_patched?(underlying)
      return unless underlying.is_a?(String) && !underlying.empty?

      return unless ASSESS.enabled?

      with_contrast_scope do
        cs__splat_tags(fiber, underlying)
        fiber.cs__properties.build_event(
            FIBER_NEW_NODE,
            fiber,
            underlying,
            fiber,
            [])
      end
    end

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

    # Some propagation occurred, but we're not sure what the
    # exact transformation was. To be safe, we just explode
    # all the tags from the source to the return.
    #
    # If the return already had that tag, the existing tag
    # range is recycled to save us an object.
    def cs__splat_tags ret, source = self
      source.cs__properties.tag_keys.each do |key|
        length = Contrast::Utils::StringUtils.ret_length(ret)
        existing = ret.cs__properties.fetch_tag(key)
        # if the tag already exists, drop all but the first range
        # then change that range to cover the entire return
        if existing&.any?
          existing.drop(existing.length - 1)
          range = existing[0]
          range.repurpose(0, length)
        else
          span = Contrast::Agent::Assess::AdjustedSpan.new(0, length)
          ret.cs__properties.add_tag(key, span)
        end
      end
    end
  end
end

Fiber.instrument_fiber_track