From ba6b242fae359e45fb9d8c41515df8e09b1bda66 Mon Sep 17 00:00:00 2001
From: Ouyang Yadong <oyydoibh@gmail.com>
Date: Sun, 7 Oct 2018 21:18:15 +0800
Subject: [PATCH] tls: make StreamWrap work correctly in "drain" callback

When an instance of StreamWrap is shutting down and a "drain" event
is emitted, the instance will abort as its
`this[kCurrentShutdownRequest]` is already set. The following test
will fail before this commit.
---
 lib/internal/wrap_js_stream.js          |  7 ++--
 test/parallel/test-stream-wrap-drain.js | 50 +++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 test/parallel/test-stream-wrap-drain.js

diff --git a/lib/internal/wrap_js_stream.js b/lib/internal/wrap_js_stream.js
index b51e008705c85e..e6de433676f7d3 100644
--- a/lib/internal/wrap_js_stream.js
+++ b/lib/internal/wrap_js_stream.js
@@ -100,9 +100,6 @@ class JSStreamWrap extends Socket {
   }
 
   doShutdown(req) {
-    assert.strictEqual(this[kCurrentShutdownRequest], null);
-    this[kCurrentShutdownRequest] = req;
-
     // TODO(addaleax): It might be nice if we could get into a state where
     // DoShutdown() is not called on streams while a write is still pending.
     //
@@ -113,8 +110,10 @@ class JSStreamWrap extends Socket {
     // so for now that is supported here.
 
     if (this[kCurrentWriteRequest] !== null)
-      return this.on('drain', () => this.doShutdown(req));
+      return this.once('drain', () => this.doShutdown(req));
     assert.strictEqual(this[kCurrentWriteRequest], null);
+    assert.strictEqual(this[kCurrentShutdownRequest], null);
+    this[kCurrentShutdownRequest] = req;
 
     const handle = this._handle;
 
diff --git a/test/parallel/test-stream-wrap-drain.js b/test/parallel/test-stream-wrap-drain.js
new file mode 100644
index 00000000000000..068e2d7fd45979
--- /dev/null
+++ b/test/parallel/test-stream-wrap-drain.js
@@ -0,0 +1,50 @@
+// Flags: --expose-internals
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const { StreamWrap } = require('_stream_wrap');
+const { Duplex } = require('stream');
+const { internalBinding } = require('internal/test/binding');
+const { ShutdownWrap } = internalBinding('stream_wrap');
+
+// This test makes sure that when an instance of JSStreamWrap is waiting for
+// a "drain" event to `doShutdown`, the instance will work correctly when a
+// "drain" event emitted.
+{
+  let resolve = null;
+
+  class TestDuplex extends Duplex {
+    _write(chunk, encoding, callback) {
+      // We will resolve the write later.
+      resolve = () => {
+        callback();
+      };
+    }
+
+    _read() {}
+  }
+
+  const testDuplex = new TestDuplex();
+  const socket = new StreamWrap(testDuplex);
+
+  socket.write(
+    // Make the buffer long enough so that the `Writable` will emit "drain".
+    Buffer.allocUnsafe(socket.writableHighWaterMark * 2),
+    common.mustCall()
+  );
+
+  // Make sure that the 'drain' events will be emitted.
+  testDuplex.on('drain', common.mustCall(() => {
+    console.log('testDuplex drain');
+  }));
+
+  assert.strictEqual(typeof resolve, 'function');
+
+  const req = new ShutdownWrap();
+  req.oncomplete = common.mustCall();
+  req.handle = socket._handle;
+  // Should not throw.
+  socket._handle.shutdown(req);
+
+  resolve();
+}