From 338d8db27437b0ab92853379023d362b5d3d1c3e Mon Sep 17 00:00:00 2001
From: John Bland <106998124+jpbland1@users.noreply.github.com>
Date: Fri, 10 Feb 2023 13:05:52 -0500
Subject: [PATCH] Ecc ctx state fix (#6077)

* set the client state correctly when wc_ecc_ctx_set_kdf_salt is called
* add test that covers wc_ecc_ctx_set_kdf_salt
* use shared key and smallstack to reduce test stack usage
---
 wolfcrypt/src/ecc.c   |   2 +-
 wolfcrypt/test/test.c | 102 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/wolfcrypt/src/ecc.c b/wolfcrypt/src/ecc.c
index 10d8849bd..67e06ae40 100644
--- a/wolfcrypt/src/ecc.c
+++ b/wolfcrypt/src/ecc.c
@@ -13053,7 +13053,7 @@ int wc_ecc_ctx_set_kdf_salt(ecEncCtx* ctx, const byte* salt, word32 len)
     ctx->kdfSaltSz = len;
 
     if (ctx->protocol == REQ_RESP_CLIENT) {
-        ctx->srvSt = ecSRV_SALT_SET;
+        ctx->cliSt = ecCLI_SALT_SET;
     }
     else if (ctx->protocol == REQ_RESP_SERVER) {
         ctx->srvSt = ecSRV_SALT_SET;
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index e3a2dcb6b..f12318bcd 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -26591,6 +26591,105 @@ done:
 #if defined(HAVE_ECC_ENCRYPT) && defined(HAVE_AES_CBC) && \
     (defined(WOLFSSL_AES_128) || defined(WOLFSSL_AES_256))
 
+static int ecc_ctx_kdf_salt_test(WC_RNG* rng, ecc_key* a, ecc_key* b)
+{
+#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
+    byte* plaintext;
+    byte* encrypted;
+    byte* decrypted;
+#else
+    byte plaintext[128];
+    byte encrypted[128];
+    byte decrypted[128];
+#endif
+    ecEncCtx* aCtx = NULL;
+    ecEncCtx* bCtx = NULL;
+    const byte salt[16] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
+        15};
+    int ret = 0, aRet = -1, bRet = -1;
+    const char* message = "Hello wolfSSL!";
+    word32 plaintextLen = sizeof(message), encryptLen = 128, decryptLen = 128;
+
+#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
+    plaintext = XMALLOC(128, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    encrypted = XMALLOC(128, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    decrypted = XMALLOC(128, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+#endif
+
+    ret = aRet = wc_ecc_init(a);
+
+    if (ret == 0)
+        ret = bRet = wc_ecc_init(b);
+
+    if (ret == 0)
+        ret = wc_ecc_make_key(rng, 32, a);
+
+    if (ret == 0)
+        ret = wc_ecc_make_key(rng, 32, b);
+
+    /* create context */
+    if (ret == 0) {
+        aCtx = wc_ecc_ctx_new(REQ_RESP_CLIENT, rng);
+
+        if (aCtx == NULL)
+            ret = -1;
+    }
+
+    if (ret == 0) {
+        bCtx = wc_ecc_ctx_new(REQ_RESP_SERVER, rng);
+
+        if (bCtx == NULL)
+            ret = -1;
+    }
+
+    /* set salt */
+    if (ret == 0)
+        ret = wc_ecc_ctx_set_kdf_salt(aCtx, salt, sizeof(salt));
+
+    if (ret == 0)
+        ret = wc_ecc_ctx_set_kdf_salt(bCtx, salt, sizeof(salt));
+
+    XMEMCPY(plaintext, message, XSTRLEN(message));
+
+    while (plaintextLen % AES_BLOCK_SIZE != 0) {
+        plaintextLen++;
+    }
+
+    /* encrypt */
+    if (ret == 0)
+        ret = wc_ecc_encrypt(a, b, plaintext, plaintextLen, encrypted,
+            &encryptLen, aCtx);
+
+    /* decrypt */
+    if (ret == 0)
+        ret = wc_ecc_decrypt(b, a, encrypted, encryptLen, decrypted,
+            &decryptLen, bCtx);
+
+    /* compare */
+    if (ret == 0 && XMEMCMP(decrypted, (byte*)message, sizeof(message)) != 0)
+        ret = -1;
+
+    if (aRet == 0)
+        wc_ecc_free(a);
+
+    if (bRet == 0)
+        wc_ecc_free(b);
+
+    if (aCtx != NULL)
+        wc_ecc_ctx_free(aCtx);
+
+    if (bCtx != NULL)
+        wc_ecc_ctx_free(bCtx);
+
+#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
+    XFREE(plaintext, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(encrypted, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+    XFREE(decrypted, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
+#endif
+
+    return ret;
+}
+
 /* ecc_encrypt_e2e_test() uses wc_ecc_ctx_set_algo(), which was added in
  * wolfFIPS 5.3.
  * ecc_encrypt_kat() is used only by ecc_encrypt_e2e_test().
@@ -27218,6 +27317,9 @@ WOLFSSL_TEST_SUBROUTINE int ecc_encrypt_test(void)
         }
     }
 #endif
+    if (ret == 0) {
+        ret = ecc_ctx_kdf_salt_test(&rng, userA, userB);
+    }
 #endif
 
 #endif /* !HAVE_FIPS || FIPS_VERSION_GE(5,3) */