instrumentation.rb in sqreen-alt-1.13.0

- old
+ new

@@ -9,10 +9,11 @@
 require 'sqreen/events/remote_exception'
 require 'sqreen/rules_signature'
 require 'sqreen/shared_storage'
 require 'sqreen/rules_callbacks/record_request_context'
 require 'sqreen/rules_callbacks/run_req_start_actions'
+require 'sqreen/rules_callbacks/run_block_user_actions'
 require 'set'
 
 # How to override a class method:
 #
 # class Cache
@@ -243,10 +244,14 @@
     end
 
     def self.define_callback_method(meth, original_meth, klass_name)
       @sqreen_multi_instr ||= nil
       proc do |*args, &block|
+        record_req_hp = @@record_request_hookpoints.include?([klass_name, meth]) &&
+          Sqreen::PerformanceNotifications.listen_for?
+        Sqreen::PerformanceNotifications::BinnedMetrics.start_request if record_req_hp
+
         budget = nil
         skip_call = Thread.current[:sqreen_in_use]
         begin
           if !skip_call && Sqreen.performance_budget
             # Not using framework here to try to get a bit more perf by not loading cbs
@@ -368,23 +373,25 @@
                                                             &block)
             returns.each do |ret|
               next unless ret.is_a? Hash
               case ret[:status]
               when :raise, 'raise'
+                raise ret[:exception] if ret.key?(:exception)
                 raise Sqreen::AttackBlocked, "Sqreen blocked a security threat (type: #{ret[:rule_name]}). No action is required."
               when :override, 'override'
                 result = ret[:new_return_value]
               else
                 next
               end
             end
             result
           ensure
-            if @@record_request_hookpoints.include?([klass_name, meth]) && Sqreen::PerformanceNotifications.listen_for?
+            if record_req_hp
               Sqreen::PerformanceNotifications.instrument('Callbacks/hooks_reporting/pre') do
                 Sqreen::PerformanceNotifications::LogPerformance.next_request
                 Sqreen::PerformanceNotifications::NewRelic.next_request
+                Sqreen::PerformanceNotifications::BinnedMetrics.finish_request
               end
             end
             Thread.current[:sqreen_in_use] = false
           end
         end
@@ -437,10 +444,12 @@
         elsif private_method_defined?(meth)
           method_kind = :private
         end
         alias_method meth, saved_meth_name
         send(method_kind, meth)
+
+        remove_method saved_meth_name
       end
     end
 
     def get_saved_method_name(meth, suffix = nil)
       "#{meth}_sq#{suffix}_not_modified".to_sym
@@ -448,10 +457,19 @@
 
     def override_instance_method(klass_name, meth)
       saved_meth_name = get_saved_method_name(meth)
       new_method      = "#{meth}_modified".to_sym
 
+      # do not include ancestors o/wise we might get a saved method from a
+      # superclass which is never called because the original name
+      # (instrumented) method of the superclass has been overridden
+      private_meths = klass_name.private_instance_methods(false)
+      if private_meths.include?(saved_meth_name)
+        Sqreen.log.debug { "#{saved_meth_name} found #{klass_name}##{meth} already instrumented" }
+        return saved_meth_name
+      end
+
       p = Instrumentation.define_callback_method(meth, saved_meth_name,
                                                  klass_name)
       method_kind = nil
       klass_name.class_eval do
         alias_method saved_meth_name, meth
@@ -662,10 +680,12 @@
     attr_accessor :metrics_engine
 
     # @return [Array<Sqreen::CB>]
     def hardcoded_callbacks(framework)
       [
-        Sqreen::Rules::RunReqStartActions.new(framework)
+        Sqreen::Rules::RunReqStartActions.new(framework),
+        Sqreen::Rules::RunBlockUserActions.new(Sqreen, :identify, 0),
+        Sqreen::Rules::RunBlockUserActions.new(Sqreen, :auth_track, 1),
       ]
     end
 
     # Instrument the application code using the rules
     # @param rules [Array<Hash>] Rules to instrument