Make it work with async hooks

HISTORY.md

@@ -1,6 +1,7 @@
 unreleased
 ==========
 
+  * Prevent loss of async hooks context
   * deps: [email protected]
     - deps: [email protected]
     - deps: [email protected]

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,33 @@ 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) {
+    /* istanbul ignore next */
+    return {}
+  }
+}
+
+/**
+ * Wrap function with async resource
+ * @private
+ */
+
+function wrap (fn) {
+  if (!asyncHooks.AsyncResource) {
+    /* istanbul ignore next */
+    return fn
+  }
+
+  // AsyncResource.bind static method backported
+  var res = new asyncHooks.AsyncResource(fn.name || 'bound-anonymous-fn')
+  return res.runInAsyncScope.bind(res, fn, null)
+}

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 {}
+  }
+}