From cfafbd96590a5e8e2bf80ff40e61f6d47c64dfe4 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Fri, 2 Nov 2018 11:01:39 -0700
Subject: [PATCH] Added the prime check to the functions wolfSSL_SetTmpDh() and
 wolfSSL_CTX_SetTmpDh().

---
 src/ssl.c | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/src/ssl.c b/src/ssl.c
index 5ec1a3297..6bcd80b44 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -1477,6 +1477,24 @@ int wolfSSL_SetTmpDH(WOLFSSL* ssl, const unsigned char* p, int pSz,
     if (ssl->options.side == WOLFSSL_CLIENT_END)
         return SIDE_ERROR;
 
+    #ifndef WOLFSSL_OLD_PRIME_CHECK
+    {
+        DhKey checkKey;
+        int error, freeKey = 0;
+
+        error = wc_InitDhKey(&checkKey);
+        if (!error) {
+            freeKey = 1;
+            error = wc_DhSetCheckKey(&checkKey,
+                                 p, pSz, g, gSz, NULL, 0, 0, ssl->rng);
+        }
+        if (freeKey)
+            wc_FreeDhKey(&checkKey);
+        if (error)
+            return error;
+    }
+    #endif
+
     if (ssl->buffers.serverDH_P.buffer && ssl->buffers.weOwnDH) {
         XFREE(ssl->buffers.serverDH_P.buffer, ssl->heap, DYNAMIC_TYPE_PUBLIC_KEY);
         ssl->buffers.serverDH_P.buffer = NULL;
@@ -1545,6 +1563,28 @@ int wolfSSL_CTX_SetTmpDH(WOLFSSL_CTX* ctx, const unsigned char* p, int pSz,
     if (pSz > ctx->maxDhKeySz)
         return DH_KEY_SIZE_E;
 
+    #ifndef WOLFSSL_OLD_PRIME_CHECK
+    {
+        DhKey checkKey;
+        WC_RNG rng;
+        int error, freeKey = 0;
+
+        error = wc_InitRng(&rng);
+        if (!error)
+            error = wc_InitDhKey(&checkKey);
+        if (!error) {
+            freeKey = 1;
+            error = wc_DhSetCheckKey(&checkKey,
+                                 p, pSz, g, gSz, NULL, 0, 0, &rng);
+        }
+        if (freeKey)
+            wc_FreeDhKey(&checkKey);
+        wc_FreeRng(&rng);
+        if (error)
+            return error;
+    }
+    #endif
+
     XFREE(ctx->serverDH_P.buffer, ctx->heap, DYNAMIC_TYPE_PUBLIC_KEY);
     ctx->serverDH_P.buffer = NULL;
     XFREE(ctx->serverDH_G.buffer, ctx->heap, DYNAMIC_TYPE_PUBLIC_KEY);