From 422740ecd34b59211a806cff14514a919856a7e5 Mon Sep 17 00:00:00 2001
From: Darshan Sen <darshan.sen@postman.com>
Date: Sat, 11 Sep 2021 20:52:01 +0530
Subject: [PATCH] src,crypto: eliminate code duplication between
 StatelessDiffieHellman*

Signed-off-by: Darshan Sen <darshan.sen@postman.com>
---
 src/crypto/crypto_dh.cc | 107 ++++++++++++++++++++++------------------
 1 file changed, 59 insertions(+), 48 deletions(-)

diff --git a/src/crypto/crypto_dh.cc b/src/crypto/crypto_dh.cc
index 1c48f98656fd21..c5c726d9e491e4 100644
--- a/src/crypto/crypto_dh.cc
+++ b/src/crypto/crypto_dh.cc
@@ -28,6 +28,7 @@ using v8::ReadOnly;
 using v8::SideEffectType;
 using v8::Signature;
 using v8::String;
+using v8::Undefined;
 using v8::Value;
 
 namespace crypto {
@@ -539,62 +540,73 @@ WebCryptoKeyExportStatus DHKeyExportTraits::DoExport(
 }
 
 namespace {
-AllocatedBuffer StatelessDiffieHellman(
-    Environment* env,
-    ManagedEVPPKey our_key,
-    ManagedEVPPKey their_key) {
-  size_t out_size;
-
-  EVPKeyCtxPointer ctx(EVP_PKEY_CTX_new(our_key.get(), nullptr));
+bool CreateStatelessDiffieHellmanContext(
+    const ManagedEVPPKey& our_key,
+    const ManagedEVPPKey& their_key,
+    char** buf,
+    size_t* out_size) {
+  const EVPKeyCtxPointer ctx(EVP_PKEY_CTX_new(our_key.get(), nullptr));
   if (!ctx ||
       EVP_PKEY_derive_init(ctx.get()) <= 0 ||
       EVP_PKEY_derive_set_peer(ctx.get(), their_key.get()) <= 0 ||
-      EVP_PKEY_derive(ctx.get(), nullptr, &out_size) <= 0)
-    return AllocatedBuffer();
+      EVP_PKEY_derive(ctx.get(), nullptr, out_size) <= 0)
+    return false;
 
-  AllocatedBuffer result = AllocatedBuffer::AllocateManaged(env, out_size);
-  CHECK_NOT_NULL(result.data());
+  *buf = MallocOpenSSL<char>(*out_size);
+  size_t remainder_size = *out_size;
+  if (EVP_PKEY_derive(ctx.get(),
+                      reinterpret_cast<unsigned char*>(*buf),
+                      &remainder_size) <= 0) {
+    OPENSSL_clear_free(buf, *out_size);
+    return false;
+  }
 
-  unsigned char* data = reinterpret_cast<unsigned char*>(result.data());
-  if (EVP_PKEY_derive(ctx.get(), data, &out_size) <= 0)
-    return AllocatedBuffer();
+  ZeroPadDiffieHellmanSecret(remainder_size, *buf, *out_size);
 
-  ZeroPadDiffieHellmanSecret(out_size, &result);
-  return result;
+  return true;
 }
 
-// The version of StatelessDiffieHellman that returns an AllocatedBuffer
-// is not threadsafe because of the AllocatedBuffer allocation of a
-// v8::BackingStore (it'll cause much crashing if we call it from a
-// libuv worker thread). This version allocates a ByteSource instead,
-// which we can convert into a v8::BackingStore later.
-// TODO(@jasnell): Eliminate the code duplication between these two
-// versions of the function.
-ByteSource StatelessDiffieHellmanThreadsafe(
+Local<Value> StatelessDiffieHellman(
     Environment* env,
-    ManagedEVPPKey our_key,
-    ManagedEVPPKey their_key) {
+    const ManagedEVPPKey& our_key,
+    const ManagedEVPPKey& their_key) {
+  char* buf = nullptr;
   size_t out_size;
-
-  EVPKeyCtxPointer ctx(EVP_PKEY_CTX_new(our_key.get(), nullptr));
-  if (!ctx ||
-      EVP_PKEY_derive_init(ctx.get()) <= 0 ||
-      EVP_PKEY_derive_set_peer(ctx.get(), their_key.get()) <= 0 ||
-      EVP_PKEY_derive(ctx.get(), nullptr, &out_size) <= 0)
-    return ByteSource();
-
-  char* buf = MallocOpenSSL<char>(out_size);
-  ByteSource out = ByteSource::Allocated(buf, out_size);
-
-  if (EVP_PKEY_derive(
-          ctx.get(),
-          reinterpret_cast<unsigned char*>(buf),
-          &out_size) <= 0) {
+  if (!CreateStatelessDiffieHellmanContext(our_key,
+                                           their_key,
+                                           &buf,
+                                           &out_size))
+    return Undefined(env->isolate());
+
+  void* hint = reinterpret_cast<void*>(static_cast<uintptr_t>(out_size));
+  return Buffer::New(env,
+                     buf,
+                     out_size,
+                     [](char* data, void* hint) {
+                       size_t out_size = static_cast<size_t>(
+                           reinterpret_cast<uintptr_t>(hint));
+                       OPENSSL_clear_free(data, out_size);
+                     },
+                     hint).FromMaybe(Local<Object>());
+}
+
+// The version of StatelessDiffieHellman that returns a Buffer
+// is not threadsafe because of the allocation of a v8::BackingStore
+// (it'll cause much crashing if we call it from a libuv worker
+// thread). This version allocates a ByteSource instead, which we can
+// convert into a v8::BackingStore later.
+ByteSource StatelessDiffieHellmanThreadsafe(
+    const ManagedEVPPKey& our_key,
+    const ManagedEVPPKey& their_key) {
+  char* buf = nullptr;
+  size_t out_size;
+  if (!CreateStatelessDiffieHellmanContext(our_key,
+                                           their_key,
+                                           &buf,
+                                           &out_size))
     return ByteSource();
-  }
 
-  ZeroPadDiffieHellmanSecret(out_size, buf, out.size());
-  return out;
+  return ByteSource::Allocated(buf, out_size);
 }
 }  // namespace
 
@@ -612,11 +624,11 @@ void DiffieHellman::Stateless(const FunctionCallbackInfo<Value>& args) {
   ManagedEVPPKey our_key = our_key_object->Data()->GetAsymmetricKey();
   ManagedEVPPKey their_key = their_key_object->Data()->GetAsymmetricKey();
 
-  AllocatedBuffer out = StatelessDiffieHellman(env, our_key, their_key);
-  if (out.size() == 0)
+  Local<Value> out = StatelessDiffieHellman(env, our_key, their_key);
+  if (out->IsUndefined() || Buffer::Length(out) == 0)
     return ThrowCryptoError(env, ERR_get_error(), "diffieHellman failed");
 
-  args.GetReturnValue().Set(out.ToBuffer().FromMaybe(Local<Value>()));
+  args.GetReturnValue().Set(out);
 }
 
 Maybe<bool> DHBitsTraits::AdditionalConfig(
@@ -661,7 +673,6 @@ bool DHBitsTraits::DeriveBits(
     const DHBitsConfig& params,
     ByteSource* out) {
   *out = StatelessDiffieHellmanThreadsafe(
-      env,
       params.private_key->GetAsymmetricKey(),
       params.public_key->GetAsymmetricKey());
   return true;