Fix AppSec ActiveRecord instrumentation for Rails 4

Author: y9v

Matrixfile

@@ -264,7 +264,7 @@
     'redis-3'      => '✅ 2.5 / ✅ 2.6 / ✅ 2.7 / ✅ 3.0 / ✅ 3.1 / ✅ 3.2 / ✅ 3.3 / ✅ 3.4 / ✅ jruby',
   },
   'appsec:active_record' => {
-    'relational_db' => '❌ 2.5 / ❌ 2.6 / ✅ 2.7 / ✅ 3.0 / ✅ 3.1 / ✅ 3.2 / ✅ 3.3 / ✅ 3.4 / ✅ jruby',
+    'relational_db' => '✅ 2.5 / ✅ 2.6 / ✅ 2.7 / ✅ 3.0 / ✅ 3.1 / ✅ 3.2 / ✅ 3.3 / ✅ 3.4 / ✅ jruby',
   },
   'appsec:rack' => {
     'rack-latest' => '✅ 2.5 / ✅ 2.6 / ✅ 2.7 / ✅ 3.0 / ✅ 3.1 / ✅ 3.2 / ✅ 3.3 / ✅ 3.4 / ✅ jruby',

lib/datadog/appsec/contrib/active_record/instrumentation.rb

@@ -61,15 +61,33 @@ def execute_and_clear(sql, *args, **rest)
             end
           end
 
+          # patch for postgres adapter in ActiveRecord 4
+          module Rails4ExecuteAndClearAdapterPatch
+            def execute_and_clear(sql, name, binds)
+              Instrumentation.detect_sql_injection(sql, adapter_name)
+
+              super
+            end
+          end
+
           # patch for mysql2 and sqlite3 adapters in ActiveRecord < 7.1
-          # this patch is also used when using JDBC adapter
+          # also used for postgres adapter in ActiveRecord >= 7.1 when used together with JDBC adapter
           module ExecQueryAdapterPatch
             def exec_query(sql, *args, **rest)
               Instrumentation.detect_sql_injection(sql, adapter_name)
 
               super
             end
           end
+
+          # patch for mysql2 and sqlite3 db adapters in ActiveRecord 4
+          module Rails4ExecQueryAdapterPatch
+            def exec_query(sql, *args)
+              Instrumentation.detect_sql_injection(sql, adapter_name)
+
+              super
+            end
+          end
         end
       end
     end

lib/datadog/appsec/contrib/active_record/patcher.rb

@@ -19,30 +19,75 @@ def target_version
           end
 
           def patch
+            # Rails 7.0 intruduced new on-load hooks for sqlite3 and postgresql adapters
+            # The load hook for mysql2 adapter was introduced in Rails 7.1
+            #
+            # If the adapter is not loaded when the :active_record load hook is called,
+            # we need to add a load hook for the adapter
             ActiveSupport.on_load :active_record do
-              instrumentation_module = if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
-                                         Instrumentation::InternalExecQueryAdapterPatch
-                                       else
-                                         Instrumentation::ExecQueryAdapterPatch
-                                       end
-
               if defined?(::ActiveRecord::ConnectionAdapters::SQLite3Adapter)
-                ::ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(instrumentation_module)
+                ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_sqlite3_adapter
+              else
+                ActiveSupport.on_load :active_record_sqlite3adapter do
+                  ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_sqlite3_adapter
+                end
               end
 
               if defined?(::ActiveRecord::ConnectionAdapters::Mysql2Adapter)
-                ::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(instrumentation_module)
+                ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_mysql2_adapter
+              else
+                ActiveSupport.on_load :active_record_mysql2adapter do
+                  ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_mysql2_adapter
+                end
               end
 
               if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
-                unless defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
-                  instrumentation_module = Instrumentation::ExecuteAndClearAdapterPatch
+                ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_postgresql_adapter
+              else
+                ActiveSupport.on_load :active_record_postgresqladapter do
+                  ::Datadog::AppSec::Contrib::ActiveRecord::Patcher.patch_postgresql_adapter
                 end
-
-                ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(instrumentation_module)
               end
             end
           end
+
+          def patch_sqlite3_adapter
+            instrumentation_module = if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
+                                       Instrumentation::InternalExecQueryAdapterPatch
+                                     elsif ::ActiveRecord.gem_version.segments.first == 4
+                                       Instrumentation::Rails4ExecQueryAdapterPatch
+                                     else
+                                       Instrumentation::ExecQueryAdapterPatch
+                                     end
+
+            ::ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(instrumentation_module)
+          end
+
+          def patch_mysql2_adapter
+            instrumentation_module = if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
+                                       Instrumentation::InternalExecQueryAdapterPatch
+                                     elsif ::ActiveRecord.gem_version.segments.first == 4
+                                       Instrumentation::Rails4ExecQueryAdapterPatch
+                                     else
+                                       Instrumentation::ExecQueryAdapterPatch
+                                     end
+
+            ::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(instrumentation_module)
+          end
+
+          def patch_postgresql_adapter
+            jdbc_defined = defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
+
+            instrumentation_module = if jdbc_defined && ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
+                                       Instrumentation::InternalExecQueryAdapterPatch
+                                     elsif jdbc_defined && ::ActiveRecord.gem_version.segments.first == 4
+                                       Instrumentation::Rails4ExecuteAndClearAdapterPatch
+                                     else
+                                       Instrumentation::ExecuteAndClearAdapterPatch
+                                     end
+
+            ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(instrumentation_module)
+          end
         end
       end
     end

sig/datadog/appsec/contrib/active_record/instrumentation.rbs

@@ -16,6 +16,14 @@ module Datadog
           module ExecQueryAdapterPatch
             def exec_query: (String sql, *untyped args, **untyped rest) -> untyped
           end
+
+          module Rails4ExecuteAndClearAdapterPatch
+            def execute_and_clear: (String sql, String name, untyped binds) -> untyped
+          end
+
+          module Rails4ExecQueryAdapterPatch
+            def exec_query: (String sql, *untyped args) -> untyped
+          end
         end
       end
     end

sig/datadog/appsec/contrib/active_record/patcher.rbs

@@ -8,6 +8,12 @@ module Datadog
           def self?.target_version: () -> Gem::Version?
 
           def self?.patch: () -> void
+
+          def self?.patch_sqlite3_adapter: () -> void
+
+          def self?.patch_mysql2_adapter: () -> void
+
+          def self?.patch_postgresql_adapter: () -> void
         end
       end
     end