From 186d3c2eb494e7799cc3b410adfb34ce094e9a3d Mon Sep 17 00:00:00 2001
From: Andras Fekete <andras@wolfssl.com>
Date: Mon, 18 Sep 2023 13:12:00 -0400
Subject: [PATCH] Fixes to various Async issues

---
 src/internal.c      | 19 +++++++++++++------
 src/tls.c           | 31 +++++++++++++++++++++----------
 wolfcrypt/src/ecc.c |  5 -----
 3 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index fadb60e54..382830b6f 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -5349,7 +5349,7 @@ int EccMakeKey(WOLFSSL* ssl, ecc_key* key, ecc_key* peer)
 
 #ifdef WOLFSSL_ASYNC_CRYPT
     /* initialize event */
-    ret = wolfSSL_AsyncInit(ssl, &key->asyncDev, WC_ASYNC_FLAG_NONE);
+    ret = wolfSSL_AsyncInit(ssl, &key->asyncDev, WC_ASYNC_FLAG_CALL_AGAIN);
     if (ret != 0)
         return ret;
 #endif
@@ -38733,18 +38733,25 @@ int wolfSSL_AsyncPop(WOLFSSL* ssl, byte* state)
 
         ret = wolfAsync_EventPop(event, WOLF_EVENT_TYPE_ASYNC_WOLFSSL);
         if (ret != WC_NO_PENDING_E && ret != WC_PENDING_E) {
-
             /* advance key share state if doesn't need called again */
             if (state && (asyncDev->event.flags & WC_ASYNC_FLAG_CALL_AGAIN) == 0) {
                 (*state)++;
             }
-
-            /* clear event */
+            /* clear event and async device */
             XMEMSET(&asyncDev->event, 0, sizeof(WOLF_EVENT));
-
-            /* clear async dev */
             ssl->asyncDev = NULL;
         }
+    #if !defined(WOLFSSL_ASYNC_CRYPT_SW) && \
+        (defined(WOLF_CRYPTO_CB) || defined(HAVE_PK_CALLBACKS))
+        else if (ret == WC_PENDING_E) {
+            /* Allow the underlying crypto API to be called again to trigger the
+             * crypto or PK callback. The actual callback must be called, since
+             * the completion is not detected in the poll like Intel QAT or
+             * Nitrox */
+            ret = wolfEventQueue_Remove(&ssl->ctx->event_queue, event);
+            printf("Queue_Remove: %d",ret);
+        }
+    #endif
     }
     else {
         ret = WC_NO_PENDING_E;
diff --git a/src/tls.c b/src/tls.c
index 92bac11b5..f719e186f 100644
--- a/src/tls.c
+++ b/src/tls.c
@@ -7392,7 +7392,7 @@ static int TLSX_KeyShare_GenEccKey(WOLFSSL *ssl, KeyShareEntry* kse)
     word16 curveId = (word16) ECC_CURVE_INVALID;
     ecc_key* eccKey = (ecc_key*)kse->key;
 
-    /* TODO: [TLS13] The key sizes should come from wolfcrypt. */
+    /* TODO: [TLS13] Get key sizes using wc_ecc_get_curve_size_from_id. */
     /* Translate named group to a curve id. */
     switch (kse->group) {
     #if (!defined(NO_ECC256)  || defined(HAVE_ALL_CURVES)) && ECC_MIN_KEY_SZ <= 256
@@ -7431,9 +7431,6 @@ static int TLSX_KeyShare_GenEccKey(WOLFSSL *ssl, KeyShareEntry* kse)
     }
 
     if (kse->key == NULL) {
-        kse->keyLen = keySize;
-        kse->pubKeyLen = keySize * 2 + 1;
-
     #if defined(WOLFSSL_RENESAS_TSIP_TLS)
         ret = tsip_Tls13GenEccKeyPair(ssl, kse);
         if (ret != CRYPTOCB_UNAVAILABLE) {
@@ -7447,9 +7444,13 @@ static int TLSX_KeyShare_GenEccKey(WOLFSSL *ssl, KeyShareEntry* kse)
             return MEMORY_E;
         }
 
-        /* Make an ECC key */
+        /* Initialize an ECC key struct for the ephemeral key */
         ret = wc_ecc_init_ex((ecc_key*)kse->key, ssl->heap, ssl->devId);
+
         if (ret == 0) {
+            kse->keyLen = keySize;
+            kse->pubKeyLen = keySize * 2 + 1;
+
             /* setting eccKey means okay to call wc_ecc_free */
             eccKey = (ecc_key*)kse->key;
 
@@ -7461,11 +7462,21 @@ static int TLSX_KeyShare_GenEccKey(WOLFSSL *ssl, KeyShareEntry* kse)
                 /* set curve info for EccMakeKey "peer" info */
                 ret = wc_ecc_set_curve(eccKey, kse->keyLen, curveId);
                 if (ret == 0) {
-                    /* Generate ephemeral ECC key */
-                    /* For async this is called once and when event is done, the
-                    *   provided buffers in key be populated.
-                    * Final processing is x963 key export below. */
-                    ret = EccMakeKey(ssl, eccKey, eccKey);
+            #ifdef WOLFSSL_ASYNC_CRYPT
+                    /* Detect when private key generation is done */
+                    if (ssl->error == WC_PENDING_E &&
+                            eccKey->type == ECC_PRIVATEKEY) {
+                        ret = 0; /* ECC Key Generation is done */
+                    }
+                    else
+            #endif
+                    {
+                        /* Generate ephemeral ECC key */
+                        /* For async this is called once and when event is done, the
+                        *   provided buffers in key be populated.
+                        * Final processing is x963 key export below. */
+                        ret = EccMakeKey(ssl, eccKey, eccKey);
+                    }
                 }
             #ifdef WOLFSSL_ASYNC_CRYPT
                 if (ret == WC_PENDING_E)
diff --git a/wolfcrypt/src/ecc.c b/wolfcrypt/src/ecc.c
index 59ff57aa6..b3d883bc2 100644
--- a/wolfcrypt/src/ecc.c
+++ b/wolfcrypt/src/ecc.c
@@ -5006,11 +5006,6 @@ int wc_ecc_shared_secret_ex(ecc_key* private_key, ecc_point* point,
             if (private_key->asyncDev.marker == WOLFSSL_ASYNC_MARKER_ECC) {
                 err = wc_ecc_shared_secret_gen_async(private_key, point,
                     out, outlen);
-                if (err == 0) {
-                    /* exit early */
-                    RESTORE_VECTOR_REGISTERS();
-                    return err;
-                }
             }
             else
         #endif