after_load_patch.rb in contrast-agent-3.12.0

- old
+ new

@@ -1,23 +1,28 @@
 # Copyright (c) 2020 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 # frozen_string_literal: true
 
+cs__scoped_require 'contrast/components/interface'
+cs__scoped_require 'contrast/extension/module'
 cs__scoped_require 'contrast/utils/class_util'
 
 module Contrast
   module Agent
     module Patching
       module Policy
         # Used to handle tracking patches that need to apply special instrumentation when a module is loaded
         class AfterLoadPatch
-          attr_reader :applied, :module_name, :instrumentation_file_path, :method_to_instrument
+          include Contrast::Components::Interface
+          access_component :scope
+          attr_reader :applied, :module_name, :instrumentation_file_path, :method_to_instrument, :instrumenting_module
 
-          def initialize module_name, instrumentation_file_path, method_to_instrument: nil
+          def initialize module_name, instrumentation_file_path, method_to_instrument: nil, instrumenting_module:
             @applied = false
             @module_name = module_name
             @method_to_instrument = method_to_instrument
             @instrumentation_file_path = instrumentation_file_path
+            @instrumenting_module = instrumenting_module
           end
 
           def applied?
             applied
           end
@@ -53,9 +58,13 @@
             Contrast::Utils::ClassUtil.truly_defined?(module_name)
           end
 
           def instrument!
             cs__scoped_require instrumentation_file_path
+            if instrumenting_module
+              mod = Module.cs__const_get(instrumenting_module)
+              with_contrast_scope { mod.instrument } if mod
+            end
             @applied = true
           end
 
           private
 