From 162f14aaf9a43d2ab1b69f64a8b7244885f039b3 Mon Sep 17 00:00:00 2001
From: Juliusz Sosinowicz <juliusz@wolfssl.com>
Date: Tue, 17 Aug 2021 19:42:16 +0200
Subject: [PATCH] Implement `DH_set_length`.

---
 src/ssl.c            | 14 ++++++++++++++
 tests/api.c          |  3 +++
 wolfssl/openssl/dh.h |  2 ++
 3 files changed, 19 insertions(+)

diff --git a/src/ssl.c b/src/ssl.c
index 7b612ffa1..5dfd13739 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -31495,6 +31495,20 @@ int wolfSSL_DH_compute_key(unsigned char* key, const WOLFSSL_BIGNUM* otherPub,
 
 
 #if defined(OPENSSL_VERSION_NUMBER) && OPENSSL_VERSION_NUMBER >= 0x10100000L
+int wolfSSL_DH_set_length(WOLFSSL_DH *dh, long len)
+{
+    WOLFSSL_ENTER("wolfSSL_DH_set_length");
+
+    /* len is checked at generation */
+    if (dh == NULL) {
+        WOLFSSL_MSG("Bad function arguments");
+        return WOLFSSL_FAILURE;
+    }
+
+    dh->length = (int)len;
+    return WOLFSSL_SUCCESS;
+}
+
 /* ownership of p,q,and g get taken over by "dh" on success and should be free'd
  * with a call to wolfSSL_DH_free -- not individually.
  *
diff --git a/tests/api.c b/tests/api.c
index 522cd4e65..0a4f4295a 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -38999,6 +38999,9 @@ static void test_wolfSSL_d2i_DHparams(void)
     AssertNotNull(dh->p);
     AssertNotNull(dh->g);
     AssertTrue(pt != buf);
+#if defined(OPENSSL_VERSION_NUMBER) && OPENSSL_VERSION_NUMBER >= 0x10100000L
+    AssertIntEQ(DH_set_length(dh, BN_num_bits(dh->p)), WOLFSSL_SUCCESS);
+#endif
     AssertIntEQ(DH_generate_key(dh), WOLFSSL_SUCCESS);
 
     /* Invalid cases */
diff --git a/wolfssl/openssl/dh.h b/wolfssl/openssl/dh.h
index b63b6c2d1..cf89067fc 100644
--- a/wolfssl/openssl/dh.h
+++ b/wolfssl/openssl/dh.h
@@ -67,6 +67,7 @@ WOLFSSL_API int wolfSSL_DH_generate_key(WOLFSSL_DH*);
 WOLFSSL_API int wolfSSL_DH_compute_key(unsigned char* key, const WOLFSSL_BIGNUM* pub,
                                      WOLFSSL_DH*);
 WOLFSSL_API int wolfSSL_DH_LoadDer(WOLFSSL_DH*, const unsigned char*, int sz);
+WOLFSSL_API int wolfSSL_DH_set_length(WOLFSSL_DH*, long);
 WOLFSSL_API int wolfSSL_DH_set0_pqg(WOLFSSL_DH*, WOLFSSL_BIGNUM*,
     WOLFSSL_BIGNUM*, WOLFSSL_BIGNUM*);
                   
@@ -80,6 +81,7 @@ WOLFSSL_API int wolfSSL_DH_set0_pqg(WOLFSSL_DH*, WOLFSSL_BIGNUM*,
 #define DH_size         wolfSSL_DH_size
 #define DH_generate_key wolfSSL_DH_generate_key
 #define DH_compute_key  wolfSSL_DH_compute_key
+#define DH_set_length   wolfSSL_DH_set_length
 #define DH_set0_pqg     wolfSSL_DH_set0_pqg
 #define DH_get0_pqg     wolfSSL_DH_get0_pqg
 #define DH_get0_key     wolfSSL_DH_get0_key