lib/contrast/agent/assess/policy/preshift.rb in contrast-agent-4.11.0 vs lib/contrast/agent/assess/policy/preshift.rb in contrast-agent-4.12.0

@@ -1,19 +1,21 @@
 # Copyright (c) 2021 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 # frozen_string_literal: true
 
 require 'contrast/components/logger'
+require 'contrast/utils/lru_cache'
 
 module Contrast
   module Agent
     module Assess
       # In order to properly shift tags to account for the changes this method
       # caused, we'll need to store the state before the change occurred.
       class PreShift
         include Contrast::Components::Logger::InstanceMethods
         extend Contrast::Components::Logger::InstanceMethods
 
+        @lru_cache = Contrast::Utils::LRUCache.new
         UNDUPLICABLE_MODULES = [
           Enumerator # dup'ing results in 'can't copy execution context'
         ].cs__freeze
 
         attr_accessor :object, :object_length, :args, :arg_lengths
@@ -68,36 +70,47 @@
             !UNDUPLICABLE_MODULES.include?(check)
           end
 
           def append_object_details preshift, initializing, object
             can = can_dup?(initializing, object)
-            preshift.object = can ? object.dup : object
+            preshift.object = if @lru_cache.key?(object.__id__) && !Contrast::Agent::Assess::Tracker.tracked?(object)
+                                @lru_cache[object.__id__]
+                              else
+                                can ? object.dup : object
+                              end
             preshift.object_length = if Contrast::Utils::DuckUtils.quacks_to?(preshift.object, :length)
                                        object.length
                                      else
                                        0
                                      end
+
             return unless can
+
             return unless Contrast::Agent::Assess::Tracker.tracked?(object)
 
             Contrast::Agent::Assess::Tracker.copy(object, preshift.object)
+            @lru_cache[object.__id__] = object
           end
 
           def append_arg_details preshift, args
             args_length = args.length
             preshift.args = Array.new(args_length)
             preshift.arg_lengths = Array.new(args_length)
             idx = 0
             while idx < args_length
-              original_arg = args[idx]
-              p_arg = can_dup?(false, original_arg) ? original_arg.dup : original_arg
+              or_arg = args[idx]
+              p_arg = if @lru_cache.key?(or_arg.__id__)
+                        @lru_cache[or_arg.__id__]
+                      else
+                        can_dup?(false, or_arg) ? or_arg.dup : or_arg
+                      end
               preshift.args[idx] = p_arg
               preshift.arg_lengths[idx] = Contrast::Utils::DuckUtils.quacks_to?(p_arg, :length) ? p_arg.length : 0
               idx += 1
-              next if p_arg.__id__ == original_arg.__id__
-              next unless Contrast::Agent::Assess::Tracker.tracked?(original_arg)
+              next if p_arg.__id__ == or_arg.__id__
 
-              Contrast::Agent::Assess::Tracker.copy(original_arg, p_arg)
+              Contrast::Agent::Assess::Tracker.copy(or_arg, p_arg)
+              @lru_cache[p_arg.__id__] = p_arg
             end
           end
         end
       end
     end