Skip to content

Commit d7bdb68

Browse files
committed
instrument fastify reply header
1 parent b053322 commit d7bdb68

File tree

9 files changed

+80
-43
lines changed

9 files changed

+80
-43
lines changed

packages/datadog-instrumentations/src/fastify.js

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ const queryParamsReadCh = channel('datadog:fastify:query-params:finish')
1111
const cookieParserReadCh = channel('datadog:fastify-cookie:read:finish')
1212
const responsePayloadReadCh = channel('datadog:fastify:response:finish')
1313
const pathParamsReadCh = channel('datadog:fastify:path-params:finish')
14+
const finishSetHeaderCh = channel('datadog:fastify:set-header:finish')
1415

1516
const parsingResources = new WeakMap()
1617
const cookiesPublished = new WeakSet()
@@ -257,3 +258,18 @@ addHook({ name: 'fastify', versions: ['2'] }, fastify => {
257258
addHook({ name: 'fastify', versions: ['1'] }, fastify => {
258259
return shimmer.wrapFunction(fastify, fastify => wrapFastify(fastify, false))
259260
})
261+
262+
addHook({ name: 'fastify', file: 'lib/reply.js', versions: ['>=3'] }, Reply => {
263+
264+
shimmer.wrap(Reply.prototype, 'header', header => function (key, value = '') {
265+
const result = header.apply(this, arguments)
266+
267+
if (finishSetHeaderCh.hasSubscribers && key && value) {
268+
finishSetHeaderCh.publish({ name: key, value, res: this.raw })
269+
}
270+
271+
return result
272+
})
273+
274+
return Reply
275+
})

packages/datadog-instrumentations/src/http/server.js

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -112,14 +112,6 @@ function wrapWriteHead (writeHead) {
112112
// this doesn't support explicit duplicate headers, but it's an edge case
113113
const responseHeaders = Object.assign(this.getHeaders(), obj)
114114

115-
// Publish all headers from writeHead to ensure frameworks like Fastify
116-
// that bypass setHeader() are still captured for analysis
117-
if (finishSetHeaderCh.hasSubscribers) {
118-
for (const [name, value] of Object.entries(responseHeaders)) {
119-
finishSetHeaderCh.publish({ name, value, res: this })
120-
}
121-
}
122-
123115
startWriteHeadCh.publish({
124116
req: this.req,
125117
res: this,

packages/dd-trace/src/appsec/iast/analyzers/hsts-header-missing-analyzer.js

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ class HstsHeaderMissingAnalyzer extends MissingHeaderAnalyzer {
1010
super(HSTS_HEADER_MISSING, HSTS_HEADER_NAME)
1111
}
1212

13-
_isVulnerableFromRequestAndResponse (req, res) {
14-
const headerValues = this._getHeaderValues(res, HSTS_HEADER_NAME)
13+
_isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
14+
const headerValues = this._getHeaderValues(res, storedHeaders, HSTS_HEADER_NAME)
1515
return this._isHttpsProtocol(req) && (
1616
headerValues.length === 0 ||
1717
headerValues.some(headerValue => !this._isHeaderValid(headerValue))

packages/dd-trace/src/appsec/iast/analyzers/missing-header-analyzer.js

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ class MissingHeaderAnalyzer extends Analyzer {
2828
}, (data) => this.analyze(data))
2929
}
3030

31-
_getHeaderValues (res, headerName) {
32-
const headerValue = res.getHeader(headerName)
31+
_getHeaderValues (res, storedHeaders, headerName) {
32+
const headerValue = res.getHeader(headerName) || storedHeaders[headerName]
3333
if (Array.isArray(headerValue)) {
3434
return headerValue
3535
}
@@ -46,8 +46,8 @@ class MissingHeaderAnalyzer extends Analyzer {
4646
return `${type}:${this.config.tracerConfig.service}`
4747
}
4848

49-
_getEvidence ({ res }) {
50-
const headerValues = this._getHeaderValues(res, this.headerName)
49+
_getEvidence ({ res, storedHeaders }) {
50+
const headerValues = this._getHeaderValues(res, storedHeaders, this.headerName)
5151
let value
5252
if (headerValues.length === 1) {
5353
value = headerValues[0]
@@ -57,19 +57,19 @@ class MissingHeaderAnalyzer extends Analyzer {
5757
return { value }
5858
}
5959

60-
_isVulnerable ({ req, res }, context) {
61-
if (!IGNORED_RESPONSE_STATUS_LIST.has(res.statusCode) && this._isResponseHtml(res)) {
62-
return this._isVulnerableFromRequestAndResponse(req, res)
60+
_isVulnerable ({ req, res, storedHeaders }, context) {
61+
if (!IGNORED_RESPONSE_STATUS_LIST.has(res.statusCode) && this._isResponseHtml(res, storedHeaders)) {
62+
return this._isVulnerableFromRequestAndResponse(req, res, storedHeaders)
6363
}
6464
return false
6565
}
6666

67-
_isVulnerableFromRequestAndResponse (req, res) {
67+
_isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
6868
return false
6969
}
7070

71-
_isResponseHtml (res) {
72-
const contentTypes = this._getHeaderValues(res, 'content-type')
71+
_isResponseHtml (res, storedHeaders) {
72+
const contentTypes = this._getHeaderValues(res, storedHeaders, 'content-type')
7373
return contentTypes.some(contentType => {
7474
return contentType && HTML_CONTENT_TYPES.some(htmlContentType => {
7575
return htmlContentType === contentType || contentType.startsWith(htmlContentType + ';')

packages/dd-trace/src/appsec/iast/analyzers/set-cookies-header-interceptor.js

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,25 +7,29 @@ class SetCookiesHeaderInterceptor extends Plugin {
77
constructor () {
88
super()
99
this.cookiesInRequest = new WeakMap()
10-
this.addSub('datadog:http:server:response:set-header:finish', ({ name, value, res }) => {
11-
if (name.toLowerCase() === 'set-cookie') {
12-
let allCookies = value
13-
if (typeof value === 'string') {
14-
allCookies = [value]
15-
}
16-
const alreadyCheckedCookies = this._getAlreadyCheckedCookiesInResponse(res)
17-
18-
let location
19-
allCookies.forEach(cookieString => {
20-
if (!alreadyCheckedCookies.includes(cookieString)) {
21-
alreadyCheckedCookies.push(cookieString)
22-
const parsedCookie = this._parseCookie(cookieString, location)
23-
setCookieChannel.publish(parsedCookie)
24-
location = parsedCookie.location
25-
}
26-
})
10+
11+
this.addSub('datadog:http:server:response:set-header:finish', ({ name, value, res }) => this._handleCookies(name, value, res))
12+
this.addSub('datadog:fastify:set-header:finish', ({ name, value, res }) => this._handleCookies(name, value, res))
13+
}
14+
15+
_handleCookies (name, value, res) {
16+
if (name.toLowerCase() === 'set-cookie') {
17+
let allCookies = value
18+
if (typeof value === 'string') {
19+
allCookies = [value]
2720
}
28-
})
21+
const alreadyCheckedCookies = this._getAlreadyCheckedCookiesInResponse(res)
22+
23+
let location
24+
allCookies.forEach(cookieString => {
25+
if (!alreadyCheckedCookies.includes(cookieString)) {
26+
alreadyCheckedCookies.push(cookieString)
27+
const parsedCookie = this._parseCookie(cookieString, location)
28+
setCookieChannel.publish(parsedCookie)
29+
location = parsedCookie.location
30+
}
31+
})
32+
}
2933
}
3034

3135
_parseCookie (cookieString, location) {

packages/dd-trace/src/appsec/iast/analyzers/unvalidated-redirect-analyzer.js

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,10 @@ const {
99
HTTP_REQUEST_PARAMETER
1010
} = require('../taint-tracking/source-types')
1111

12-
const EXCLUDED_PATHS = getNodeModulesPaths('express/lib/response.js')
12+
const EXCLUDED_PATHS = [
13+
getNodeModulesPaths('express/lib/response.js'),
14+
getNodeModulesPaths('fastify/lib/reply.js'),
15+
]
1316

1417
const VULNERABLE_SOURCE_TYPES = new Set([
1518
HTTP_REQUEST_BODY,
@@ -23,6 +26,7 @@ class UnvalidatedRedirectAnalyzer extends InjectionAnalyzer {
2326

2427
onConfigure () {
2528
this.addSub('datadog:http:server:response:set-header:finish', ({ name, value }) => this.analyze(name, value))
29+
this.addSub('datadog:fastify:set-header:finish', ({ name, value }) => this.analyze(name, value))
2630
}
2731

2832
analyze (name, value) {

packages/dd-trace/src/appsec/iast/analyzers/xcontenttype-header-missing-analyzer.js

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ class XcontenttypeHeaderMissingAnalyzer extends MissingHeaderAnalyzer {
1010
super(XCONTENTTYPE_HEADER_MISSING, XCONTENTTYPEOPTIONS_HEADER_NAME)
1111
}
1212

13-
_isVulnerableFromRequestAndResponse (req, res) {
14-
const headerValues = this._getHeaderValues(res, XCONTENTTYPEOPTIONS_HEADER_NAME)
13+
_isVulnerableFromRequestAndResponse (req, res, storedHeaders) {
14+
const headerValues = this._getHeaderValues(res, storedHeaders, XCONTENTTYPEOPTIONS_HEADER_NAME)
1515
return headerValues.length === 0 || headerValues.some(headerValue => headerValue.trim().toLowerCase() !== 'nosniff')
1616
}
1717
}

packages/dd-trace/src/appsec/iast/index.js

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ const { IAST_ENABLED_TAG_KEY } = require('./tags')
1616
const iastTelemetry = require('./telemetry')
1717
const { enable: enableFsPlugin, disable: disableFsPlugin, IAST_MODULE } = require('../rasp/fs-plugin')
1818
const securityControls = require('./security-controls')
19+
const { responseWriteHead } = require('../channels')
20+
21+
// Map raw response -> accumulated headers set via writeHead
22+
const collectedResponseHeaders = new WeakMap()
1923

2024
// TODO Change to `apm:http:server:request:[start|close]` when the subscription
2125
// order of the callbacks can be enforce
@@ -33,6 +37,7 @@ function enable (config, _tracer) {
3337
enableTaintTracking(config.iast, iastTelemetry.verbosity)
3438
requestStart.subscribe(onIncomingHttpRequestStart)
3539
requestClose.subscribe(onIncomingHttpRequestEnd)
40+
responseWriteHead.subscribe(onResponseWriteHeadCollect)
3641
overheadController.configure(config.iast)
3742
overheadController.startGlobalContext()
3843
securityControls.configure(config.iast)
@@ -53,6 +58,7 @@ function disable () {
5358
overheadController.finishGlobalContext()
5459
if (requestStart.hasSubscribers) requestStart.unsubscribe(onIncomingHttpRequestStart)
5560
if (requestClose.hasSubscribers) requestClose.unsubscribe(onIncomingHttpRequestEnd)
61+
if (responseWriteHead.hasSubscribers) responseWriteHead.unsubscribe(onResponseWriteHeadCollect)
5662
vulnerabilityReporter.stop()
5763
}
5864

@@ -87,7 +93,13 @@ function onIncomingHttpRequestEnd (data) {
8793
const topContext = web.getContext(data.req)
8894
const iastContext = iastContextFunctions.getIastContext(store, topContext)
8995
if (iastContext?.rootSpan) {
90-
iastResponseEnd.publish(data)
96+
const storedHeaders = collectedResponseHeaders.get(data.res) || {}
97+
98+
iastResponseEnd.publish({ ...data, storedHeaders })
99+
100+
if (storedHeaders) {
101+
collectedResponseHeaders.delete(data.res)
102+
}
91103

92104
const vulnerabilities = iastContext.vulnerabilities
93105
const rootSpan = iastContext.rootSpan
@@ -103,4 +115,13 @@ function onIncomingHttpRequestEnd (data) {
103115
}
104116
}
105117

118+
function onResponseWriteHeadCollect ({ res, responseHeaders = {} }) {
119+
if (!res) return
120+
121+
if (Object.keys(responseHeaders).length) {
122+
collectedResponseHeaders.set(res, responseHeaders)
123+
}
124+
}
125+
126+
106127
module.exports = { enable, disable, onIncomingHttpRequestEnd, onIncomingHttpRequestStart }

packages/dd-trace/test/appsec/index.fastify.plugin.spec.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -515,7 +515,7 @@ withVersions('fastify', 'fastify', '5.4.0', version => {
515515
return agent.close({ ritmReset: false })
516516
})
517517

518-
it.only('should trigger HTTP header instrumentation when setting cookies', async () => {
518+
it('should trigger HTTP header instrumentation when setting cookies', async () => {
519519
const res = await axios.get('/test-cookie')
520520

521521
console.log('Response headers:', res.headers)

0 commit comments

Comments
 (0)