contrast-agent 6.2.0 → 6.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b285299c51379c206ff8c7a64426ef5207594391d59c1ce284bae3a899d6d56f
-  data.tar.gz: f7847b7523600ebad30ba380cd338d36f6ed2e8136e86d2fdb71f4c7953f530a
+  metadata.gz: bdb8bbd08206dafbb18eee844e1487701c419d27c1df8e4b58c4c226b1995d5c
+  data.tar.gz: 57bbfe5e3dca05729f95f2994246e74d59f7eecc2f7fe434e265fa5c60cc3734
 SHA512:
-  metadata.gz: 3edc5557919437f5c7f76b741ed4133263bd975baebb6c81d8ac416713571c6ef871e9ab4d3b1ddf34aa0ff548242e71081443bc5c3cd04c2f4bae35b030e866
-  data.tar.gz: 21c33b917d5f9c0f3b8070bde64b42c8d800dd73c952b0d45830b9d13eb6135ba086c9abf43b3003168f5c3902753cf1a480444bbfca506f00d7bf51b5039d03
+  metadata.gz: 5f6d80b48529b52719fb7aa4608e0cd42f8ec613e61138d7f7a977652054f9a24f05a78c0d79821749bc4cc1d86b0878009e93f65be47155e47c9e1ec084780d
+  data.tar.gz: ad3d855319ebf15f1af37cac9507c2774af76a6151d7e9a4559c9de23891bcc41f0642729087165b19d20ee4891e327895f9e0d9c9c5d4e402e3a4b901b44484

data/.gitignore

@@ -56,8 +56,5 @@ contrast-agent-*.gem
 .ruby-gemset
 service_executables/*-*
 
-# Generated Protobuf files
-/lib/contrast/api/*.pb.rb
-
 # IDE stuff
 tags

data/.simplecov

@@ -4,5 +4,6 @@
 SimpleCov.minimum_coverage(line: 94)
 SimpleCov.start do
   add_filter '/spec/'
+  add_filter '/lib/protobuf/'
   enable_coverage :branch
 end

data/Rakefile

@@ -17,30 +17,3 @@ Dir['ext/cs__*'].each do |extension|
     ext.lib_dir = "lib/#{ name }"
   end
 end
-
-desc 'compile the protobuf files for the agent, translating them to .rb classes'
-task :contrast_pb_compile do
-  # do some stuff before compile
-
-  # Invoke the protobuf compile task with your sensible defaults
-  ::Rake::Task['protobuf:compile'].invoke('lib', './agent-service-api/protobuf ./agent-service-api/protobuf/dtm.proto',
-                                          'lib/contrast/api',
-                                          nil)
-
-  ::Rake::Task['protobuf:compile'].reenable
-
-  ::Rake::Task['protobuf:compile'].invoke('lib',
-                                          './agent-service-api/protobuf ./agent-service-api/protobuf/settings.proto',
-                                          'lib/contrast/api',
-                                          nil)
-
-  ['dtm.pb.rb', 'settings.pb.rb'].each do |target_file|
-    target_path = File.absolute_path(File.join(__dir__, "./lib/contrast/api/#{ target_file }"))
-    unless File.exist?(target_path)
-      puts "File not found #{ target_path }"
-      exit 1
-    end
-  end
-
-  puts 'Protobuf copied successfully'
-end

data/ext/cs__assess_basic_object/cs__assess_basic_object.c

@@ -17,9 +17,10 @@
  *  }
  */
 
-VALUE contrast_check_and_register_instance_patch(
-    const char *module_name, const char *method_name,
-    VALUE(c_fn)(const int, VALUE *, const VALUE));
+VALUE contrast_check_and_register_instance_patch(const char *module_name,
+                                                 const char *method_name,
+                                                 VALUE(c_fn)(const int, VALUE *,
+                                                             const VALUE));
 
 void contrast_assess_instance_eval_trigger_check(VALUE self, VALUE source,
                                                  VALUE ret) {
@@ -65,6 +66,7 @@ void Init_cs__assess_basic_object(void) {
      * but if someone else patched BasicObject#instance_eval,
      * IDK if this is intentional... noting it.  -ajm
      */
-    contrast_check_and_register_instance_patch("BasicObject", "instance_eval",
-                            contrast_assess_basic_object_instance_eval);
+    contrast_check_and_register_instance_patch(
+        "BasicObject", "instance_eval",
+        contrast_assess_basic_object_instance_eval);
 }

data/ext/cs__assess_kernel/cs__assess_kernel.c

@@ -26,14 +26,25 @@ contrast_patched_kernel_exec(const int argc, const VALUE *argv,
     return rb_funcall(self, rb_sym_assess_kernel_exec, argc, *argv);
 }
 
+/* Check and see if the Kernel#exec is already prepended */
+VALUE contrast_is_kernel_exec_prepended() {
+    return contrast_check_prepended(
+        rb_const_get(rb_cObject, rb_intern("Kernel")), rb_sym_kernel_exec,
+        Qfalse);
+}
+
 void Init_cs__assess_kernel(void) {
+    VALUE rb_sym_kernel_exec = ID2SYM(rb_intern("exec"));
     kernel_propagator = rb_define_module_under(core_assess, "KernelPropagator");
     exec_apply_trigger = rb_intern("apply_trigger");
 
     rb_sym_assess_kernel_exec =
         contrast_register_patch("Kernel", "exec", contrast_patched_kernel_exec);
 
-    /* should return the same value as above */
-    rb_sym_assess_kernel_exec = contrast_register_singleton_patch(
-        "Kernel", "exec", contrast_patched_kernel_exec);
+    /* check if prepended and register singleton patch if true */
+    if (contrast_is_kernel_exec_prepended() == Qtrue) {
+        /* should return the same value as above */
+        rb_sym_assess_kernel_exec = contrast_register_singleton_patch(
+            "Kernel", "exec", contrast_patched_kernel_exec);
+    }
 }

data/ext/cs__assess_kernel/cs__assess_kernel.h

@@ -3,6 +3,8 @@
 static VALUE exec_apply_trigger;
 static VALUE kernel_propagator;
 static VALUE rb_sym_assess_kernel_exec;
+static VALUE rb_sym_kernel_exec;
+VALUE contrast_is_kernel_exec_prepended();
 
 VALUE
 contrast_patched_kernel_exec(const int argc, const VALUE *argv,

data/ext/cs__assess_marshal_module/cs__assess_marshal_module.c

@@ -18,7 +18,7 @@
  *     return rb_marshal_load_with_proc(port, proc);
  * }
  */
-static VALUE contrast_marshal_module_load(const int argc, const VALUE *argv) {
+static VALUE contrast_marshal_module_prepend(const int argc, const VALUE *argv) {
     VALUE result;
     VALUE source_string;
 
@@ -67,6 +67,13 @@ void Init_cs__assess_marshal_module(void) {
     rb_sym_assess_marshal_load = rb_intern("cs__load_assess");
     rb_sym_protect_marshal_load = rb_intern("cs__load_protect");
 
-    contrast_register_singleton_prepend_patch("Marshal", "load",
-                                              &contrast_marshal_module_load);
+    /* check if prepended called form the after_load_patcher */
+    VALUE marshal = rb_const_get(rb_cObject, rb_intern("Marshal"));
+    VALUE marshal_load = ID2SYM(rb_intern("load"));
+    VALUE is_prepended = contrast_check_prepended(marshal, marshal_load, Qfalse);
+
+    if (is_prepended == Qtrue) {
+        contrast_register_singleton_prepend_patch("Marshal", "load",
+                                              &contrast_marshal_module_prepend);
+    }
 }

data/ext/cs__assess_marshal_module/cs__assess_marshal_module.h

@@ -5,6 +5,7 @@ static VALUE marshal_propagator;
 static VALUE rb_sym_assess_marshal_load;
 static VALUE rb_sym_protect_marshal_load;
 static VALUE properties_hash;
+static VALUE marshal, marshal_load, is_prepended;
 
 /*
  * Rails is a jerk. In Rails 5, they decided to do away with the alias chaining
@@ -14,7 +15,7 @@ static VALUE properties_hash;
  * special case this for now.
  * -HM (shamelessly commenting on DP's work)
  */
-static VALUE contrast_marshal_module_load(const int argc,
+static VALUE contrast_marshal_module_prepend(const int argc,
                                                  const VALUE *argv);
 
 void Init_cs__assess_marshal_module(void);

data/ext/cs__assess_regexp/cs__assess_regexp.c

@@ -7,15 +7,15 @@
 #include <ruby.h>
 
 extern VALUE contrast_force_patch(const int argc, VALUE *argv) {
-  return contrast_check_and_register_instance_patch(
+    return contrast_check_and_register_instance_patch(
         "Regexp", "=~", contrast_assess_regexp_equal_squiggle);
 }
 
 /* check if method is prepended and register instance alias or prepend patch */
 VALUE contrast_check_and_register_instance_patch(const char *module_name,
-                                                  const char *method_name,
-                                                  VALUE(c_fn)(const int, VALUE *,
-                                                              const VALUE));
+                                                 const char *method_name,
+                                                 VALUE(c_fn)(const int, VALUE *,
+                                                             const VALUE));
 
 void contrast_alias_method(const VALUE target, const char *to,
                            const char *from);
@@ -58,8 +58,10 @@ void Init_cs__assess_regexp(void) {
     rb_global_variable(&rb_sym_string);
     rb_sym_back_ref = ID2SYM(rb_intern("back_ref"));
     rb_global_variable(&rb_sym_back_ref);
-    rb_define_singleton_method(assess, "contrast_force_repatch_regexp", contrast_force_patch, 0);
+    rb_define_singleton_method(assess, "contrast_force_repatch_regexp",
+                               contrast_force_patch, 0);
 
-    rb_sym_assess_regexp_equal_squiggle = contrast_check_and_register_instance_patch(
-      "Regexp", "=~", contrast_assess_regexp_equal_squiggle);
+    rb_sym_assess_regexp_equal_squiggle =
+        contrast_check_and_register_instance_patch(
+            "Regexp", "=~", contrast_assess_regexp_equal_squiggle);
 }

data/ext/{cs__assess_string_interpolation26/cs__assess_string_interpolation26.c → cs__assess_string_interpolation/cs__assess_string_interpolation.c}

@@ -1,14 +1,25 @@
 /* Copyright (c) 2022 Contrast Security, Inc.  See
  * https://www.contrastsecurity.com/enduser-terms-0317a for more details. */
 
-#include "cs__assess_string_interpolation26.h"
+#include "cs__assess_string_interpolation.h"
 #include "../cs__common/cs__common.h"
+#include "../cs__scope/cs__scope.h"
 #include <ruby.h>
 
 static VALUE rb_str_concat_literals_hook(size_t num, VALUE *strary) {
     VALUE result = rb_str_concat_literals_original(num, strary);
     VALUE rb_params = rb_ary_new_from_values((int)num, strary);
-    rb_funcall(string_propagator, track_interpolation, 2, rb_params, result);
+    VALUE in_contrast_scope = inst_methods_in_cntr_scope(contrast_patcher(), 0);
+
+    if (in_contrast_scope == Qfalse) {
+        /* enter scope */
+        inst_methods_enter_cntr_scope(contrast_patcher(), 0);
+        rb_funcall(string_propagator, track_interpolation, 2, rb_params,
+                   result);
+        /* exit scope */
+        inst_methods_exit_cntr_scope(contrast_patcher(), 0);
+    }
+
     return result;
 }
 
@@ -20,7 +31,7 @@ static int install_hooks() {
     return 0;
 }
 
-void Init_cs__assess_string_interpolation26(void) {
+void Init_cs__assess_string_interpolation(void) {
     string_propagator =
         rb_define_class_under(core_assess, "StringPropagator", rb_cObject);
     track_interpolation = rb_intern("track_interpolation");

data/ext/{cs__assess_string_interpolation26/cs__assess_string_interpolation26.h → cs__assess_string_interpolation/cs__assess_string_interpolation.h}

@@ -10,4 +10,4 @@ static VALUE rb_str_concat_literals_hook(size_t num, VALUE *strary);
 
 static int install_hooks();
 
-void Init_cs__assess_string_interpolation26(void);
+void Init_cs__assess_string_interpolation(void);

data/ext/cs__common/cs__common.c

@@ -97,9 +97,10 @@ VALUE contrast_register_singleton_prepend_patch(const char *module_name,
 /* module name c_char "Module"; */
 /* method name c_char  "method"; */
 /* c_func => pointer */
-VALUE contrast_check_and_register_instance_patch(
-    const char *module_name, const char *method_name,
-    VALUE(c_fn)(const int, VALUE *, const VALUE)) {
+VALUE contrast_check_and_register_instance_patch(const char *module_name,
+                                                 const char *method_name,
+                                                 VALUE(c_fn)(const int, VALUE *,
+                                                             const VALUE)) {
 
     VALUE object, method, is_prepended, patch_type;
     /* check if method is prepended */
@@ -291,5 +292,5 @@ void Init_cs__common(void) {
                                contrast_check_prepended, 2);
     /* defined for object lookout */
     rb_define_singleton_method(assess, "cs__object_method_prepended?",
-                               contrast_lookout_prepended, 4);
+                               contrast_lookout_prepended, 3);
 }

data/ext/cs__contrast_patch/cs__contrast_patch.c

@@ -101,7 +101,18 @@ VALUE rescue_func(VALUE arg1) {
     return Qnil;
 }
 
+/**
+  * In the event that the original_method call throws an exception we need to ensure that contrast_post_patch is called
+  * to report that error. However, if there is no error we will call post_patch with the original_return instead of
+  * Qnil.
+  *
+**/
 VALUE contrast_patch_call_ensure(const VALUE *args) {
+    // we do not need to ensure that post patch is called if no error was thrown
+    if(!RTEST(rb_errinfo())) {
+        return Qnil;
+    }
+
     int argc;
     VALUE object, preshift, method_policy, method;
     VALUE *argv;
@@ -125,6 +136,7 @@ VALUE ensure_wrapper(const VALUE *args) {
     original_args = (VALUE)args[1];
     ensure_args = (VALUE)args[2];
 
+    //this ensure if being treated as a rescue due to issues surrounding Kernel#throw
     return rb_ensure(original_method, original_args, contrast_patch_call_ensure,
                      (VALUE)ensure_args);
 }
@@ -220,21 +232,15 @@ VALUE contrast_run_patches(const VALUE *wrapped_args) {
      * If the original method threw an exception, contrast_patch_call_rescue
      * re-raises the original exception, which unwinds the stack back to the
      * call site.  This means the rest of this function is not executed.
+     * post_patch is called in the ensure_wrapper on exception.  rb_rescue
+     * raises the exception so the below will not be executed in that event.
      */
 
-    /* Invoke Contrast post-call patching.
-     * Post-call patching may transform the return value,
-     * hence the assignment.
-     */
-    transformed_ret = contrast_call_post_patch(method_policy, preshift, object,
+    /* Invoke Contrast post-call patching. */
+    contrast_call_post_patch(method_policy, preshift, object,
                                                original_ret, argc, argv);
 
-    /* Special case for tracking frozen sources */
-    if (transformed_ret != Qnil) {
-        return transformed_ret;
-    } else {
-        return original_ret;
-    }
+    return original_ret;
 }
 
 VALUE contrast_ensure_function(const VALUE method_policy) {

data/lib/contrast/agent/assess/events/source_event.rb

@@ -56,18 +56,8 @@ module Contrast
 
           # Probably only for source events, but we'll go with source_type instead. java & .net support source_type
           # in propagation events, so we'll future proof this
-          def build_event_source_dtm
-            # You can have a source w/o a name, but not w/o a type
-            return unless source_type
-
-            dtm = Contrast::Api::Dtm::TraceEventSource.new
-            dtm.type = forced_source_type
-            dtm.name = forced_source_name
-            dtm
-          end
-
-          # Probably only for source events, but we'll go with source_type instead. java & .net support source_type
-          # in propagation events, so we'll future proof this
+          #
+          # @return [Contrast::Agent::Reporting::TraceEventSource, nil]
           def build_event_source
             # You can have a source w/o a name, but not w/o a type
             return unless source_type
@@ -89,6 +79,20 @@ module Contrast
             event_dtm.target = @policy_node.target_string
             @policy_node.targets[0]
           end
+
+          private
+
+          # Probably only for source events, but we'll go with source_type instead. java & .net support source_type
+          # in propagation events, so we'll future proof this
+          def build_event_source_dtm
+            # You can have a source w/o a name, but not w/o a type
+            return unless source_type
+
+            dtm = Contrast::Api::Dtm::TraceEventSource.new
+            dtm.type = forced_source_type
+            dtm.name = forced_source_name
+            dtm
+          end
         end
       end
     end

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

@@ -14,6 +14,7 @@ module Contrast
           FROZEN_FINALIZED_IDS = Set.new
 
           def []= key, obj
+            return unless 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

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

@@ -5,6 +5,7 @@ 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'
+require 'contrast/components/logger'
 
 module Contrast
   module Agent
@@ -13,6 +14,7 @@ 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 Contrast::Components::Logger::InstanceMethods
           include PolicyNodeUtils
           JSON_TAGS = 'tags'
           JSON_DATAFLOW = 'dataflow'
@@ -107,12 +109,16 @@ module Contrast
           def validate
             super
             validate_tags
+          rescue ArgumentError => e
+            logger.debug('Validation of policy node failed with: ', e)
+            nil
           end
 
           # TeamServer is picky. The tags here match to ENUMs there. If there
           # isn't a matching ENUM in TS land, the database gets got. We really
           # don't want to get them, so we're going to prevent the node from being
           # made.
+          # @raise[ArgumentError] raises if any of the tags is invalid
           def validate_tags
             return unless tags
 

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

@@ -8,6 +8,7 @@ require 'contrast/components/logger'
 require 'contrast/utils/object_share'
 require 'contrast/utils/sha256_builder'
 require 'contrast/utils/assess/propagation_method_utils'
+require 'contrast/utils/assess/event_limit_utils'
 require 'contrast/agent/assess/events/event_data'
 require 'contrast/utils/assess/object_store'
 
@@ -21,6 +22,7 @@ module Contrast
         module PropagationMethod
           extend Contrast::Components::Logger::InstanceMethods
           extend Contrast::Utils::Assess::PropagationMethodUtils
+          extend Contrast::Utils::Assess::EventLimitUtils
 
           @properties = Contrast::Utils::Assess::ObjectStore.new
 
@@ -38,6 +40,7 @@ module Contrast
             def apply_propagation method_policy, preshift, object, ret, args, block
               return unless (propagation_node = method_policy.propagation_node)
               return unless propagation_node.use_original_object? || preshift
+              return if event_limit?(method_policy)
 
               target = determine_target(propagation_node, ret, object, args)
               propagation_data = Contrast::Agent::Assess::Events::EventData.new(nil, nil, object, ret, args)
@@ -58,8 +61,6 @@ module Contrast
             #                         ret [Object] the Return of the invoked method
             #                         args [Array<Object>] the Arguments with which the method was invoked
             # @param block [Block] the Block passed to the original method
-            # @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_propagator propagation_node, preshift, target, propagation_data, block
               return unless propagation_possible?(propagation_node, target)
 
@@ -133,16 +134,6 @@ module Contrast
               true
             end
 
-            # Safely duplicate the target, or return nil
-            #
-            # @param target [Object] the thing to check for duplication
-            # @return [Object, nil]
-            def safe_dup target
-              target.dup
-            rescue StandardError => _e
-              nil
-            end
-
             # Iterate over each key and value in a hash to allow for propagation to each.
             #
             # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
@@ -196,28 +187,15 @@ module Contrast
             # @param _block [Block] the Block passed to the original method
             def handle_cs_properties_propagation propagation_node, preshift, target, propagation_data, _block
               return if propagation_node.action == NOOP_ACTION
-              return unless can_propagate?(propagation_node, preshift, target)
+              return unless can_propagate?(propagation_node, preshift, target, propagation_data)
               return unless (propagation_class = find_propagation_class(propagation_node))
 
-              restore_frozen_state = false
-              if target.cs__frozen? && !Contrast::Agent::Assess::Tracker.trackable?(target)
-                return unless can_handle_frozen?(propagation_node)
-                return unless (dup = safe_dup(propagation_data.ret))
-
-                restore_frozen_state = true
-                ret = dup
-                target = ret
-                Contrast::Agent::Assess::Tracker.pre_freeze(ret)
-                ret.cs__freeze
-                # double check that we were able to finalize the replaced return
-                return unless Contrast::Agent::Assess::Tracker.trackable?(target)
-              end
-
               # 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)
+              update_properties(propagation_node, target, propagation_data)
+              increment_event_count(propagation_node)
             end
 
             def handle_propagation propagation_class, propagation_node, source, target
@@ -228,7 +206,7 @@ module Contrast
               end
             end
 
-            def update_properties restore_frozen_state, propagation_node, target, propagation_data, ret
+            def update_properties propagation_node, target, propagation_data
               if propagation_node.use_original_on_bang_method?
                 properties = use_original_object_properties(propagation_data)
 
@@ -248,11 +226,10 @@ module Contrast
               event_data = Contrast::Agent::Assess::Events::EventData.new(propagation_node,
                                                                           target,
                                                                           propagation_data.object,
-                                                                          ret,
+                                                                          propagation_data.ret,
                                                                           propagation_data.args)
               properties.build_event(event_data)
               logger.trace('Propagation detected', node_id: propagation_node.id, target_id: target.__id__)
-              restore_frozen_state ? ret : nil
             end
 
             # Find the propagation class from the given node, if one exists.
@@ -269,17 +246,6 @@ module Contrast
               propagation_class
             end
 
-            # We can handle frozen propagation iff we're allowed to, as determined by configuration, and the target of
-            # the propagation is a return, as that's a replaceable value.
-            #
-            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs a
-            #   propagation event.
-            # @return [Boolean]
-            def can_handle_frozen? propagation_node
-              ::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

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

@@ -3,6 +3,7 @@
 
 require 'contrast/agent/assess/policy/policy_node'
 require 'contrast/api/decorators/trace_taint_range_tags'
+require 'contrast/components/logger'
 
 module Contrast
   module Agent
@@ -14,6 +15,8 @@ module Contrast
         # untrusted data (indicate points in the application where user
         # controlled input is modified).
         class PropagationNode < PolicyNode
+          include Contrast::Components::Logger::InstanceMethods
+
           JSON_ACTION = 'action'
           JSON_UNTAGS = 'untags'
           JSON_PATCH_CLASS = 'patch_class'
@@ -41,6 +44,9 @@ module Contrast
             @patch_method = propagation_hash[JSON_PATCH_METHOD]
             @patch_method = @patch_method.to_sym if @patch_method
             validate
+          rescue ArgumentError => e
+            logger.error('Propagation Node Initialization failed with: ', e)
+            nil
           end
 
           def node_class
@@ -56,6 +62,7 @@ module Contrast
 
           # Standard validation + TS trace version two rules:
           # Must have source, target, and action
+          # @raise[ArgumentError] raises if any of the required propagation node field is not valid, or is missing
           def validate
             super
             raise(ArgumentError, "Propagator #{ id } did not have a proper action. Unable to create.") unless action
@@ -78,6 +85,7 @@ module Contrast
             validate_untags
           end
 
+          # @raise[ArgumentError] raises if any of the tags is invalid
           def validate_untags
             return unless untags
 

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

@@ -25,6 +25,8 @@ module Contrast
                 Contrast::Agent::Assess::Tracker.tracked?(value)
               end
 
+              # @raise [NoMethodError] This is being raised if any of the implementing subclasses does not have
+              #     that method implemented, but is being called on.
               def propagate _propagation_node, _preshift, _target
                 raise(NoMethodError("Expected Base propagator subclass: #{ cs__class } to implement #propagate"))
               end

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

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require 'contrast/extension/module'
+require 'contrast/utils/assess/event_limit_utils'
 
 module Contrast
   module Agent
@@ -13,6 +14,8 @@ module Contrast
           # action knows the class and method it should call to preform this
           # action.
           module Custom
+            extend Contrast::Utils::Assess::EventLimitUtils
+
             class << self
               def propagate propagation_node, preshift, ret, block
                 clazz = propagation_node.patch_class
@@ -26,6 +29,7 @@ module Contrast
                   propagation_node.patch_class = clazz
                 end
                 clazz.send(method, propagation_node, preshift, ret, block)
+                increment_event_count(propagation_node)
               end
             end
           end

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

@@ -1,6 +1,8 @@
 # Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 # frozen_string_literal: true
 
+require 'contrast/utils/assess/event_limit_utils'
+
 module Contrast
   module Agent
     module Assess
@@ -11,6 +13,8 @@ module Contrast
           # results in new source nodes to track which columns in the database
           # have been tainted.
           class DatabaseWrite < Contrast::Agent::Assess::Policy::Propagator::Base
+            extend Contrast::Utils::Assess::EventLimitUtils
+
             class << self
               def propagate propagation_node, preshift, target
                 return unless Contrast::ASSESS.require_dynamic_sources?
@@ -22,6 +26,7 @@ module Contrast
                 known_tainted = ::Contrast::ASSESS.tainted_columns[class_name]
                 propagation_node.sources.each do |source|
                   handle_write(propagation_node, source, preshift, target, known_tainted, tainted_columns)
+                  increment_event_count(propagation_node)
                 end
                 return if tainted_columns.empty?