fix tests error

Author: Yu-Ma28051503

tests/api/test_mldsa.c

@@ -730,7 +730,9 @@ int test_wc_dilithium_pub_from_priv(void)
 {
     EXPECT_DECLS;
 #if defined(HAVE_DILITHIUM) && defined(WOLFSSL_WC_DILITHIUM) && \
-    defined(WOLFSSL_DILITHIUM_PRIVATE_KEY) && defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
+    defined(WOLFSSL_DILITHIUM_PRIVATE_KEY) && \
+    defined(WOLFSSL_DILITHIUM_PUBLIC_KEY) && \
+    !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
     dilithium_key* key = NULL;
     dilithium_key* importKey = NULL;
     WC_RNG rng;
@@ -739,7 +741,9 @@ int test_wc_dilithium_pub_from_priv(void)
     byte* pubKey = NULL;
     word32 pubKeyLen = DILITHIUM_MAX_PUB_KEY_SIZE;
     byte* origPub = NULL;
+#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
     word32 origPubLen = DILITHIUM_MAX_PUB_KEY_SIZE;
+#endif
     int ret;
 
     key = (dilithium_key*)XMALLOC(sizeof(*key), NULL, DYNAMIC_TYPE_TMP_BUFFER);
@@ -761,7 +765,9 @@ int test_wc_dilithium_pub_from_priv(void)
     if (importKey != NULL) XMEMSET(importKey, 0, sizeof(*importKey));
     XMEMSET(&rng, 0, sizeof(WC_RNG));
 
+#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
     ExpectIntEQ(wc_InitRng(&rng), 0);
+#endif
     ExpectIntEQ(wc_dilithium_init(key), 0);
 
 #ifndef WOLFSSL_NO_ML_DSA_44
@@ -823,7 +829,9 @@ int test_wc_dilithium_pub_from_priv(void)
 
     wc_dilithium_free(importKey);
     wc_dilithium_free(key);
+#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
     wc_FreeRng(&rng);
+#endif
 
     XFREE(pubKey, NULL, DYNAMIC_TYPE_TMP_BUFFER);
     XFREE(origPub, NULL, DYNAMIC_TYPE_TMP_BUFFER);

wolfcrypt/src/dilithium.c

@@ -8003,29 +8003,37 @@ static int dilithium_pub_from_priv(dilithium_key* key)
     int ret = 0;
     const wc_dilithium_params* params = key->params;
     const byte* pub_seed = key->k;
-    const byte* s1p = pub_seed + DILITHIUM_PUB_SEED_SZ + DILITHIUM_K_SZ + DILITHIUM_TR_SZ;
-    const byte* s2p = s1p + params->s1EncSz;
+    const byte* s1p = NULL;
+    const byte* s2p = NULL;
     sword32* a  = NULL;
     sword32* s1 = NULL;
     sword32* s2 = NULL;
     sword32* t  = NULL;
     byte* t0 = NULL;
     byte* t1 = key->p + DILITHIUM_PUB_SEED_SZ;
 
-    /* Allocate and create cached values. */
-#ifndef WC_DILITHIUM_CACHE_MATRIX_A
-    a = (sword32*)XMALLOC(params->aSz, key->heap,
-        DYNAMIC_TYPE_DILITHIUM);
-    if (a == NULL) {
-        ret = MEMORY_E;
+    if ((params == NULL) || (!key->prvKeySet)) {
+        ret = BAD_STATE_E;
     }
-    else {
-        XMEMSET(a, 0, params->aSz);
+
+    if (ret == 0) {
+        s1p = pub_seed + DILITHIUM_PUB_SEED_SZ + DILITHIUM_K_SZ + DILITHIUM_TR_SZ;
+        s2p = s1p + params->s1EncSz;
     }
 
+    /* Allocate and create cached values. */
+#ifndef WC_DILITHIUM_CACHE_MATRIX_A
     if (ret == 0) {
-        ret = dilithium_expand_a(&key->shake, pub_seed, params->k, params->l,
-            a, key->heap);
+        a = (sword32*)XMALLOC(params->aSz, key->heap,
+            DYNAMIC_TYPE_DILITHIUM);
+        if (a == NULL) {
+            ret = MEMORY_E;
+        }
+        else {
+            XMEMSET(a, 0, params->aSz);
+            ret = dilithium_expand_a(&key->shake, pub_seed, params->k,
+                params->l, a, key->heap);
+        }
     }
 #else
     if (ret == 0) {
@@ -8061,6 +8069,7 @@ static int dilithium_pub_from_priv(dilithium_key* key)
         else {
             s1 = key->s1;
             s2 = key->s2;
+	    dilithium_vec_invntt_full(s2, params->k);
         }
     }
 #endif
@@ -8088,28 +8097,38 @@ static int dilithium_pub_from_priv(dilithium_key* key)
         dilithium_matrix_mul(t, a, s1, params->k, params->l);
         dilithium_vec_invntt_full(t, params->k);
         dilithium_vec_add(t, s2, params->k);
-        /* NTT s2 */
-        dilithium_vec_ntt_small(s2, params->k);
 
         /* Make positive for decomposing. */
         dilithium_vec_make_pos(t, params->k);
         /* Decompose t in t0 and t1 and encode into public and private key. */
         dilithium_vec_encode_t0_t1(t, params->k, t0, t1);
     }
-
+#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
+    if (s2 != NULL) {
+    	/* NTT s2 */
+    	dilithium_vec_ntt_small(s2, params->k);
+    }
+#endif
 #ifndef WC_DILITHIUM_CACHE_MATRIX_A
-    XMEMSET(a, 0, params->aSz);
-    XFREE(a, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    if (a != NULL) {
+    	XMEMSET(a, 0, params->aSz);
+    	XFREE(a, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
 #endif
 #ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
-    XMEMSET(s1, 0, params->s1Sz + params->s2Sz);
-    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    if (s1 != NULL) {
+    	XMEMSET(s1, 0, params->s1Sz + params->s2Sz);
+    	XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
 #endif
-    XMEMSET(t0, 0, params->s2Sz);
-    XMEMSET(t, 0, params->s2Sz);
-    XFREE(t0, key->heap, DYNAMIC_TYPE_DILITHIUM);
-    XFREE(t, key->heap, DYNAMIC_TYPE_DILITHIUM);
-
+    if (t0 != NULL) {
+    	XMEMSET(t0, 0, params->s2Sz);
+	XFREE(t0, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
+    if (t != NULL) {
+    	XMEMSET(t, 0, params->s2Sz);
+    	XFREE(t, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
     if (ret == 0) {
 #ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
     #ifndef WC_DILITHIUM_FIXED_ARRAY
@@ -10311,7 +10330,12 @@ int wc_dilithium_pub_from_priv(dilithium_key* key)
 
     if (ret == 0) {
 #ifdef WOLFSSL_WC_DILITHIUM
-        ret = dilithium_pub_from_priv(key);
+        if ((key->params == NULL) || (!key->prvKeySet)) {
+            ret = BAD_STATE_E;
+        }
+        else {
+            ret = dilithium_pub_from_priv(key);
+        }
 #elif defined(HAVE_LIBOQS)
         ret = NOT_COMPILED_IN;
 #endif