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

require 'contrast/utils/duck_utils'

module Contrast
  module Agent
    module Assess
      module Finalizers
        # An extension of Hash that doesn't impact GC of the object being stored by storing its ID as a Key to lookup
        # and registering a finalizer on the object to remove its entry from the Hash immediately after it's GC'd.
        class Hash < Hash
          FROZEN_FINALIZED_IDS = Set.new

          def []= key, obj
            return unless ::Contrast::AGENT.enabled? && ::Contrast::ASSESS.enabled?

            # We can't finalize frozen things, so only act on those that went through .pre_freeze
            if key.cs__frozen?
              return unless FROZEN_FINALIZED_IDS.include?(key.__id__)
            else
              ObjectSpace.define_finalizer(key, finalizing_proc)
            end
            super key.__id__, obj
          end

          def [] key
            super key.__id__
          end

          # Something is trackable if it is not a collection and either not frozen or it was frozen after we put a
          # finalizer on it.
          #
          # @param key [Object] the thing to determine if trackable
          # @return [Boolean]
          def trackable? key
            return false unless key
            # Track things in these, not them themselves.
            return false if Contrast::Utils::DuckUtils.iterable_hash?(key)
            return false if Contrast::Utils::DuckUtils.iterable_enumerable?(key)
            # If it's not frozen, we can finalize/ track it.
            return true unless key.cs__frozen?

            # Otherwise, we can only track it if we've finalized it in our freeze patch.
            FROZEN_FINALIZED_IDS.include?(key.__id__)
          end

          # Determine if the given Object is tracked, meaning it has a known set of properties and those properties are
          # tracked.
          #
          # @param key [Object] the Object whose properties, by id, we want to check for tracked status
          # @return [Boolean]
          def tracked? key
            key?(key.__id__) && fetch(key.__id__, nil)&.tracked?
          end

          # Create a Proc to remove the given key from our frozen and properties tracking during finalization of the
          # Object to which the given key_id pertains. The ObjectSpace's finalizer mechanism will handle passing this
          # ID in, so we only need to define the Proc once.
          #
          # NOTE: by necessity, this is the only method which takes the __id__, not the Object itself. You CANNOT pass
          # the Object to this as a finalizer cannot hold reference to the Object being finalized; that prevents GC,
          # which introduces a memory leak and defeats the entire purpose of this.
          #
          # @return [Proc] the Proc to remove references to the ID of the GC'd object from our Hash
          def finalizing_proc
            @_finalizing_proc ||= proc do |key_id|
              FROZEN_FINALIZED_IDS.delete(key_id)
              delete(key_id)
            end
          end

          # Frozen things cannot be finalized. To avoid any issue here, we intercept the #freeze call and set
          # finalizers on the Object. To ensure later we know it's been pre-finalized, we add it's __id__ to our
          # tracking.
          #
          # @param key [Object] the Object on which we need to pre-define finalizers
          def pre_freeze key
            return unless ::Contrast::AGENT.enabled? && ::Contrast::ASSESS.enabled?
            return if key.cs__frozen?
            return if FROZEN_FINALIZED_IDS.include?(key.__id__)

            ObjectSpace.define_finalizer(key, finalizing_proc)

            FROZEN_FINALIZED_IDS << key.__id__
          rescue StandardError => _e
            nil
          end
        end
      end
    end
  end
end