dd-trace 4.21.0 → 4.23.0

package/index.d.ts

@@ -567,6 +567,11 @@ export declare interface TracerOptions {
      */
     blockedTemplateJson?: string,
 
+    /**
+     * Specifies a path to a custom blocking template json file for graphql requests
+     */
+    blockedTemplateGraphql?: string,
+
     /**
      * Controls the automated user event tracking configuration
      */

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "dd-trace",
-  "version": "4.21.0",
+  "version": "4.23.0",
   "description": "Datadog APM tracing client for JavaScript",
   "main": "index.js",
   "typings": "index.d.ts",
@@ -68,8 +68,8 @@
     "node": ">=16"
   },
   "dependencies": {
-    "@datadog/native-appsec": "5.0.0",
-    "@datadog/native-iast-rewriter": "2.2.1",
+    "@datadog/native-appsec": "6.0.0",
+    "@datadog/native-iast-rewriter": "2.2.2",
     "@datadog/native-iast-taint-tracking": "1.6.4",
     "@datadog/native-metrics": "^2.0.0",
     "@datadog/pprof": "4.1.0",
@@ -79,7 +79,7 @@
     "crypto-randomuuid": "^1.0.0",
     "dc-polyfill": "^0.1.2",
     "ignore": "^5.2.4",
-    "import-in-the-middle": "^1.4.2",
+    "import-in-the-middle": "^1.7.1",
     "int64-buffer": "^0.1.9",
     "ipaddr.js": "^2.1.0",
     "istanbul-lib-coverage": "3.2.0",
@@ -99,9 +99,9 @@
     "path-to-regexp": "^0.1.2",
     "pprof-format": "^2.0.7",
     "protobufjs": "^7.2.5",
-    "tlhunter-sorted-set": "^0.1.0",
     "retry": "^0.13.1",
-    "semver": "^7.5.4"
+    "semver": "^7.5.4",
+    "tlhunter-sorted-set": "^0.1.0"
   },
   "devDependencies": {
     "@types/node": ">=16",

package/packages/datadog-instrumentations/src/apollo-server-core.js

@@ -0,0 +1,41 @@
+'use strict'
+
+const { AbortController } = require('node-abort-controller')
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const dc = require('dc-polyfill')
+
+const requestChannel = dc.tracingChannel('datadog:apollo-server-core:request')
+
+addHook({ name: 'apollo-server-core', file: 'dist/runHttpQuery.js', versions: ['>3.0.0'] }, runHttpQueryModule => {
+  const HttpQueryError = runHttpQueryModule.HttpQueryError
+
+  shimmer.wrap(runHttpQueryModule, 'runHttpQuery', function wrapRunHttpQuery (originalRunHttpQuery) {
+    return async function runHttpQuery () {
+      if (!requestChannel.start.hasSubscribers) {
+        return originalRunHttpQuery.apply(this, arguments)
+      }
+
+      const abortController = new AbortController()
+      const abortData = {}
+
+      const runHttpQueryResult = requestChannel.tracePromise(
+        originalRunHttpQuery,
+        { abortController, abortData },
+        this,
+        ...arguments)
+
+      const abortPromise = new Promise((resolve, reject) => {
+        abortController.signal.addEventListener('abort', (event) => {
+          // runHttpQuery callbacks are writing the response on resolve/reject.
+          // We should return blocking data in the apollo-server-core HttpQueryError object
+          reject(new HttpQueryError(abortData.statusCode, abortData.message, true, abortData.headers))
+        }, { once: true })
+      })
+
+      return Promise.race([runHttpQueryResult, abortPromise])
+    }
+  })
+
+  return runHttpQueryModule
+})

package/packages/datadog-instrumentations/src/apollo-server.js

@@ -0,0 +1,83 @@
+'use strict'
+
+const { AbortController } = require('node-abort-controller')
+const dc = require('dc-polyfill')
+
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const graphqlMiddlewareChannel = dc.tracingChannel('datadog:apollo:middleware')
+
+const requestChannel = dc.tracingChannel('datadog:apollo:request')
+
+let HeaderMap
+
+function wrapExecuteHTTPGraphQLRequest (originalExecuteHTTPGraphQLRequest) {
+  return async function executeHTTPGraphQLRequest () {
+    if (!HeaderMap || !requestChannel.start.hasSubscribers) {
+      return originalExecuteHTTPGraphQLRequest.apply(this, arguments)
+    }
+
+    const abortController = new AbortController()
+    const abortData = {}
+
+    const graphqlResponseData = requestChannel.tracePromise(
+      originalExecuteHTTPGraphQLRequest,
+      { abortController, abortData },
+      this,
+      ...arguments)
+
+    const abortPromise = new Promise((resolve, reject) => {
+      abortController.signal.addEventListener('abort', (event) => {
+        // This method is expected to return response data
+        // with headers, status and body
+        const headers = new HeaderMap()
+        Object.keys(abortData.headers).forEach(key => {
+          headers.set(key, abortData.headers[key])
+        })
+
+        resolve({
+          headers: headers,
+          status: abortData.statusCode,
+          body: {
+            kind: 'complete',
+            string: abortData.message
+          }
+        })
+      }, { once: true })
+    })
+
+    return Promise.race([abortPromise, graphqlResponseData])
+  }
+}
+
+function apolloExpress4Hook (express4) {
+  shimmer.wrap(express4, 'expressMiddleware', function wrapExpressMiddleware (originalExpressMiddleware) {
+    return function expressMiddleware (server, options) {
+      const originalMiddleware = originalExpressMiddleware.apply(this, arguments)
+
+      return shimmer.wrap(originalMiddleware, function (req, res, next) {
+        if (!graphqlMiddlewareChannel.start.hasSubscribers) {
+          return originalMiddleware.apply(this, arguments)
+        }
+
+        return graphqlMiddlewareChannel.traceSync(originalMiddleware, { req }, this, ...arguments)
+      })
+    }
+  })
+  return express4
+}
+
+function apolloHeaderMapHook (headerMap) {
+  HeaderMap = headerMap.HeaderMap
+  return headerMap
+}
+
+function apolloServerHook (apolloServer) {
+  shimmer.wrap(apolloServer.ApolloServer.prototype, 'executeHTTPGraphQLRequest', wrapExecuteHTTPGraphQLRequest)
+  return apolloServer
+}
+
+addHook({ name: '@apollo/server', file: 'dist/cjs/ApolloServer.js', versions: ['>=4.0.0'] }, apolloServerHook)
+addHook({ name: '@apollo/server', file: 'dist/cjs/express4/index.js', versions: ['>=4.0.0'] }, apolloExpress4Hook)
+addHook({ name: '@apollo/server', file: 'dist/cjs/utils/HeaderMap.js', versions: ['>=4.0.0'] }, apolloHeaderMapHook)

package/packages/datadog-instrumentations/src/child-process.js

@@ -9,11 +9,10 @@ const shimmer = require('../../datadog-shimmer')
 const childProcessChannel = channel('datadog:child_process:execution:start')
 const execMethods = ['exec', 'execFile', 'fork', 'spawn', 'execFileSync', 'execSync', 'spawnSync']
 const names = ['child_process', 'node:child_process']
-names.forEach(name => {
-  addHook({ name }, childProcess => {
-    shimmer.massWrap(childProcess, execMethods, wrapChildProcessMethod())
-    return childProcess
-  })
+
+addHook({ name: names }, childProcess => {
+  shimmer.massWrap(childProcess, execMethods, wrapChildProcessMethod())
+  return childProcess
 })
 
 function wrapChildProcessMethod () {

package/packages/datadog-instrumentations/src/couchbase.js

@@ -252,9 +252,10 @@ addHook({ name: 'couchbase', file: 'lib/cluster.js', versions: ['^3.0.7', '^3.1.
   return Cluster
 })
 
-// semver >=3.2.0
+// semver >=3.2.2
+// NOTE: <3.2.2 segfaults on cluster.close() https://issues.couchbase.com/browse/JSCBC-936
 
-addHook({ name: 'couchbase', file: 'dist/collection.js', versions: ['>=3.2.0'] }, collection => {
+addHook({ name: 'couchbase', file: 'dist/collection.js', versions: ['>=3.2.2'] }, collection => {
   const Collection = collection.Collection
 
   wrapAllNames(['upsert', 'insert', 'replace'], name => {
@@ -264,7 +265,7 @@ addHook({ name: 'couchbase', file: 'dist/collection.js', versions: ['>=3.2.0'] }
   return collection
 })
 
-addHook({ name: 'couchbase', file: 'dist/bucket.js', versions: ['>=3.2.0'] }, bucket => {
+addHook({ name: 'couchbase', file: 'dist/bucket.js', versions: ['>=3.2.2'] }, bucket => {
   const Bucket = bucket.Bucket
   shimmer.wrap(Bucket.prototype, 'collection', getCollection => {
     return function () {
@@ -278,7 +279,7 @@ addHook({ name: 'couchbase', file: 'dist/bucket.js', versions: ['>=3.2.0'] }, bu
   return bucket
 })
 
-addHook({ name: 'couchbase', file: 'dist/cluster.js', versions: ['3.2.0 - 3.2.1', '>=3.2.2'] }, (cluster) => {
+addHook({ name: 'couchbase', file: 'dist/cluster.js', versions: ['>=3.2.2'] }, (cluster) => {
   const Cluster = cluster.Cluster
 
   shimmer.wrap(Cluster.prototype, 'query', wrapV3Query)

package/packages/datadog-instrumentations/src/crypto.js

@@ -11,8 +11,9 @@ const cryptoCipherCh = channel('datadog:crypto:cipher:start')
 
 const hashMethods = ['createHash', 'createHmac', 'createSign', 'createVerify', 'sign', 'verify']
 const cipherMethods = ['createCipheriv', 'createDecipheriv']
+const names = ['crypto', 'node:crypto']
 
-addHook({ name: 'crypto' }, crypto => {
+addHook({ name: names }, crypto => {
   shimmer.massWrap(crypto, hashMethods, wrapCryptoMethod(cryptoHashCh))
   shimmer.massWrap(crypto, cipherMethods, wrapCryptoMethod(cryptoCipherCh))
   return crypto

package/packages/datadog-instrumentations/src/dns.js

@@ -18,8 +18,9 @@ const rrtypes = {
 }
 
 const rrtypeMap = new WeakMap()
+const names = ['dns', 'node:dns']
 
-addHook({ name: 'dns' }, dns => {
+addHook({ name: names }, dns => {
   dns.lookup = wrap('apm:dns:lookup', dns.lookup, 2)
   dns.lookupService = wrap('apm:dns:lookup_service', dns.lookupService, 3)
   dns.resolve = wrap('apm:dns:resolve', dns.resolve, 2)

package/packages/datadog-instrumentations/src/graphql.js

@@ -1,5 +1,7 @@
 'use strict'
 
+const { AbortController } = require('node-abort-controller')
+
 const {
   addHook,
   channel,
@@ -37,6 +39,13 @@ const validateStartCh = channel('apm:graphql:validate:start')
 const validateFinishCh = channel('apm:graphql:validate:finish')
 const validateErrorCh = channel('apm:graphql:validate:error')
 
+class AbortError extends Error {
+  constructor (message) {
+    super(message)
+    this.name = 'AbortError'
+  }
+}
+
 function getOperation (document, operationName) {
   if (!document || !Array.isArray(document.definitions)) {
     return
@@ -175,11 +184,11 @@ function wrapExecute (execute) {
           docSource: documentSources.get(document)
         })
 
-        const context = { source, asyncResource, fields: {} }
+        const context = { source, asyncResource, fields: {}, abortController: new AbortController() }
 
         contexts.set(contextValue, context)
 
-        return callInAsyncScope(exe, asyncResource, this, arguments, (err, res) => {
+        return callInAsyncScope(exe, asyncResource, this, arguments, context.abortController, (err, res) => {
           if (finishResolveCh.hasSubscribers) finishResolvers(context)
 
           const error = err || (res && res.errors && res.errors[0])
@@ -207,7 +216,7 @@ function wrapResolve (resolve) {
 
     const field = assertField(context, info, args)
 
-    return callInAsyncScope(resolve, field.asyncResource, this, arguments, (err) => {
+    return callInAsyncScope(resolve, field.asyncResource, this, arguments, context.abortController, (err) => {
       updateFieldCh.publish({ field, info, err })
     })
   }
@@ -217,10 +226,15 @@ function wrapResolve (resolve) {
   return resolveAsync
 }
 
-function callInAsyncScope (fn, aR, thisArg, args, cb) {
+function callInAsyncScope (fn, aR, thisArg, args, abortController, cb) {
   cb = cb || (() => {})
 
   return aR.runInAsyncScope(() => {
+    if (abortController?.signal.aborted) {
+      cb(null, null)
+      throw new AbortError('Aborted')
+    }
+
     try {
       const result = fn.apply(thisArg, args)
       if (result && typeof result.then === 'function') {

package/packages/datadog-instrumentations/src/helpers/hooks.js

@@ -1,6 +1,8 @@
 'use strict'
 
 module.exports = {
+  '@apollo/server': () => require('../apollo-server'),
+  'apollo-server-core': () => require('../apollo-server-core'),
   '@aws-sdk/smithy-client': () => require('../aws-sdk'),
   '@cucumber/cucumber': () => require('../cucumber'),
   '@playwright/test': () => require('../playwright'),
@@ -29,7 +31,6 @@ module.exports = {
   'bunyan': () => require('../bunyan'),
   'cassandra-driver': () => require('../cassandra-driver'),
   'child_process': () => require('../child-process'),
-  'node:child_process': () => require('../child-process'),
   'connect': () => require('../connect'),
   'cookie': () => require('../cookie'),
   'cookie-parser': () => require('../cookie-parser'),
@@ -43,7 +44,6 @@ module.exports = {
   'fastify': () => require('../fastify'),
   'find-my-way': () => require('../find-my-way'),
   'fs': () => require('../fs'),
-  'node:fs': () => require('../fs'),
   'generic-pool': () => require('../generic-pool'),
   'graphql': () => require('../graphql'),
   'grpc': () => require('../grpc'),
@@ -77,6 +77,13 @@ module.exports = {
   'mysql2': () => require('../mysql2'),
   'net': () => require('../net'),
   'next': () => require('../next'),
+  'node:child_process': () => require('../child-process'),
+  'node:crypto': () => require('../crypto'),
+  'node:dns': () => require('../dns'),
+  'node:http': () => require('../http'),
+  'node:http2': () => require('../http2'),
+  'node:https': () => require('../http'),
+  'node:net': () => require('../net'),
   'oracledb': () => require('../oracledb'),
   'openai': () => require('../openai'),
   'paperplane': () => require('../paperplane'),

package/packages/datadog-instrumentations/src/helpers/instrument.js

@@ -21,11 +21,16 @@ exports.channel = function (name) {
  * @param Function hook
  */
 exports.addHook = function addHook ({ name, versions, file }, hook) {
-  if (!instrumentations[name]) {
-    instrumentations[name] = []
+  if (typeof name === 'string') {
+    name = [name]
   }
 
-  instrumentations[name].push({ name, versions, file, hook })
+  for (const val of name) {
+    if (!instrumentations[val]) {
+      instrumentations[val] = []
+    }
+    instrumentations[val].push({ name: val, versions, file, hook })
+  }
 }
 
 // AsyncResource.bind exists and binds `this` properly only from 17.8.0 and up.

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -24,6 +24,7 @@ if (!disabledInstrumentations.has('fetch')) {
   require('../fetch')
 }
 
+const HOOK_SYMBOL = Symbol('hookExportsMap')
 // TODO: make this more efficient
 
 for (const packageName of names) {
@@ -42,14 +43,29 @@ for (const packageName of names) {
     for (const { name, file, versions, hook } of instrumentations[packageName]) {
       const fullFilename = filename(name, file)
 
+      // Create a WeakMap associated with the hook function so that patches on the same moduleExport only happens once
+      // for example by instrumenting both dns and node:dns double the spans would be created
+      // since they both patch the same moduleExport, this WeakMap is used to mitigate that
+      if (!hook[HOOK_SYMBOL]) {
+        hook[HOOK_SYMBOL] = new WeakMap()
+      }
+
       if (moduleName === fullFilename) {
         const version = moduleVersion || getVersion(moduleBaseDir)
 
         if (matchVersion(version, versions)) {
+          // Check if the hook already has a set moduleExport
+          if (hook[HOOK_SYMBOL].has(moduleExports)) {
+            return moduleExports
+          }
+
           try {
             loadChannel.publish({ name, version, file })
-
-            moduleExports = hook(moduleExports, version)
+            // Send the name and version of the module back to the callback because now addHook
+            // takes in an array of names so by passing the name the callback will know which module name is being used
+            moduleExports = hook(moduleExports, version, name)
+            // Set the moduleExports in the hooks weakmap
+            hook[HOOK_SYMBOL].set(moduleExports, name)
           } catch (e) {
             log.error(e)
           }

package/packages/datadog-instrumentations/src/http/client.js

@@ -14,9 +14,9 @@ const endChannel = channel('apm:http:client:request:end')
 const asyncStartChannel = channel('apm:http:client:request:asyncStart')
 const errorChannel = channel('apm:http:client:request:error')
 
-addHook({ name: 'https' }, hookFn)
+const names = ['http', 'https', 'node:http', 'node:https']
 
-addHook({ name: 'http' }, hookFn)
+addHook({ name: names }, hookFn)
 
 function hookFn (http) {
   patch(http, 'request')
@@ -69,29 +69,17 @@ function patch (http, methodName) {
         try {
           const req = request.call(this, options, callback)
           const emit = req.emit
-
-          const requestSetTimeout = req.setTimeout
+          const setTimeout = req.setTimeout
 
           ctx.req = req
 
           // tracked to accurately discern custom request socket timeout
           let customRequestTimeout = false
-
           req.setTimeout = function () {
             customRequestTimeout = true
-            return requestSetTimeout.apply(this, arguments)
+            return setTimeout.apply(this, arguments)
           }
 
-          req.on('socket', socket => {
-            if (socket) {
-              const socketSetTimeout = socket.setTimeout
-              socket.setTimeout = function () {
-                customRequestTimeout = true
-                return socketSetTimeout.apply(this, arguments)
-              }
-            }
-          })
-
           req.emit = function (eventName, arg) {
             switch (eventName) {
               case 'response': {

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

@@ -15,14 +15,17 @@ const finishSetHeaderCh = channel('datadog:http:server:response:set-header:finis
 
 const requestFinishedSet = new WeakSet()
 
-addHook({ name: 'https' }, http => {
-  // http.ServerResponse not present on https
+const httpNames = ['http', 'node:http']
+const httpsNames = ['https', 'node:https']
+
+addHook({ name: httpNames }, http => {
+  shimmer.wrap(http.ServerResponse.prototype, 'emit', wrapResponseEmit)
   shimmer.wrap(http.Server.prototype, 'emit', wrapEmit)
   return http
 })
 
-addHook({ name: 'http' }, http => {
-  shimmer.wrap(http.ServerResponse.prototype, 'emit', wrapResponseEmit)
+addHook({ name: httpsNames }, http => {
+  // http.ServerResponse not present on https
   shimmer.wrap(http.Server.prototype, 'emit', wrapEmit)
   return http
 })

package/packages/datadog-instrumentations/src/http2/client.js

@@ -10,6 +10,8 @@ const asyncStartChannel = channel('apm:http2:client:request:asyncStart')
 const asyncEndChannel = channel('apm:http2:client:request:asyncEnd')
 const errorChannel = channel('apm:http2:client:request:error')
 
+const names = ['http2', 'node:http2']
+
 function createWrapEmit (ctx) {
   return function wrapEmit (emit) {
     return function (event, arg1) {
@@ -66,7 +68,7 @@ function wrapConnect (connect) {
   }
 }
 
-addHook({ name: 'http2' }, http2 => {
+addHook({ name: names }, http2 => {
   shimmer.wrap(http2, 'connect', wrapConnect)
 
   return http2

package/packages/datadog-instrumentations/src/http2/server.js

@@ -14,7 +14,9 @@ const startServerCh = channel('apm:http2:server:request:start')
 const errorServerCh = channel('apm:http2:server:request:error')
 const finishServerCh = channel('apm:http2:server:request:finish')
 
-addHook({ name: 'http2' }, http2 => {
+const names = ['http2', 'node:http2']
+
+addHook({ name: names }, http2 => {
   shimmer.wrap(http2, 'createSecureServer', wrapCreateServer)
   shimmer.wrap(http2, 'createServer', wrapCreateServer)
   return http2

package/packages/datadog-instrumentations/src/jest.js

@@ -403,7 +403,7 @@ function jestAdapterWrapper (jestAdapter, jestVersion) {
           const coverageFiles = getCoveredFilenamesFromCoverage(environment.global.__coverage__)
             .map(filename => getTestSuitePath(filename, environment.rootDir))
           asyncResource.runInAsyncScope(() => {
-            testSuiteCodeCoverageCh.publish([...coverageFiles, environment.testSuite])
+            testSuiteCodeCoverageCh.publish({ coverageFiles, testSuite: environment.testSuite })
           })
         }
         testSuiteFinishCh.publish({ status, errorMessage })

package/packages/datadog-instrumentations/src/net.js

@@ -17,8 +17,16 @@ const errorTCPCh = channel('apm:net:tcp:error')
 
 const connectionCh = channel(`apm:net:tcp:connection`)
 
-addHook({ name: 'net' }, net => {
-  require('dns')
+const names = ['net', 'node:net']
+
+addHook({ name: names }, (net, version, name) => {
+  // explicitly require dns so that net gets an instrumented instance
+  // so that we don't miss the dns calls
+  if (name === 'net') {
+    require('dns')
+  } else {
+    require('node:dns')
+  }
 
   shimmer.wrap(net.Socket.prototype, 'connect', connect => function () {
     if (!startICPCh.hasSubscribers || !startTCPCh.hasSubscribers) {

package/packages/datadog-instrumentations/src/next.js

@@ -65,7 +65,7 @@ function wrapRenderToHTML (renderToHTML) {
 
 function wrapRenderErrorToHTML (renderErrorToHTML) {
   return function (err, req, res, pathname, query) {
-    return instrument(req, res, () => renderErrorToHTML.apply(this, arguments))
+    return instrument(req, res, err, () => renderErrorToHTML.apply(this, arguments))
   }
 }
 
@@ -76,8 +76,8 @@ function wrapRenderToResponse (renderToResponse) {
 }
 
 function wrapRenderErrorToResponse (renderErrorToResponse) {
-  return function (ctx) {
-    return instrument(ctx.req, ctx.res, () => renderErrorToResponse.apply(this, arguments))
+  return function (ctx, err) {
+    return instrument(ctx.req, ctx.res, err, () => renderErrorToResponse.apply(this, arguments))
   }
 }
 
@@ -111,13 +111,23 @@ function getPageFromPath (page, dynamicRoutes = []) {
   return getPagePath(page)
 }
 
-function instrument (req, res, handler) {
+function instrument (req, res, error, handler) {
+  if (typeof error === 'function') {
+    handler = error
+    error = null
+  }
+
   req = req.originalRequest || req
   res = res.originalResponse || res
 
   // TODO support middleware properly in the future?
   const isMiddleware = req.headers[MIDDLEWARE_HEADER]
-  if (isMiddleware || requests.has(req)) return handler()
+  if (isMiddleware || requests.has(req)) {
+    if (error) {
+      errorChannel.publish({ error })
+    }
+    return handler()
+  }
 
   requests.add(req)
 

package/packages/datadog-instrumentations/src/rhea.js

@@ -22,7 +22,7 @@ const dispatchReceiveCh = channel('apm:rhea:receive:dispatch')
 const errorReceiveCh = channel('apm:rhea:receive:error')
 const finishReceiveCh = channel('apm:rhea:receive:finish')
 
-const contexts = new WeakMap()
+const contexts = new WeakMap() // key: delivery Fn, val: context
 
 addHook({ name: 'rhea', versions: ['>=1'] }, rhea => {
   shimmer.wrap(rhea.message, 'encode', encode => function (msg) {
@@ -52,7 +52,8 @@ addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/link.js' }, obj => {
       startSendCh.publish({ targetAddress, host, port, msg })
       const delivery = send.apply(this, arguments)
       const context = {
-        asyncResource
+        asyncResource,
+        connection: this.connection
       }
       contexts.set(delivery, context)
 
@@ -80,7 +81,8 @@ addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/link.js' }, obj => {
 
         if (msgObj.delivery) {
           const context = {
-            asyncResource
+            asyncResource,
+            connection: this.connection
           }
           contexts.set(msgObj.delivery, context)
           msgObj.delivery.update = wrapDeliveryUpdate(msgObj.delivery, msgObj.delivery.update)
@@ -114,7 +116,7 @@ addHook({ name: 'rhea', versions: ['>=1'], file: 'lib/connection.js' }, Connecti
 
           asyncResource.runInAsyncScope(() => {
             errorReceiveCh.publish(error)
-            beforeFinish(delivery, null)
+            exports.beforeFinish(delivery, null)
             finishReceiveCh.publish()
           })
         })
@@ -187,7 +189,7 @@ function patchCircularBuffer (proto, Session) {
                 const state = remoteState && remoteState.constructor
                   ? entry.remote_state.constructor.composite_type : undefined
                 asyncResource.runInAsyncScope(() => {
-                  beforeFinish(entry, state)
+                  exports.beforeFinish(entry, state)
                   finishSendCh.publish()
                 })
               }
@@ -217,13 +219,13 @@ function addToInFlightDeliveries (connection, delivery) {
 }
 
 function beforeFinish (delivery, state) {
-  const obj = contexts.get(delivery)
-  if (obj) {
+  const context = contexts.get(delivery)
+  if (context) {
     if (state) {
       dispatchReceiveCh.publish({ state })
     }
-    if (obj.connection && obj.connection[inFlightDeliveries]) {
-      obj.connection[inFlightDeliveries].delete(delivery)
+    if (context.connection && context.connection[inFlightDeliveries]) {
+      context.connection[inFlightDeliveries].delete(delivery)
     }
   }
 }
@@ -238,3 +240,7 @@ function getStateFromData (stateData) {
     }
   }
 }
+
+module.exports.inFlightDeliveries = inFlightDeliveries
+module.exports.beforeFinish = beforeFinish
+module.exports.contexts = contexts