Index: azure_iot_hub_f767zi/trunk/wolfssl-4.7.0/wolfcrypt/src/dh.c
===================================================================
--- azure_iot_hub_f767zi/trunk/wolfssl-4.4.0/wolfcrypt/src/dh.c	(revision 457)
+++ azure_iot_hub_f767zi/trunk/wolfssl-4.7.0/wolfcrypt/src/dh.c	(revision 464)
@@ -931,8 +931,8 @@
     key->heap = heap; /* for XMALLOC/XFREE in future */
 
-#if !defined(WOLFSSL_QT) && !defined(OPENSSL_ALL)
+#ifdef WOLFSSL_DH_EXTRA
+    if (mp_init_multi(&key->p, &key->g, &key->q, &key->pub, &key->priv, NULL) != MP_OKAY)
+#else
     if (mp_init_multi(&key->p, &key->g, &key->q, NULL, NULL, NULL) != MP_OKAY)
-#else
-    if (mp_init_multi(&key->p,&key->g,&key->q,&key->pub,&key->priv,NULL) != MP_OKAY)
 #endif
         return MEMORY_E;
@@ -961,4 +961,8 @@
         mp_clear(&key->g);
         mp_clear(&key->q);
+    #ifdef WOLFSSL_DH_EXTRA
+        mp_clear(&key->pub);
+        mp_forcezero(&key->priv);
+    #endif
 
     #if defined(WOLFSSL_ASYNC_CRYPT) && defined(WC_ASYNC_ENABLE_DH)
@@ -1149,5 +1153,4 @@
 
     mp_forcezero(tmpX);
-    mp_clear(tmpX);
     mp_clear(tmpQ);
 #ifdef WOLFSSL_SMALL_STACK
@@ -1207,5 +1210,9 @@
         }
 
-        ret = wc_RNG_GenerateBlock(rng, priv, sz);
+        if (sz > *privSz)
+            ret = WC_KEY_SIZE_E;
+
+        if (ret == 0)
+            ret = wc_RNG_GenerateBlock(rng, priv, sz);
 
         if (ret == 0) {
@@ -1231,4 +1238,5 @@
     int ret = 0;
 #ifndef WOLFSSL_SP_MATH
+    word32 binSz = 0;
 #ifdef WOLFSSL_SMALL_STACK
     mp_int* x;
@@ -1255,5 +1263,5 @@
 #endif
 
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
 #ifdef WOLFSSL_SMALL_STACK
     x = (mp_int*)XMALLOC(sizeof(mp_int), key->heap, DYNAMIC_TYPE_DH);
@@ -1280,9 +1288,16 @@
         ret = MP_EXPTMOD_E;
 
+    if (ret == 0) {
+        binSz = mp_unsigned_bin_size(y);
+        if (binSz > *pubSz) {
+            ret = WC_KEY_SIZE_E;
+        }
+    }
+
     if (ret == 0 && mp_to_unsigned_bin(y, pub) != MP_OKAY)
         ret = MP_TO_E;
 
     if (ret == 0)
-        *pubSz = mp_unsigned_bin_size(y);
+        *pubSz = binSz;
 
     mp_clear(y);
@@ -1468,4 +1483,5 @@
     }
 
+    /* SP 800-56Ar3, section 5.6.2.3.1, process step 2 */
     if (ret == 0 && prime != NULL) {
 #ifdef WOLFSSL_HAVE_SP_DH
@@ -1486,5 +1502,5 @@
         else
 #endif
-#ifdef WOLFSSL_SP_NO_4096
+#ifdef WOLFSSL_SP_4096
         if (mp_count_bits(&key->p) == 4096) {
             ret = sp_ModExp_4096(y, q, p, y);
@@ -1497,6 +1513,5 @@
 
         {
-    /* SP 800-56Ar3, section 5.6.2.3.1, process step 2 */
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
             /* calculate (y^q) mod(p), store back into y */
             if (mp_exptmod(y, q, p, y) != MP_OKAY)
@@ -1783,5 +1798,5 @@
 #endif
         {
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
             if (mp_exptmod(&key->g, privateKey, &key->p, checkKey) != MP_OKAY)
                 ret = MP_EXPTMOD_E;
@@ -1799,5 +1814,4 @@
 
     mp_forcezero(privateKey);
-    mp_clear(privateKey);
     mp_clear(publicKey);
     mp_clear(checkKey);
@@ -1835,5 +1849,4 @@
 }
 
-
 static int wc_DhAgree_Sync(DhKey* key, byte* agree, word32* agreeSz,
     const byte* priv, word32 privSz, const byte* otherPub, word32 pubSz)
@@ -1841,12 +1854,12 @@
     int ret = 0;
 #ifdef WOLFSSL_SMALL_STACK
-    mp_int* y;
-#ifndef WOLFSSL_SP_MATH
-    mp_int* x;
-    mp_int* z;
+    mp_int* y = NULL;
+#if !defined(WOLFSSL_SP_MATH)
+    mp_int* x = NULL;
+    mp_int* z = NULL;
 #endif
 #else
     mp_int y[1];
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
     mp_int x[1];
     mp_int z[1];
@@ -1870,5 +1883,5 @@
     if (y == NULL)
         return MEMORY_E;
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
     x = (mp_int*)XMALLOC(sizeof(mp_int), key->heap, DYNAMIC_TYPE_DH);
     if (x == NULL) {
@@ -1899,5 +1912,5 @@
         mp_clear(y);
     #ifdef WOLFSSL_SMALL_STACK
-    #ifndef WOLFSSL_SP_MATH
+    #if !defined(WOLFSSL_SP_MATH)
         XFREE(z, key->heap, DYNAMIC_TYPE_DH);
         XFREE(x, key->heap, DYNAMIC_TYPE_DH);
@@ -1921,5 +1934,5 @@
         mp_clear(y);
     #ifdef WOLFSSL_SMALL_STACK
-    #ifndef WOLFSSL_SP_MATH
+    #if !defined(WOLFSSL_SP_MATH)
         XFREE(z, key->heap, DYNAMIC_TYPE_DH);
         XFREE(x, key->heap, DYNAMIC_TYPE_DH);
@@ -1943,5 +1956,5 @@
         mp_clear(y);
     #ifdef WOLFSSL_SMALL_STACK
-    #ifndef WOLFSSL_SP_MATH
+    #if !defined(WOLFSSL_SP_MATH)
         XFREE(z, key->heap, DYNAMIC_TYPE_DH);
         XFREE(x, key->heap, DYNAMIC_TYPE_DH);
@@ -1954,5 +1967,5 @@
 #endif
 
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
     if (mp_init_multi(x, y, z, 0, 0, 0) != MP_OKAY) {
     #ifdef WOLFSSL_SMALL_STACK
@@ -1986,8 +1999,10 @@
     mp_clear(y);
     mp_forcezero(x);
+#else
+    ret = WC_KEY_SIZE_E;
 #endif
 
 #ifdef WOLFSSL_SMALL_STACK
-#ifndef WOLFSSL_SP_MATH
+#if !defined(WOLFSSL_SP_MATH)
     XFREE(z, key->heap, DYNAMIC_TYPE_DH);
     XFREE(x, key->heap, DYNAMIC_TYPE_DH);
@@ -2066,15 +2081,51 @@
 }
 
-#if defined(WOLFSSL_QT) || defined(OPENSSL_ALL)
+#ifdef WOLFSSL_DH_EXTRA
+WOLFSSL_LOCAL int wc_DhKeyCopy(DhKey* src, DhKey* dst)
+{
+    int ret;
+
+    if (!src || !dst || src == dst) {
+        WOLFSSL_MSG("Parameters not provided or are the same");
+        return BAD_FUNC_ARG;
+    }
+
+    if ((ret = mp_copy(&src->p, &dst->p)) != MP_OKAY) {
+        WOLFSSL_MSG("mp_copy error");
+        return ret;
+    }
+
+    if ((ret = mp_copy(&src->g, &dst->g)) != MP_OKAY) {
+        WOLFSSL_MSG("mp_copy error");
+        return ret;
+    }
+
+    if ((ret = mp_copy(&src->q, &dst->q)) != MP_OKAY) {
+        WOLFSSL_MSG("mp_copy error");
+        return ret;
+    }
+
+    if ((ret = mp_copy(&src->pub, &dst->pub)) != MP_OKAY) {
+        WOLFSSL_MSG("mp_copy error");
+        return ret;
+    }
+
+    if ((ret = mp_copy(&src->priv, &dst->priv)) != MP_OKAY) {
+        WOLFSSL_MSG("mp_copy error");
+        return ret;
+    }
+
+    dst->heap = src->heap;
+
+    return MP_OKAY;
+}
+
 /* Sets private and public key in DhKey if both are available, otherwise sets
-    either private or public key, depending on which is available.
-    Returns WOLFSSL_SUCCESS if at least one of the keys was set. */
-WOLFSSL_LOCAL int wc_DhSetFullKeys(DhKey* key,const byte* priv_key,word32 privSz,
-                                   const byte* pub_key, word32 pubSz)
-{
-    byte havePriv = 0;
-    byte havePub = 0;
-    mp_int* keyPriv = NULL;
-    mp_int* keyPub  = NULL;
+    either private or public key, depending on which is available. */
+int wc_DhImportKeyPair(DhKey* key, const byte* priv, word32 privSz,
+                       const byte* pub, word32 pubSz)
+{
+    byte havePriv, havePub;
+    mp_int *keyPriv = NULL, *keyPub  = NULL;
 
     if (key == NULL) {
@@ -2082,6 +2133,6 @@
     }
 
-    havePriv = ( (priv_key != NULL) && (privSz > 0) );
-    havePub  = ( (pub_key  != NULL) && (pubSz  > 0) );
+    havePriv = ( (priv != NULL) && (privSz > 0) );
+    havePub  = ( (pub  != NULL) && (pubSz  > 0) );
 
     if (!havePub && !havePriv) {
@@ -2089,51 +2140,91 @@
         return BAD_FUNC_ARG;
     }
+
     /* Set Private Key */
-    if (havePriv == TRUE) {
+    if (havePriv) {
         /* may have leading 0 */
-        if (priv_key[0] == 0) {
-            privSz--; priv_key++;
+        if (priv[0] == 0) {
+            privSz--; priv++;
         }
         if (mp_init(&key->priv) != MP_OKAY)
-            havePriv = FALSE;
-    }
-
-    if (havePriv == TRUE) {
-        if (mp_read_unsigned_bin(&key->priv, priv_key, privSz) != MP_OKAY) {
-            havePriv = FALSE;
+            havePriv = 0;
+    }
+    if (havePriv) {
+        if (mp_read_unsigned_bin(&key->priv, priv, privSz) != MP_OKAY) {
+            mp_clear(&key->priv);
+            havePriv = 0;
         } else {
             keyPriv = &key->priv;
-            WOLFSSL_MSG("DH Private Key Set.");
+            WOLFSSL_MSG("DH Private Key Set");
         }
     }
 
     /* Set Public Key */
-    if (havePub == TRUE) {
+    if (havePub) {
         /* may have leading 0 */
-        if (pub_key[0] == 0) {
-            pubSz--; pub_key++;
+        if (pub[0] == 0) {
+            pubSz--; pub++;
         }
         if (mp_init(&key->pub) != MP_OKAY)
-            havePub = FALSE;
-    }
-
-    if (havePub == TRUE) {
-        if (mp_read_unsigned_bin(&key->pub, pub_key, pubSz) != MP_OKAY) {
-            havePub = FALSE;
+            havePub = 0;
+    }
+    if (havePub) {
+        if (mp_read_unsigned_bin(&key->pub, pub, pubSz) != MP_OKAY) {
+            mp_clear(&key->pub);
+            havePub = 0;
         } else {
             keyPub = &key->pub;
-            WOLFSSL_MSG("DH Public Key Set.");
-        }
-    }
-    /* Free Memory if error occured */
-    if (havePriv == FALSE && keyPriv != NULL)
+            WOLFSSL_MSG("DH Public Key Set");
+        }
+    }
+    /* Free Memory if error occurred */
+    if (havePriv == 0 && keyPriv != NULL)
         mp_clear(keyPriv);
-    if (havePub == FALSE && keyPub != NULL)
+    if (havePub == 0 && keyPub != NULL)
         mp_clear(keyPub);
 
-    /* WOLFSSL_SUCCESS if private or public was set else WOLFSSL_FAILURE */
-    return havePriv || havePub;
-}
-#endif
+    if (havePriv == 0 && havePub == 0) {
+        return MEMORY_E;
+    }
+
+    return 0;
+}
+
+/* Can be used with WOLFSSL_DH_EXTRA when key is loaded with 
+    wc_DhKeyDecode or wc_DhImportKeyPair */
+int wc_DhExportKeyPair(DhKey* key, byte* priv, word32* pPrivSz, 
+    byte* pub, word32* pPubSz)
+{
+    int ret = 0;
+    word32 pubSz, privSz;
+
+    if (key == NULL || (priv && pPrivSz == NULL) || (pub && pPubSz == NULL)) {
+        return BAD_FUNC_ARG;
+    }
+
+    if (priv) {
+        privSz = mp_unsigned_bin_size(&key->priv);
+        if (privSz > *pPrivSz) {
+            return BUFFER_E;
+        }
+        *pPrivSz = privSz;
+        ret |= mp_to_unsigned_bin(&key->priv, priv);
+    }
+
+    if (pub) {
+        pubSz = mp_unsigned_bin_size(&key->pub);
+        if (pubSz > *pPubSz) {
+            return BUFFER_E;
+        }
+        *pPubSz = pubSz;
+        ret |= mp_to_unsigned_bin(&key->pub,  pub);
+    }
+
+    if (ret != 0)
+        ret = ASN_DH_KEY_E;
+    return ret;
+}
+
+#endif /* WOLFSSL_DH_EXTRA */
 
 static int _DhSetKey(DhKey* key, const byte* p, word32 pSz, const byte* g,
@@ -2368,5 +2459,10 @@
     if (ret == 0) {
         /* at this point tmp generates a group of order q mod p */
+#ifndef USE_FAST_MATH
+        /* Exchanging is quick when the data pointer can be copied. */
         mp_exch(&tmp, &dh->g);
+#else
+        mp_copy(&tmp, &dh->g);
+#endif
     }
 
