@@ -2274,14 +2274,31 @@ void mlkem_decapsulate(const sword16* s, sword16* w, sword16* u,
22742274static int mlkem_gen_matrix_k2_avx2 (sword16 * a , byte * seed , int transposed )
22752275{
22762276 int i ;
2277+ #ifdef WOLFSSL_SMALL_STACK
2278+ byte * rand = NULL ;
2279+ word64 * state = NULL ;
2280+ #else
22772281 byte rand [4 * GEN_MATRIX_SIZE + 2 ];
22782282 word64 state [25 * 4 ];
2283+ #endif
22792284 unsigned int ctr0 ;
22802285 unsigned int ctr1 ;
22812286 unsigned int ctr2 ;
22822287 unsigned int ctr3 ;
22832288 byte * p ;
22842289
2290+ #ifdef WOLFSSL_SMALL_STACK
2291+ rand = (byte * )XMALLOC (4 * GEN_MATRIX_SIZE + 2 , NULL ,
2292+ DYNAMIC_TYPE_TMP_BUFFER );
2293+ state = (word64 * )XMALLOC (sizeof (word64 ) * 25 * 4 , NULL ,
2294+ DYNAMIC_TYPE_TMP_BUFFER );
2295+ if ((rand == NULL ) || (state == NULL )) {
2296+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2297+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2298+ return MEMORY_E ;
2299+ }
2300+ #endif
2301+
22852302 /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
22862303 rand [4 * GEN_MATRIX_SIZE + 0 ] = 0xff ;
22872304 rand [4 * GEN_MATRIX_SIZE + 1 ] = 0xff ;
@@ -2345,6 +2362,11 @@ static int mlkem_gen_matrix_k2_avx2(sword16* a, byte* seed, int transposed)
23452362 p , XOF_BLOCK_SIZE );
23462363 }
23472364
2365+ #ifdef WOLFSSL_SMALL_STACK
2366+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2367+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2368+ #endif
2369+
23482370 return 0 ;
23492371}
23502372#endif
@@ -2365,14 +2387,31 @@ static int mlkem_gen_matrix_k3_avx2(sword16* a, byte* seed, int transposed)
23652387{
23662388 int i ;
23672389 int k ;
2390+ #ifdef WOLFSSL_SMALL_STACK
2391+ byte * rand = NULL ;
2392+ word64 * state = NULL ;
2393+ #else
23682394 byte rand [4 * GEN_MATRIX_SIZE + 2 ];
23692395 word64 state [25 * 4 ];
2396+ #endif
23702397 unsigned int ctr0 ;
23712398 unsigned int ctr1 ;
23722399 unsigned int ctr2 ;
23732400 unsigned int ctr3 ;
23742401 byte * p ;
23752402
2403+ #ifdef WOLFSSL_SMALL_STACK
2404+ rand = (byte * )XMALLOC (4 * GEN_MATRIX_SIZE + 2 , NULL ,
2405+ DYNAMIC_TYPE_TMP_BUFFER );
2406+ state = (word64 * )XMALLOC (sizeof (word64 ) * 25 * 4 , NULL ,
2407+ DYNAMIC_TYPE_TMP_BUFFER );
2408+ if ((rand == NULL ) || (state == NULL )) {
2409+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2410+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2411+ return MEMORY_E ;
2412+ }
2413+ #endif
2414+
23762415 /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
23772416 rand [4 * GEN_MATRIX_SIZE + 0 ] = 0xff ;
23782417 rand [4 * GEN_MATRIX_SIZE + 1 ] = 0xff ;
@@ -2473,6 +2512,11 @@ static int mlkem_gen_matrix_k3_avx2(sword16* a, byte* seed, int transposed)
24732512 XOF_BLOCK_SIZE );
24742513 }
24752514
2515+ #ifdef WOLFSSL_SMALL_STACK
2516+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2517+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2518+ #endif
2519+
24762520 return 0 ;
24772521}
24782522#endif
@@ -2492,14 +2536,31 @@ static int mlkem_gen_matrix_k4_avx2(sword16* a, byte* seed, int transposed)
24922536{
24932537 int i ;
24942538 int k ;
2539+ #ifdef WOLFSSL_SMALL_STACK
2540+ byte * rand = NULL ;
2541+ word64 * state = NULL ;
2542+ #else
24952543 byte rand [4 * GEN_MATRIX_SIZE + 2 ];
24962544 word64 state [25 * 4 ];
2545+ #endif
24972546 unsigned int ctr0 ;
24982547 unsigned int ctr1 ;
24992548 unsigned int ctr2 ;
25002549 unsigned int ctr3 ;
25012550 byte * p ;
25022551
2552+ #ifdef WOLFSSL_SMALL_STACK
2553+ rand = (byte * )XMALLOC (4 * GEN_MATRIX_SIZE + 2 , NULL ,
2554+ DYNAMIC_TYPE_TMP_BUFFER );
2555+ state = (word64 * )XMALLOC (sizeof (word64 ) * 25 * 4 , NULL ,
2556+ DYNAMIC_TYPE_TMP_BUFFER );
2557+ if ((rand == NULL ) || (state == NULL )) {
2558+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2559+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2560+ return MEMORY_E ;
2561+ }
2562+ #endif
2563+
25032564 /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
25042565 rand [4 * GEN_MATRIX_SIZE + 0 ] = 0xff ;
25052566 rand [4 * GEN_MATRIX_SIZE + 1 ] = 0xff ;
@@ -2563,6 +2624,11 @@ static int mlkem_gen_matrix_k4_avx2(sword16* a, byte* seed, int transposed)
25632624 a += 4 * MLKEM_N ;
25642625 }
25652626
2627+ #ifdef WOLFSSL_SMALL_STACK
2628+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2629+ XFREE (state , NULL , DYNAMIC_TYPE_TMP_BUFFER );
2630+ #endif
2631+
25662632 return 0 ;
25672633}
25682634#endif /* WOLFSSL_KYBER1024 || WOLFSSL_WC_ML_KEM_1024 */
@@ -4120,7 +4186,17 @@ static int mlkem_get_noise_k2_avx2(MLKEM_PRF_T* prf, sword16* vec1,
41204186 sword16 * vec2 , sword16 * poly , byte * seed )
41214187{
41224188 int ret = 0 ;
4189+ #ifdef WOLFSSL_SMALL_STACK
4190+ byte * rand ;
4191+ #else
41234192 byte rand [4 * PRF_RAND_SZ ];
4193+ #endif
4194+
4195+ #ifdef WOLFSSL_SMALL_STACK
4196+ rand = (byte * )XMALLOC (4 * PRF_RAND_SZ , NULL , DYNAMIC_TYPE_TMP_BUFFER );
4197+ if (rand == NULL )
4198+ return MEMORY_E ;
4199+ #endif
41244200
41254201 mlkem_get_noise_x4_eta3_avx2 (rand , seed );
41264202 mlkem_cbd_eta3_avx2 (vec1 , rand + 0 * PRF_RAND_SZ );
@@ -4137,6 +4213,10 @@ static int mlkem_get_noise_k2_avx2(MLKEM_PRF_T* prf, sword16* vec1,
41374213 ret = mlkem_get_noise_eta2_avx2 (prf , poly , seed );
41384214 }
41394215
4216+ #ifdef WOLFSSL_SMALL_STACK
4217+ XFREE (rand , NULL , DYNAMIC_TYPE_TMP_BUFFER );
4218+ #endif
4219+
41404220 return ret ;
41414221}
41424222#endif
0 commit comments