Skip to content

Commit e32ebd7

Browse files
uurienrochdev
uurien
authored and
rochdev
committed
Remove Async Resource from knex (#6241)
* Remove Async Resource from knex * Tiny improvements
1 parent 6d107fd commit e32ebd7

File tree

5 files changed

+86
-28
lines changed

5 files changed

+86
-28
lines changed

‎packages/datadog-instrumentations/src/knex.js

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
'use strict'
22

3-
const { addHook, channel, AsyncResource } = require('./helpers/instrument')
3+
const { addHook, channel } = require('./helpers/instrument')
44
const { wrapThen } = require('./helpers/promise')
55
const shimmer = require('../../datadog-shimmer')
66

77
const startRawQueryCh = channel('datadog:knex:raw:start')
8+
const rawQuerySubscribes = channel('datadog:knex:raw:subscribes')
89
const finishRawQueryCh = channel('datadog:knex:raw:finish')
910

1011
patch('lib/query/builder.js')
@@ -22,8 +23,8 @@ function patch (file) {
2223
})
2324
}
2425

25-
function finish () {
26-
finishRawQueryCh.publish()
26+
function finish (context, cb) {
27+
finishRawQueryCh.runStores(context, cb)
2728
}
2829

2930
addHook({
@@ -43,21 +44,18 @@ addHook({
4344
return raw.apply(this, arguments)
4445
}
4546

46-
const asyncResource = new AsyncResource('bound-anonymous-fn')
47-
48-
return asyncResource.runInAsyncScope(() => {
49-
startRawQueryCh.publish({ sql, dialect: this.dialect })
50-
47+
const context = { sql, dialect: this.dialect }
48+
return startRawQueryCh.runStores(context, () => {
5149
const rawResult = raw.apply(this, arguments)
5250
shimmer.wrap(rawResult, 'then', originalThen => function () {
53-
return asyncResource.runInAsyncScope(() => {
54-
arguments[0] = wrapCallbackWithFinish(arguments[0], finish)
55-
if (arguments[1]) arguments[1] = wrapCallbackWithFinish(arguments[1], finish)
51+
return rawQuerySubscribes.runStores(context, () => {
52+
arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
53+
if (arguments[1]) arguments[1] = wrapCallbackWithFinish(arguments[1], finish, context)
5654

5755
const originalThenResult = originalThen.apply(this, arguments)
5856

5957
shimmer.wrap(originalThenResult, 'catch', originalCatch => function () {
60-
arguments[0] = wrapCallbackWithFinish(arguments[0], finish)
58+
arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
6159
return originalCatch.apply(this, arguments)
6260
})
6361

@@ -66,23 +64,23 @@ addHook({
6664
})
6765

6866
shimmer.wrap(rawResult, 'asCallback', originalAsCallback => function () {
69-
return asyncResource.runInAsyncScope(() => {
70-
arguments[0] = wrapCallbackWithFinish(arguments[0], finish)
67+
return rawQuerySubscribes.runStores(context, () => {
68+
arguments[0] = wrapCallbackWithFinish(arguments[0], finish, context)
7169
return originalAsCallback.apply(this, arguments)
7270
})
7371
})
7472

7573
return rawResult
7674
})
7775
})
76+
7877
return Knex
7978
})
8079

81-
function wrapCallbackWithFinish (callback, finish) {
80+
function wrapCallbackWithFinish (callback, finish, context) {
8281
if (typeof callback !== 'function') return callback
8382

8483
return shimmer.wrapFunction(callback, callback => function () {
85-
finish()
86-
callback.apply(this, arguments)
84+
finish(context, () => callback.apply(this, arguments))
8785
})
8886
}

‎packages/dd-trace/src/appsec/iast/analyzers/sql-injection-analyzer.js

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,16 @@ class SqlInjectionAnalyzer extends StoredInjectionAnalyzer {
3333
this.addSub('datadog:mysql:pool:query:start', ({ sql }) => this.setStoreAndAnalyze(sql, 'MYSQL'))
3434
this.addSub('datadog:mysql:pool:query:finish', () => this.returnToParentStore())
3535

36-
this.addSub('datadog:knex:raw:start', ({ sql, dialect: knexDialect }) => {
36+
this.addBind('datadog:knex:raw:start', (context) => {
37+
const { sql, dialect: knexDialect } = context
3738
const dialect = this.normalizeKnexDialect(knexDialect)
38-
this.setStoreAndAnalyze(sql, dialect)
39+
const currentStore = this.getStoreAndAnalyze(sql, dialect)
40+
context.currentStore = currentStore
41+
return currentStore
3942
})
40-
this.addSub('datadog:knex:raw:finish', () => this.returnToParentStore())
43+
44+
this.addBind('datadog:knex:raw:subscribes', ({ currentStore }) => currentStore)
45+
this.addBind('datadog:knex:raw:finish', ({ currentStore }) => currentStore?.sqlParentStore)
4146
}
4247

4348
setStoreAndAnalyze (query, dialect) {
@@ -57,8 +62,7 @@ class SqlInjectionAnalyzer extends StoredInjectionAnalyzer {
5762
}
5863
}
5964

60-
returnToParentStore () {
61-
const store = storage('legacy').getStore()
65+
returnToParentStore (store = storage('legacy').getStore()) {
6266
if (store && store.sqlParentStore) {
6367
storage('legacy').enterWith(store.sqlParentStore)
6468
}

‎packages/dd-trace/test/appsec/iast/analyzers/resources/knex-sql-injection-methods.js

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,25 @@ function executeKnexNestedRawQueryAsCallback (knex, taintedSql, sqlToFail, cb) {
3535
})
3636
}
3737

38+
async function executeKnexAsyncNestedRawQuery (knex, taintedSql, notTaintedSql) {
39+
await knex.raw(notTaintedSql)
40+
await knex.raw(taintedSql)
41+
}
42+
43+
async function executeKnexAsyncNestedRawQueryAsAsyncTryCatch (knex, taintedSql, sqlToFail) {
44+
try {
45+
await knex.raw(sqlToFail)
46+
} catch (e) {
47+
await knex.raw(taintedSql)
48+
}
49+
}
50+
3851
module.exports = {
3952
executeKnexRawQuery,
4053
executeKnexNestedRawQuery,
54+
executeKnexAsyncNestedRawQuery,
4155
executeKnexNestedRawQueryOnRejectedInThen,
4256
executeKnexNestedRawQueryWitCatch,
43-
executeKnexNestedRawQueryAsCallback
57+
executeKnexNestedRawQueryAsCallback,
58+
executeKnexAsyncNestedRawQueryAsAsyncTryCatch
4459
}

‎packages/dd-trace/test/appsec/iast/analyzers/sql-injection-analyzer.knex.plugin.spec.js

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ describe('sql-injection-analyzer with knex', () => {
1414
withVersions('knex', 'knex', knexVersion => {
1515
if (!semver.satisfies(knexVersion, '>=2')) return
1616

17-
withVersions('pg', 'pg', pgVersion => {
17+
withVersions('pg', 'pg', () => {
1818
let knex
1919

2020
prepareTestServerForIast('knex + pg',
@@ -88,6 +88,26 @@ describe('sql-injection-analyzer with knex', () => {
8888
})
8989
})
9090

91+
describe('nested raw query - using async instead of then', () => {
92+
testThatRequestHasVulnerability(() => {
93+
const store = storage('legacy').getStore()
94+
const iastCtx = iastContextFunctions.getIastContext(store)
95+
96+
let taintedSql = 'SELECT 1'
97+
taintedSql = newTaintedString(iastCtx, taintedSql, 'param', 'Request')
98+
99+
const notTaintedSql = 'SELECT 1'
100+
101+
return queryMethods.executeKnexAsyncNestedRawQuery(knex, taintedSql, notTaintedSql)
102+
}, 'SQL_INJECTION', {
103+
occurrences: 1,
104+
location: {
105+
path: 'knex-sql-injection-methods.js',
106+
line: 40
107+
}
108+
})
109+
})
110+
91111
describe('nested raw query - onRejected as then argument', () => {
92112
testThatRequestHasVulnerability(() => {
93113
const store = storage('legacy').getStore()
@@ -128,6 +148,26 @@ describe('sql-injection-analyzer with knex', () => {
128148
})
129149
})
130150

151+
describe('nested raw query - async try catch', () => {
152+
testThatRequestHasVulnerability(() => {
153+
const store = storage('legacy').getStore()
154+
const iastCtx = iastContextFunctions.getIastContext(store)
155+
156+
let taintedSql = 'SELECT 1'
157+
taintedSql = newTaintedString(iastCtx, taintedSql, 'param', 'Request')
158+
159+
const sqlToFail = 'SELECT * FROM NON_EXISTSING_TABLE'
160+
161+
return queryMethods.executeKnexAsyncNestedRawQueryAsAsyncTryCatch(knex, taintedSql, sqlToFail)
162+
}, 'SQL_INJECTION', {
163+
occurrences: 1,
164+
location: {
165+
path: 'knex-sql-injection-methods.js',
166+
line: 47
167+
}
168+
})
169+
})
170+
131171
describe('nested raw query - asCallback', () => {
132172
testThatRequestHasVulnerability(() => {
133173
return new Promise((resolve, reject) => {

‎packages/dd-trace/test/appsec/iast/analyzers/sql-injection-analyzer.spec.js

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,21 @@ describe('sql-injection-analyzer', () => {
6464
sqlInjectionAnalyzer.configure(true)
6565

6666
it('should subscribe to mysql, mysql2 and pg start query channel', () => {
67-
expect(sqlInjectionAnalyzer._subscriptions).to.have.lengthOf(9)
67+
expect(sqlInjectionAnalyzer._subscriptions).to.have.lengthOf(7)
6868
expect(sqlInjectionAnalyzer._subscriptions[0]._channel.name).to.equals('apm:mysql:query:start')
6969
expect(sqlInjectionAnalyzer._subscriptions[1]._channel.name).to.equals('datadog:mysql2:outerquery:start')
7070
expect(sqlInjectionAnalyzer._subscriptions[2]._channel.name).to.equals('apm:pg:query:start')
7171
expect(sqlInjectionAnalyzer._subscriptions[3]._channel.name).to.equals('datadog:sequelize:query:finish')
7272
expect(sqlInjectionAnalyzer._subscriptions[4]._channel.name).to.equals('datadog:pg:pool:query:finish')
7373
expect(sqlInjectionAnalyzer._subscriptions[5]._channel.name).to.equals('datadog:mysql:pool:query:start')
7474
expect(sqlInjectionAnalyzer._subscriptions[6]._channel.name).to.equals('datadog:mysql:pool:query:finish')
75-
expect(sqlInjectionAnalyzer._subscriptions[7]._channel.name).to.equals('datadog:knex:raw:start')
76-
expect(sqlInjectionAnalyzer._subscriptions[8]._channel.name).to.equals('datadog:knex:raw:finish')
7775

78-
expect(sqlInjectionAnalyzer._bindings).to.have.lengthOf(2)
76+
expect(sqlInjectionAnalyzer._bindings).to.have.lengthOf(5)
7977
expect(sqlInjectionAnalyzer._bindings[0]._channel.name).to.equals('datadog:sequelize:query:start')
8078
expect(sqlInjectionAnalyzer._bindings[1]._channel.name).to.equals('datadog:pg:pool:query:start')
79+
expect(sqlInjectionAnalyzer._bindings[2]._channel.name).to.equals('datadog:knex:raw:start')
80+
expect(sqlInjectionAnalyzer._bindings[3]._channel.name).to.equals('datadog:knex:raw:subscribes')
81+
expect(sqlInjectionAnalyzer._bindings[4]._channel.name).to.equals('datadog:knex:raw:finish')
8182
})
8283

8384
it('should not detect vulnerability when no query', () => {

0 commit comments

Comments
 (0)