From bced81d2348474e01149d2dc03af6baac68f25fa Mon Sep 17 00:00:00 2001
From: David Garske <david@wolfssl.com>
Date: Tue, 31 Jan 2017 16:42:06 -0800
Subject: [PATCH] Improve handling of mp_init / mp_clear for DH and DSA after
 speed-up.

---
 wolfcrypt/src/dh.c  | 19 ++++++++++---------
 wolfcrypt/src/dsa.c | 37 +++++++++++++++++++++----------------
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/wolfcrypt/src/dh.c b/wolfcrypt/src/dh.c
index 15b557a76..c617fd823 100644
--- a/wolfcrypt/src/dh.c
+++ b/wolfcrypt/src/dh.c
@@ -51,20 +51,21 @@
 
 void wc_InitDhKey(DhKey* key)
 {
-    (void)key;
-/* TomsFastMath doesn't use memory allocation */
-#ifndef USE_FAST_MATH
-    key->p.dp = NULL;
-    key->g.dp = NULL;
-#endif
+    if (key) {
+        mp_init(&key->p);
+        mp_init(&key->g);
+    }
 }
 
 
 void wc_FreeDhKey(DhKey* key)
 {
-    (void)key;
-    mp_clear(&key->p);
-    mp_clear(&key->g);
+    if (key) {
+    #ifndef USE_FAST_MATH
+        mp_clear(&key->p);
+        mp_clear(&key->g);
+    #endif
+    }
 }
 
 
diff --git a/wolfcrypt/src/dsa.c b/wolfcrypt/src/dsa.c
index eaac64346..dd30d1bf7 100644
--- a/wolfcrypt/src/dsa.c
+++ b/wolfcrypt/src/dsa.c
@@ -52,18 +52,20 @@ enum {
 
 void wc_InitDsaKey(DsaKey* key)
 {
+    if (key == NULL)
+        return;
+
     key->type = -1;  /* haven't decided yet */
     key->heap = NULL;
 
-/* TomsFastMath doesn't use memory allocation */
-#ifndef USE_FAST_MATH
-    key->p.dp = 0;   /* public  alloc parts */
-    key->q.dp = 0;    
-    key->g.dp = 0;    
-    key->y.dp = 0;    
+    /* public  alloc parts */
+    mp_init(&key->p);
+    mp_init(&key->q);
+    mp_init(&key->g);
+    mp_init(&key->y);
 
-    key->x.dp = 0;   /* private alloc parts */
-#endif
+    /* private alloc parts */
+    mp_init(&key->x);
 }
 
 
@@ -78,11 +80,14 @@ int wc_InitDsaKey_h(DsaKey* key, void* h)
 
 void wc_FreeDsaKey(DsaKey* key)
 {
-    (void)key;
-/* TomsFastMath doesn't use memory allocation */
-#ifndef USE_FAST_MATH
+    if (key == NULL)
+        return;
+
     if (key->type == DSA_PRIVATE)
-        mp_clear(&key->x);
+        mp_forcezero(&key->x);
+
+#ifndef USE_FAST_MATH
+    mp_clear(&key->x);
     mp_clear(&key->y);
     mp_clear(&key->g);
     mp_clear(&key->q);
@@ -148,7 +153,7 @@ int wc_MakeDsaKey(WC_RNG *rng, DsaKey *dsa)
     }
 
     dsa->type = DSA_PRIVATE;
-    
+
     return MP_OKAY;
 }
 
@@ -356,7 +361,7 @@ int wc_DsaSign(const byte* digest, byte* out, DsaKey* key, WC_RNG* rng)
     byte*  tmp = out;  /* initial output pointer */
 
     sz = min((int)sizeof(buffer), mp_unsigned_bin_size(&key->q));
-    
+
     if (mp_init_multi(&k, &kInv, &r, &s, &H, 0) != MP_OKAY)
         return MP_INIT_E;
 
@@ -370,12 +375,12 @@ int wc_DsaSign(const byte* digest, byte* out, DsaKey* key, WC_RNG* rng)
 
         if (mp_read_unsigned_bin(&k, buffer, sz) != MP_OKAY)
             ret = MP_READ_E;
-            
+
         /* k is a random numnber and it should be less than q
          * if k greater than repeat
          */
     } while (mp_cmp(&k, &key->q) != MP_LT);
-    
+
     if (ret == 0 && mp_cmp_d(&k, 1) != MP_GT)
         ret = MP_CMP_E;