1 files changed, 11 insertions, 13 deletions
diff --git a/security/nss/lib/freebl/dh.c b/security/nss/lib/freebl/dh.c
index 97025c7e2..6f2bafda2 100644
--- a/security/nss/lib/freebl/dh.c
+++ b/security/nss/lib/freebl/dh.c
@@ -14,9 +14,9 @@
 #include "secerr.h"
 
 #include "blapi.h"
+#include "blapii.h"
 #include "secitem.h"
 #include "mpi.h"
-#include "mpprime.h"
 #include "secmpi.h"
 
 #define KEA_DERIVED_SECRET_LEN 128
@@ -46,9 +46,7 @@ DH_GenParam(int primeLen, DHParams **params)
 {
     PLArenaPool *arena;
     DHParams *dhparams;
-    unsigned char *pb = NULL;
     unsigned char *ab = NULL;
-    unsigned long counter = 0;
     mp_int p, q, a, h, psub1, test;
     mp_err err = MP_OKAY;
     SECStatus rv = SECSuccess;
@@ -81,17 +79,17 @@ DH_GenParam(int primeLen, DHParams **params)
     CHECK_MPI_OK(mp_init(&psub1));
     CHECK_MPI_OK(mp_init(&test));
     /* generate prime with MPI, uses Miller-Rabin to generate strong prime. */
-    pb = PORT_Alloc(primeLen);
-    CHECK_SEC_OK(RNG_GenerateGlobalRandomBytes(pb, primeLen));
-    pb[0] |= 0x80;            /* set high-order bit */
-    pb[primeLen - 1] |= 0x01; /* set low-order bit  */
-    CHECK_MPI_OK(mp_read_unsigned_octets(&p, pb, primeLen));
-    CHECK_MPI_OK(mpp_make_prime(&p, primeLen * 8, PR_TRUE, &counter));
+    CHECK_SEC_OK(generate_prime(&p, primeLen));
     /* construct Sophie-Germain prime q = (p-1)/2. */
     CHECK_MPI_OK(mp_sub_d(&p, 1, &psub1));
     CHECK_MPI_OK(mp_div_2(&psub1, &q));
     /* construct a generator from the prime. */
     ab = PORT_Alloc(primeLen);
+    if (!ab) {
+        PORT_SetError(SEC_ERROR_NO_MEMORY);
+        rv = SECFailure;
+        goto cleanup;
+    }
     /* generate a candidate number a in p's field */
     CHECK_SEC_OK(RNG_GenerateGlobalRandomBytes(ab, primeLen));
     CHECK_MPI_OK(mp_read_unsigned_octets(&a, ab, primeLen));
@@ -121,16 +119,16 @@ cleanup:
     mp_clear(&h);
     mp_clear(&psub1);
     mp_clear(&test);
-    if (pb)
-        PORT_ZFree(pb, primeLen);
-    if (ab)
+    if (ab) {
         PORT_ZFree(ab, primeLen);
+    }
     if (err) {
         MP_TO_SEC_ERROR(err);
         rv = SECFailure;
     }
-    if (rv)
+    if (rv != SECSuccess) {
         PORT_FreeArena(arena, PR_TRUE);
+    }
     return rv;
 }