contrast-agent 4.1.0 → 4.4.0

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

@@ -12,8 +12,7 @@ module Contrast
           class Center < Contrast::Agent::Assess::Policy::Propagator::Base
             class << self
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 sources = propagation_node.sources
                 source1 = find_source(sources[0], preshift)

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

@@ -12,7 +12,7 @@ module Contrast
           # of tags from the source to the target. Each node with the CUSTOM
           # action knows the class and method it should call to preform this
           # action.
-          class Custom
+          module Custom
             class << self
               def propagate propagation_node, preshift, ret, block
                 clazz = propagation_node.patch_class

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

@@ -24,25 +24,8 @@ module Contrast
 
                 known_tainted = ASSESS.tainted_columns[class_name]
                 propagation_node.sources.each do |source|
-                  arg = preshift.args[source]
-                  next unless arg.cs__respond_to?(:each_pair)
-
-                  arg.each_pair do |key, value|
-                    next unless value
-                    next if known_tainted&.include?(key)
-
-                    properties = Contrast::Agent::Assess::Tracker.properties(value)
-                    next unless properties
-
-                    # TODO: RUBY-540 handle sanitization, handle nested objects
-                    Contrast::Agent::Assess::Policy::PropagationMethod.apply_tags(propagation_node, value)
-                    properties.build_event(propagation_node, value, preshift.object, target, preshift.args)
-                    next unless tracked_value?(value)
-
-                    tainted_columns[key] = properties
-                  end
+                  handle_write(propagation_node, source, preshift, target, known_tainted, tainted_columns)
                 end
-
                 return if tainted_columns.empty?
 
                 if known_tainted
@@ -53,6 +36,26 @@ module Contrast
 
                 Contrast::Agent::Assess::Policy::DynamicSourceFactory.create_sources class_type, tainted_columns
               end
+
+              private
+
+              def handle_write propagation_node, source, preshift, target, known_tainted, tainted_columns
+                arg = preshift.args[source]
+                return unless arg.cs__respond_to?(:each_pair)
+
+                arg.each_pair do |key, value|
+                  next unless value
+                  next if known_tainted&.include?(key)
+                  next unless (properties = Contrast::Agent::Assess::Tracker.properties!(value))
+
+                  # TODO: RUBY-540 handle sanitization, handle nested objects
+                  Contrast::Agent::Assess::Policy::PropagationMethod.apply_tags(propagation_node, value)
+                  properties.build_event(propagation_node, value, preshift.object, target, preshift.args)
+                  next unless tracked_value?(value)
+
+                  tainted_columns[key] = properties
+                end
+              end
             end
           end
         end

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

@@ -16,8 +16,7 @@ module Contrast
               # Unlike additive propagation, this currently only supports one source
               # We assume that insert changes the preshift target
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[1], preshift)
 

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

@@ -14,8 +14,7 @@ module Contrast
               # Keep means the tags just pass from the source to the target
               # as is.
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[0], preshift)
                 properties.copy_from(source, target, 0, propagation_node.untags)

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

@@ -67,11 +67,12 @@ module Contrast
               def square_bracket_single argument_index, preshift, return_value, propagation_node
                 original_start_index = preshift.object.begin(argument_index)
                 original_end_index = preshift.object.end(argument_index)
-                original_properties = Contrast::Agent::Assess::Tracker.properties(preshift.object)
+                return unless (original_properties = Contrast::Agent::Assess::Tracker.properties(preshift.object))
+
                 applicable_tags = original_properties.tags_at_range(original_start_index...original_end_index)
                 return if applicable_tags.empty?
+                return unless (return_properties = Contrast::Agent::Assess::Tracker.properties!(return_value))
 
-                return_properties = Contrast::Agent::Assess::Tracker.properties(return_value)
                 applicable_tags.each do |tag_name, tag_ranges|
                   return_properties.set_tags(tag_name, tag_ranges)
                 end

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

@@ -15,8 +15,7 @@ module Contrast
               # String has some silly methods like next. Basically, this flips a
               # character in a predictable manner
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[0], preshift)
                 properties.copy_from(source, target, 0, propagation_node.untags)

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

@@ -14,8 +14,7 @@ module Contrast
               # For the source, prepend its tags to the target. It's basically the
               # opposite of append. :-P
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 sources = propagation_node.sources
                 # source1 is the copy of the thing being prepended to

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

@@ -17,8 +17,7 @@ module Contrast
               # 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
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[0], preshift)
                 properties.copy_from(source, target, 0, propagation_node.untags)
@@ -27,8 +26,7 @@ module Contrast
               end
 
               def handle_removal source_chars, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source_idx = 0
 

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

@@ -14,8 +14,7 @@ module Contrast
               # Replace means we're replacing the target w/ the source. Anything
               # on the source should be passed to the target.
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[0], preshift)
                 properties.clear_tags

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

@@ -13,8 +13,7 @@ module Contrast
           class Reverse < Contrast::Agent::Assess::Policy::Propagator::Base
             class << self
               def propagate propagation_node, preshift, target
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 source = find_source(propagation_node.sources[0], preshift)
                 properties.copy_from(source, target, 0, propagation_node.untags)

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

@@ -14,9 +14,6 @@ module Contrast
           class Select
             class << self
               def select_tagger patcher, preshift, ret, _block
-                properties = Contrast::Agent::Assess::Tracker.properties(ret)
-                return unless properties
-
                 source = preshift.object
                 args = preshift.args
 
@@ -31,7 +28,9 @@ module Contrast
                   return
                 end
 
-                source_properties = Contrast::Agent::Assess::Tracker.properties(source)
+                return unless (source_properties = Contrast::Agent::Assess::Tracker.properties(source))
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(ret))
+
                 properties.build_event(
                     patcher,
                     ret,

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

@@ -19,31 +19,17 @@ module Contrast
                   when Contrast::Utils::ObjectShare::OBJECT_KEY
                     tracked_inputs << preshift.object if Contrast::Agent::Assess::Tracker.tracked?(preshift.object)
                   else
-                    arg = preshift.args[source]
-                    if arg.is_a?(String)
-                      tracked_inputs << arg if Contrast::Agent::Assess::Tracker.tracked?(arg)
-                    elsif Contrast::Utils::DuckUtils.iterable_hash?(arg)
-                      arg.each_pair do |key, value|
-                        tracked_inputs << key if tracked_value?(key)
-                        tracked_inputs << value if tracked_value?(value)
-                      end
-                    elsif Contrast::Utils::DuckUtils.iterable_enumerable?(arg)
-                      arg.each do |value|
-                        tracked_inputs << value if tracked_value?(value)
-                      end
-                    end
+                    find_argument_inputs(tracked_inputs, preshift.args[source])
                   end
                 end
 
                 splat_tags(tracked_inputs, target)
-                Contrast::Agent::Assess::Tracker.properties(target).cleanup_tags
+                Contrast::Agent::Assess::Tracker.properties(target)&.cleanup_tags
               end
 
               def splat_tags tracked_inputs, target
                 return if tracked_inputs.empty?
-
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 tracked_inputs.each do |input|
                   input_properties = Contrast::Agent::Assess::Tracker.properties(input)
@@ -52,6 +38,28 @@ module Contrast
                   properties.splat_from(input, target)
                 end
               end
+
+              private
+
+              # The arguments to the splat method are complex and of multiple types. As such, we need to handle
+              # Strings and iterables to determine the tracked inputs on which to act.
+              #
+              # @param tracked_inputs [Array] storage for the inputs to act on later
+              # @param arg [Object] an input to the method which act as sources for this propagation.
+              def find_argument_inputs tracked_inputs, arg
+                if arg.is_a?(String)
+                  tracked_inputs << arg if Contrast::Agent::Assess::Tracker.tracked?(arg)
+                elsif Contrast::Utils::DuckUtils.iterable_hash?(arg)
+                  arg.each_pair do |key, value|
+                    tracked_inputs << key if tracked_value?(key)
+                    tracked_inputs << value if tracked_value?(value)
+                  end
+                elsif Contrast::Utils::DuckUtils.iterable_enumerable?(arg)
+                  arg.each do |value|
+                    tracked_inputs << value if tracked_value?(value)
+                  end
+                end
+              end
             end
           end
         end

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

@@ -15,13 +15,12 @@ module Contrast
           class Split < Contrast::Agent::Assess::Policy::Propagator::Base
             include Contrast::Components::Interface
 
-            access_component :agent, :logging
+            access_component :agent, :logging, :scope
 
             SPLIT_TRACKER = Contrast::Utils::ThreadTracker.new
+
             class << self
-              # Propagate taint from a source as it is split into composite
-              # sections. This method MUST return nil, otherwise it risks
-              # changing the result of of the propagation.
+              # Propagate taint from a source as it is split into composite sections.
               #
               # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode]
               #   the node that governs this propagation event.
@@ -29,171 +28,139 @@ module Contrast
               #   of the state of the code just prior to the invocation of the
               #   patched method.
               # @param target [Array, String] the target to which to propagate.
-              # @return [nil]
+              # @return [nil] so as not to risk changing the result of the propagation.
+
               def propagate propagation_node, preshift, target
                 logger.trace('Propagation detected',
                              node_id: propagation_node.id,
                              target_id: target.__id__)
-                unless target.is_a?(Array)
-                  Contrast::Agent::Assess::Policy::Propagator::Keep.propagate(propagation_node, preshift, target)
-                  properties = Contrast::Agent::Assess::Tracker.properties(target)
-                  properties.build_event(propagation_node, target, object, ret, args)
-                  return
-                end
 
                 source = find_source(propagation_node.sources[0], preshift)
+                return unless (source_properties = Contrast::Agent::Assess::Tracker.properties(source))
 
-                separator_length = if propagation_node.method_name == :grapheme_clusters
-                                     # grapheme_clusters break the string apart based on each "user-perceived" character
-                                     0
-                                   else
-                                     # The default for String#split is to use a single whitespace
-                                     preshift&.args&.first&.to_s&.length || $FIELD_SEPARATOR&.to_s&.length || 1
-                                   end
+                separator_length = find_separator_length(propagation_node, preshift)
 
                 current_index = 0
-                target.each do |elem|
-                  elem_length = elem.length
-                  range = current_index...(current_index + elem_length)
-                  elem_properties = Contrast::Agent::Assess::Tracker.properties(elem)
-                  next unless elem_properties
+                target.each do |target_elem|
+                  next unless (elem_properties = Contrast::Agent::Assess::Tracker.properties!(target_elem))
 
-                  source_properties = Contrast::Agent::Assess::Tracker.properties(source)
+                  # Get tags for element from source by element range.
+                  range = current_index...(current_index + target_elem.length)
                   tags = source_properties.tags_at_range(range)
+
+                  # Set element properties accordingly.
                   elem_properties.clear_tags
-                  tags.each_pair do |key, value|
-                    elem_properties.set_tags(key, value)
-                  end
-                  elem_properties.build_event(propagation_node, elem, preshift.object, target, preshift.args, 0)
+                  tags.each_pair { |key, value| elem_properties.set_tags(key, value) }
+                  elem_properties.build_event(propagation_node, target_elem, preshift.object, target, preshift.args, 0)
                   elem_properties.add_properties(propagation_node.properties)
-                  current_index = current_index + elem_length + separator_length
+
+                  current_index = range.end + separator_length
                 end
                 nil
               end
 
-              # Marks the point in which the String#split method is called.
-              # Responsible for building the context required to propagate when
-              # the results of #split are yielded directly to a block
+              # Context for block split execution.
               #
-              # @param string [String] the String on which split is invoked
-              # @param args [Array<Object>] the arguments passed to the
-              #   original split call
-              def begin_split string, args
-                save_split_depth!
-                depth = SPLIT_TRACKER.get(:split_depth)
-                save_split_index!(depth)
-                save_split_value!(depth, string, args)
-              rescue Exception => e # rubocop:disable Lint/RescueException
-                # don't let our errors propagate and disable String#split for
-                # this since we're in an error state
-                logger.warn('Unable to record split context', e)
-                end_split
-              end
+              # @param string [String] the String on which split is invoked.
+              # @param args [Array<Object>] the arguments passed to the original split call.
+              def wrap_split string, args
+                # String#split start. Build context and yield.
+                begin
+                  enter_split_scope!
+                  save_split_index!
+                  save_split_value!(string, args)
+                rescue Exception => e # rubocop:disable Lint/RescueException
+                  logger.warn('Unable to record split context', e)
+                end
 
-              # Marks the point in which the String#split method is exited.
-              # Responsible for removing the context required to propagate when
-              # the results of #split are yielded directly to a block
-              def end_split
-                depth = SPLIT_TRACKER.get(:split_depth)
-                return unless depth
-
-                depth -= 1
-                if depth.negative?
-                  SPLIT_TRACKER.delete(:split_depth)
-                  SPLIT_TRACKER.delete(:split_index)
-                  SPLIT_TRACKER.delete(:split_value)
-                else
-                  SPLIT_TRACKER.set(:split_depth, depth)
+                yield
+              ensure
+                # String#split exit. Remove propagation context.
+                begin
+                  exit_split_scope!
+                  unless in_split_scope?
+                    SPLIT_TRACKER.delete(:split_index)
+                    SPLIT_TRACKER.delete(:split_value)
+                  end
+                rescue StandardError => e
+                  logger.warn('Unable to remove split context', e)
                 end
-              rescue StandardError => e
-                logger.warn('Unable to remove split context', e)
               end
 
-              # This method is called whenever an rb_yield is called. We need
-              # to leave it as soon as possible with as little work as
-              # possible.
+              # This method is called whenever an rb_yield is called.
+              # We need to leave it as soon as possible with as little work as possible.
               #
-              # @param target [String] the entity being passed to the yield
-              #   block
+              # @param target [String] the entity being passed to the yield block
               def propagate_yield target
-                depth, index = nil
-
-                depth = SPLIT_TRACKER.get(:split_depth)
-                return unless depth
-
-                source = SPLIT_TRACKER.get(:split_value)&.fetch(depth)
-                return unless source
-
-                index = SPLIT_TRACKER.get(:split_index)&.fetch(depth)
-                return unless index
-
-                properties = Contrast::Agent::Assess::Tracker.properties(target)
-                return unless properties
+                return unless (source = SPLIT_TRACKER.get(:split_value)&.fetch(split_scope_depth))
+                return unless (index = SPLIT_TRACKER.get(:split_index)&.fetch(split_scope_depth))
+                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))
 
                 true_source = source[index]
                 properties.copy_from(true_source, target)
               rescue StandardError => e
                 logger.warn('Unable to track within split context', e)
               ensure
-                if depth && index
+                if in_split_scope? && index
                   idx = SPLIT_TRACKER.get(:split_index)
-                  idx[depth] = index + 1 if defined?(idx) && idx.is_a?(Array)
+                  idx[split_scope_depth] = index + 1 if defined?(idx) && idx.is_a?(Array)
                 end
               end
 
+              # Load patch.
               def instrument_string_split
-                if @_instrument_string_split.nil?
-                  @_instrument_string_split = begin
-                    require 'cs__assess_yield_track/cs__assess_yield_track' if AGENT.patch_yield? && Funchook.available?
-                    true
-                  rescue StandardError => e
-                    logger.error('Error loading split rb_yield patch', e)
-                    false
-                  end
+                @_instrument_string_split ||= begin
+                  require 'cs__assess_yield_track/cs__assess_yield_track' if AGENT.patch_yield? && Funchook.available?
+                  true
+                rescue StandardError => e
+                  logger.error('Error loading split rb_yield patch', e)
+                  false
                 end
-                @_instrument_string_split
               end
 
               private
 
-              def save_split_depth!
-                depth = SPLIT_TRACKER.get(:split_depth)
-                if depth
-                  depth += 1
-                  SPLIT_TRACKER.set(:split_depth, depth)
-                else
-                  SPLIT_TRACKER.set(:split_depth, 0)
-                end
+              # grapheme_clusters break the string apart based on each "user-perceived" character. Otherwise, the
+              # default for String#split is to use a single whitespace.
+              #
+              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
+              #   propagation event.
+              # @param preshift [Contrast::Agent::Assess::PreShift] The capture of the state of the code just prior to
+              #   the invocation of the patched method.
+              def find_separator_length propagation_node, preshift
+                return 0 if propagation_node.method_name == :grapheme_clusters
+
+                preshift&.args&.first&.to_s&.length || $FIELD_SEPARATOR&.to_s&.length || 1
               end
 
-              def save_split_index! depth
-                split_index = SPLIT_TRACKER.get(:split_index)
-                unless split_index
+              # Save index of the current split object.
+              # Create index tracking array as needed.
+              def save_split_index!
+                unless (split_index = SPLIT_TRACKER.get(:split_index))
                   split_index = []
                   SPLIT_TRACKER.set(:split_index, split_index)
                 end
-                # save the index to the ThreadLocal; not useless
-                split_index[depth] = 0 # rubocop:disable Lint/UselessSetterCall
+                # save the index to the ThreadLocal; not useless.
+                split_index[split_scope_depth] = 0 # rubocop:disable Lint/UselessSetterCall
               end
 
-              def save_split_value! depth, string, args
+              # Save value of the current split object.
+              # Create value tracking array as needed.
+              def save_split_value! string, args
                 preshift = Contrast::Agent::Assess::PreShift.build_preshift(split_node, string, args)
                 target = string.split
                 propagate(split_node, preshift, target)
-                split_value = SPLIT_TRACKER.get(:split_value)
-                unless split_value
+                unless (split_value = SPLIT_TRACKER.get(:split_value))
                   split_value = []
                   SPLIT_TRACKER.set(:split_value, split_value)
                 end
-                # save the target to the ThreadLocal; not useless
-                split_value[depth] = target # rubocop:disable Lint/UselessSetterCall
+                # Save the target to the ThreadLocal; not useless.
+                split_value[split_scope_depth] = target # rubocop:disable Lint/UselessSetterCall
               end
 
-              # Quick hook to the String#split propagation node in our Assess
-              # policy
+              # Quick hook to the String#split propagation node in our Assess policy
               #
-              # @return [Contrast::Agent::Assess::Policy::PropagationNode]
-              #   String#split node
+              # @return [Contrast::Agent::Assess::Policy::PropagationNode] String#split node
               def split_node
                 @_split_node ||= begin
                   Contrast::Agent::Assess::Policy::Policy.instance.propagators.find do |node|
@@ -215,19 +182,15 @@ if RUBY_VERSION >= '2.6.0'
   class String
     alias_method :cs__patched_string_split_special, :split
 
-    # override of the the standard split method to handle the 2.6 direct
-    # yield case.
+    # Override of the the standard split method to handle the 2.6 direct yield case.
     #
     # Note: because this patch is applied before our standard propagation, this
-    # call wrapped in it. As such, any call here happens in scope, so there is
+    # call is wrapped in it. As such, any call here happens in scope, so there is
     # no need to manage it on our own.
     def split *args, &block
       if block
-        Contrast::Agent::Assess::Policy::Propagator::Split.begin_split(self, args)
-        begin
+        Contrast::Agent::Assess::Policy::Propagator::Split.wrap_split(self, args) do
           cs__patched_string_split_special(*args, &block)
-        ensure
-          Contrast::Agent::Assess::Policy::Propagator::Split.end_split
         end
       else
         cs__patched_string_split_special(*args, &block)