2 files changed, 27 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_;