From e6b511932555baf1a485891cacb939d4cd3379a4 Mon Sep 17 00:00:00 2001
From: Trevor Norris <trev.norris@gmail.com>
Date: Tue, 20 Oct 2015 12:18:15 -0600
Subject: [PATCH] async_wrap: allow some hooks to be optional

Only enforce that the init callback is passed to setupHooks(). The
remaining hooks can be optionally passed.

Throw if async_wrap.enable() runs before setting the init callback or if
setupHooks() is called while async wrap is enabled.

Add test to verify calls throw appropriately.

PR-URL: https://github.com/nodejs/node/pull/3461
Reviewed-By: Fedor Indutny <fedor@indutny.com>
---
 src/async-wrap.cc                             | 18 ++++++++++-----
 .../parallel/test-async-wrap-throw-no-init.js | 22 +++++++++++++++++++
 2 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 test/parallel/test-async-wrap-throw-no-init.js

diff --git a/src/async-wrap.cc b/src/async-wrap.cc
index 4f27e5116dca8d..767e76831c3d04 100644
--- a/src/async-wrap.cc
+++ b/src/async-wrap.cc
@@ -103,6 +103,9 @@ RetainedObjectInfo* WrapperInfo(uint16_t class_id, Local<Value> wrapper) {
 
 static void EnableHooksJS(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
+  Local<Function> init_fn = env->async_hooks_init_function();
+  if (init_fn.IsEmpty() || !init_fn->IsFunction())
+    return env->ThrowTypeError("init callback is not assigned to a function");
   env->async_hooks()->set_enable_callbacks(1);
 }
 
@@ -116,13 +119,18 @@ static void DisableHooksJS(const FunctionCallbackInfo<Value>& args) {
 static void SetupHooks(const FunctionCallbackInfo<Value>& args) {
   Environment* env = Environment::GetCurrent(args);
 
-  CHECK(args[0]->IsFunction());
-  CHECK(args[1]->IsFunction());
-  CHECK(args[2]->IsFunction());
+  if (env->async_hooks()->callbacks_enabled())
+    return env->ThrowError("hooks should not be set while also enabled");
+
+  if (!args[0]->IsFunction())
+    return env->ThrowTypeError("init callback must be a function");
 
   env->set_async_hooks_init_function(args[0].As<Function>());
-  env->set_async_hooks_pre_function(args[1].As<Function>());
-  env->set_async_hooks_post_function(args[2].As<Function>());
+
+  if (args[1]->IsFunction())
+    env->set_async_hooks_pre_function(args[1].As<Function>());
+  if (args[2]->IsFunction())
+    env->set_async_hooks_post_function(args[2].As<Function>());
 }
 
 
diff --git a/test/parallel/test-async-wrap-throw-no-init.js b/test/parallel/test-async-wrap-throw-no-init.js
new file mode 100644
index 00000000000000..b2f60f3215e850
--- /dev/null
+++ b/test/parallel/test-async-wrap-throw-no-init.js
@@ -0,0 +1,22 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const async_wrap = process.binding('async_wrap');
+
+
+assert.throws(function() {
+  async_wrap.setupHooks(null);
+}, /init callback must be a function/);
+
+assert.throws(function() {
+  async_wrap.enable();
+}, /init callback is not assigned to a function/);
+
+// Should not throw
+async_wrap.setupHooks(() => {});
+async_wrap.enable();
+
+assert.throws(function() {
+  async_wrap.setupHooks(() => {});
+}, /hooks should not be set while also enabled/);