DataDog · uurien · Nov 8, 2023 · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023
@@ -29,6 +29,7 @@ jobs:
  uses: actions/checkout@v3
  with:
  repository: 'DataDog/system-tests'
+ ref: 'ugaitz/modify-nextjs-initialization'
 
  - name: Checkout dd-trace-js
  uses: actions/checkout@v3

@@ -0,0 +1,50 @@
+let i = 0
+const loadMap = new Map()
+const MAX_REWRITE_TIME_MS = 100
+let communicationPort
+
+export async function initialize({ port }) {
+ communicationPort = port
+ communicationPort.on('message', (msg) => {
+ if (msg.id !== undefined) {
+ const cb = loadMap.get(msg.id)
+ if (cb) {
+ cb({ id: msg.id, source: msg.source })
+ }
+ }
+ })
+}
+
+export async function load(url, context, nextLoad) {
+ const nextLoadResult = await nextLoad(url, context)
+ if (nextLoadResult.source) {
+ const id = i++
+ return new Promise((resolve) => {
+ let resolved = false
+ const timeout = setTimeout(() => {
+ if (!resolved) {
+ resolved = true
+ loadMap.delete(id)
+ resolve(nextLoadResult)
+ }
+ }, MAX_REWRITE_TIME_MS)
+ timeout.unref && timeout.unref()
+
+ loadMap.set(id, ({ source }) => {
+ if (!resolved) {
+ resolved = true
+ nextLoadResult.source = Buffer.from(source)
+ clearTimeout(timeout)
+ resolve(nextLoadResult)
+ }
+ })
+ communicationPort.postMessage({
+ id,
+ url,
+ source: nextLoadResult.source
+ })
+ })
+
+ }
+ return nextLoadResult
+}
@@ -1,9 +1,15 @@
+'use strict'
+
+const { isMainThread, MessageChannel } = require('worker_threads')
+const { pathToFileURL } = require('url')
+const dc = require('dc-polyfill')
+const { register } = require('module')
+
 const vulnerabilityReporter = require('./vulnerability-reporter')
 const { enableAllAnalyzers, disableAllAnalyzers } = require('./analyzers')
 const web = require('../../plugins/util/web')
 const { storage } = require('../../../../datadog-core')
 const overheadController = require('./overhead-controller')
-const dc = require('dc-polyfill')
 const iastContextFunctions = require('./iast-context')
 const {
  enableTaintTracking,
@@ -14,15 +20,21 @@ const {
 } = require('./taint-tracking')
 const { IAST_ENABLED_TAG_KEY } = require('./tags')
 const iastTelemetry = require('./telemetry')
+const iastLog = require('./iast-log')
 
 // TODO Change to `apm:http:server:request:[start|close]` when the subscription
 // order of the callbacks can be enforce
 const requestStart = dc.channel('dd-trace:incomingHttpRequestStart')
 const requestClose = dc.channel('dd-trace:incomingHttpRequestEnd')
+const sourcePreloadChannel = dc.channel('datadog:esm:source:preload')
 const iastResponseEnd = dc.channel('datadog:iast:response-end')
 
+let registered = false
+let esmRewriterEnabled = false
+
 function enable (config, _tracer) {
  iastTelemetry.configure(config, config.iast && config.iast.telemetryVerbosity)
+ enableEsmRewriter()
  enableAllAnalyzers(config)
  enableTaintTracking(config.iast, iastTelemetry.verbosity)
  requestStart.subscribe(onIncomingHttpRequestStart)
@@ -32,8 +44,47 @@ function enable (config, _tracer) {
  vulnerabilityReporter.start(config, _tracer)
 }
 
+function enableEsmRewriter () {
+ if (register && isMainThread) {
+ // When we are not in the main app thread it could be an ESM loaders thread,
+ // we don't want to register the rewriter twice
+ esmRewriterEnabled = true
+
+ if (registered) return
+
+ try {
+ const { port1, port2 } = new MessageChannel()
+ registered = true
+
+ port1.on('message', (originalData) => {
+ if (esmRewriterEnabled && sourcePreloadChannel.hasSubscribers) {
+ const data = { ...originalData }
+ sourcePreloadChannel.publish(data)
+ port1.postMessage(data)
+ } else {
+ port1.postMessage(originalData)
+ }
+ })
+ port1.unref()
+
+ register('./hooks.mjs', {
+ parentURL: pathToFileURL(__filename),
+ data: { port: port2 },
+ transferList: [port2]
+ })
+ } catch (e) {
+ iastLog.errorAndPublish(e)
+ }
+ }
+}
+
+function disableEsmRewriter () {
+ esmRewriterEnabled = false
+}
+
 function disable () {
  iastTelemetry.stop()
+ disableEsmRewriter()
  disableAllAnalyzers()
  disableTaintTracking()
  overheadController.finishGlobalContext()

@@ -10,6 +10,8 @@ const { getRewriteFunction } = require('./rewriter-telemetry')
 const dc = require('dc-polyfill')
 
 const hardcodedSecretCh = dc.channel('datadog:secrets:result')
+const sourcePreloadChannel = dc.channel('datadog:esm:source:preload')
+
 let rewriter
 let getPrepareStackTrace
 
@@ -81,19 +83,12 @@ function getPrepareStackTraceAccessor () {
 }
 
 function getCompileMethodFn (compileMethod) {
- const rewriteFn = getRewriteFunction(rewriter)
  return function (content, filename) {
  try {
- if (isPrivateModule(filename) && isNotLibraryFile(filename)) {
- const rewritten = rewriteFn(content, filename)
+ const rewritten = rewriteContent(content, filename)
 
- if (rewritten?.literalsResult && hardcodedSecretCh.hasSubscribers) {
- hardcodedSecretCh.publish(rewritten.literalsResult)
- }
-
- if (rewritten?.content) {
- return compileMethod.apply(this, [rewritten.content, filename])
- }
+ if (rewritten?.content) {
+ return compileMethod.apply(this, [rewritten.content, filename])
  }
  } catch (e) {
  iastLog.error(`Error rewriting ${filename}`)
@@ -103,6 +98,30 @@ function getCompileMethodFn (compileMethod) {
  }
 }
 
+function rewriteContent (content, filename) {
+ const rewriteFn = getRewriteFunction(rewriter)
+
+ if (isPrivateModule(filename) && isNotLibraryFile(filename)) {
+ if (typeof content !== 'string') {
+ content = content.toString()
+ }
+ const rewritten = rewriteFn(content, filename)
+
+ if (rewritten?.literalsResult && hardcodedSecretCh.hasSubscribers) {
+ hardcodedSecretCh.publish(rewritten.literalsResult)
+ }
+ return rewritten
+ }
+ return undefined
+}
+
+function rewriteESMModule (data) {
+ const rewritten = rewriteContent(Buffer.from(data.source), data.url)
+ if (rewritten) {
+ data.source = Buffer.from(rewritten.content)
+ }
+}
+
 function enableRewriter (telemetryVerbosity) {
  try {
  const rewriter = getRewriter(telemetryVerbosity)
@@ -112,6 +131,7 @@ function enableRewriter (telemetryVerbosity) {
  Object.defineProperty(global.Error, 'prepareStackTrace', getPrepareStackTraceAccessor())
  }
  shimmer.wrap(Module.prototype, '_compile', compileMethod => getCompileMethodFn(compileMethod))
+ sourcePreloadChannel.subscribe(rewriteESMModule)
  }
  } catch (e) {
  iastLog.error('Error enabling TaintTracking Rewriter')
@@ -121,6 +141,7 @@ function enableRewriter (telemetryVerbosity) {
 
 function disableRewriter () {
  shimmer.unwrap(Module.prototype, '_compile')
+ sourcePreloadChannel.unsubscribe(rewriteESMModule)
  Error.prepareStackTrace = originalPrepareStackTrace
 }
 

@@ -1,5 +1,6 @@
 'use strict'
 
+const { isMainThread } = require('worker_threads')
 const log = require('../log')
 const RuleManager = require('./rule_manager')
 const remoteConfig = require('./remote_config')
@@ -29,7 +30,7 @@ let isEnabled = false
 let config
 
 function enable (_config) {
- if (isEnabled) return
+ if (isEnabled || !isMainThread) return
 
  try {
  appsecTelemetry.enable(_config.telemetry)

@@ -127,7 +127,10 @@ describe('IAST Index', () => {
  }
  mockIast = proxyquire('../../../src/appsec/iast', {
  './vulnerability-reporter': mockVulnerabilityReporter,
- './overhead-controller': mockOverheadController
+ './overhead-controller': mockOverheadController,
+ 'module': {
+ register: sinon.stub()
+ }
  })
  })
 

@@ -0,0 +1,3 @@
+import { sum } from './esm-two.mjs'
+
+sum()
@@ -0,0 +1,3 @@
+export function sum (a, b) {
+ return a + b
+}
@@ -1,7 +1,13 @@
 'use strict'
-
+const path = require('path')
+const os = require('os')
+const { pathToFileURL } = require('node:url')
 const { expect } = require('chai')
+const dc = require('dc-polyfill')
 const proxyquire = require('proxyquire')
+const semver = require('semver')
+const iast = require('../../../../src/appsec/iast')
+const Config = require('../../../../src/config')
 
 describe('IAST Rewriter', () => {
  it('Addon should return a rewritter instance', () => {
@@ -13,12 +19,7 @@ describe('IAST Rewriter', () => {
  })
 
  describe('Enabling rewriter', () => {
- let rewriter, iastTelemetry
-
- const shimmer = {
- wrap: sinon.spy(),
- unwrap: sinon.spy()
- }
+ let rewriter, iastTelemetry, shimmer
 
  class Rewriter {
  rewrite (content, filename) {
@@ -33,7 +34,11 @@ describe('IAST Rewriter', () => {
 
  beforeEach(() => {
  iastTelemetry = {
- add: sinon.spy()
+ add: sinon.stub()
+ }
+ shimmer = {
+ wrap: sinon.stub(),
+ unwrap: sinon.stub()
  }
  rewriter = proxyquire('../../../../src/appsec/iast/taint-tracking/rewriter', {
  '@datadog/native-iast-rewriter': { Rewriter, getPrepareStackTrace: function () {} },
@@ -43,7 +48,7 @@ describe('IAST Rewriter', () => {
  })
 
  afterEach(() => {
- sinon.restore()
+ rewriter.disableRewriter()
  })
 
  it('Should wrap module compile method on taint tracking enable', () => {
@@ -117,4 +122,62 @@ describe('IAST Rewriter', () => {
  process.env.NODE_OPTIONS = origNodeOptions
  })
  })
+
+ describe('ESM rewriter hooks', () => {
+ const sourcePreloadChannel = dc.channel('datadog:esm:source:preload')
+
+ beforeEach(() => {
+ iast.enable(new Config({
+ experimental: {
+ iast: {
+ enabled: true
+ }
+ }
+ }))
+ })
+
+ it('should rewrite in channel object data', () => {
+ const source = Buffer.from(`export function test(b,c) { return b + c }`)
+ const preloadData = {
+ source,
+ url: pathToFileURL(path.join(os.tmpdir(), 'test1.mjs')).toString()
+ }
+
+ sourcePreloadChannel.publish(preloadData)
+
+ expect(preloadData.source.toString()).to.contain('_ddiast.plusOperator(')
+ })
+
+ it('should not rewrite in node_modules file', () => {
+ const source = Buffer.from(`export function test(b,c) { return b + c }`)
+ const preloadData = {
+ source,
+ url: pathToFileURL(path.join(os.tmpdir(), 'node_modules', 'test1.mjs')).toString()
+ }
+
+ sourcePreloadChannel.publish(preloadData)
+
+ expect(preloadData.source).to.be.equal(source)
+ })
+
+ if (semver.satisfies(process.versions.node, '>=20.6.0')) {
+ it('should publish events when module is imported', (done) => {
+ const channelCallback = sinon.stub()
+ sourcePreloadChannel.subscribe(channelCallback)
+
+ import('./resources/esm-one.mjs').then(() => {
+ expect(channelCallback).to.have.been.calledTwice
+ expect(channelCallback.firstCall.args[0].url).to.contain('esm-one.mjs')
+ expect(channelCallback.secondCall.args[0].url).to.contain('esm-two.mjs')
+
+ sourcePreloadChannel.unsubscribe(channelCallback)
+ done()
+ }).catch(done)
+ })
+ }
+
+ afterEach(() => {
+ iast.disable()
+ })
+ })
 })