From f9d5327987fb353ede87667aab8bc10116781f42 Mon Sep 17 00:00:00 2001
From: Justin Smith <justsmth@amazon.com>
Date: Wed, 25 Sep 2024 17:15:09 -0400
Subject: [PATCH] Implement more EVP_PKEY_DH functionality

---
 crypto/evp_extra/p_dh.c          | 52 ++++++++++++++++++++++++++++++--
 crypto/fipsmodule/evp/internal.h |  2 ++
 include/openssl/evp.h            |  2 ++
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/crypto/evp_extra/p_dh.c b/crypto/evp_extra/p_dh.c
index 2ae75f9293..68cc98288d 100644
--- a/crypto/evp_extra/p_dh.c
+++ b/crypto/evp_extra/p_dh.c
@@ -21,6 +21,9 @@
 
 typedef struct dh_pkey_ctx_st {
   int pad;
+  /* Parameter gen parameters */
+  int prime_len;
+  int generator;
 } DH_PKEY_CTX;
 
 static int pkey_dh_init(EVP_PKEY_CTX *ctx) {
@@ -104,7 +107,7 @@ static int pkey_dh_derive(EVP_PKEY_CTX *ctx, uint8_t *out, size_t *out_len) {
   return 1;
 }
 
-static int pkey_dh_ctrl(EVP_PKEY_CTX *ctx, int type, int p1, void *p2) {
+static int pkey_dh_ctrl(EVP_PKEY_CTX *ctx, int type, int p1, void *_p2) {
   DH_PKEY_CTX *dctx = ctx->data;
   switch (type) {
     case EVP_PKEY_CTRL_PEER_KEY:
@@ -116,21 +119,54 @@ static int pkey_dh_ctrl(EVP_PKEY_CTX *ctx, int type, int p1, void *p2) {
       dctx->pad = p1;
       return 1;
 
+    case EVP_PKEY_CTRL_DH_PARAMGEN_PRIME_LEN:
+      if(p1 < 256) {
+        return -2;
+      }
+      dctx->prime_len = p1;
+      return 1;
+
+    case EVP_PKEY_CTRL_DH_PARAMGEN_GENERATOR:
+      dctx->generator = p1;
+      return 1;
+
     default:
       OPENSSL_PUT_ERROR(EVP, EVP_R_COMMAND_NOT_SUPPORTED);
       return 0;
   }
 }
 
+static int pkey_dh_paramgen(EVP_PKEY_CTX *ctx, EVP_PKEY *pkey) {
+  DH_PKEY_CTX *dctx = ctx->data;
+
+  DH *dh = DH_new();
+  if (dh == NULL) {
+    return 0;
+  }
+  int ret = DH_generate_parameters_ex(dh, dctx->prime_len, dctx->generator, NULL);
+  if (ret) {
+    EVP_PKEY_assign_DH(pkey, dh);
+  } else {
+    DH_free(dh);
+  }
+  return ret;
+}
+
+
 static int pkey_dh_ctrl_str(EVP_PKEY_CTX *ctx, const char *type,
                             const char *value) {
   // We don't support:
-  // * dh_paramgen_prime_len
   // * dh_rfc5114
   // * dh_param
-  // * dh_paramgen_generator
   // * dh_paramgen_subprime_len
   // * dh_paramgen_type
+  if (strcmp(type, "dh_paramgen_prime_len") == 0) {
+
+  }
+
+  if (strcmp(type, "dh_paramgen_generator") == 0) {
+  }
+
 
   if (strcmp(type, "dh_pad") == 0) {
     char* str_end = NULL;
@@ -152,6 +188,7 @@ const EVP_PKEY_METHOD dh_pkey_meth = {
     .cleanup = pkey_dh_cleanup,
     .keygen = pkey_dh_keygen,
     .derive = pkey_dh_derive,
+    .paramgen = pkey_dh_paramgen,
     .ctrl = pkey_dh_ctrl,
     .ctrl_str = pkey_dh_ctrl_str
 };
@@ -160,3 +197,12 @@ int EVP_PKEY_CTX_set_dh_pad(EVP_PKEY_CTX *ctx, int pad) {
   return EVP_PKEY_CTX_ctrl(ctx, EVP_PKEY_DH, EVP_PKEY_OP_DERIVE,
                            EVP_PKEY_CTRL_DH_PAD, pad, NULL);
 }
+
+int EVP_PKEY_CTX_set_dh_paramgen_prime_len(EVP_PKEY_CTX *ctx, int pbits) {
+  return EVP_PKEY_CTX_ctrl(ctx, EVP_PKEY_DH, EVP_PKEY_OP_PARAMGEN,
+                    EVP_PKEY_CTRL_DH_PARAMGEN_PRIME_LEN, pbits, NULL);
+}
+int EVP_PKEY_CTX_set_dh_paramgen_generator(EVP_PKEY_CTX *ctx, int gen) {
+  return EVP_PKEY_CTX_ctrl(ctx, EVP_PKEY_DH, EVP_PKEY_OP_PARAMGEN,
+                   EVP_PKEY_CTRL_DH_PARAMGEN_GENERATOR, gen, NULL);
+}
diff --git a/crypto/fipsmodule/evp/internal.h b/crypto/fipsmodule/evp/internal.h
index 9aa81882ea..004476e79c 100644
--- a/crypto/fipsmodule/evp/internal.h
+++ b/crypto/fipsmodule/evp/internal.h
@@ -241,6 +241,8 @@ int EVP_RSA_PKEY_CTX_ctrl(EVP_PKEY_CTX *ctx, int optype, int cmd, int p1, void *
 #define EVP_PKEY_CTRL_HKDF_SALT (EVP_PKEY_ALG_CTRL + 17)
 #define EVP_PKEY_CTRL_HKDF_INFO (EVP_PKEY_ALG_CTRL + 18)
 #define EVP_PKEY_CTRL_DH_PAD (EVP_PKEY_ALG_CTRL + 19)
+#define EVP_PKEY_CTRL_DH_PARAMGEN_PRIME_LEN (EVP_PKEY_ALG_CTRL + 20)
+#define EVP_PKEY_CTRL_DH_PARAMGEN_GENERATOR (EVP_PKEY_ALG_CTRL + 21)
 
 struct evp_pkey_ctx_st {
   // Method associated with this operation
diff --git a/include/openssl/evp.h b/include/openssl/evp.h
index d45f4fb282..656f961146 100644
--- a/include/openssl/evp.h
+++ b/include/openssl/evp.h
@@ -180,6 +180,8 @@ OPENSSL_EXPORT int EVP_PKEY_set1_DH(EVP_PKEY *pkey, DH *key);
 OPENSSL_EXPORT int EVP_PKEY_assign_DH(EVP_PKEY *pkey, DH *key);
 OPENSSL_EXPORT DH *EVP_PKEY_get0_DH(const EVP_PKEY *pkey);
 OPENSSL_EXPORT DH *EVP_PKEY_get1_DH(const EVP_PKEY *pkey);
+OPENSSL_EXPORT int EVP_PKEY_CTX_set_dh_paramgen_prime_len(EVP_PKEY_CTX *ctx, int pbits);
+OPENSSL_EXPORT int EVP_PKEY_CTX_set_dh_paramgen_generator(EVP_PKEY_CTX *ctx, int gen);
 
 #define EVP_PKEY_NONE NID_undef
 #define EVP_PKEY_RSA NID_rsaEncryption