3 files changed, 53 insertions, 29 deletions
diff --git a/src/node_crypto.cc b/src/node_crypto.cc
index 02cf4ef518..84f006b00f 100644
--- a/src/node_crypto.cc
+++ b/src/node_crypto.cc
@@ -4881,44 +4881,40 @@ void DiffieHellman::ComputeSecret(const FunctionCallbackInfo<Value>& args) {
 }
 
 
-void DiffieHellman::SetPublicKey(const FunctionCallbackInfo<Value>& args) {
-  DiffieHellman* diffieHellman;
-  ASSIGN_OR_RETURN_UNWRAP(&diffieHellman, args.Holder());
-  Environment* env = diffieHellman->env();
+void DiffieHellman::SetKey(const v8::FunctionCallbackInfo<v8::Value>& args,
+                           BIGNUM* (DH::*field), const char* what) {
+  Environment* env = Environment::GetCurrent(args);
 
-  if (!diffieHellman->initialised_) {
-    return ThrowCryptoError(env, ERR_get_error(), "Not initialized");
-  }
+  DiffieHellman* dh;
+  ASSIGN_OR_RETURN_UNWRAP(&dh, args.Holder());
+  if (!dh->initialised_) return env->ThrowError("Not initialized");
+
+  BIGNUM** num = &((dh->dh)->*field);
+  char errmsg[64];
 
   if (args.Length() == 0) {
-    return env->ThrowError("Public key argument is mandatory");
-  } else {
-    THROW_AND_RETURN_IF_NOT_BUFFER(args[0], "Public key");
-    diffieHellman->dh->pub_key = BN_bin2bn(
-        reinterpret_cast<unsigned char*>(Buffer::Data(args[0])),
-        Buffer::Length(args[0]), 0);
+    snprintf(errmsg, sizeof(errmsg), "%s argument is mandatory", what);
+    return env->ThrowError(errmsg);
+  }
+
+  if (!Buffer::HasInstance(args[0])) {
+    snprintf(errmsg, sizeof(errmsg), "%s must be a buffer", what);
+    return env->ThrowTypeError(errmsg);
   }
+
+  *num = BN_bin2bn(reinterpret_cast<unsigned char*>(Buffer::Data(args[0])),
+                   Buffer::Length(args[0]), *num);
+  CHECK_NE(*num, nullptr);
 }
 
 
-void DiffieHellman::SetPrivateKey(const FunctionCallbackInfo<Value>& args) {
-  DiffieHellman* diffieHellman;
-  ASSIGN_OR_RETURN_UNWRAP(&diffieHellman, args.Holder());
-  Environment* env = diffieHellman->env();
+void DiffieHellman::SetPublicKey(const FunctionCallbackInfo<Value>& args) {
+  SetKey(args, &DH::pub_key, "Public key");
+}
 
-  if (!diffieHellman->initialised_) {
-    return ThrowCryptoError(env, ERR_get_error(), "Not initialized");
-  }
 
-  if (args.Length() == 0) {
-    return env->ThrowError("Private key argument is mandatory");
-  } else {
-    THROW_AND_RETURN_IF_NOT_BUFFER(args[0], "Private key");
-    diffieHellman->dh->priv_key = BN_bin2bn(
-        reinterpret_cast<unsigned char*>(Buffer::Data(args[0])),
-        Buffer::Length(args[0]),
-        0);
-  }
+void DiffieHellman::SetPrivateKey(const FunctionCallbackInfo<Value>& args) {
+  SetKey(args, &DH::priv_key, "Private key");
 }
 
 
diff --git a/src/node_crypto.h b/src/node_crypto.h
index 01c799a8b1..1d823bcb35 100644
--- a/src/node_crypto.h
+++ b/src/node_crypto.h
@@ -692,6 +692,8 @@ class DiffieHellman : public BaseObject {
  private:
   static void GetField(const v8::FunctionCallbackInfo<v8::Value>& args,
                        BIGNUM* (DH::*field), const char* err_if_null);
+  static void SetKey(const v8::FunctionCallbackInfo<v8::Value>& args,
+                     BIGNUM* (DH::*field), const char* what);
   bool VerifyContext();
 
   bool initialised_;
diff --git a/test/parallel/test-crypto-dh-leak.js b/test/parallel/test-crypto-dh-leak.js
new file mode 100644
index 0000000000..2ca26d3e9b
--- /dev/null
+++ b/test/parallel/test-crypto-dh-leak.js
@@ -0,0 +1,26 @@
+// Flags: --expose-gc
+'use strict';
+
+const common = require('../common');
+if (!common.hasCrypto)
+  common.skip('missing crypto');
+
+const assert = require('assert');
+const crypto = require('crypto');
+
+const before = process.memoryUsage().rss;
+{
+  const dh = crypto.createDiffieHellman(common.hasFipsCrypto ? 1024 : 256);
+  const publicKey = dh.generateKeys();
+  const privateKey = dh.getPrivateKey();
+  for (let i = 0; i < 5e4; i += 1) {
+    dh.setPublicKey(publicKey);
+    dh.setPrivateKey(privateKey);
+  }
+}
+global.gc();
+const after = process.memoryUsage().rss;
+
+// RSS should stay the same, ceteris paribus, but allow for
+// some slop because V8 mallocs memory during execution.
+assert(after - before < 5 << 20);