From 869f42a90198d3afd201b0cc3cd4c0e9761c8fa9 Mon Sep 17 00:00:00 2001
From: Darshan Sen <darshan.sen@postman.com>
Date: Thu, 14 Oct 2021 19:26:36 +0530
Subject: [PATCH] src,crypto: use `std::variant` in DH params

Signed-off-by: Darshan Sen <darshan.sen@postman.com>
---
 src/crypto/crypto_dh.cc | 26 ++++++++++++++++----------
 src/crypto/crypto_dh.h  |  5 ++---
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/crypto/crypto_dh.cc b/src/crypto/crypto_dh.cc
index 86475e3b1b1e017..a5a87b3edf8e6ac 100644
--- a/src/crypto/crypto_dh.cc
+++ b/src/crypto/crypto_dh.cc
@@ -9,6 +9,8 @@
 #include "threadpoolwork-inl.h"
 #include "v8.h"
 
+#include <variant>
+
 namespace node {
 
 using v8::ArrayBuffer;
@@ -437,7 +439,7 @@ Maybe<bool> DhKeyGenTraits::AdditionalConfig(
       return Nothing<bool>();
     }
 
-    params->params.prime_fixed_value = BignumPointer(
+    params->params.prime = BignumPointer(
         BN_bin2bn(reinterpret_cast<const unsigned char*>(group->prime),
                   group->prime_size, nullptr));
     params->params.generator = group->gen;
@@ -449,14 +451,14 @@ Maybe<bool> DhKeyGenTraits::AdditionalConfig(
         THROW_ERR_OUT_OF_RANGE(env, "Invalid prime size");
         return Nothing<bool>();
       }
-      params->params.prime_size = size;
+      params->params.prime = size;
     } else {
       ArrayBufferOrViewContents<unsigned char> input(args[*offset]);
       if (UNLIKELY(!input.CheckSizeInt32())) {
         THROW_ERR_OUT_OF_RANGE(env, "prime is too big");
         return Nothing<bool>();
       }
-      params->params.prime_fixed_value = BignumPointer(
+      params->params.prime = BignumPointer(
           BN_bin2bn(input.data(), input.size(), nullptr));
     }
 
@@ -470,31 +472,33 @@ Maybe<bool> DhKeyGenTraits::AdditionalConfig(
 
 EVPKeyCtxPointer DhKeyGenTraits::Setup(DhKeyPairGenConfig* params) {
   EVPKeyPointer key_params;
-  if (params->params.prime_fixed_value) {
+  if (BignumPointer* prime_fixed_value =
+          std::get_if<BignumPointer>(&params->params.prime)) {
     DHPointer dh(DH_new());
     if (!dh)
       return EVPKeyCtxPointer();
 
-    BIGNUM* prime = params->params.prime_fixed_value.get();
+    BIGNUM* prime = prime_fixed_value->get();
     BignumPointer bn_g(BN_new());
     if (!BN_set_word(bn_g.get(), params->params.generator) ||
-        !DH_set0_pqg(dh.get(), prime, nullptr, bn_g.get()))
+        !DH_set0_pqg(dh.get(), prime, nullptr, bn_g.get())) {
       return EVPKeyCtxPointer();
+    }
 
-    params->params.prime_fixed_value.release();
+    prime_fixed_value->release();
     bn_g.release();
 
     key_params = EVPKeyPointer(EVP_PKEY_new());
     CHECK(key_params);
-    EVP_PKEY_assign_DH(key_params.get(), dh.release());
-  } else {
+    CHECK_EQ(EVP_PKEY_assign_DH(key_params.get(), dh.release()), 1);
+  } else if (int* prime_size = std::get_if<int>(&params->params.prime)) {
     EVPKeyCtxPointer param_ctx(EVP_PKEY_CTX_new_id(EVP_PKEY_DH, nullptr));
     EVP_PKEY* raw_params = nullptr;
     if (!param_ctx ||
         EVP_PKEY_paramgen_init(param_ctx.get()) <= 0 ||
         EVP_PKEY_CTX_set_dh_paramgen_prime_len(
             param_ctx.get(),
-            params->params.prime_size) <= 0 ||
+            *prime_size) <= 0 ||
         EVP_PKEY_CTX_set_dh_paramgen_generator(
             param_ctx.get(),
             params->params.generator) <= 0 ||
@@ -503,6 +507,8 @@ EVPKeyCtxPointer DhKeyGenTraits::Setup(DhKeyPairGenConfig* params) {
     }
 
     key_params = EVPKeyPointer(raw_params);
+  } else {
+    UNREACHABLE();
   }
 
   EVPKeyCtxPointer ctx(EVP_PKEY_CTX_new(key_params.get(), nullptr));
diff --git a/src/crypto/crypto_dh.h b/src/crypto/crypto_dh.h
index fecbf41070bda4a..722bb1dad81a1f2 100644
--- a/src/crypto/crypto_dh.h
+++ b/src/crypto/crypto_dh.h
@@ -10,6 +10,7 @@
 #include "memory_tracker.h"
 #include "v8.h"
 
+#include <variant>
 #include <vector>
 
 namespace node {
@@ -58,12 +59,10 @@ class DiffieHellman : public BaseObject {
 };
 
 struct DhKeyPairParams final : public MemoryRetainer {
-  // TODO(tniessen): Use std::variant instead.
   // Diffie-Hellman can either generate keys using a fixed prime, or by first
   // generating a random prime of a given size (in bits). Only one of both
   // options may be specified.
-  BignumPointer prime_fixed_value;
-  unsigned int prime_size;
+  std::variant<BignumPointer, int> prime;
   unsigned int generator;
   SET_NO_MEMORY_INFO()
   SET_MEMORY_INFO_NAME(DhKeyPairParams)