contrast-agent 5.1.0 → 6.0.0

data/ext/cs__scope/cs__scope.h

@@ -0,0 +1,88 @@
+#include <ruby.h>
+
+/* Calls to Contrast modules and classes */
+VALUE scope_interface;
+VALUE scope_inst_methods, scope_mod, scope_klass;
+
+/* IVs */
+static VALUE rb_iv_cntr_scope;
+static VALUE rb_iv_dslr_scope;
+static VALUE rb_iv_split_scope;
+
+/* Constants */
+static VALUE rb_const_ec;
+static VALUE rb_const_mon;
+static VALUE rb_const_ec_keys;
+
+/* Symbols */
+static VALUE rb_sym_scope_mod;
+static VALUE rb_sym_contrast;
+static VALUE rb_sym_deserialization;
+static VALUE rb_sym_split;
+
+/* methods names - only reusable ones */
+VALUE rb_method_name_init;
+VALUE rb_method_name_in_scope;
+VALUE rb_method_name_enter_scope;
+VALUE rb_method_name_exit_scope;
+VALUE rb_method_name_scope_for_current_ec;
+VALUE rb_method_name_in_cntr_scope;
+VALUE rb_method_name_enter_cntr_scope;
+VALUE rb_method_name_exit_cntr_scope;
+VALUE rb_method_name_in_dslr_scope;
+VALUE rb_method_name_enter_dslr_scope;
+VALUE rb_method_name_exit_dslr_scope;
+VALUE rb_method_name_in_split_scope;
+VALUE rb_method_name_enter_split_scope;
+VALUE rb_method_name_exit_split_scope;
+VALUE rb_method_name_split_scope_depth;
+
+/* Scope class */
+VALUE contrast_scope_klass_init(VALUE self, VALUE args);
+VALUE in_cntr_scope(VALUE self, VALUE args);
+VALUE enter_cntr_scope(VALUE self, VALUE args);
+VALUE exit_cntr_scope(VALUE self, VALUE args);
+VALUE in_split_scope(VALUE self, VALUE args);
+VALUE enter_split_scope(VALUE self, VALUE args);
+VALUE exit_split_scope(VALUE self, VALUE args);
+VALUE split_scope_depth(VALUE self, VALUE args);
+VALUE in_dsrl_scope_scope(VALUE self, VALUE args);
+VALUE enter_dsrl_scope(VALUE self, VALUE args);
+VALUE exit_dsrl_scope(VALUE self, VALUE args);
+VALUE scope_klass_in_scope(VALUE self, VALUE method_scope_sym);
+VALUE scope_klass_enter_scope(VALUE self, VALUE method_scope_sym);
+VALUE scope_klass_exit_scope(VALUE self, VALUE method_scope_sym);
+
+/* Scope interface */
+VALUE contrast_scope_interface_init(VALUE self, VALUE args);
+VALUE contrast_scope_for_current_ec(VALUE self, VALUE args);
+
+/* Scope instance methods */
+VALUE inst_methods_in_cntr_scope(VALUE self, VALUE args);
+VALUE inst_methods_enter_cntr_scope(VALUE self, VALUE args);
+VALUE inst_methods_exit_cntr_scope(VALUE self, VALUE args);
+VALUE inst_methods_in_split_scope(VALUE self, VALUE args);
+VALUE inst_methods_enter_split_scope(VALUE self, VALUE args);
+VALUE inst_methods_exit_split_scope(VALUE self, VALUE args);
+VALUE inst_methods_split_scope_depth(VALUE self, VALUE args);
+VALUE inst_methods_in_dsrl_scope(VALUE self, VALUE args);
+VALUE inst_methods_enter_dsrl_scope(VALUE self, VALUE args);
+VALUE inst_methods_exit_dsrl_scope(VALUE self, VALUE args);
+VALUE inst_methods_in_scope(VALUE self, VALUE method_scope_sym);
+VALUE inst_methods_enter_scope(VALUE self, VALUE method_scope_sym);
+VALUE inst_methods_exit_scope(VALUE self, VALUE method_scope_sym);
+
+/* Scope components module */
+VALUE scope_mod_sweep_dead_ecs(VALUE self, VALUE args);
+VALUE inst_methods_enter_method_scope(VALUE self, VALUE scopes_to_enter);
+VALUE inst_methods_exit_method_scope(VALUE self, VALUE scopes_to_exit);
+
+/* Helpers */
+VALUE is_in_scope(int scope);
+VALUE get_ec();
+VALUE rb_new_c_scope();
+void rb_raise_scope_no_method_err(const VALUE method_scope_sym);
+int scope_increase(int scope);
+int scope_decrease(int scope);
+
+void Init_cs__scope(void);

data/ext/cs__scope/extconf.rb

@@ -0,0 +1,5 @@
+# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
+# frozen_string_literal: true
+
+$TO_MAKE = File.basename(__dir__)
+require_relative '../extconf_common'

data/lib/contrast/agent/assess/contrast_event.rb

@@ -15,20 +15,27 @@ module Contrast
     module Assess
       # This class holds the data about an event in the application We'll use it to build an event that TeamServer can
       # consume if the object to which this event belongs ends in a trigger.
-      #
-      # @attr_reader event_id [Integer] the atomic id of this event
-      # @attr_reader policy_node [Contrast::Agent::Assess::Policy::PolicyNode] the node that governs this event.
-      # @attr_reader stack_trace [Array<String>] the execution stack at the time the method for this event was invoked
-      # @attr_reader time [Integer] the time, in epoch ms, when this event was created
-      # @attr_reader thread [Integer] the object id of the thread on which this event was generated
-      # @attr_reader object [Contrast::Agent::Assess::ContrastObject] the safe representation of the Object on which
-      #   the method was invoked
-      # @attr_reader ret [Contrast::Agent::Assess::ContrastObject] the safe representation of the Return of the invoked
-      #   method
-      # @attr_reader args [Array<Contrast::Agent::Assess::ContrastObject>] the safe representation of the Arguments
-      #   with which the method was invoked
       class ContrastEvent
-        attr_reader :event_id, :policy_node, :stack_trace, :time, :thread, :object, :ret, :args, :tags
+        # @return [Integer] the atomic id of this event
+        attr_reader :event_id
+        # @return [Contrast::Agent::Assess::Policy::PolicyNode] the node that governs this event.
+        attr_reader :policy_node
+        # @return [Array<String>] the execution stack at the time the method for this event was invoked
+        attr_reader :stack_trace
+        # @return [Integer] the time, in epoch ms, when this event was created
+        attr_reader :time
+        # @return [Integer] the object id of the thread on which this event was generated
+        attr_reader :thread
+        # @return [Contrast::Agent::Assess::ContrastObject] the safe representation of the Object on which the method
+        #   was invoked
+        attr_reader :object
+        # @return [Contrast::Agent::Assess::ContrastObject] the safe representation of the Return of the invoked method
+        attr_reader :ret
+        # @return [Array<Contrast::Agent::Assess::ContrastObject, nil>] the safe representation of the Arguments with
+        #   which the method was invoked
+        attr_reader :args
+        # @return [Hash<Contrast::Agent::Assess::Tag>]
+        attr_reader :tags
 
         # We need this to track the parent id's of events to build up a flow chart of the finding
         @atomic_id = 0

data/lib/contrast/agent/assess/contrast_object.rb

@@ -42,8 +42,11 @@ module Contrast
           end
         end
 
+        # Is the object this represents tracked or not?
+        #
+        # @return [Boolean]
         def tracked?
-          tags&.any?
+          !!tags&.any?
         end
       end
     end

data/lib/contrast/agent/assess/finalizers/hash.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require 'contrast/utils/duck_utils'
+require 'contrast/agent/assess/policy/propagation_method'
 
 module Contrast
   module Agent
@@ -66,6 +67,7 @@ module Contrast
           def finalizing_proc
             @_finalizing_proc ||= proc do |key_id|
               FROZEN_FINALIZED_IDS.delete(key_id)
+              Contrast::Agent::Assess::Policy::PropagationMethod.instance_variable_get(:@properties).delete(key_id)
               delete(key_id)
             end
           end

data/lib/contrast/agent/assess/policy/policy_node.rb

@@ -3,6 +3,8 @@
 
 require 'contrast/agent/patching/policy/policy_node'
 require 'contrast/api/decorators/trace_taint_range_tags'
+require 'contrast/utils/object_share'
+require 'contrast/agent/assess/policy/policy_node_utils'
 
 module Contrast
   module Agent
@@ -11,9 +13,30 @@ module Contrast
         # This class functions to translate our policy.json into an actionable
         # Ruby object, allowing for dynamic patching over hardcoded patching.
         class PolicyNode < Contrast::Agent::Patching::Policy::PolicyNode
+          include PolicyNodeUtils
+
           attr_accessor :tags, :type
           attr_reader   :sources, :targets, :source_string, :target_string
 
+          # Here are all methods that can use original objects without mutation the source.
+          # For methods with REMOVE action, their (!) bang alternative is not listed, since
+          # this methods tends to mutate the original (shortens it's length, remove special
+          # symbols..) and causes mismatch in tags range representation of the target. In
+          # short preshift is needed.
+          ORIGINAL_OBJECT_METHODS = %i[
+            capitalize
+            capitalize!
+            chomp
+            to_s
+            downcase
+            downcase!
+            lstrip
+            strip
+            upcase!
+            upcase
+          ].cs__freeze
+          TO_S = 'to_s'
+
           def initialize policy_hash = {}
             super(policy_hash)
             @source_string = policy_hash[JSON_SOURCE]
@@ -21,6 +44,20 @@ module Contrast
             @tags = Set.new(policy_hash[JSON_TAGS])
             @sources = convert_policy_markers(source_string)
             @targets = convert_policy_markers(target_string)
+            @_use_original_object = ORIGINAL_OBJECT_METHODS.include?(@method_name)
+            @_use_original_on_bang_method = assign_on_bang_check policy_hash
+          end
+
+          def assign_on_bang_check policy_hash
+            # If we have KEEP action on String, and the method is to_s, that method would return self:
+            # String#to_s => self or string. This method is included here to cover the situations such as
+            # String.to_s.html_safe, where normally the dynamic sources properties get lost. To solve this
+            # we will simply return the original object here.
+            return true if @_use_original_object && policy_hash[JSON_METHOD_NAME] == TO_S
+
+            @_use_original_object &&
+                # Check if method name ends with a (!) bang unless is the to_s method:
+                policy_hash[JSON_METHOD_NAME].end_with?(Contrast::Utils::ObjectShare::BANG)
           end
 
           def feature
@@ -119,36 +156,22 @@ module Contrast
           JSON_TAGS = 'tags'
           JSON_DATAFLOW = 'dataflow'
 
-          private
-
-          # Given a policy string in the format A,B,C, populate the given array
-          # 1) Split on ','
-          # 2) If 'O' or 'R', add the array, else it's P and needs to be
-          #    converted. P type will either be P# where # is the index
-          #    of the parameter. Drop the P and store the # as an int.
+          # This method will check if a method is fit to use it's original object and
+          # that the method is without bang - it does not change the source, but rather
+          # creates a copy of it.
           #
-          # @param markers [String] the String from the policy to parse
-          # @return [Array] the array generated by converting the marker string
-          def convert_policy_markers markers
-            return Contrast::Utils::ObjectShare::EMPTY_ARRAY unless markers
-            return Contrast::Utils::ObjectShare::EMPTY_ARRAY if markers.empty?
-
-            converted = []
-            markers.split(Contrast::Utils::ObjectShare::COMMA).each do |t|
-              converted << case t
-                           when Contrast::Utils::ObjectShare::OBJECT_KEY,
-                   Contrast::Utils::ObjectShare::RETURN_KEY
-
-                             t
-                           else
-                             Integer(t[1..])
-                           end
-            end
-            converted
+          # @return true | false
+          def use_original_object?
+            @_use_original_object && Contrast::ASSESS.track_original_object?
           end
 
-          def all_type? marker
-            marker.include?(Contrast::Utils::ObjectShare::COMMA)
+          # This method will check if a method is fit to use it's original object and
+          # that the target return is the same as object - a bang method modifying the
+          # source.
+          #
+          # @return true | false
+          def use_original_on_bang_method?
+            @_use_original_on_bang_method && Contrast::ASSESS.track_original_object?
           end
         end
       end

data/lib/contrast/agent/assess/policy/policy_node_utils.rb

@@ -0,0 +1,51 @@
+# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
+# frozen_string_literal: true
+
+require 'contrast/agent/patching/policy/policy_node'
+require 'contrast/api/decorators/trace_taint_range_tags'
+require 'contrast/utils/object_share'
+
+module Contrast
+  module Agent
+    module Assess
+      module Policy
+        # This module will extend the PolicyNode, which functions is to translate
+        # our policy.json into an actionable Ruby object, allowing for dynamic
+        # patching over hardcoded patching.
+        module PolicyNodeUtils
+          private
+
+          # Given a policy string in the format A,B,C, populate the given array
+          # 1) Split on ','
+          # 2) If 'O' or 'R', add the array, else it's P and needs to be
+          #    converted. P type will either be P# where # is the index
+          #    of the parameter. Drop the P and store the # as an int.
+          #
+          # @param markers [String] the String from the policy to parse
+          # @return [Array] the array generated by converting the marker string
+          def convert_policy_markers markers
+            return Contrast::Utils::ObjectShare::EMPTY_ARRAY unless markers
+            return Contrast::Utils::ObjectShare::EMPTY_ARRAY if markers.empty?
+
+            converted = []
+            markers.split(Contrast::Utils::ObjectShare::COMMA).each do |maker|
+              converted << case maker
+                           when Contrast::Utils::ObjectShare::OBJECT_KEY,
+                             Contrast::Utils::ObjectShare::RETURN_KEY
+
+                             maker
+                           else
+                             Integer(maker[1..])
+                           end
+            end
+            converted
+          end
+
+          def all_type? marker
+            marker.include?(Contrast::Utils::ObjectShare::COMMA)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/contrast/agent/assess/policy/preshift.rb

@@ -14,9 +14,12 @@ module Contrast
         extend Contrast::Components::Logger::InstanceMethods
 
         @lru_cache = Contrast::Utils::LRUCache.new
-        UNDUPLICABLE_MODULES = [
+        UNDUPLICABLE_MODULES = [ # rubocop: disable Style/MutableConstant
           Enumerator # dup'ing results in 'can't copy execution context'
-        ].cs__freeze
+        ]
+        # dup'ing IO::Buffer results in null buffer.
+        UNDUPLICABLE_MODULES << IO::Buffer if RUBY_VERSION >= '3.1.0'
+        UNDUPLICABLE_MODULES.cs__freeze
 
         attr_accessor :object, :object_length, :args, :arg_lengths
 
@@ -41,6 +44,9 @@ module Contrast
             return unless propagation_node
             return unless ::Contrast::ASSESS.enabled?
 
+            # don't do preshift just use the original object
+            return if propagation_node.use_original_object?
+
             initializing = propagation_node.method_name == :initialize
             return if unsafe_io_object?(object, initializing)
 

data/lib/contrast/agent/assess/policy/propagation_method.rb

@@ -9,6 +9,7 @@ require 'contrast/utils/object_share'
 require 'contrast/utils/sha256_builder'
 require 'contrast/utils/assess/propagation_method_utils'
 require 'contrast/agent/assess/events/event_data'
+require 'contrast/utils/assess/object_store'
 
 module Contrast
   module Agent
@@ -21,6 +22,8 @@ module Contrast
           extend Contrast::Components::Logger::InstanceMethods
           extend Contrast::Utils::Assess::PropagationMethodUtils
 
+          @properties = Contrast::Utils::Assess::ObjectStore.new
+
           class << self
             # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy] the policy that governs the
             #   patches to this method
@@ -33,10 +36,8 @@ module Contrast
             # @return [Object, nil] the tracked Return or nil if no changes were made; will replace the return of the
             #   original function if not nil
             def apply_propagation method_policy, preshift, object, ret, args, block
-              return unless method_policy.propagation_node
-              return unless preshift
-
-              propagation_node = method_policy.propagation_node
+              return unless (propagation_node = method_policy.propagation_node)
+              return unless propagation_node.use_original_object? || preshift
 
               target = determine_target(propagation_node, ret, object, args)
               propagation_data = Contrast::Agent::Assess::Events::EventData.new(nil, nil, object, ret, args)
@@ -211,20 +212,39 @@ module Contrast
                 # double check that we were able to finalize the replaced return
                 return unless Contrast::Agent::Assess::Tracker.trackable?(target)
               end
-              propagation_class.propagate(propagation_node, preshift, target)
+
+              # If we are using the original object tracking, the preshift object is not created.
+              # Instead identify the source as the original object itself and propagate with it.
+              source = propagation_node.use_original_object? ? propagation_data.object : preshift
+              handle_propagation propagation_class, propagation_node, source, target
               update_properties restore_frozen_state, propagation_node, target, propagation_data, ret
             end
 
+            def handle_propagation propagation_class, propagation_node, source, target
+              if propagation_node.patch_method
+                propagation_class.send(propagation_node.patch_method, propagation_node, source, target)
+              else
+                propagation_class.propagate(propagation_node, source, target)
+              end
+            end
+
             def update_properties restore_frozen_state, propagation_node, target, propagation_data, ret
-              # Once we've propagated, attempt to tag the target if there is a tag(s) to be applied
-              apply_tags(propagation_node, target)
-              # Even though we skipped propagating tags from the source if they were included in untags, the target may
-              # have already had some on it. Let's go ahead and remove them. In this order, untags takes precedent over
-              # tags; but we control both and there should never be a propagator that has a tag in its untag.
-              apply_untags(propagation_node, target)
-              return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
+              if propagation_node.use_original_on_bang_method?
+                properties = use_original_object_properties propagation_data
 
-              properties.add_properties(propagation_node.properties)
+                return unless properties
+              else
+                # Once we've propagated, attempt to tag the target if there is a tag(s) to be applied
+                apply_tags(propagation_node, target)
+                # Even though we skipped propagating tags from the source if they were included in untags, the target
+                # may have already had some on it. Let's go ahead and remove them. In this order, untags takes
+                # precedent over tags; but we control both and there should never be a propagator that has a
+                # tag in its untag.
+                apply_untags(propagation_node, target)
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
+
+                properties.add_properties(propagation_node.properties)
+              end
               event_data = Contrast::Agent::Assess::Events::EventData.new(propagation_node,
                                                                           target,
                                                                           propagation_data.object,
@@ -259,6 +279,20 @@ module Contrast
               ::Contrast::ASSESS.track_frozen_sources? &&
                   propagation_node.targets[0] == Contrast::Utils::ObjectShare::RETURN_KEY
             end
+
+            # For certain bang methods we return the same object, no need to create expensive new properties.
+            #
+            # @param propagation_data [Contrast::Agent::Assess::Events::EventData] used to hold object, args
+            # and ret.
+            # @return properties [Contrast::Agent::Assess::Properties] the original properties transfered to
+            # target.
+            def use_original_object_properties propagation_data
+              unless @properties[propagation_data.object.__id__]
+                @properties[propagation_data.object.__id__] = Contrast::Agent::Assess::Tracker.properties!(
+                    propagation_data.object)
+              end
+              @properties[propagation_data.object.__id__]
+            end
           end
         end
         # rubocop:enable Metrics/ModuleLength

data/lib/contrast/agent/assess/policy/propagation_node.rb

@@ -118,11 +118,8 @@ module Contrast
           # it is an interesting security event that has a meaningful
           # change.
           def tagger?
-            @_tagger ||= begin
-              has_tags = tags&.any?
-              has_untags = untags&.any?
-              has_tags || has_untags
-            end
+            @_tagger = tags&.any? || untags&.any? if @_tagger.nil?
+            @_tagger
           end
         end
       end

data/lib/contrast/agent/assess/policy/propagator/buffer.rb

@@ -0,0 +1,118 @@
+# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
+# frozen_string_literal: true
+
+require 'contrast/agent/assess/policy/propagator/select'
+
+module Contrast
+  module Agent
+    module Assess
+      module Policy
+        module Propagator
+          # Propagation that results in all the tags of the source being
+          # applied to the target at the point of insertion. The target's
+          # preexisting tags are shifted to account for this insertion.
+          class Buffer < Contrast::Agent::Assess::Policy::Propagator::Base
+            class << self
+              # For the source, append its tags to the target.
+              # Once the tag is applied, shift it to the location of the insert
+              # Unlike additive propagation, this currently only supports one source
+              # We assume that insert changes the preshift target
+              def propagate_insert propagation_node, preshift, target
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
+
+                source = find_source(propagation_node.sources[0], preshift)
+
+                patcher_target = propagation_node.targets[0]
+                preshift_target = case patcher_target
+                                  when Contrast::Utils::ObjectShare::OBJECT_KEY
+                                    preshift.object
+                                  else
+                                    # this is hardly reached b/c insert supports only one source
+                                    # changed to second argument: string.insert[idx, string_to_insert]
+                                    # previously was args[int] => produces exception
+                                    preshift.args[1]
+                                  end
+
+                # in case of situations, where buffer is initialized with X size and it's not fully
+                preshift_target = preshift_target.get_string
+                target = target.get_string
+                # Find the first difference between the source to which
+                # we inserted and the result. That is the insertion
+                # point on which all tags need to be adjusted
+                # If the insertion point is the end of the string, preshift length is returned
+                # https://stackoverflow.com/questions/31714522/find-the-first-differing-character-between-two-strings-in-ruby
+                insert_point = (0...preshift_target.length).find do |i|
+                  preshift_target[i] != target[i]
+                end || preshift_target.length
+                # Depending what's inserted, we might be wrong. For instance, inserting 'foo'
+                # into 'asdfasdf' could result in 'asdfoofasdf'. we'd be off by one b/c of the 'f'
+                insert_point = target.rindex(source, insert_point)
+                overflow = insert_point...(insert_point + source.length)
+
+                # handle shifting the inserted range
+                properties.shift_tags([overflow])
+
+                properties.copy_from(source, target, insert_point, propagation_node.untags)
+                properties.cleanup_tags
+              end
+
+              # This method will implement both KEEP and SPLAT actions.
+              #
+              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode]
+              # @param preshift [Contrast::Agent::Assess::Preshift]
+              # @param ret [Object] Return target from method invocation.
+              def buffer_keep_splat propagation_node, preshift, ret
+                # If IO::Buffer gets duplicated the dup buffer is not allocated,
+                # it's empty, So a way of doing it is to use the original object
+                # itself. The class is included in the preshift's UNDUPLICABLE_MODULES.
+                return unless preshift.object
+                # We need the return in both propagation cases. So this propagation is
+                # applied in the Contrast::Utils::Patching::PatchUtils.apply_post_patch.
+                # we use this for the IO::Buffer#get_value and it requires 2 args always.
+                return unless ret || Contrast::Agent::Assess::Tracker.tracked?(preshift.object)
+
+                if preshift.args[1] == 0 # rubocop:disable Style/NumericPredicate
+                  # KEEP
+                  Contrast::Agent::Assess::Policy::Propagator::Keep.propagate(propagation_node, preshift, ret)
+                else
+                  # SPLAT
+                  Contrast::Agent::Assess::Policy::Propagator::Splat.propagate(propagation_node, preshift, ret)
+                end
+              end
+
+              # This method will implement both SELECT and KEEP actions for
+              # IO::Buffer
+              #
+              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode]
+              # @param preshift [Contrast::Agent::Assess::Preshift]
+              # @param ret [Object] Return targer from method invocation.
+              # @param _block [nil, {}] block passed.
+              def propagate_keep_select propagation_node, preshift, ret, _block
+                # We need the return in both propagation cases. So this propagation is
+                # applied in the Contrast::Utils::Patching::PatchUtils.apply_post_patch.
+                # If IO::Buffer gets duplicated the dup buffer is not allocated,
+                # it's empty, So a way of doing it is to use the original object
+                # itself. The class is handled with no dup in preshift.
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(ret)) &&
+                    (source = preshift.object)
+
+                # KEEP means the tags just pass from the source to the target as is.
+                # There is a case with IO::Buffer use, when we don't pass any args,
+                # and the target (ret) in this case is the whole strings. We should
+                # use Keep action. Otherwise the SELECT action covers all other cases.
+                if preshift.args.empty?
+                  # KEEP
+                  properties.copy_from(source, ret, 0, propagation_node.untags)
+                  ret
+                else
+                  # SELECT
+                  Contrast::Agent::Assess::Policy::Propagator::Select.select_tagger propagation_node, preshift, ret, nil
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/contrast/agent/assess/policy/propagator/keep.rb

@@ -11,12 +11,27 @@ module Contrast
           # are unaffected beyond any merging of overlapping tags.
           class Keep < Contrast::Agent::Assess::Policy::Propagator::Base
             class << self
-              # Keep means the tags just pass from the source to the target
-              # as is.
-              def propagate propagation_node, preshift, target
+              # Keep means the tags just pass from the source to the target as is.
+              #
+              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] Current node
+              # governing the propagation.
+              # @param preshift_or_object [Contrast::Agent::Assess::Preshift] current preshift state or
+              # source object passed if track_original_object flag is active.
+              # @param target [Object] To where the tainted data flows.
+              def propagate propagation_node, preshift_or_object, target
                 return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
-                source = find_source(propagation_node.sources[0], preshift)
+                source = if propagation_node.use_original_object?
+                           # for now this is used with string based object and with methods not mutating
+                           # original object, so we will always return 'O' key which is the object itself.
+                           # Because preshift is disabled for this methods we use it's parameter to pass
+                           # in the original object, preshift = object. To see list of methods in which
+                           # this is used:
+                           # Contrast::Utils::MethodCheck::ORIGINAL_OBJECT_METHODS
+                           preshift_or_object
+                         else
+                           find_source(propagation_node.sources[0], preshift_or_object)
+                         end
                 properties.copy_from(source, target, 0, propagation_node.untags)
               end
             end

data/lib/contrast/agent/assess/policy/propagator/match_data.rb

@@ -37,6 +37,8 @@ module Contrast
               end
 
               def captures_tagger propagation_node, preshift, ret, _block
+                return unless ret
+
                 idx = 0
                 while idx < ret.length
                   return_value = ret[idx]

data/lib/contrast/agent/assess/policy/propagator/remove.rb

@@ -13,10 +13,26 @@ module Contrast
             class << self
               # For the source, append its tags to the target. Once the tag is applied, remove the section that was
               # removed by the delete. Unlike additive propagation, this currently only supports one source.
-              def propagate propagation_node, preshift, target
+              #
+              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] Current node
+              # governing the propagation.
+              # @param preshift_or_object [Contrast::Agent::Assess::Preshift] current preshift state or
+              # source object passed if track_original_object flag is active.
+              # @param target [Object] To where the tainted data flows.
+              def propagate propagation_node, preshift_or_object, target
                 return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
-                source = find_source(propagation_node.sources[0], preshift)
+                source = if propagation_node.use_original_object?
+                           # for now this is used with string based object and with methods not mutating
+                           # original object (clearing empty spaces, \n..) so we will always return 'O' key
+                           # which is the object itself. Because preshift is disabled for this methods
+                           # we use it's parameter to pass in the original object( in this case just a unchained
+                           # source), preshift = object. To see list of methods in which this is used:
+                           # Contrast::Utils::MethodCheck::ORIGINAL_OBJECT_METHODS
+                           preshift_or_object
+                         else
+                           find_source(propagation_node.sources[0], preshift_or_object)
+                         end
                 properties.copy_from(source, target, 0, propagation_node.untags)
                 handle_removal(propagation_node, source, target)
               end