Fix ActiveRecord adapter patching for Postgres with JDBC

Author: y9v

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

@@ -43,7 +43,7 @@ def detect_sql_injection(sql, adapter_name)
             end
           end
 
-          # patch for all adapters in ActiveRecord >= 7.1
+          # patch for mysql2, sqlite3, and postgres+jdbc adapters in ActiveRecord >= 7.1
           module InternalExecQueryAdapterPatch
             def internal_exec_query(sql, *args, **rest)
               Instrumentation.detect_sql_injection(sql, adapter_name)
@@ -52,37 +52,36 @@ def internal_exec_query(sql, *args, **rest)
             end
           end
 
-          # patch for postgres adapter in ActiveRecord < 7.1
-          module ExecuteAndClearAdapterPatch
-            def execute_and_clear(sql, *args, **rest)
+          # patch for mysql2, sqlite3, and postgres+jdbc adapters in ActiveRecord < 7.1
+          module ExecQueryAdapterPatch
+            def exec_query(sql, *args, **rest)
               Instrumentation.detect_sql_injection(sql, adapter_name)
 
               super
             end
           end
 
-          # patch for postgres adapter in ActiveRecord 4
-          module Rails4ExecuteAndClearAdapterPatch
-            def execute_and_clear(sql, name, binds)
+          # patch for mysql2, sqlite3, and postgres+jdbc db adapters in ActiveRecord 4
+          module Rails4ExecQueryAdapterPatch
+            def exec_query(sql, *args)
               Instrumentation.detect_sql_injection(sql, adapter_name)
 
               super
             end
           end
 
-          # patch for mysql2 and sqlite3 adapters in ActiveRecord < 7.1
-          # also used for postgres adapter in ActiveRecord >= 7.1 when used together with JDBC adapter
-          module ExecQueryAdapterPatch
-            def exec_query(sql, *args, **rest)
+          # patch for non-jdbc postgres adapter in ActiveRecord > 4
+          module ExecuteAndClearAdapterPatch
+            def execute_and_clear(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)
+          # patch for non-jdbc postgres adapter in ActiveRecord 4
+          module Rails4ExecuteAndClearAdapterPatch
+            def execute_and_clear(sql, name, binds)
               Instrumentation.detect_sql_injection(sql, adapter_name)
 
               super

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

@@ -76,16 +76,22 @@ def patch_mysql2_adapter
           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_module = if ::ActiveRecord.gem_version.segments.first == 4
                                        Instrumentation::Rails4ExecuteAndClearAdapterPatch
                                      else
                                        Instrumentation::ExecuteAndClearAdapterPatch
                                      end
 
+            if defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
+              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
+            end
+
             ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(instrumentation_module)
           end
         end