more work on DECLARE_VAR -- proper handling of failed allocations. WIP.

wolfcrypt/benchmark/benchmark.c

@@ -2224,6 +2224,12 @@ static void bench_aesgcm_internal(int doAsync, const byte* key, word32 keySz,
 
     DECLARE_VAR(bench_additional, byte, AES_AUTH_ADD_SZ, HEAP_HINT);
     DECLARE_VAR(bench_tag, byte, AES_AUTH_TAG_SZ, HEAP_HINT);
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((bench_additional == NULL) || (bench_tag == NULL)) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
 
     /* clear for done cleanup */
     XMEMSET(enc, 0, sizeof(enc));
@@ -2663,13 +2669,19 @@ void bench_aesccm(void)
 
     DECLARE_VAR(bench_additional, byte, AES_AUTH_ADD_SZ, HEAP_HINT);
     DECLARE_VAR(bench_tag, byte, AES_AUTH_TAG_SZ, HEAP_HINT);
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((bench_additional == NULL) || (bench_tag == NULL)) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
 
     XMEMSET(bench_tag, 0, AES_AUTH_TAG_SZ);
     XMEMSET(bench_additional, 0, AES_AUTH_ADD_SZ);
 
     if ((ret = wc_AesCcmSetKey(&enc, bench_key, 16)) != 0) {
         printf("wc_AesCcmSetKey failed, ret = %d\n", ret);
-        return;
+        goto exit;
     }
 
     bench_stats_start(&count, &start);
@@ -2694,6 +2706,7 @@ void bench_aesccm(void)
     } while (bench_stats_sym_check(start));
     bench_stats_sym_finish("AES-CCM-Dec", 0, count, bench_size, start, ret);
 
+  exit:
 
     FREE_VAR(bench_additional, HEAP_HINT);
     FREE_VAR(bench_tag, HEAP_HINT);
@@ -4472,6 +4485,12 @@ static void bench_rsa_helper(int doAsync, RsaKey rsaKey[BENCH_MAX_PENDING],
     const char**desc = bench_desc_words[lng_index];
 #ifndef WOLFSSL_RSA_VERIFY_ONLY
     DECLARE_VAR_INIT(message, byte, len, messageStr, HEAP_HINT);
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if (message == NULL) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
 #endif
     #if !defined(WOLFSSL_MDK5_COMPLv5)
     /* MDK5 compiler regard this as a executable statement, and does not allow declarations after the line. */
@@ -4493,9 +4512,17 @@ static void bench_rsa_helper(int doAsync, RsaKey rsaKey[BENCH_MAX_PENDING],
     #endif
 
     DECLARE_ARRAY_DYNAMIC_EXE(enc, byte, BENCH_MAX_PENDING, rsaKeySz, HEAP_HINT);
+    if (enc[0] == NULL) {
+        printf("malloc failed\n");
+        goto exit;
+    }
     #if !defined(WOLFSSL_RSA_VERIFY_INLINE) && \
                     !defined(WOLFSSL_RSA_PUBLIC_ONLY)
         DECLARE_ARRAY_DYNAMIC_EXE(out, byte, BENCH_MAX_PENDING, rsaKeySz, HEAP_HINT);
+        if (out[0] == NULL) {
+            printf("malloc failed\n");
+            goto exit;
+        }
     #endif
 
     if (!rsa_sign_verify) {
@@ -4842,6 +4869,12 @@ void bench_dh(int doAsync)
     DECLARE_ARRAY(agree, byte, BENCH_MAX_PENDING, BENCH_DH_KEY_SIZE, HEAP_HINT);
     DECLARE_ARRAY(priv, byte, BENCH_MAX_PENDING, BENCH_DH_PRIV_SIZE, HEAP_HINT);
     DECLARE_VAR(priv2, byte, BENCH_DH_PRIV_SIZE, HEAP_HINT);
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((pub == NULL) || (pub2 == NULL) || (agree == NULL) || (priv == NULL) || (priv2 == NULL)) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
 
     (void)tmp;
 
@@ -5279,6 +5312,21 @@ void bench_ecc(int doAsync)
     DECLARE_ARRAY(digest, byte, BENCH_MAX_PENDING, BENCH_ECC_SIZE, HEAP_HINT);
 #endif
 
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+#ifdef HAVE_ECC_DHE
+    if (shared == NULL) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
+#if !defined(NO_ASN) && defined(HAVE_ECC_SIGN)
+    if ((sig == NULL) || (digest == NULL)) {
+        printf("malloc failed\n");
+        goto exit;
+    }
+#endif
+#endif
+
     /* clear for done cleanup */
     XMEMSET(&genKey, 0, sizeof(genKey));
 #ifdef HAVE_ECC_DHE

wolfcrypt/test/test.c

@@ -1149,7 +1149,7 @@ initDefaultName();
         test_pass("mp       test passed!\n");
 #endif
 
-#if defined(WOLFSSL_PUBLIC_MP) && defined(WOLFSSL_KEY_GEN)
+#if defined(WOLFSSL_PUBLIC_MP) && defined(WOLFSSL_KEY_GEN) && !defined(WOLFSSL_OLD_PRIME_CHECK)
     if ( (ret = prime_test()) != 0)
         return err_sys("prime    test failed!\n", ret);
     else
@@ -11952,6 +11952,11 @@ static int rsa_pss_test(WC_RNG* rng, RsaKey* key)
     DECLARE_VAR(out, byte, RSA_TEST_BYTES, HEAP_HINT);
     DECLARE_VAR(sig, byte, RSA_TEST_BYTES, HEAP_HINT);
 
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((in == NULL) || (out == NULL) || (sig == NULL))
+        ERROR_OUT(MEMORY_E, exit_rsa_pss);
+#endif
+
     /* Test all combinations of hash and MGF. */
     for (j = 0; j < (int)(sizeof(hash)/sizeof(*hash)); j++) {
         /* Calculate hash of message. */
@@ -12237,6 +12242,11 @@ static int rsa_no_pad_test(void)
     DECLARE_VAR(out, byte, RSA_TEST_BYTES, HEAP_HINT);
     DECLARE_VAR(plain, byte, RSA_TEST_BYTES, HEAP_HINT);
 
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((out == NULL) || (plain == NULL))
+        ERROR_OUT(MEMORY_E, exit_rsa_nopadding);
+#endif
+
     /* initialize stack structures */
     XMEMSET(&rng, 0, sizeof(rng));
     XMEMSET(&key, 0, sizeof(key));
@@ -13050,9 +13060,9 @@ static int rsa_test(void)
     DECLARE_VAR(plain, byte, RSA_TEST_BYTES, HEAP_HINT);
 #endif
 
-#ifdef WOLFSSL_ASYNC_CRYPT
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
-    if (in == NULL)
+    if ((in == NULL) || (out == NULL) || (plain == NULL))
-        return MEMORY_E;
+        ERROR_OUT(MEMORY_E, exit_rsa);
 #endif
 
     /* initialize stack structures */
@@ -13068,7 +13078,7 @@ static int rsa_test(void)
 #if !defined(HAVE_USER_RSA) && !defined(NO_ASN)
     ret = rsa_decode_test(&key);
     if (ret != 0)
-        return ret;
+        ERROR_OUT(ret, exit_rsa);
 #endif
 
 #ifdef USE_CERT_BUFFERS_1024
@@ -13092,13 +13102,8 @@ static int rsa_test(void)
 #endif
 
     tmp = (byte*)XMALLOC(bytes, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
-    if (tmp == NULL
+    if (tmp == NULL)
-    #ifdef WOLFSSL_ASYNC_CRYPT
+        ERROR_OUT(-7700, exit_rsa);
-        || out == NULL || plain == NULL
-    #endif
-    ) {
-        return -7700;
-    }
 
 #ifdef USE_CERT_BUFFERS_1024
     XMEMCPY(tmp, client_key_der_1024, (size_t)sizeof_client_key_der_1024);
@@ -18051,11 +18056,18 @@ static int ecc_test_vector_item(const eccVector* vector)
     DECLARE_VAR(sigRaw, byte, ECC_SIG_SIZE, HEAP_HINT);
 #endif
 
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
+    if ((sig == NULL)
+#if !defined(NO_ASN) && !defined(HAVE_SELFTEST)
+        || (sigRaw == NULL)
+#endif
+        )
+        ERROR_OUT(MEMORY_E, done);
+#endif
+
     ret = wc_ecc_init_ex(&userA, HEAP_HINT, devId);
-    if (ret != 0) {
+    if (ret != 0)
-        FREE_VAR(sig, HEAP_HINT);
+        goto done;
-        return ret;
-    }
 
     ret = wc_ecc_import_raw(&userA, vector->Qx, vector->Qy,
                                                   vector->d, vector->curveName);
@@ -18897,7 +18909,7 @@ static int ecc_test_curve_size(WC_RNG* rng, int keySize, int testVerifyCount,
     ecc_key userA, userB, pubKey;
     int     curveSize;
 
-#ifdef WOLFSSL_SMALL_STACK
+#ifdef DECLARE_VAR_IS_HEAP_ALLOC
 #if (defined(HAVE_ECC_DHE) || defined(HAVE_ECC_CDH)) && \
     !defined(WOLFSSL_ATECC508A) && !defined(WOLFSSL_ATECC608A)
     if ((sharedA == NULL) || (sharedB == NULL))

wolfssl/wolfcrypt/types.h

@@ -395,8 +395,9 @@ decouple library dependencies with standard string, memory and so on.
         #endif /* WOLFSSL_STATIC_MEMORY */
     #endif
 
-    /* declare/free variable handling for async */
+    /* declare/free variable handling for async and smallstack */
     #if defined(WOLFSSL_ASYNC_CRYPT) || defined(WOLFSSL_SMALL_STACK)
+        #define DECLARE_VAR_IS_HEAP_ALLOC
         #define DECLARE_VAR(VAR_NAME, VAR_TYPE, VAR_SIZE, HEAP) \
             VAR_TYPE* VAR_NAME = (VAR_TYPE*)XMALLOC(sizeof(VAR_TYPE) * VAR_SIZE, (HEAP), DYNAMIC_TYPE_WOLF_BIGINT);
         #define DECLARE_VAR_INIT(VAR_NAME, VAR_TYPE, VAR_SIZE, INIT_VALUE, HEAP) \
@@ -409,9 +410,19 @@ decouple library dependencies with standard string, memory and so on.
             })
         #define DECLARE_ARRAY(VAR_NAME, VAR_TYPE, VAR_ITEMS, VAR_SIZE, HEAP) \
             VAR_TYPE* VAR_NAME[VAR_ITEMS]; \
-            int idx##VAR_NAME; \
+            int idx##VAR_NAME, inner_idx_##VAR_NAME; \
             for (idx##VAR_NAME=0; idx##VAR_NAME<VAR_ITEMS; idx##VAR_NAME++) { \
                 VAR_NAME[idx##VAR_NAME] = (VAR_TYPE*)XMALLOC(VAR_SIZE, (HEAP), DYNAMIC_TYPE_WOLF_BIGINT); \
+                if (VAR_NAME[idx##VAR_NAME] == NULL) { \
+                    for (inner_idx_##VAR_NAME = 0; inner_idx_##VAR_NAME < idx##VAR_NAME; inner_idx_##VAR_NAME++) { \
+                        XFREE(VAR_NAME[inner_idx_##VAR_NAME], HEAP, DYNAMIC_TYPE_WOLF_BIGINT); \
+                        VAR_NAME[inner_idx_##VAR_NAME] = NULL; \
+                    } \
+                    for (inner_idx_##VAR_NAME = idx##VAR_NAME + 1; inner_idx_##VAR_NAME < VAR_ITEMS; inner_idx_##VAR_NAME++) { \
+                        VAR_NAME[inner_idx_##VAR_NAME] = NULL; \
+                    } \
+                    break; \
+                } \
             }
         #define FREE_VAR(VAR_NAME, HEAP) \
             XFREE(VAR_NAME, (HEAP), DYNAMIC_TYPE_WOLF_BIGINT);
@@ -426,6 +437,7 @@ decouple library dependencies with standard string, memory and so on.
         #define FREE_ARRAY_DYNAMIC(VAR_NAME, VAR_ITEMS, HEAP) \
             FREE_ARRAY(VAR_NAME, VAR_ITEMS, HEAP)
     #else
+        #undef DECLARE_VAR_IS_HEAP_ALLOC
         #define DECLARE_VAR(VAR_NAME, VAR_TYPE, VAR_SIZE, HEAP) \
             VAR_TYPE VAR_NAME[VAR_SIZE]
         #define DECLARE_VAR_INIT(VAR_NAME, VAR_TYPE, VAR_SIZE, INIT_VALUE, HEAP) \
@@ -437,10 +449,20 @@ decouple library dependencies with standard string, memory and so on.
 
         #define DECLARE_ARRAY_DYNAMIC_DEC(VAR_NAME, VAR_TYPE, VAR_ITEMS, VAR_SIZE, HEAP) \
             VAR_TYPE* VAR_NAME[VAR_ITEMS]; \
-            int idx##VAR_NAME;
+            int idx##VAR_NAME, inner_idx_##VAR_NAME;
         #define DECLARE_ARRAY_DYNAMIC_EXE(VAR_NAME, VAR_TYPE, VAR_ITEMS, VAR_SIZE, HEAP) \
             for (idx##VAR_NAME=0; idx##VAR_NAME<VAR_ITEMS; idx##VAR_NAME++) { \
                 VAR_NAME[idx##VAR_NAME] = (VAR_TYPE*)XMALLOC(VAR_SIZE, (HEAP), DYNAMIC_TYPE_TMP_BUFFER); \
+                if (VAR_NAME[idx##VAR_NAME] == NULL) { \
+                    for (inner_idx_##VAR_NAME = 0; inner_idx_##VAR_NAME < idx##VAR_NAME; inner_idx_##VAR_NAME++) { \
+                        XFREE(VAR_NAME[inner_idx_##VAR_NAME], HEAP, DYNAMIC_TYPE_TMP_BUFFER); \
+                        VAR_NAME[inner_idx_##VAR_NAME] = NULL; \
+                    } \
+                    for (inner_idx_##VAR_NAME = idx##VAR_NAME + 1; inner_idx_##VAR_NAME < VAR_ITEMS; inner_idx_##VAR_NAME++) { \
+                        VAR_NAME[inner_idx_##VAR_NAME] = NULL; \
+                    } \
+                    break; \
+                } \
             }
         #define FREE_ARRAY_DYNAMIC(VAR_NAME, VAR_ITEMS, HEAP) \
             for (idx##VAR_NAME=0; idx##VAR_NAME<VAR_ITEMS; idx##VAR_NAME++) { \