From 8599540d3c37cb19c76e0c0839174a057abbcada Mon Sep 17 00:00:00 2001
From: Krystian Jarmicki <jarmicki@gmail.com>
Date: Sat, 10 Oct 2020 21:48:46 +0200
Subject: [PATCH] Prevent loss of async hooks context

closes #71
---
 HISTORY.md    |  5 +++++
 index.js      | 31 ++++++++++++++++++++++++++++++-
 test/index.js | 44 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/HISTORY.md b/HISTORY.md
index 4bac5e80..3ac9fd82 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,3 +1,8 @@
+unreleased
+==========
+
+  * Prevent loss of async hooks context
+
 2.4.2 / 2021-11-16
 ==================
 
diff --git a/index.js b/index.js
index 7fe81860..425c2bea 100644
--- a/index.js
+++ b/index.js
@@ -12,6 +12,7 @@
  * @private
  */
 
+var asyncHooks = tryRequireAsyncHooks()
 var bytes = require('bytes')
 var createError = require('http-errors')
 var iconv = require('iconv-lite')
@@ -105,7 +106,7 @@ function getRawBody (stream, options, callback) {
 
   if (done) {
     // classic callback style
-    return readStream(stream, encoding, length, limit, done)
+    return readStream(stream, encoding, length, limit, wrap(done))
   }
 
   return new Promise(function executor (resolve, reject) {
@@ -284,3 +285,31 @@ function readStream (stream, encoding, length, limit, callback) {
     stream.removeListener('close', cleanup)
   }
 }
+
+/**
+ * Try to require async_hooks
+ * @private
+ */
+
+function tryRequireAsyncHooks () {
+  try {
+    return require('async_hooks')
+  } catch (e) {
+    return {}
+  }
+}
+
+/**
+ * Wrap function with async resource
+ * @private
+ */
+
+function wrap (fn) {
+  if (!asyncHooks.AsyncResource) {
+    return fn
+  }
+
+  // AsyncResource.bind static method backported
+  var res = new asyncHooks.AsyncResource(fn.name || 'bound-anonymous-fn')
+  return res.runInAsyncScope.bind(res, fn, null)
+}
diff --git a/test/index.js b/test/index.js
index 244f03be..3e8c3819 100644
--- a/test/index.js
+++ b/test/index.js
@@ -1,4 +1,5 @@
 var assert = require('assert')
+var asyncHooks = tryRequire('async_hooks')
 var fs = require('fs')
 var getRawBody = require('..')
 var path = require('path')
@@ -8,6 +9,10 @@ var EventEmitter = require('events').EventEmitter
 var Promise = global.Promise || require('bluebird')
 var Readable = require('readable-stream').Readable
 
+var describeAsyncHooks = typeof asyncHooks.AsyncLocalStorage === 'function'
+  ? describe
+  : describe.skip
+
 var file = path.join(__dirname, 'index.js')
 var length = fs.statSync(file).size
 var string = fs.readFileSync(file, 'utf8')
@@ -272,6 +277,37 @@ describe('Raw Body', function () {
     })
   })
 
+  describeAsyncHooks('with async local storage', function () {
+    it('should presist store in callback', function (done) {
+      var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+      var store = { foo: 'bar' }
+      var stream = createStream()
+
+      asyncLocalStorage.run(store, function () {
+        getRawBody(stream, function (err, buf) {
+          if (err) return done(err)
+          assert.ok(buf.length > 0)
+          assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+          done()
+        })
+      })
+    })
+
+    it('should presist store in promise', function (done) {
+      var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+      var store = { foo: 'bar' }
+      var stream = createStream()
+
+      asyncLocalStorage.run(store, function () {
+        getRawBody(stream).then(function (buf) {
+          assert.ok(buf.length > 0)
+          assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+          done()
+        }, done)
+      })
+    })
+  })
+
   describe('when an encoding is set', function () {
     it('should return a string', function (done) {
       getRawBody(createStream(), {
@@ -411,3 +447,11 @@ function createStream (buf) {
 function throwExpectedError () {
   throw new Error('expected error')
 }
+
+function tryRequire (name) {
+  try {
+    return require(name)
+  } catch (e) {
+    return {}
+  }
+}