From 51957bedb950c67b8c039a455e2ae54a16248a2a Mon Sep 17 00:00:00 2001
From: Justin Smith <justsmth@amazon.com>
Date: Thu, 17 Oct 2024 15:22:31 -0400
Subject: [PATCH] DH paramgen callback

---
 crypto/evp_extra/evp_extra_test.cc | 54 +++++++++++++++++++++++++++++-
 crypto/evp_extra/p_dh.c            | 15 ++++++++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/crypto/evp_extra/evp_extra_test.cc b/crypto/evp_extra/evp_extra_test.cc
index c312d30626..079d19f3e0 100644
--- a/crypto/evp_extra/evp_extra_test.cc
+++ b/crypto/evp_extra/evp_extra_test.cc
@@ -2983,7 +2983,7 @@ static int dummy_gen_cb(EVP_PKEY_CTX *ctx) {
   return 1;  // Return success (1).
 }
 
-TEST(EVPExtraTest, Callbacks) {
+TEST(EVPExtraTest, KeygenCallbacks) {
   bssl::UniquePtr<EVP_PKEY_CTX> ctx(EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr));
   ASSERT_TRUE(ctx);
 
@@ -3025,3 +3025,55 @@ TEST(EVPExtraTest, Callbacks) {
   }
 }
 
+
+TEST(EVPExtraTest, ParamgenCallbacks) {
+  bssl::UniquePtr<EVP_PKEY_CTX> ctx(EVP_PKEY_CTX_new_id(EVP_PKEY_DH, nullptr));
+  ASSERT_TRUE(ctx);
+
+  // Check the initial values of |ctx->keygen_info|.
+  int keygen_info = EVP_PKEY_CTX_get_keygen_info(ctx.get(), -1);
+  ASSERT_EQ(keygen_info, EVP_PKEY_CTX_KEYGEN_INFO_COUNT);
+  for (int i = 0; i < keygen_info; i++) {
+    EXPECT_EQ(EVP_PKEY_CTX_get_keygen_info(ctx.get(), i), 0);
+  }
+
+
+  // Generating an DH params will trigger the callback.
+  EVP_PKEY *pkey = EVP_PKEY_new();
+  ASSERT_EQ(EVP_PKEY_paramgen_init(ctx.get()), 1);
+
+  ASSERT_TRUE(EVP_PKEY_CTX_set_dh_paramgen_prime_len(ctx.get(), 512));
+  ASSERT_TRUE(EVP_PKEY_paramgen(ctx.get(), &pkey));
+  ASSERT_TRUE(pkey);
+
+  // Verify that |ctx->keygen_info| has not been updated since a callback hasn't
+  // been set.
+  for (int i = 0; i < keygen_info; i++) {
+    EXPECT_EQ(EVP_PKEY_CTX_get_keygen_info(ctx.get(), i), 0);
+  }
+
+  // Now we set the application data and callback.
+  dummy_cb_app_data app_data{false};
+  EVP_PKEY_CTX_set_app_data(ctx.get(), &app_data);
+  EVP_PKEY_CTX_set_cb(ctx.get(), dummy_gen_cb);
+  EXPECT_FALSE(app_data.state);
+
+  // Call key generation again to trigger the callback.
+  ASSERT_TRUE(EVP_PKEY_paramgen(ctx.get(), &pkey));
+  ASSERT_TRUE(pkey);
+  bssl::UniquePtr<EVP_PKEY> ptr(pkey);
+
+  // The callback function should set the state to true. The contents of
+  // |ctx->keygen_info| will only be populated once the callback has been set.
+  EXPECT_TRUE(app_data.state);
+
+  for (int i = 0; i < keygen_info; i++) {
+    // DH_generate_parameters_ex makes a final call to `BN_GENCB_call(cb, 3, 0)`
+    if(i == 0) {
+      EXPECT_EQ(EVP_PKEY_CTX_get_keygen_info(ctx.get(), i), 3);
+    } else {
+      EXPECT_EQ(EVP_PKEY_CTX_get_keygen_info(ctx.get(), i), 0);
+    }
+  }
+}
+
diff --git a/crypto/evp_extra/p_dh.c b/crypto/evp_extra/p_dh.c
index 71d7c10e41..e0f194023d 100644
--- a/crypto/evp_extra/p_dh.c
+++ b/crypto/evp_extra/p_dh.c
@@ -143,18 +143,31 @@ static int pkey_dh_ctrl(EVP_PKEY_CTX *ctx, int type, int p1, void *_p2) {
 }
 
 static int pkey_dh_paramgen(EVP_PKEY_CTX *ctx, EVP_PKEY *pkey) {
+  int ret = 0;
   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);
+
+  BN_GENCB *pkey_ctx_cb = NULL;
+  if (ctx->pkey_gencb) {
+    pkey_ctx_cb = BN_GENCB_new();
+    if (pkey_ctx_cb == NULL) {
+      goto end;
+    }
+    evp_pkey_set_cb_translate(pkey_ctx_cb, ctx);
+  }
+
+  ret = DH_generate_parameters_ex(dh, dctx->prime_len, dctx->generator, pkey_ctx_cb);
+end:
   if (ret == 1) {
     EVP_PKEY_assign_DH(pkey, dh);
   } else {
     ret = 0;
     DH_free(dh);
   }
+  BN_GENCB_free(pkey_ctx_cb);
   return ret;
 }