Skip to content

Commit ee4c3a2

Browse files
committed
Simplify AppSec ActiveRecord patcher and split specs by adapter
1 parent 776d405 commit ee4c3a2

File tree

10 files changed

+346
-262
lines changed

10 files changed

+346
-262
lines changed

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,14 @@ def detect_sql_injection(sql, adapter_name)
1212
scope = AppSec.active_scope
1313
return unless scope
1414

15+
# libddwaf expects db system to be lowercase,
16+
# in case of sqlite adapter, libddwaf expects 'sqlite' as db system
17+
db_system = adapter_name.downcase
18+
db_system = 'sqlite' if db_system == 'sqlite3'
19+
1520
ephemeral_data = {
1621
'server.db.statement' => sql,
17-
'server.db.system' => adapter_name.downcase.gsub(/\d{1}\z/, '')
22+
'server.db.system' => db_system
1823
}
1924

2025
waf_timeout = Datadog.configuration.appsec.waf_timeout

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ module Datadog
77
module AppSec
88
module Contrib
99
module ActiveRecord
10-
# Description of ActiveRecord integration
10+
# This class provides helper methods that are used when patching ActiveRecord
1111
class Integration
1212
include Datadog::AppSec::Contrib::Integration
1313

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

+16-16
Original file line numberDiff line numberDiff line change
@@ -23,27 +23,27 @@ def target_version
2323

2424
def patch
2525
ActiveSupport.on_load :active_record do
26-
if defined? ::ActiveRecord::ConnectionAdapters::SQLite3Adapter
27-
::ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(Patcher.prepended_class_name(:sqlite3))
26+
instrumentation_module = if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
27+
Instrumentation::InternalExecQueryAdapterPatch
28+
else
29+
Instrumentation::ExecQueryAdapterPatch
30+
end
31+
32+
if defined?(::ActiveRecord::ConnectionAdapters::SQLite3Adapter)
33+
::ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(instrumentation_module)
2834
end
2935

30-
if defined? ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
31-
::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Patcher.prepended_class_name(:postgresql))
36+
if defined?(::ActiveRecord::ConnectionAdapters::Mysql2Adapter)
37+
::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(instrumentation_module)
3238
end
3339

34-
if defined? ::ActiveRecord::ConnectionAdapters::Mysql2Adapter
35-
::ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(Patcher.prepended_class_name(:mysql2))
36-
end
37-
end
38-
end
40+
if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
41+
unless defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
42+
instrumentation_module = Instrumentation::ExecuteAndClearAdapterPatch
43+
end
3944

40-
def prepended_class_name(adapter_name)
41-
if ::ActiveRecord.gem_version >= Gem::Version.new('7.1')
42-
Instrumentation::InternalExecQueryAdapterPatch
43-
elsif adapter_name == :postgresql && !defined?(::ActiveRecord::ConnectionAdapters::JdbcAdapter)
44-
Instrumentation::ExecuteAndClearAdapterPatch
45-
else
46-
Instrumentation::ExecQueryAdapterPatch
45+
::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(instrumentation_module)
46+
end
4747
end
4848
end
4949
end

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module Datadog
33
module Contrib
44
module ActiveRecord
55
module Instrumentation
6-
def self?.detect_sqli: (String sql, String adapter_name) -> void
6+
def self?.detect_sql_injection: (String sql, String adapter_name) -> void
77

88
module InternalExecQueryAdapterPatch
99
def internal_exec_query: (String sql, *untyped args, **untyped rest) -> untyped

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

-2
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@ module Datadog
1010
def self?.target_version: () -> Gem::Version?
1111

1212
def self?.patch: () -> void
13-
14-
def self?.prepended_class_name: (Symbol adapter_name) -> class
1513
end
1614
end
1715
end

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

-171
This file was deleted.
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
# frozen_string_literal: true
2+
3+
require 'datadog/appsec/spec_helper'
4+
require 'active_record'
5+
6+
require 'spec/datadog/tracing/contrib/rails/support/deprecation'
7+
8+
if PlatformHelpers.jruby?
9+
require 'activerecord-jdbc-adapter'
10+
else
11+
require 'mysql2'
12+
end
13+
14+
RSpec.describe 'AppSec ActiveRecord integration for Mysql2 adapter' do
15+
let(:telemetry) { instance_double(Datadog::Core::Telemetry::Component) }
16+
let(:ruleset) { Datadog::AppSec::Processor::RuleLoader.load_rules(ruleset: :recommended, telemetry: telemetry) }
17+
let(:processor) { Datadog::AppSec::Processor.new(ruleset: ruleset, telemetry: telemetry) }
18+
let(:context) { processor.new_context }
19+
20+
let(:span) { Datadog::Tracing::SpanOperation.new('root') }
21+
let(:trace) { Datadog::Tracing::TraceOperation.new }
22+
23+
let!(:user_class) do
24+
stub_const('User', Class.new(ActiveRecord::Base)).tap do |klass|
25+
klass.establish_connection(db_config)
26+
27+
klass.connection.create_table 'users', force: :cascade do |t|
28+
t.string :name, null: false
29+
t.string :email, null: false
30+
t.timestamps
31+
end
32+
33+
# prevent internal sql requests from showing up
34+
klass.count
35+
klass.first
36+
end
37+
end
38+
39+
let(:db_config) do
40+
{
41+
adapter: 'mysql2',
42+
database: ENV.fetch('TEST_MYSQL_DB', 'mysql'),
43+
host: ENV.fetch('TEST_MYSQL_HOST', '127.0.0.1'),
44+
password: ENV.fetch('TEST_MYSQL_ROOT_PASSWORD', 'root'),
45+
port: ENV.fetch('TEST_MYSQL_PORT', '3306')
46+
}
47+
end
48+
49+
before do
50+
Datadog.configure do |c|
51+
c.appsec.enabled = true
52+
c.appsec.instrument :active_record
53+
end
54+
55+
Datadog::AppSec::Scope.activate_scope(trace, span, processor)
56+
57+
raise_on_rails_deprecation!
58+
end
59+
60+
after do
61+
Datadog.configuration.reset!
62+
63+
Datadog::AppSec::Scope.deactivate_scope
64+
processor.finalize
65+
end
66+
67+
it 'calls waf with correct arguments when querying using .where' do
68+
expect(Datadog::AppSec.active_scope.processor_context).to(
69+
receive(:run).with(
70+
{},
71+
{
72+
'server.db.statement' => "SELECT `users`.* FROM `users` WHERE `users`.`name` = 'Bob'",
73+
'server.db.system' => 'mysql2'
74+
},
75+
Datadog.configuration.appsec.waf_timeout
76+
).and_call_original
77+
)
78+
79+
User.where(name: 'Bob').to_a
80+
end
81+
82+
it 'calls waf with correct arguments when querying using .find_by_sql' do
83+
expect(Datadog::AppSec.active_scope.processor_context).to(
84+
receive(:run).with(
85+
{},
86+
{
87+
'server.db.statement' => "SELECT * FROM users WHERE name = 'Bob'",
88+
'server.db.system' => 'mysql2'
89+
},
90+
Datadog.configuration.appsec.waf_timeout
91+
).and_call_original
92+
)
93+
94+
User.find_by_sql("SELECT * FROM users WHERE name = 'Bob'").to_a
95+
end
96+
97+
it 'adds an event to processor context if waf status is :match' do
98+
expect(Datadog::AppSec.active_scope.processor_context).to(
99+
receive(:run).and_return(instance_double(Datadog::AppSec::WAF::Result, status: :match, actions: {}))
100+
)
101+
102+
expect(Datadog::AppSec.active_scope.processor_context.events).to receive(:<<).and_call_original
103+
104+
User.where(name: 'Bob').to_a
105+
end
106+
end

0 commit comments

Comments
 (0)