From ca09364b8b97f809a9d23ddc124560f12cb89ed7 Mon Sep 17 00:00:00 2001
From: Matteo Collina <hello@matteocollina.com>
Date: Mon, 17 Feb 2020 13:56:56 +0100
Subject: [PATCH] Avoid deadlocks when using then and returning self.

Fixes #94
---
 boot.js                      | 14 +++++++++++---
 test/after-and-ready.test.js | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/boot.js b/boot.js
index de0a5cc..44afe35 100644
--- a/boot.js
+++ b/boot.js
@@ -6,6 +6,7 @@ const inherits = require('util').inherits
 const TimeTree = require('./time-tree')
 const Plugin = require('./plugin')
 const debug = require('debug')('avvio')
+const kAvvio = Symbol('kAvvio')
 
 function wrap (server, opts, instance) {
   const expose = opts.expose || {}
@@ -33,6 +34,7 @@ function wrap (server, opts, instance) {
   }
 
   Object.defineProperty(server, 'then', { get: thenify.bind(instance) })
+  server[kAvvio] = true
 
   server[afterKey] = function (func) {
     if (typeof func !== 'function') {
@@ -188,6 +190,8 @@ function assertPlugin (plugin) {
   }
 }
 
+Boot.prototype[kAvvio] = true
+
 // load a plugin
 Boot.prototype.use = function (plugin, opts) {
   this._lastUsed = this._addPlugin(plugin, opts, false)
@@ -344,7 +348,11 @@ function thenify () {
   // await server.ready() as ready() resolves
   // with the server, end we will end up here
   // because of automatic promise chaining.
-  if (this.booted) return
+  if (this.booted) {
+    debug('thenify returning null because we are already booted')
+    return
+  }
+  debug('thenify')
   const p = this._loadRegistered()
   return p.then.bind(p)
 }
@@ -359,14 +367,14 @@ function callWithCbOrNextTick (func, cb, context) {
   if (func.length === 0) {
     this._error = err
     res = func()
-    if (res && typeof res.then === 'function') {
+    if (res && !res[kAvvio] && typeof res.then === 'function') {
       res.then(() => process.nextTick(cb), (e) => process.nextTick(cb, e))
     } else {
       process.nextTick(cb)
     }
   } else if (func.length === 1) {
     res = func(err)
-    if (res && typeof res.then === 'function') {
+    if (res && !res[kAvvio] && typeof res.then === 'function') {
       res.then(() => process.nextTick(cb), (e) => process.nextTick(cb, e))
     } else {
       process.nextTick(cb)
diff --git a/test/after-and-ready.test.js b/test/after-and-ready.test.js
index 7eadc58..fd8d23d 100644
--- a/test/after-and-ready.test.js
+++ b/test/after-and-ready.test.js
@@ -753,3 +753,37 @@ test('preReady event (errored)', (t) => {
     t.is(order.shift(), 3)
   })
 })
+
+test('after return self', (t) => {
+  t.plan(6)
+
+  const app = boot()
+  let pluginLoaded = false
+  let afterCalled = false
+  let second = false
+
+  app.use(function (s, opts, done) {
+    t.notOk(afterCalled, 'after not called')
+    pluginLoaded = true
+    done()
+  })
+
+  app.after(function () {
+    t.ok(pluginLoaded, 'afterred!')
+    afterCalled = true
+    // happens with after(() => app.use(..))
+    return app
+  })
+
+  app.use(function (s, opts, done) {
+    t.ok(afterCalled, 'after called')
+    second = true
+    done()
+  })
+
+  app.on('start', () => {
+    t.ok(afterCalled, 'after called')
+    t.ok(pluginLoaded, 'plugin loaded')
+    t.ok(second, 'second plugin loaded')
+  })
+})