From f776371874db55119c0f53466f02aa14dbf5eb98 Mon Sep 17 00:00:00 2001
From: Daniel Pouzzner <douzzer@wolfssl.com>
Date: Thu, 26 Jan 2023 15:00:56 -0600
Subject: [PATCH] wolfcrypt/src/hpke.c: add PRIVATE_KEY_{UNLOCK,LOCK}()
 wrappers in wc_HpkeSealBase() and wc_HpkeOpenBase();

wolfcrypt/test/test.c: remove PRIVATE_KEY_{UNLOCK,LOCK}() wrappers from hpke_test_single(), and do a smallstack refactor.
---
 wolfcrypt/src/hpke.c  | 11 ++++++++++-
 wolfcrypt/test/test.c | 23 +++++++++++++++++++----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/wolfcrypt/src/hpke.c b/wolfcrypt/src/hpke.c
index 604f9fb0b..b9e3c3b2c 100644
--- a/wolfcrypt/src/hpke.c
+++ b/wolfcrypt/src/hpke.c
@@ -942,14 +942,19 @@ int wc_HpkeSealBase(Hpke* hpke, void* ephemeralKey, void* receiverKey,
     }
 #endif
 
+    PRIVATE_KEY_UNLOCK();
+
     /* setup the context and pubKey */
     ret = wc_HpkeSetupBaseSender(hpke, context, ephemeralKey, receiverKey, info,
         infoSz);
 
     /* run seal using the context */
-    if (ret == 0)
+    if (ret == 0) {
         ret = wc_HpkeContextSealBase(hpke, context, aad, aadSz, plaintext,
             ptSz, ciphertext);
+    }
+
+    PRIVATE_KEY_LOCK();
 
 #ifdef WOLFSSL_SMALL_STACK
     XFREE(context, hpke->heap, DYNAMIC_TYPE_TMP_BUFFER);
@@ -1174,6 +1179,8 @@ int wc_HpkeOpenBase(Hpke* hpke, void* receiverKey, const byte* pubKey,
     }
 #endif
 
+    PRIVATE_KEY_UNLOCK();
+
     /* setup receiver */
     ret = wc_HpkeSetupBaseReceiver(hpke, context, receiverKey, pubKey,
         pubKeySz, info, infoSz);
@@ -1184,6 +1191,8 @@ int wc_HpkeOpenBase(Hpke* hpke, void* receiverKey, const byte* pubKey,
             ctSz, plaintext);
     }
 
+    PRIVATE_KEY_LOCK();
+
 #ifdef WOLFSSL_SMALL_STACK
     XFREE(context, hpke->heap, DYNAMIC_TYPE_TMP_BUFFER);
 #endif
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 339e50918..0453f089a 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -22524,14 +22524,28 @@ static int hpke_test_single(Hpke* hpke)
     byte plaintext[MAX_HPKE_LABEL_SZ];
     void* receiverKey = NULL;
     void* ephemeralKey = NULL;
+#ifdef WOLFSSL_SMALL_STACK
+    uint8_t *pubKey = NULL; /* public key */
+    word16 pubKeySz = (word16)HPKE_Npk_MAX;
+#else
     uint8_t pubKey[HPKE_Npk_MAX]; /* public key */
     word16 pubKeySz = (word16)sizeof(pubKey);
+#endif
 
     rngRet = ret = wc_InitRng(rng);
 
     if (ret != 0)
         return ret;
 
+#ifdef WOLFSSL_SMALL_STACK
+    if (ret == 0) {
+        pubKey = (uint8_t *)XMALLOC(pubKeySz, HEAP_HINT,
+                                    DYNAMIC_TYPE_TMP_BUFFER);
+        if (pubKey == NULL)
+            ret = MEMORY_E;
+    }
+#endif
+
     /* generate the keys */
     if (ret == 0)
         ret = wc_HpkeGenerateKeyPair(hpke, &ephemeralKey, rng);
@@ -22541,13 +22555,11 @@ static int hpke_test_single(Hpke* hpke)
 
     /* seal */
     if (ret == 0) {
-        PRIVATE_KEY_UNLOCK();
         ret = wc_HpkeSealBase(hpke, ephemeralKey, receiverKey,
             (byte*)info_text, (word32)XSTRLEN(info_text),
             (byte*)aad_text, (word32)XSTRLEN(aad_text),
             (byte*)start_text, (word32)XSTRLEN(start_text),
             ciphertext);
-        PRIVATE_KEY_LOCK();
     }
 
     /* export ephemeral key */
@@ -22556,13 +22568,11 @@ static int hpke_test_single(Hpke* hpke)
 
     /* open with exported ephemeral key */
     if (ret == 0) {
-        PRIVATE_KEY_UNLOCK();
         ret = wc_HpkeOpenBase(hpke, receiverKey, pubKey, pubKeySz,
             (byte*)info_text, (word32)XSTRLEN(info_text),
             (byte*)aad_text, (word32)XSTRLEN(aad_text),
             ciphertext, (word32)XSTRLEN(start_text),
             plaintext);
-        PRIVATE_KEY_LOCK();
     }
 
     if (ret == 0)
@@ -22574,6 +22584,11 @@ static int hpke_test_single(Hpke* hpke)
     if (receiverKey != NULL)
         wc_HpkeFreeKey(hpke, hpke->kem, receiverKey, hpke->heap);
 
+#ifdef WOLFSSL_SMALL_STACK
+    if (pubKey != NULL)
+        XFREE(pubKey, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+#endif
+
     if (rngRet == 0)
         wc_FreeRng(rng);