src,crypto: use std::variant in DH params

Signed-off-by: Darshan Sen <darshan.sen@postman.com>
nodejs · Oct 14, 2021 · 29a2559 · 29a2559
1 parent b80b85e
commit 29a2559
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git 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
@@ -10,7 +10,7 @@
 #include "memory_tracker.h"
 #include "v8.h"
 
-#include <vector>
+#include <variant>
 
 namespace node {
 namespace crypto {
@@ -58,12 +58,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)