Simplify AppSec ActiveRecord patcher and split specs by adapter

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

@@ -12,9 +12,14 @@ def detect_sql_injection(sql, adapter_name)
             scope = AppSec.active_scope
             return unless scope
 
+            # libddwaf expects db system to be lowercase,
+            # in case of sqlite adapter, libddwaf expects 'sqlite' as db system
+            db_system = adapter_name.downcase
+            db_system = 'sqlite' if db_system == 'sqlite3'
+
             ephemeral_data = {
               'server.db.statement' => sql,
-              'server.db.system' => adapter_name.downcase.gsub(/\d{1}\z/, '')
+              'server.db.system' => db_system
             }
 
             waf_timeout = Datadog.configuration.appsec.waf_timeout

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

@@ -7,7 +7,7 @@ module Datadog
   module AppSec
     module Contrib
       module ActiveRecord
-        # Description of ActiveRecord integration
+        # This class provides helper methods that are used when patching ActiveRecord
         class Integration
           include Datadog::AppSec::Contrib::Integration
 

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

@@ -23,27 +23,27 @@ def target_version
 
           def patch
             ActiveSupport.on_load :active_record do
-              if defined? ::ActiveRecord::ConnectionAdapters::SQLite3Adapter
-                ::ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(Patcher.prepended_class_name(:sqlite3))
+              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)
               end
 
-              if defined? ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-                ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Patcher.prepended_class_name(:postgresql))
+              if defined?(::ActiveRecord::ConnectionAdapters::Mysql2Adapter)
+                ::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(instrumentation_module)
               end
 
-              if defined? ::ActiveRecord::ConnectionAdapters::Mysql2Adapter
-                ::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(Patcher.prepended_class_name(:mysql2))
-              end
-            end
-          end
+              if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+                unless defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
+                  instrumentation_module = Instrumentation::ExecuteAndClearAdapterPatch
+                end
 
-          def prepended_class_name(adapter_name)
-            if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
-              Instrumentation::InternalExecQueryAdapterPatch
-            elsif adapter_name == :postgresql && !defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
-              Instrumentation::ExecuteAndClearAdapterPatch
-            else
-              Instrumentation::ExecQueryAdapterPatch
+                ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(instrumentation_module)
+              end
             end
           end
         end

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

@@ -3,7 +3,7 @@ module Datadog
     module Contrib
       module ActiveRecord
         module Instrumentation
-          def self?.detect_sqli: (String sql, String adapter_name) -> void
+          def self?.detect_sql_injection: (String sql, String adapter_name) -> void
 
           module InternalExecQueryAdapterPatch
             def internal_exec_query: (String sql, *untyped args, **untyped rest) -> untyped

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

@@ -10,8 +10,6 @@ module Datadog
           def self?.target_version: () -> Gem::Version?
 
           def self?.patch: () -> void
-
-          def self?.prepended_class_name: (Symbol adapter_name) -> class
         end
       end
     end

spec/datadog/appsec/contrib/active_record/mysql2_adapter_spec.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require 'datadog/appsec/spec_helper'
+require 'active_record'
+
+require 'spec/datadog/tracing/contrib/rails/support/deprecation'
+
+if PlatformHelpers.jruby?
+  require 'activerecord-jdbc-adapter'
+else
+  require 'mysql2'
+end
+
+RSpec.describe 'AppSec ActiveRecord integration for Mysql2 adapter' do
+  let(:telemetry) { instance_double(Datadog::Core::Telemetry::Component) }
+  let(:ruleset) { Datadog::AppSec::Processor::RuleLoader.load_rules(ruleset: :recommended, telemetry: telemetry) }
+  let(:processor) { Datadog::AppSec::Processor.new(ruleset: ruleset, telemetry: telemetry) }
+  let(:context) { processor.new_context }
+
+  let(:span) { Datadog::Tracing::SpanOperation.new('root') }
+  let(:trace) { Datadog::Tracing::TraceOperation.new }
+
+  let!(:user_class) do
+    stub_const('User', Class.new(ActiveRecord::Base)).tap do |klass|
+      klass.establish_connection(db_config)
+
+      klass.connection.create_table 'users', force: :cascade do |t|
+        t.string :name, null: false
+        t.string :email, null: false
+        t.timestamps
+      end
+
+      # prevent internal sql requests from showing up
+      klass.count
+      klass.first
+    end
+  end
+
+  let(:db_config) do
+    {
+      adapter: 'mysql2',
+      database: ENV.fetch('TEST_MYSQL_DB', 'mysql'),
+      host: ENV.fetch('TEST_MYSQL_HOST', '127.0.0.1'),
+      password: ENV.fetch('TEST_MYSQL_ROOT_PASSWORD', 'root'),
+      port: ENV.fetch('TEST_MYSQL_PORT', '3306')
+    }
+  end
+
+  before do
+    Datadog.configure do |c|
+      c.appsec.enabled = true
+      c.appsec.instrument :active_record
+    end
+
+    Datadog::AppSec::Scope.activate_scope(trace, span, processor)
+
+    raise_on_rails_deprecation!
+  end
+
+  after do
+    Datadog.configuration.reset!
+
+    Datadog::AppSec::Scope.deactivate_scope
+    processor.finalize
+  end
+
+  it 'calls waf with correct arguments when querying using .where' do
+    expect(Datadog::AppSec.active_scope.processor_context).to(
+      receive(:run).with(
+        {},
+        {
+          'server.db.statement' => "SELECT `users`.* FROM `users` WHERE `users`.`name` = 'Bob'",
+          'server.db.system' => 'mysql2'
+        },
+        Datadog.configuration.appsec.waf_timeout
+      ).and_call_original
+    )
+
+    User.where(name: 'Bob').to_a
+  end
+
+  it 'calls waf with correct arguments when querying using .find_by_sql' do
+    expect(Datadog::AppSec.active_scope.processor_context).to(
+      receive(:run).with(
+        {},
+        {
+          'server.db.statement' => "SELECT * FROM users WHERE name = 'Bob'",
+          'server.db.system' => 'mysql2'
+        },
+        Datadog.configuration.appsec.waf_timeout
+      ).and_call_original
+    )
+
+    User.find_by_sql("SELECT * FROM users WHERE name = 'Bob'").to_a
+  end
+
+  it 'adds an event to processor context if waf status is :match' do
+    expect(Datadog::AppSec.active_scope.processor_context).to(
+      receive(:run).and_return(instance_double(Datadog::AppSec::WAF::Result, status: :match, actions: {}))
+    )
+
+    expect(Datadog::AppSec.active_scope.processor_context.events).to receive(:<<).and_call_original
+
+    User.where(name: 'Bob').to_a
+  end
+end