luna - wolfssl/wolfcrypt/src/dilithium.c

    1/* dilithium.c
    2 *
    3 * Copyright (C) 2006-2026 wolfSSL Inc.
    4 *
    5 * This file is part of wolfSSL.
    6 *
    7 * wolfSSL is free software; you can redistribute it and/or modify
    8 * it under the terms of the GNU General Public License as published by
    9 * the Free Software Foundation; either version 3 of the License, or
   10 * (at your option) any later version.
   11 *
   12 * wolfSSL is distributed in the hope that it will be useful,
   13 * but WITHOUT ANY WARRANTY; without even the implied warranty of
   14 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   15 * GNU General Public License for more details.
   16 *
   17 * You should have received a copy of the GNU General Public License
   18 * along with this program; if not, write to the Free Software
   19 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
   20 */
   21
   22/* Based on ed448.c and Reworked for Dilithium by Anthony Hu.
   23 * WolfSSL implementation by Sean Parkinson.
   24 */
   25
   26/* Possible Dilithium/ML-DSA options:
   27 *
   28 * HAVE_DILITHIUM                                             Default: OFF
   29 *   Enables the code in this file to be compiled.
   30 *
   31 * WOLFSSL_NO_ML_DSA_44                                       Default: OFF
   32 *   Does not compile in parameter set ML-DSA-44 and any code specific to that
   33 *   parameter set.
   34 * WOLFSSL_NO_ML_DSA_65                                       Default: OFF
   35 *   Does not compile in parameter set ML-DSA-65 and any code specific to that
   36 *   parameter set.
   37 * WOLFSSL_NO_ML_DSA_87                                       Default: OFF
   38 *   Does not compile in parameter set ML-DSA-87 and any code specific to that
   39 *   parameter set.
   40 *
   41 * WOLFSSL_DILITHIUM_NO_LARGE_CODE                            Default: OFF
   42 *   Compiles smaller, fast code with speed trade-off.
   43 * WOLFSSL_DILITHIUM_SMALL                                    Default: OFF
   44 *   Compiles to small code size with a speed trade-off.
   45 * WOLFSSL_DILITHIUM_VERIFY_ONLY                              Default: OFF
   46 *   Compiles in only the verification and public key operations.
   47 * WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM                         Default: OFF
   48 *   Compiles verification implementation that uses smaller amounts of memory.
   49 * WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC                         Default: OFF
   50 *   Only works with WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM.
   51 *   Don't allocate memory with XMALLOC. Memory is pinned against key.
   52 * WOLFSSL_DILITHIUM_ASSIGN_KEY                               Default: OFF
   53 *   Key data is assigned into Dilithium key rather than copied.
   54 *   Life of key data passed in is tightly coupled to life of Dilithium key.
   55 *   Cannot be used when make key is enabled.
   56 * WOLFSSL_DILITHIUM_DYNAMIC_KEYS                                Default: OFF
   57 *   Key buffers (public and private) are dynamically allocated on the heap
   58 *   instead of being static arrays in the key struct. Buffers are right-sized
   59 *   for the key's ML-DSA level and only allocated when needed (e.g. no private
   60 *   key buffer for verify-only keys). Reduces memory footprint significantly.
   61 *   Cannot be used with WOLFSSL_DILITHIUM_ASSIGN_KEY.
   62 * WOLFSSL_DILITHIUM_SIGN_SMALL_MEM                           Default: OFF
   63 *   Compiles signature implementation that uses smaller amounts of memory but
   64 *   is considerably slower.
   65 * WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC                   Default: OFF
   66 *   Compiles signature implementation that uses smaller amounts of memory but
   67 *   is considerably slower. Allocates vectors and decodes private key data
   68 *   into them upfront.
   69 * WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A                 Default: OFF
   70 *   Compiles signature implementation that uses smaller amounts of memory but
   71 *   is slower. Allocates matrix A and calculates it upfront.
   72 * WOLFSSL_DILITHIUM_MAKE_KEY_SMALL_MEM                       Default: OFF
   73 *   Compiles key generation implementation that uses smaller amounts of memory
   74 *   but is slower.
   75 * WOLFSSL_DILITHIUM_SMALL_MEM_POLY64                         Default: OFF
   76 *   Compiles the small memory implementations to use a 64-bit polynomial.
   77 *   Uses 2KB of memory but is slightly quicker (2.75-7%).
   78 *
   79 * WOLFSSL_DILITHIUM_ALIGNMENT                                Default: 8
   80 *   Use to indicate whether loading and storing of words needs to be aligned.
   81 *   Default is to use WOLFSSL_GENERAL_ALIGNMENT - should be 4 on some ARM CPUs.
   82 *   Set this value explicitly if specific Dilithium implementation alignment is
   83 *   needed.
   84 *
   85 * WOLFSSL_DILITHIUM_NO_ASN1                                  Default: OFF
   86 *   Disables any ASN.1 encoding or decoding code.
   87 * WOLFSSL_DILITHIUM_REVERSE_HASH_OID                         Default: OFF
   88 *   Reverse the DER encoded hash oid when signing and verifying a pre-hashed
   89 *   message.
   90 *
   91 * WC_DILITHIUM_CACHE_MATRIX_A                                Default: OFF
   92 *   Enable caching of the A matrix on import.
   93 *   Less work is required in sign and verify operations.
   94 * WC_DILITHIUM_CACHE_PRIV_VECTORS                            Default: OFF
   95 *   Enable caching of private key vectors on import.
   96 *   Enables WC_DILITHIUM_CACHE_MATRIX_A.
   97 *   Less work is required in sign operations.
   98 * WC_DILITHIUM_CACHE_PUB_VECTORS                             Default: OFF
   99 *   Enable caching of public key vectors on import.
  100 *   Enables WC_DILITHIUM_CACHE_MATRIX_A.
  101 *   Less work is required in sign operations.
  102 * WC_DILITHIUM_FIXED_ARRAY                                   Default: OFF
  103 *   Make the matrix and vectors of cached data fixed arrays that have
  104 *   maximumal sizes for the configured parameters.
  105 *   Useful in low dynamic memory situations.
  106 *
  107 * WOLFSSL_DILITHIUM_SIGN_CHECK_Y                             Default: OFF
  108 *   Check vector y is in required range as an early check on valid vector z.
  109 *   Falsely reports invalid in approximately 1-2% of checks.
  110 *   All valid reports are true.
  111 *   Fast fail gives faster signing times on average.
  112 *   DO NOT enable this if implementation must be conformant to FIPS 204.
  113 * WOLFSSL_DILITHIUM_SIGN_CHECK_W0                            Default: OFF
  114 *   Check vector w0 is in required range as an early check on valid vector r0.
  115 *   Falsely reports invalid in approximately 3-5% of checks.
  116 *   All valid reports are true.
  117 *   Fast fail gives faster signing times on average.
  118 *   DO NOT enable this if implementation must be conformant to FIPS 204.
  119 *
  120 * DILITHIUM_MUL_SLOW                                         Default: OFF
  121 *   Define when multiplying by Q / 44 is slower than masking.
  122 *   Only applies to ML-DSA-44.
  123 * DILITHIUM_MUL_44_SLOW                                      Default: OFF
  124 *   Define when multiplying by 44 is slower than by 11.
  125 *   Only applies to ML-DSA-44.
  126 * DILITHIUM_MUL_11_SLOW                                      Default: OFF
  127 *   Define when multiplying by 11 is slower than adding and shifting.
  128 *   Only applies to ML-DSA-44.
  129 * DILITHIUM_MUL_QINV_SLOW                                    Default: OFF
  130 *   Define when multiplying by QINV 0x3802001 is slower than add, subtract and
  131 *   shift equivalent.
  132 * DILITHIUM_MUL_Q_SLOW                                       Default: OFF
  133 *   Define when multiplying by Q 0x7fe001 is slower than add, subtract and
  134 *   shift equivalent.
  135 */
  136
  137#include <wolfssl/wolfcrypt/libwolfssl_sources.h>
  138
  139#if FIPS_VERSION3_GE(2,0,0)
  140    /* set NO_WRAPPERS before headers, use direct internal f()s not wrappers */
  141    #define FIPS_NO_WRAPPERS
  142#endif
  143
  144#ifndef WOLFSSL_DILITHIUM_NO_ASN1
  145#include <wolfssl/wolfcrypt/asn.h>
  146#endif
  147
  148#if defined(HAVE_DILITHIUM)
  149
  150#include <wolfssl/wolfcrypt/dilithium.h>
  151#include <wolfssl/wolfcrypt/hash.h>
  152#include <wolfssl/wolfcrypt/sha3.h>
  153#include <wolfssl/wolfcrypt/cpuid.h>
  154#include <wolfssl/wolfcrypt/error-crypt.h>
  155#ifdef NO_INLINE
  156    #include <wolfssl/wolfcrypt/misc.h>
  157#else
  158    #define WOLFSSL_MISC_INCLUDED
  159    #include <wolfcrypt/src/misc.c>
  160#endif
  161
  162#if defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC) && \
  163        !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)
  164    #define WOLFSSL_DILITHIUM_SIGN_SMALL_MEM
  165#endif
  166#if defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A) && \
  167        !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)
  168    #define WOLFSSL_DILITHIUM_SIGN_SMALL_MEM
  169    #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
  170        #error "PRECALC and PRECALC_A are equivalent to non small mem"
  171    #endif
  172#endif
  173
  174#if defined(USE_INTEL_SPEEDUP)
  175static cpuid_flags_t cpuid_flags = WC_CPUID_INITIALIZER;
  176#endif
  177
  178#ifdef DEBUG_DILITHIUM
  179void print_polys(const char* name, const sword32* a, int d1, int d2);
  180void print_polys(const char* name, const sword32* a, int d1, int d2)
  181{
  182    int i;
  183    int j;
  184    int k;
  185
  186    fprintf(stderr, "%s: %d %d\n", name, d1, d2);
  187    for (i = 0; i < d1; i++) {
  188        for (j = 0; j < d2; j++) {
  189            for (k = 0; k < 256; k++) {
  190                fprintf(stderr, "%9d,", a[(i*d2*256) + (j*256) + k]);
  191                if ((k % 8) == 7) fprintf(stderr, "\n");
  192            }
  193            fprintf(stderr, "\n");
  194        }
  195    }
  196}
  197#endif
  198
  199#ifdef DEBUG_DILITHIUM
  200void print_data(const char* name, const byte* d, int len);
  201void print_data(const char* name, const byte* d, int len)
  202{
  203    int i;
  204
  205    fprintf(stderr, "%s\n", name);
  206    for (i = 0; i < len; i++) {
  207        fprintf(stderr, "0x%02x,", d[i]);
  208        if ((i % 16) == 15) fprintf(stderr, "\n");
  209    }
  210    fprintf(stderr, "\n");
  211}
  212#endif
  213
  214#if defined(WOLFSSL_NO_ML_DSA_44) && defined(WOLFSSL_NO_ML_DSA_65) && \
  215    defined(WOLFSSL_NO_ML_DSA_87)
  216    #error "No Dilithium parameters chosen"
  217#endif
  218
  219#if defined(WOLFSSL_DILITHIUM_ASSIGN_KEY) && \
  220    !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
  221    #error "Cannot use assign key when making keys"
  222#endif
  223
  224#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
  225    defined(WOLFSSL_DILITHIUM_ASSIGN_KEY)
  226    #error "Cannot use both WOLFSSL_DILITHIUM_DYNAMIC_KEYS and WOLFSSL_DILITHIUM_ASSIGN_KEY"
  227#endif
  228
  229
  230/* Number of bytes from first block to use for sign. */
  231#define DILITHIUM_SIGN_BYTES            8
  232
  233
  234/* Length of seed in bytes when generating y. */
  235#define DILITHIUM_Y_SEED_SZ             (DILITHIUM_PRIV_RAND_SEED_SZ + 2)
  236
  237
  238/* Length of seed in bytes used in generating matrix a. */
  239#define DILITHIUM_GEN_A_SEED_SZ         (DILITHIUM_PUB_SEED_SZ + 2)
  240/* Length of seed in bytes used in generating vectors s1 and s2. */
  241#define DILITHIUM_GEN_S_SEED_SZ         (DILITHIUM_PRIV_SEED_SZ + 2)
  242
  243
  244/* MAX: (256 * 8 / (17 + 1)) = 576, or ((256 * 8 / (19 + 1)) = 640
  245 * but need blocks of 17 * 8 bytes: 5 * 17 * 8 = 680 */
  246#define DILITHIUM_MAX_V_BLOCKS          5
  247/* Maximum number of bytes to generate into v to make y. */
  248#define DILITHIUM_MAX_V                 (DILITHIUM_MAX_V_BLOCKS * 8 * 17)
  249
  250
  251/* 2 blocks, each block 136 bytes = 272 bytes.
  252 * ETA 2: Min req is 128 but reject rate is 2 in 16 so we need 146.3 on average.
  253 * ETA 4: Min req is 128 but reject rate is 7 in 16 so we need 227.6 on average.
  254 */
  255#define DILITHIUM_GEN_S_NBLOCKS         2
  256/* Number of bytes to a block of SHAKE-256 when generating s1 and s2. */
  257#define DILITHIUM_GEN_S_BLOCK_BYTES    (WC_SHA3_256_COUNT * 8)
  258/* Number of bytes to generate with SHAKE-256 when generating s1 and s2. */
  259#define DILITHIUM_GEN_S_BYTES           \
  260    (DILITHIUM_GEN_S_NBLOCKS * DILITHIUM_GEN_S_BLOCK_BYTES)
  261
  262/* Length of the hash OID to include in pre-hash message. */
  263#define DILITHIUM_HASH_OID_LEN         11
  264
  265
  266/* The ML-DSA parameters sets. */
  267static const wc_dilithium_params dilithium_params[] = {
  268#ifndef WOLFSSL_NO_ML_DSA_44
  269    { WC_ML_DSA_44, PARAMS_ML_DSA_44_K, PARAMS_ML_DSA_44_L,
  270      PARAMS_ML_DSA_44_ETA, PARAMS_ML_DSA_44_ETA_BITS,
  271      PARAMS_ML_DSA_44_TAU, PARAMS_ML_DSA_44_BETA, PARAMS_ML_DSA_44_OMEGA,
  272      PARAMS_ML_DSA_44_LAMBDA,
  273      PARAMS_ML_DSA_44_GAMMA1_BITS, PARAMS_ML_DSA_44_GAMMA2,
  274      PARAMS_ML_DSA_44_W1_ENC_SZ, PARAMS_ML_DSA_44_A_SIZE,
  275      PARAMS_ML_DSA_44_S1_SIZE, PARAMS_ML_DSA_44_S1_ENC_SIZE,
  276      PARAMS_ML_DSA_44_S2_SIZE, PARAMS_ML_DSA_44_S2_ENC_SIZE,
  277      PARAMS_ML_DSA_44_Z_ENC_SIZE,
  278      PARAMS_ML_DSA_44_PK_SIZE, PARAMS_ML_DSA_44_SIG_SIZE },
  279#endif
  280#ifndef WOLFSSL_NO_ML_DSA_65
  281    { WC_ML_DSA_65, PARAMS_ML_DSA_65_K, PARAMS_ML_DSA_65_L,
  282      PARAMS_ML_DSA_65_ETA, PARAMS_ML_DSA_65_ETA_BITS,
  283      PARAMS_ML_DSA_65_TAU, PARAMS_ML_DSA_65_BETA, PARAMS_ML_DSA_65_OMEGA,
  284      PARAMS_ML_DSA_65_LAMBDA,
  285      PARAMS_ML_DSA_65_GAMMA1_BITS, PARAMS_ML_DSA_65_GAMMA2,
  286      PARAMS_ML_DSA_65_W1_ENC_SZ, PARAMS_ML_DSA_65_A_SIZE,
  287      PARAMS_ML_DSA_65_S1_SIZE, PARAMS_ML_DSA_65_S1_ENC_SIZE,
  288      PARAMS_ML_DSA_65_S2_SIZE, PARAMS_ML_DSA_65_S2_ENC_SIZE,
  289      PARAMS_ML_DSA_65_Z_ENC_SIZE,
  290      PARAMS_ML_DSA_65_PK_SIZE, PARAMS_ML_DSA_65_SIG_SIZE },
  291#endif
  292#ifndef WOLFSSL_NO_ML_DSA_87
  293    { WC_ML_DSA_87, PARAMS_ML_DSA_87_K, PARAMS_ML_DSA_87_L,
  294      PARAMS_ML_DSA_87_ETA, PARAMS_ML_DSA_87_ETA_BITS,
  295      PARAMS_ML_DSA_87_TAU, PARAMS_ML_DSA_87_BETA, PARAMS_ML_DSA_87_OMEGA,
  296      PARAMS_ML_DSA_87_LAMBDA,
  297      PARAMS_ML_DSA_87_GAMMA1_BITS, PARAMS_ML_DSA_87_GAMMA2,
  298      PARAMS_ML_DSA_87_W1_ENC_SZ, PARAMS_ML_DSA_87_A_SIZE,
  299      PARAMS_ML_DSA_87_S1_SIZE, PARAMS_ML_DSA_87_S1_ENC_SIZE,
  300      PARAMS_ML_DSA_87_S2_SIZE, PARAMS_ML_DSA_87_S2_ENC_SIZE,
  301      PARAMS_ML_DSA_87_Z_ENC_SIZE,
  302      PARAMS_ML_DSA_87_PK_SIZE, PARAMS_ML_DSA_87_SIG_SIZE },
  303#endif
  304#if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
  305#ifndef WOLFSSL_NO_ML_DSA_44
  306    { WC_ML_DSA_44_DRAFT, PARAMS_ML_DSA_44_K, PARAMS_ML_DSA_44_L,
  307      PARAMS_ML_DSA_44_ETA, PARAMS_ML_DSA_44_ETA_BITS,
  308      PARAMS_ML_DSA_44_TAU, PARAMS_ML_DSA_44_BETA, PARAMS_ML_DSA_44_OMEGA,
  309      PARAMS_ML_DSA_44_LAMBDA,
  310      PARAMS_ML_DSA_44_GAMMA1_BITS, PARAMS_ML_DSA_44_GAMMA2,
  311      PARAMS_ML_DSA_44_W1_ENC_SZ, PARAMS_ML_DSA_44_A_SIZE,
  312      PARAMS_ML_DSA_44_S1_SIZE, PARAMS_ML_DSA_44_S1_ENC_SIZE,
  313      PARAMS_ML_DSA_44_S2_SIZE, PARAMS_ML_DSA_44_S2_ENC_SIZE,
  314      PARAMS_ML_DSA_44_Z_ENC_SIZE,
  315      PARAMS_ML_DSA_44_PK_SIZE, PARAMS_ML_DSA_44_SIG_SIZE },
  316#endif
  317#ifndef WOLFSSL_NO_ML_DSA_65
  318    { WC_ML_DSA_65_DRAFT, PARAMS_ML_DSA_65_K, PARAMS_ML_DSA_65_L,
  319      PARAMS_ML_DSA_65_ETA, PARAMS_ML_DSA_65_ETA_BITS,
  320      PARAMS_ML_DSA_65_TAU, PARAMS_ML_DSA_65_BETA, PARAMS_ML_DSA_65_OMEGA,
  321      PARAMS_ML_DSA_65_LAMBDA,
  322      PARAMS_ML_DSA_65_GAMMA1_BITS, PARAMS_ML_DSA_65_GAMMA2,
  323      PARAMS_ML_DSA_65_W1_ENC_SZ, PARAMS_ML_DSA_65_A_SIZE,
  324      PARAMS_ML_DSA_65_S1_SIZE, PARAMS_ML_DSA_65_S1_ENC_SIZE,
  325      PARAMS_ML_DSA_65_S2_SIZE, PARAMS_ML_DSA_65_S2_ENC_SIZE,
  326      PARAMS_ML_DSA_65_Z_ENC_SIZE,
  327      PARAMS_ML_DSA_65_PK_SIZE, PARAMS_ML_DSA_65_SIG_SIZE },
  328#endif
  329#ifndef WOLFSSL_NO_ML_DSA_87
  330    { WC_ML_DSA_87_DRAFT, PARAMS_ML_DSA_87_K, PARAMS_ML_DSA_87_L,
  331      PARAMS_ML_DSA_87_ETA, PARAMS_ML_DSA_87_ETA_BITS,
  332      PARAMS_ML_DSA_87_TAU, PARAMS_ML_DSA_87_BETA, PARAMS_ML_DSA_87_OMEGA,
  333      PARAMS_ML_DSA_87_LAMBDA,
  334      PARAMS_ML_DSA_87_GAMMA1_BITS, PARAMS_ML_DSA_87_GAMMA2,
  335      PARAMS_ML_DSA_87_W1_ENC_SZ, PARAMS_ML_DSA_87_A_SIZE,
  336      PARAMS_ML_DSA_87_S1_SIZE, PARAMS_ML_DSA_87_S1_ENC_SIZE,
  337      PARAMS_ML_DSA_87_S2_SIZE, PARAMS_ML_DSA_87_S2_ENC_SIZE,
  338      PARAMS_ML_DSA_87_Z_ENC_SIZE,
  339      PARAMS_ML_DSA_87_PK_SIZE, PARAMS_ML_DSA_87_SIG_SIZE },
  340#endif
  341#endif
  342};
  343/* Number of ML-DSA parameter sets compiled in. */
  344#define DILITHIUM_PARAMS_CNT \
  345    ((unsigned int)(sizeof(dilithium_params) / sizeof(wc_dilithium_params)))
  346
  347/* Get the ML-DSA parameters that match the level.
  348 *
  349 * @param [in]  level   Level required.
  350 * @param [out] params  Parameter set.
  351 * @return  0 on success.
  352 * @return  NOT_COMPILED_IN when parameters at level are not compiled in.
  353 */
  354static int dilithium_get_params(int level, const wc_dilithium_params** params)
  355{
  356    unsigned int i;
  357    int ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);
  358
  359    for (i = 0; i < DILITHIUM_PARAMS_CNT; i++) {
  360        if (dilithium_params[i].level == level) {
  361            *params = &dilithium_params[i];
  362            ret = 0;
  363        }
  364    }
  365
  366    return ret;
  367}
  368
  369#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
  370    defined(WOLFSSL_DILITHIUM_PRIVATE_KEY)
  371/* Allocate the private key buffer for the current level if not already
  372 * allocated. Buffer is sized via wc_dilithium_size(key) and the allocated size
  373 * is stored in key->kSz for later use (ForceZero, free). On failure key->k may
  374 * remain NULL; callers must not inspect it. */
  375static int dilithium_alloc_priv_buf(dilithium_key* key)
  376{
  377    int ret = 0;
  378
  379    if (key->k == NULL) {
  380        int secSz = wc_dilithium_size(key);
  381        if (secSz < 0) {
  382            /* Should not happen, as the level checks have already been
  383             * performed, but defense-in-depth. */
  384            ret = BAD_STATE_E;
  385        }
  386        else {
  387        #ifdef USE_INTEL_SPEEDUP
  388            secSz += 8;
  389        #endif
  390            key->k = (byte*)XMALLOC((word32)secSz, key->heap,
  391                                    DYNAMIC_TYPE_DILITHIUM);
  392            if (key->k == NULL) {
  393                ret = MEMORY_E;
  394            }
  395            else {
  396                key->kSz = (word32)secSz;
  397            }
  398        }
  399    }
  400    return ret;
  401}
  402#endif
  403
  404#if defined(WOLFSSL_DILITHIUM_DYNAMIC_KEYS) && \
  405    defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
  406/* Allocate the public key buffer for the current level if not already
  407 * allocated. Buffer is sized via wc_dilithium_pub_size(key). On failure,
  408 * key->p may remain NULL; callers must not inspect it. */
  409static int dilithium_alloc_pub_buf(dilithium_key* key)
  410{
  411    int ret = 0;
  412
  413    if (key->p == NULL) {
  414        int pubSz = wc_dilithium_pub_size(key);
  415        if (pubSz < 0) {
  416            /* Should not happen, as the level checks have already been
  417             * performed, but defense-in-depth. */
  418            ret = BAD_STATE_E;
  419        }
  420        else {
  421        #ifdef USE_INTEL_SPEEDUP
  422            pubSz += 8;
  423        #endif
  424            key->p = (byte*)XMALLOC((word32)pubSz, key->heap,
  425                                    DYNAMIC_TYPE_DILITHIUM);
  426            if (key->p == NULL) {
  427                ret = MEMORY_E;
  428            }
  429        }
  430    }
  431    return ret;
  432}
  433#endif
  434
  435/******************************************************************************
  436 * Hash operations
  437 ******************************************************************************/
  438
  439/* 256-bit hash using SHAKE-256.
  440 *
  441 * FIPS 204. 8.3: H(v,d) <- SHAKE256(v,d)
  442 *
  443 * @param [in, out] shake256  SHAKE-256 object.
  444 * @param [in]      data      Buffer holding data to hash.
  445 * @param [in]      dataLen   Length of data to hash in bytes.
  446 * @param [out]     hash      Buffer to hold hash result.
  447 * @param [in]      hashLen   Number of bytes of hash to return.
  448 * @return  0 on success.
  449 * @return  Negative on error.
  450 */
  451static int dilithium_shake256(wc_Shake* shake256, const byte* data,
  452    word32 dataLen, byte* hash, word32 hashLen)
  453{
  454    int ret;
  455#ifdef USE_INTEL_SPEEDUP
  456    word64* state = shake256->s;
  457    word8 *state8 = (word8*)state;
  458
  459    if (dataLen >= WC_SHA3_256_COUNT * 8) {
  460        XMEMCPY(state, data,  WC_SHA3_256_COUNT * 8);
  461        XMEMSET(state + WC_SHA3_256_COUNT, 0,
  462            sizeof(shake256->s) - WC_SHA3_256_COUNT * 8);
  463        dataLen -= WC_SHA3_256_COUNT * 8;
  464        data    += WC_SHA3_256_COUNT * 8;
  465#ifndef WC_SHA3_NO_ASM
  466        if (IS_INTEL_AVX2(cpuid_flags) &&
  467                 (SAVE_VECTOR_REGISTERS2() == 0)) {
  468            sha3_block_avx2(state);
  469            RESTORE_VECTOR_REGISTERS();
  470        }
  471        else if (IS_INTEL_BMI2(cpuid_flags)) {
  472            sha3_block_bmi2(state);
  473        }
  474        else
  475#endif
  476        {
  477            BlockSha3(state);
  478        }
  479        if (dataLen >= WC_SHA3_256_COUNT * 8) {
  480#ifndef WC_SHA3_NO_ASM
  481            word32 n = dataLen / (WC_SHA3_256_COUNT * 8);
  482            if (IS_INTEL_AVX2(cpuid_flags) &&
  483                     (SAVE_VECTOR_REGISTERS2() == 0)) {
  484                sha3_block_n_avx2(state, data, n, WC_SHA3_256_COUNT * 8);
  485                RESTORE_VECTOR_REGISTERS();
  486                n *= WC_SHA3_256_COUNT * 8;
  487                dataLen -= n;
  488                data    += n;
  489            }
  490            else if (IS_INTEL_BMI2(cpuid_flags)) {
  491                sha3_block_n_bmi2(state, data, n, WC_SHA3_256_COUNT * 8);
  492                n *= WC_SHA3_256_COUNT * 8;
  493                dataLen -= n;
  494                data    += n;
  495            }
  496            else
  497#endif
  498            {
  499                while (dataLen >= WC_SHA3_256_COUNT * 8) {
  500                    xorbuf(state, data, WC_SHA3_256_COUNT * 8);
  501                    dataLen -= WC_SHA3_256_COUNT * 8;
  502                    data    += WC_SHA3_256_COUNT * 8;
  503                    BlockSha3(state);
  504                }
  505            }
  506        }
  507        if (dataLen > 0) {
  508            xorbuf(state, data, dataLen);
  509        }
  510        state8[dataLen] ^= 0x1f;
  511        state8[WC_SHA3_256_COUNT * 8 - 1] ^= 0x80;
  512    }
  513    else {
  514        XMEMCPY(state, data, dataLen);
  515        state8[dataLen] = 0x1f;
  516        XMEMSET(state8 + dataLen + 1, 0, sizeof(shake256->s) - (dataLen + 1));
  517        state8[WC_SHA3_256_COUNT * 8 - 1] ^= 0x80;
  518    }
  519
  520#ifndef WC_SHA3_NO_ASM
  521    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
  522        sha3_block_avx2(state);
  523        RESTORE_VECTOR_REGISTERS();
  524    }
  525    else if (IS_INTEL_BMI2(cpuid_flags)) {
  526        sha3_block_bmi2(state);
  527    }
  528    else
  529#endif
  530    {
  531        BlockSha3(state);
  532    }
  533    if (hash != (byte*)shake256->s) {
  534        XMEMCPY(hash, shake256->s, hashLen);
  535    }
  536    ret = 0;
  537#else
  538    /* Initialize SHAKE-256 operation. */
  539    ret = wc_InitShake256(shake256, NULL, INVALID_DEVID);
  540    if (ret == 0) {
  541        /* Update with data. */
  542        ret = wc_Shake256_Update(shake256, data, dataLen);
  543    }
  544    if (ret == 0) {
  545        /* Compute hash of data. */
  546        ret = wc_Shake256_Final(shake256, hash, hashLen);
  547    }
  548#endif
  549
  550    return ret;
  551}
  552
  553/* 256-bit hash using SHAKE-256.
  554 *
  555 * This is the domain-separated version of the hash.
  556 * See FIPS 204. D.3.
  557 *
  558 * FIPS 204. 8.3: H(v,d) <- SHAKE256(v,d)
  559 *
  560 * @param [in, out] shake256  SHAKE-256 object.
  561 * @param [in]      data1     First block of data to hash.
  562 * @param [in]      data1Len  Length of first block in bytes.
  563 * @param [in]      data2     Second block of data to hash.
  564 * @param [in]      data2Len  Length of second block in bytes.
  565 * @param [out]     hash      Buffer to hold hash result.
  566 * @param [in]      hashLen   Number of bytes of hash to return.
  567 * @return  0 on success.
  568 * @return  Negative on error.
  569 */
  570static int dilithium_hash256(wc_Shake* shake256, const byte* data1,
  571    word32 data1Len, const byte* data2, word32 data2Len, byte* hash,
  572    word32 hashLen)
  573{
  574    int ret;
  575
  576#ifdef USE_INTEL_SPEEDUP
  577    word64* state = shake256->s;
  578    word8 *state8 = (word8*)state;
  579
  580    if (data2Len > (WOLFSSL_MAX_32BIT - data1Len)) {
  581        return BAD_FUNC_ARG;
  582    }
  583    if (data1Len + data2Len >= WC_SHA3_256_COUNT * 8) {
  584        XMEMCPY(state8, data1, data1Len);
  585        XMEMCPY(state8 + data1Len, data2,  WC_SHA3_256_COUNT * 8 - data1Len);
  586        XMEMSET(state + WC_SHA3_256_COUNT, 0,
  587            sizeof(shake256->s) - WC_SHA3_256_COUNT * 8);
  588        data2Len -= WC_SHA3_256_COUNT * 8 - data1Len;
  589        data2    += WC_SHA3_256_COUNT * 8 - data1Len;
  590#ifndef WC_SHA3_NO_ASM
  591        if (IS_INTEL_AVX2(cpuid_flags) &&
  592                 (SAVE_VECTOR_REGISTERS2() == 0)) {
  593            sha3_block_avx2(state);
  594            RESTORE_VECTOR_REGISTERS();
  595        }
  596        else if (IS_INTEL_BMI2(cpuid_flags)) {
  597            sha3_block_bmi2(state);
  598        }
  599        else
  600#endif
  601        {
  602            BlockSha3(state);
  603        }
  604
  605        if (data2Len >= WC_SHA3_256_COUNT * 8) {
  606#ifndef WC_SHA3_NO_ASM
  607            word32 n = data2Len / (WC_SHA3_256_COUNT * 8);
  608            if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
  609                sha3_block_n_avx2(state, data2, n, WC_SHA3_256_COUNT * 8);
  610                RESTORE_VECTOR_REGISTERS();
  611                n *= WC_SHA3_256_COUNT * 8;
  612                data2Len -= n;
  613                data2    += n;
  614            }
  615            else if (IS_INTEL_BMI2(cpuid_flags)) {
  616                sha3_block_n_bmi2(state, data2, n, WC_SHA3_256_COUNT * 8);
  617                n *= WC_SHA3_256_COUNT * 8;
  618                data2Len -= n;
  619                data2    += n;
  620            }
  621            else
  622#endif
  623            {
  624                while (data2Len >= WC_SHA3_256_COUNT * 8) {
  625                    xorbuf(state, data2, WC_SHA3_256_COUNT * 8);
  626                    data2Len -= WC_SHA3_256_COUNT * 8;
  627                    data2    += WC_SHA3_256_COUNT * 8;
  628                    BlockSha3(state);
  629                }
  630            }
  631        }
  632
  633        if (data2Len > 0) {
  634            xorbuf(state, data2, data2Len);
  635        }
  636        state8[data2Len] ^= 0x1f;
  637        state8[WC_SHA3_256_COUNT * 8 - 1] ^= 0x80;
  638    }
  639    else {
  640        word32 dataLen = data1Len + data2Len;
  641
  642        XMEMCPY(state8, data1, data1Len);
  643        XMEMCPY(state8 + data1Len, data2, data2Len);
  644        state8[dataLen] = 0x1f;
  645        XMEMSET(state8 + dataLen + 1, 0, sizeof(shake256->s) - (dataLen + 1));
  646        state8[WC_SHA3_256_COUNT * 8 - 1] = 0x80;
  647    }
  648
  649#ifndef WC_SHA3_NO_ASM
  650    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
  651        sha3_block_avx2(state);
  652        RESTORE_VECTOR_REGISTERS();
  653    }
  654    else if (IS_INTEL_BMI2(cpuid_flags)) {
  655        sha3_block_bmi2(state);
  656    }
  657    else
  658#endif
  659    {
  660        BlockSha3(state);
  661    }
  662    XMEMCPY(hash, shake256->s, hashLen);
  663    ret = 0;
  664#else
  665    /* Initialize SHAKE-256 operation. */
  666    ret = wc_InitShake256(shake256, NULL, INVALID_DEVID);
  667    if (ret == 0) {
  668        /* Update with first data. */
  669        ret = wc_Shake256_Update(shake256, data1, data1Len);
  670    }
  671    if (ret == 0) {
  672        /* Update with second data. */
  673        ret = wc_Shake256_Update(shake256, data2, data2Len);
  674    }
  675    if (ret == 0) {
  676        /* Compute hash of data. */
  677        ret = wc_Shake256_Final(shake256, hash, hashLen);
  678    }
  679#endif
  680
  681    return ret;
  682}
  683
  684#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
  685/* 256-bit hash of context and message using SHAKE-256.
  686 *
  687 * FIPS 204. 5.2: Algorithm 2 ML-DSA.Sign(sk, M, ctx)
  688 *   ...
  689 *   10: M' <- BytesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
  690 *                         ctx) || M
  691 *   ...
  692 *
  693 * FIPS 204. 6.2: Algorithm 7 ML-DSA.Sign_internal(sk, M', rnd)
  694 *   ...
  695 *   6: mu <- H(BytesToBits(tr)||M', 64))
  696 *   ...
  697 *
  698 * @param [in, out] shake256  SHAKE-256 object.
  699 * @param [in]      tr        Public key hash.
  700 * @param [in]      trLen     Length of public key hash in bytes.
  701 * @param [in]      preHash   0 when message was not hashed,
  702 *                            1 when message was hashed.
  703 * @param [in]      ctx       Context of signature.
  704 * @param [in]      ctxLen    Length of context of signature in bytes.
  705 * @param [in]      msg       Message to sign.
  706 * @param [in]      msgLen    Length of message to sign in bytes.
  707 * @param [out]     hash      Buffer to hold hash result.
  708 * @param [in]      hashLen   Number of bytes of hash to return.
  709 * @return  0 on success.
  710 * @return  Negative on error.
  711 */
  712static int dilithium_hash256_ctx_msg(wc_Shake* shake256, const byte* tr,
  713    byte trLen, byte preHash, const byte* ctx, byte ctxLen, const byte* msg,
  714    word32 msgLen, byte* hash, word32 hashLen)
  715{
  716    int ret;
  717    byte prefix[2];
  718
  719    prefix[0] = preHash;
  720    prefix[1] = ctxLen;
  721
  722    /* Initialize SHAKE-256 operation. */
  723    ret = wc_InitShake256(shake256, NULL, INVALID_DEVID);
  724    if (ret == 0) {
  725        /* Update with public key hash. */
  726        ret = wc_Shake256_Update(shake256, tr, trLen);
  727    }
  728    if (ret == 0) {
  729        /* Update with context prefix - 0 | ctxLen. */
  730        ret = wc_Shake256_Update(shake256, prefix, (word32)sizeof(prefix));
  731    }
  732    if (ret == 0) {
  733        /* Update with context. */
  734        ret = wc_Shake256_Update(shake256, ctx, ctxLen);
  735    }
  736    if (ret == 0) {
  737        /* Update with message. */
  738        ret = wc_Shake256_Update(shake256, msg, msgLen);
  739    }
  740    if (ret == 0) {
  741        /* Compute hash of data. */
  742        ret = wc_Shake256_Final(shake256, hash, hashLen);
  743    }
  744
  745    return ret;
  746}
  747
  748/* Get the OID for the digest hash.
  749 *
  750 * @param [in]  hash         Hash algorithm.
  751 * @param [out] oidBuffer   Buffer to hold OID.
  752 * @param [out] oidLen      Length of OID in buffer.
  753 * @return  0 on success.
  754 * @return  BAD_FUNC_ARG if hash algorithm not known.
  755 */
  756static int dilithium_get_hash_oid(int hash, byte* oidBuffer, word32* oidLen)
  757{
  758    int ret = 0;
  759    const byte* oid;
  760
  761#ifndef WOLFSSL_DILITHIUM_NO_ASN1
  762
  763    oid = OidFromId((word32)wc_HashGetOID((enum wc_HashType)hash), oidHashType,
  764        oidLen);
  765    if ((oid != NULL) && (*oidLen <= DILITHIUM_HASH_OID_LEN - 2)) {
  766#ifndef WOLFSSL_DILITHIUM_REVERSE_HASH_OID
  767        oidBuffer[0] = 0x06;             /* ObjectID */
  768        oidBuffer[1] = (byte)*oidLen;   /* ObjectID */
  769        oidBuffer += 2;
  770        XMEMCPY(oidBuffer, oid, *oidLen);
  771#else
  772        int i;
  773        for (i = (int)*oidLen - 1; i >= 0; i--) {
  774            *(oidBuffer++) = oid[i];
  775        }
  776        *(oidBuffer++) = *oidLen;   /* ObjectID */
  777        * oidBuffer    = 0x06;   /* ObjectID */
  778#endif
  779        *oidLen += 2;
  780     }
  781     else {
  782        ret = BAD_FUNC_ARG;
  783    }
  784
  785#else
  786
  787    *oidLen = DILITHIUM_HASH_OID_LEN;
  788#ifndef NO_SHA256
  789    if (hash == WC_HASH_TYPE_SHA256) {
  790        static byte sha256Oid[DILITHIUM_HASH_OID_LEN] = {
  791            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01
  792        };
  793        oid = sha256Oid;
  794    }
  795    else
  796#endif
  797#ifdef WOLFSSL_SHA384
  798    if (hash == WC_HASH_TYPE_SHA384) {
  799        static byte sha384Oid[DILITHIUM_HASH_OID_LEN] = {
  800            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02
  801        };
  802        oid = sha384Oid;
  803    }
  804    else
  805#endif
  806#ifdef WOLFSSL_SHA512
  807    if (hash == WC_HASH_TYPE_SHA512) {
  808        static byte sha512Oid[DILITHIUM_HASH_OID_LEN] = {
  809            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03
  810        };
  811        oid = sha512Oid;
  812    }
  813    else
  814#ifndef WOLFSSL_NOSHA512_224
  815    if (hash == WC_HASH_TYPE_SHA512_224) {
  816        static byte sha512_224Oid[DILITHIUM_HASH_OID_LEN] = {
  817            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x05
  818        };
  819        oid = sha512_224Oid;
  820    }
  821    else
  822#endif
  823#ifndef WOLFSSL_NOSHA512_256
  824    if (hash == WC_HASH_TYPE_SHA512_256) {
  825        static byte sha512_256Oid[DILITHIUM_HASH_OID_LEN] = {
  826            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x06
  827        };
  828        oid = sha512_256Oid;
  829    }
  830    else
  831#endif
  832#endif
  833    if (hash == WC_HASH_TYPE_SHAKE128) {
  834        static byte shake128Oid[DILITHIUM_HASH_OID_LEN] = {
  835            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0B
  836        };
  837        oid = shake128Oid;
  838    }
  839    else if (hash == WC_HASH_TYPE_SHAKE256) {
  840        static byte shake256Oid[DILITHIUM_HASH_OID_LEN] = {
  841            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0C
  842        };
  843        oid = shake256Oid;
  844    }
  845    else if (hash == WC_HASH_TYPE_SHA3_256) {
  846        static byte sha3_256Oid[DILITHIUM_HASH_OID_LEN] = {
  847            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x08
  848        };
  849        oid = sha3_256Oid;
  850    }
  851    else if (hash == WC_HASH_TYPE_SHA3_384) {
  852        static byte sha3_384Oid[DILITHIUM_HASH_OID_LEN] = {
  853            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x09
  854        };
  855        oid = sha3_384Oid;
  856    }
  857    else if (hash == WC_HASH_TYPE_SHA3_512) {
  858        static byte sha3_512Oid[DILITHIUM_HASH_OID_LEN] = {
  859            0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0A
  860        };
  861        oid = sha3_512Oid;
  862    }
  863    else {
  864        oid = NULL;
  865        ret = BAD_FUNC_ARG;
  866    }
  867
  868    if ((oid != NULL) && (*oidLen <= DILITHIUM_HASH_OID_LEN)) {
  869#ifndef WOLFSSL_DILITHIUM_REVERSE_HASH_OID
  870        XMEMCPY(oidBuffer, oid, *oidLen);
  871#else
  872        int i;
  873        for (i = (int)*oidLen - 1; i >= 0; i--) {
  874            *(oidBuffer++) = oid[i];
  875        }
  876#endif
  877    }
  878#endif
  879
  880    return ret;
  881}
  882#endif
  883
  884#ifndef WOLFSSL_DILITHIUM_SMALL
  885/* 128-bit hash using SHAKE-128.
  886 *
  887 * FIPS 204. 8.3: H128(v,d) <- SHAKE128(v,d)
  888 *
  889 * @param [in, out] shake128  SHAKE-128 object.
  890 * @param [in]      in        Block of data to hash.
  891 * @param [in]      inLen     Length of data in bytes.
  892 * @param [out]     out       Buffer to hold hash result.
  893 * @param [in]      outLen    Number of hash blocks to return.
  894 * @return  0 on success.
  895 * @return  Negative on error.
  896 */
  897static int dilithium_squeeze128(wc_Shake* shake128, const byte* in,
  898    word32 inLen, byte* out, word32 outBlocks)
  899{
  900    int ret;
  901
  902    /* Initialize SHAKE-128 operation. */
  903    ret = wc_InitShake128(shake128, NULL, INVALID_DEVID);
  904    if (ret == 0) {
  905        /* Absorb data - update plus final. */
  906        ret = wc_Shake128_Absorb(shake128, in, inLen);
  907    }
  908    if (ret == 0) {
  909        /* Squeeze out hash data. */
  910        ret = wc_Shake128_SqueezeBlocks(shake128, out, outBlocks);
  911    }
  912
  913    return ret;
  914}
  915#endif /* WOLFSSL_DILITHIUM_SMALL */
  916
  917#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
  918    (!defined(WOLFSSL_DILITHIUM_SMALL) && \
  919     !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY))
  920/* 256-bit hash using SHAKE-256.
  921 *
  922 * FIPS 204. 8.3: H(v,d) <- SHAKE256(v,d)
  923 * Using SqueezeBlocks interface to get larger amounts of output.
  924 *
  925 * @param [in, out] shake256  SHAKE-256 object.
  926 * @param [in]      in        Block of data to hash.
  927 * @param [in]      inLen     Length of data in bytes.
  928 * @param [out]     out       Buffer to hold hash result.
  929 * @param [in]      outLen    Number of hash blocks to return.
  930 * @return  0 on success.
  931 * @return  Negative on hash error.
  932 */
  933static int dilithium_squeeze256(wc_Shake* shake256, const byte* in,
  934    word32 inLen, byte* out, word32 outBlocks)
  935{
  936    int ret;
  937#ifdef USE_INTEL_SPEEDUP
  938    word64* state = shake256->s;
  939    word8* state8 = (word8*)state;
  940
  941    XMEMCPY(state, in, inLen);
  942    state8[inLen] = 0x1f;
  943    XMEMSET(state8 + inLen + 1, 0, sizeof(shake256->s) - (inLen + 1));
  944    state8[WC_SHA3_256_COUNT * 8 - 1] ^= 0x80;
  945
  946    for (; outBlocks > 0; outBlocks--) {
  947#ifndef WC_SHA3_NO_ASM
  948        if (IS_INTEL_AVX2(cpuid_flags) &&
  949                 (SAVE_VECTOR_REGISTERS2() == 0)) {
  950            sha3_block_avx2(state);
  951            RESTORE_VECTOR_REGISTERS();
  952        }
  953        else if (IS_INTEL_BMI2(cpuid_flags)) {
  954            sha3_block_bmi2(state);
  955        }
  956        else
  957#endif
  958        {
  959            BlockSha3(state);
  960        }
  961        XMEMCPY(out, shake256->s, WC_SHA3_256_COUNT * 8);
  962        out += WC_SHA3_256_COUNT * 8;
  963    }
  964    ret = 0;
  965#else
  966    /* Initialize SHAKE-256 operation. */
  967    ret = wc_InitShake256(shake256, NULL, INVALID_DEVID);
  968    if (ret == 0) {
  969        /* Absorb data - update plus final. */
  970        ret = wc_Shake256_Absorb(shake256, in, inLen);
  971    }
  972    if (ret == 0) {
  973        /* Squeeze out hash data. */
  974        ret = wc_Shake256_SqueezeBlocks(shake256, out, outBlocks);
  975    }
  976#endif
  977
  978    return ret;
  979}
  980#endif
  981
  982/******************************************************************************
  983 * Encode/Decode operations
  984 ******************************************************************************/
  985
  986#ifndef WOLFSSL_DILITHIUM_NO_MAKE_KEY
  987/* Encode vector of polynomials with range -ETA..ETA.
  988 *
  989 * FIPS 204. 8.2: Algorithm 18 skEncode(rho, K, tr, s1, s2, t0)
  990 *   ...
  991 *   2: for i from 0 to l - 1 do
  992 *   3:     sk <- sk || BitPack(s1[i], eta, eta)
  993 *   4: end for
  994 *   ...
  995 *   OR
  996 *   ...
  997 *   5: for i from 0 to k - 1 do
  998 *   6:     sk <- sk || BitPack(s2[i], eta, eta)
  999 *   7: end for
 1000 *   ...
 1001 *
 1002 * FIPS 204. 8.2: Algorithm 11 BitPack(w, a, b)
 1003 *   1: z <- ()
 1004 *   2: for i from 0 to 255 do
 1005 *   3:     z <- z||IntegerToBits(b - wi, bitlen(a + b))
 1006 *   4: end for
 1007 *   5: return BitsToBytes(z)
 1008 *
 1009 *   IntegerToBits makes bit array with width specified from integer.
 1010 *   BitToBytes make a byte array from a bit array.
 1011 *
 1012 * @param [in]  s    Vector of polynomials to encode.
 1013 * @param [in]  d    Dimension of vector.
 1014 * @param [in]  eta  Range specifier of each value.
 1015 * @param [out] p    Buffer to encode into.
 1016 */
 1017static void dilthium_vec_encode_eta_bits_c(const sword32* s, byte d, byte eta,
 1018    byte* p)
 1019{
 1020    unsigned int i;
 1021    unsigned int j;
 1022
 1023#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 1024    /* -2..2 */
 1025    if (eta == DILITHIUM_ETA_2) {
 1026        /* Setp 2 or 5: For each polynomial of vector. */
 1027        for (i = 0; i < d; i++) {
 1028            /* Step 3 or 6.
 1029             * 3 bits to encode each number.
 1030             * 8 numbers become 3 bytes. (8 * 3 bits = 3 * 8 bits) */
 1031            for (j = 0; j < DILITHIUM_N; j += 8) {
 1032                /* Make value a positive integer. */
 1033                byte s0 = (byte)(2 - s[j + 0]);
 1034                byte s1 = (byte)(2 - s[j + 1]);
 1035                byte s2 = (byte)(2 - s[j + 2]);
 1036                byte s3 = (byte)(2 - s[j + 3]);
 1037                byte s4 = (byte)(2 - s[j + 4]);
 1038                byte s5 = (byte)(2 - s[j + 5]);
 1039                byte s6 = (byte)(2 - s[j + 6]);
 1040                byte s7 = (byte)(2 - s[j + 7]);
 1041
 1042                /* Pack 8 3-bit values into 3 bytes. */
 1043                p[0] = (byte)((s0 >> 0) | (s1 << 3) | (s2 << 6));
 1044                p[1] = (byte)((s2 >> 2) | (s3 << 1) | (s4 << 4) | (s5 << 7));
 1045                p[2] = (byte)((s5 >> 1) | (s6 << 2) | (s7 << 5));
 1046                /* Move to next place to encode into. */
 1047                p += DILITHIUM_ETA_2_BITS;
 1048            }
 1049            /* Next polynomial. */
 1050            s += DILITHIUM_N;
 1051        }
 1052    }
 1053#endif
 1054#ifndef WOLFSSL_NO_ML_DSA_65
 1055    /* -4..4 */
 1056    if (eta == DILITHIUM_ETA_4) {
 1057        for (i = 0; i < d; i++) {
 1058        #ifdef WOLFSSL_DILITHIUM_SMALL
 1059            /* Step 3 or 6.
 1060             * 4 bits to encode each number.
 1061             * 2 numbers become 1 bytes. (2 * 4 bits = 1 * 8 bits) */
 1062            for (j = 0; j < DILITHIUM_N / 2; j++) {
 1063                /* Make values positive and pack 2 4-bit values into 1 byte. */
 1064                p[j] = (byte)((((byte)(4 - s[j * 2 + 0])) << 0) |
 1065                              (((byte)(4 - s[j * 2 + 1])) << 4));
 1066            }
 1067        #else
 1068            /* Step 3 or 6.
 1069             * 4 bits to encode each number.
 1070             * 8 numbers become 4 bytes. (8 * 4 bits = 4 * 8 bits) */
 1071            for (j = 0; j < DILITHIUM_N / 2; j += 4) {
 1072                /* Make values positive and pack 2 4-bit values into 1 byte. */
 1073                p[j + 0] = (byte)((((byte)(4 - s[j * 2 + 0])) << 0) |
 1074                                  (((byte)(4 - s[j * 2 + 1])) << 4));
 1075                p[j + 1] = (byte)((((byte)(4 - s[j * 2 + 2])) << 0) |
 1076                                  (((byte)(4 - s[j * 2 + 3])) << 4));
 1077                p[j + 2] = (byte)((((byte)(4 - s[j * 2 + 4])) << 0) |
 1078                                  (((byte)(4 - s[j * 2 + 5])) << 4));
 1079                p[j + 3] = (byte)((((byte)(4 - s[j * 2 + 6])) << 0) |
 1080                                  (((byte)(4 - s[j * 2 + 7])) << 4));
 1081            }
 1082        #endif
 1083            /* Move to next place to encode into. */
 1084            p += DILITHIUM_N / 2;
 1085            /* Next polynomial. */
 1086            s += DILITHIUM_N;
 1087        }
 1088    }
 1089#endif
 1090}
 1091
 1092/* Encode vector of polynomials with range -ETA..ETA.
 1093 *
 1094 * @param [in]  s    Vector of polynomials to encode.
 1095 * @param [in]  d    Dimension of vector.
 1096 * @param [in]  eta  Range specifier of each value.
 1097 * @param [out] p    Buffer to encode into.
 1098 */
 1099static void dilthium_vec_encode_eta_bits(const sword32* s, byte d, byte eta,
 1100    byte* p)
 1101{
 1102#ifdef USE_INTEL_SPEEDUP
 1103    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1104    #if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 1105        /* -2..2 */
 1106        if (eta == DILITHIUM_ETA_2) {
 1107            wc_mldsa_vec_encode_eta_2_avx2(s, d, p);
 1108        }
 1109    #endif
 1110    #ifndef WOLFSSL_NO_ML_DSA_65
 1111        if (eta == DILITHIUM_ETA_4) {
 1112            wc_mldsa_vec_encode_eta_4_avx2(s, p);
 1113        }
 1114    #endif
 1115        RESTORE_VECTOR_REGISTERS();
 1116    }
 1117    else
 1118#endif
 1119    {
 1120        dilthium_vec_encode_eta_bits_c(s, d, eta, p);
 1121    }
 1122}
 1123#endif /* !WOLFSSL_DILITHIUM_NO_MAKE_KEY */
 1124
 1125#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 1126
 1127#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 1128/* Decode polynomial with range -2..2.
 1129 *
 1130 * FIPS 204. 7.2: Algorithm 25 skDecode(sk)
 1131 *   ...
 1132 *   5: for i from 0 to l - 1 do
 1133 *   6:     s1[i] <- BitUnpack(yi, eta, eta)
 1134 *   7: end for
 1135 *   ...
 1136 *   OR
 1137 *   ...
 1138 *   8: for i from 0 to k - 1 do
 1139 *   9:     s2[i] <- BitUnpack(zi, eta, eta)
 1140 *  10: end for
 1141 *   ...
 1142 *  Where y and z are arrays of bit arrays.
 1143 *
 1144 * @param [in]  p    Buffer of data to decode.
 1145 * @param [in]  s    Vector of decoded polynomials.
 1146 */
 1147static void dilithium_decode_eta_2_bits_c(const byte* p, sword32* s)
 1148{
 1149    unsigned int j;
 1150
 1151    /* Step 6 or 9.
 1152     * 3 bits to encode each number.
 1153     * 8 numbers from 3 bytes. (8 * 3 bits = 3 * 8 bits) */
 1154    for (j = 0; j < DILITHIUM_N; j += 8) {
 1155        /* Get 3 bits and put in range of -2..2. */
 1156        s[j + 0] = 2 - ((p[0] >> 0) & 0x7                      );
 1157        s[j + 1] = 2 - ((p[0] >> 3) & 0x7                      );
 1158        s[j + 2] = 2 - ((p[0] >> 6)       | ((p[1] << 2) & 0x7));
 1159        s[j + 3] = 2 - ((p[1] >> 1) & 0x7                      );
 1160        s[j + 4] = 2 - ((p[1] >> 4) & 0x7                      );
 1161        s[j + 5] = 2 - ((p[1] >> 7)       | ((p[2] << 1) & 0x7));
 1162        s[j + 6] = 2 - ((p[2] >> 2) & 0x7                      );
 1163        s[j + 7] = 2 - ((p[2] >> 5) & 0x7                      );
 1164        /* Move to next place to decode from. */
 1165        p += DILITHIUM_ETA_2_BITS;
 1166    }
 1167}
 1168
 1169/* Decode polynomial with range -2..2.
 1170 *
 1171 * @param [in]  p    Buffer of data to decode.
 1172 * @param [in]  s    Vector of decoded polynomials.
 1173 */
 1174static void dilithium_decode_eta_2_bits(const byte* p, sword32* s)
 1175{
 1176#ifdef USE_INTEL_SPEEDUP
 1177    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1178        wc_mldsa_decode_eta_2_avx2(p, s);
 1179        RESTORE_VECTOR_REGISTERS();
 1180    }
 1181    else
 1182#endif
 1183    {
 1184        dilithium_decode_eta_2_bits_c(p, s);
 1185    }
 1186}
 1187#endif
 1188#ifndef WOLFSSL_NO_ML_DSA_65
 1189/* Decode polynomial with range -4..4.
 1190 *
 1191 * FIPS 204. 7.2: Algorithm 25 skDecode(sk)
 1192 *   ...
 1193 *   5: for i from 0 to l - 1 do
 1194 *   6:     s1[i] <- BitUnpack(yi, eta, eta)
 1195 *   7: end for
 1196 *   ...
 1197 *   OR
 1198 *   ...
 1199 *   8: for i from 0 to k - 1 do
 1200 *   9:     s2[i] <- BitUnpack(zi, eta, eta)
 1201 *  10: end for
 1202 *   ...
 1203 *  Where y and z are arrays of bit arrays.
 1204 *
 1205 * @param [in]  p    Buffer of data to decode.
 1206 * @param [in]  s    Vector of decoded polynomials.
 1207 */
 1208static void dilithium_decode_eta_4_bits_c(const byte* p, sword32* s)
 1209{
 1210    unsigned int j;
 1211
 1212#ifdef WOLFSSL_DILITHIUM_SMALL
 1213    /* Step 6 or 9.
 1214     * 4 bits to encode each number.
 1215     * 2 numbers from 1 bytes. (2 * 4 bits = 1 * 8 bits) */
 1216    for (j = 0; j < DILITHIUM_N / 2; j++) {
 1217        /* Get 4 bits and put in range of -4..4. */
 1218        s[j * 2 + 0] = 4 - (p[j] & 0xf);
 1219        s[j * 2 + 1] = 4 - (p[j] >> 4);
 1220    }
 1221#else
 1222    /* Step 6 or 9.
 1223     * 4 bits to encode each number.
 1224     * 8 numbers from 4 bytes. (8 * 4 bits = 4 * 8 bits) */
 1225    for (j = 0; j < DILITHIUM_N / 2; j += 4) {
 1226        /* Get 4 bits and put in range of -4..4. */
 1227        s[j * 2 + 0] = 4 - (p[j + 0] & 0xf);
 1228        s[j * 2 + 1] = 4 - (p[j + 0] >> 4);
 1229        s[j * 2 + 2] = 4 - (p[j + 1] & 0xf);
 1230        s[j * 2 + 3] = 4 - (p[j + 1] >> 4);
 1231        s[j * 2 + 4] = 4 - (p[j + 2] & 0xf);
 1232        s[j * 2 + 5] = 4 - (p[j + 2] >> 4);
 1233        s[j * 2 + 6] = 4 - (p[j + 3] & 0xf);
 1234        s[j * 2 + 7] = 4 - (p[j + 3] >> 4);
 1235    }
 1236#endif /* WOLFSSL_DILITHIUM_SMALL */
 1237}
 1238
 1239/* Decode polynomial with range -4..4.
 1240 *
 1241 * @param [in]  p    Buffer of data to decode.
 1242 * @param [in]  s    Vector of decoded polynomials.
 1243 */
 1244static void dilithium_decode_eta_4_bits(const byte* p, sword32* s)
 1245{
 1246#ifdef USE_INTEL_SPEEDUP
 1247    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1248        wc_mldsa_decode_eta_4_avx2(p, s);
 1249        RESTORE_VECTOR_REGISTERS();
 1250    }
 1251    else
 1252#endif
 1253    {
 1254        dilithium_decode_eta_4_bits_c(p, s);
 1255    }
 1256}
 1257#endif
 1258
 1259#if defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 1260    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 1261     (defined(WC_DILITHIUM_CACHE_PRIV_VECTORS) || \
 1262      !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)))
 1263/* Decode vector of polynomials with range -ETA..ETA.
 1264 *
 1265 * FIPS 204. 7.2: Algorithm 25 skDecode(sk)
 1266 *   ...
 1267 *   5: for i from 0 to l - 1 do
 1268 *   6:     s1[i] <- BitUnpack(yi, eta, eta)
 1269 *   7: end for
 1270 *   ...
 1271 *   OR
 1272 *   ...
 1273 *   8: for i from 0 to k - 1 do
 1274 *   9:     s2[i] <- BitUnpack(zi, eta, eta)
 1275 *  10: end for
 1276 *   ...
 1277 *  Where y and z are arrays of bit arrays.
 1278 *
 1279 * @param [in]  p    Buffer of data to decode.
 1280 * @param [in]  eta  Range specifier of each value.
 1281 * @param [in]  s    Vector of decoded polynomials.
 1282 * @param [in]  d    Dimension of vector.
 1283 */
 1284static void dilithium_vec_decode_eta_bits(const byte* p, byte eta, sword32* s,
 1285    byte d)
 1286{
 1287    unsigned int i;
 1288
 1289#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 1290    /* -2..2 */
 1291    if (eta == DILITHIUM_ETA_2) {
 1292        /* Step 5 or 8: For each polynomial of vector */
 1293        for (i = 0; i < d; i++) {
 1294            dilithium_decode_eta_2_bits(p, s);
 1295            /* Move to next place to decode from. */
 1296            p += DILITHIUM_ETA_2_BITS * DILITHIUM_N / 8;
 1297            /* Next polynomial. */
 1298            s += DILITHIUM_N;
 1299        }
 1300    }
 1301#endif
 1302#ifndef WOLFSSL_NO_ML_DSA_65
 1303    /* -4..4 */
 1304    if (eta == DILITHIUM_ETA_4) {
 1305        /* Step 5 or 8: For each polynomial of vector */
 1306        for (i = 0; i < d; i++) {
 1307            dilithium_decode_eta_4_bits(p, s);
 1308            /* Move to next place to decode from. */
 1309            p += DILITHIUM_N / 2;
 1310            /* Next polynomial. */
 1311            s += DILITHIUM_N;
 1312        }
 1313    }
 1314#endif
 1315}
 1316#endif
 1317#endif /* !WOLFSSL_DILITHIUM_NO_SIGN || WOLFSSL_DILITHIUM_CHECK_KEY */
 1318
 1319#ifndef WOLFSSL_DILITHIUM_NO_MAKE_KEY
 1320/* Encode t into t0 and t1.
 1321 *
 1322 * FIPS 204. 8.4: Algorithm 29 Power2Round(r)
 1323 *   1: r+ <- r mod q
 1324 *   2: r0 <- r+ mod +/- 2^d
 1325 *   3: return ((r+ - r0) / 2^d, r0)
 1326 *
 1327 * FIPS 204. 7.2: Algorithm 24 skEncode(rho, K, tr, s1, s2, t0)
 1328 *   ...
 1329 *   8: for i form 0 to k - 1 do
 1330 *   9:     sk <- sk || BitPack(t0[i], s^(d-1) - 1, 2^(d-1))
 1331 *  10: end for
 1332 *
 1333 * FIPS 204. 7.2: Algorithm 22 pkEncode(rho, t1)
 1334 *   ...
 1335 *   2: for i from 0 to k - 1 do
 1336 *   3:     pk <- pk || SimpleBitPack(t1[i], 2^bitlen(q-1) - d - 1)
 1337 *   4: end for
 1338 *
 1339 * @param [in]  t    Vector of polynomials.
 1340 * @param [in]  d    Dimension of vector.
 1341 * @param [out] t0   Buffer to encode bottom part of value of t into.
 1342 * @param [out] t1   Buffer to encode top part of value of t into.
 1343 */
 1344static void dilithium_vec_encode_t0_t1_c(const sword32* t, byte d, byte* t0,
 1345    byte* t1)
 1346{
 1347    unsigned int i;
 1348    unsigned int j;
 1349
 1350    /* Alg 24, Step 8 and Alg 22, Step 2. For each polynomial of vector. */
 1351    for (i = 0; i < d; i++) {
 1352        /* Alg 24, Step 9 and Alg 22, Step 3.
 1353         * Do all polynomial values - 8 at a time. */
 1354        for (j = 0; j < DILITHIUM_N; j += 8) {
 1355            /* Take 8 values of t and take top bits and make positive. */
 1356            word16 n1_0 = (word16)((t[j + 0] + DILITHIUM_D_MAX_HALF - 1) >>
 1357                DILITHIUM_D);
 1358            word16 n1_1 = (word16)((t[j + 1] + DILITHIUM_D_MAX_HALF - 1) >>
 1359                DILITHIUM_D);
 1360            word16 n1_2 = (word16)((t[j + 2] + DILITHIUM_D_MAX_HALF - 1) >>
 1361                DILITHIUM_D);
 1362            word16 n1_3 = (word16)((t[j + 3] + DILITHIUM_D_MAX_HALF - 1) >>
 1363                DILITHIUM_D);
 1364            word16 n1_4 = (word16)((t[j + 4] + DILITHIUM_D_MAX_HALF - 1) >>
 1365                DILITHIUM_D);
 1366            word16 n1_5 = (word16)((t[j + 5] + DILITHIUM_D_MAX_HALF - 1) >>
 1367                DILITHIUM_D);
 1368            word16 n1_6 = (word16)((t[j + 6] + DILITHIUM_D_MAX_HALF - 1) >>
 1369                DILITHIUM_D);
 1370            word16 n1_7 = (word16)((t[j + 7] + DILITHIUM_D_MAX_HALF - 1) >>
 1371                DILITHIUM_D);
 1372            /* Take 8 values of t and take bottom bits and make positive. */
 1373            word16 n0_0 = (word16)(DILITHIUM_D_MAX_HALF -
 1374                                   (t[j + 0] - (n1_0 << DILITHIUM_D)));
 1375            word16 n0_1 = (word16)(DILITHIUM_D_MAX_HALF -
 1376                                   (t[j + 1] - (n1_1 << DILITHIUM_D)));
 1377            word16 n0_2 = (word16)(DILITHIUM_D_MAX_HALF -
 1378                                   (t[j + 2] - (n1_2 << DILITHIUM_D)));
 1379            word16 n0_3 = (word16)(DILITHIUM_D_MAX_HALF -
 1380                                   (t[j + 3] - (n1_3 << DILITHIUM_D)));
 1381            word16 n0_4 = (word16)(DILITHIUM_D_MAX_HALF -
 1382                                   (t[j + 4] - (n1_4 << DILITHIUM_D)));
 1383            word16 n0_5 = (word16)(DILITHIUM_D_MAX_HALF -
 1384                                   (t[j + 5] - (n1_5 << DILITHIUM_D)));
 1385            word16 n0_6 = (word16)(DILITHIUM_D_MAX_HALF -
 1386                                   (t[j + 6] - (n1_6 << DILITHIUM_D)));
 1387            word16 n0_7 = (word16)(DILITHIUM_D_MAX_HALF -
 1388                                   (t[j + 7] - (n1_7 << DILITHIUM_D)));
 1389
 1390            /* 13 bits per number.
 1391             * 8 numbers become 13 bytes. (8 * 13 bits = 13 * 8 bits) */
 1392        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 2)
 1393            word32* tp;
 1394        #endif
 1395        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1396            tp = (word32*)t0;
 1397            tp[0] =  (n0_0      ) | ((word32)n0_1 << 13) | ((word32)n0_2 << 26);
 1398            tp[1] =  (n0_2 >>  6) | ((word32)n0_3 <<  7) | ((word32)n0_4 << 20);
 1399            tp[2] =  (n0_4 >> 12) | ((word32)n0_5 <<  1) |
 1400                                    ((word32)n0_6 << 14) | ((word32)n0_7 << 27);
 1401        #else
 1402            t0[ 0] = (byte)(               (n0_0 <<  0));
 1403            t0[ 1] = (byte)((n0_0 >>  8) | (n0_1 <<  5));
 1404            t0[ 2] = (byte)((n0_1 >>  3)               );
 1405            t0[ 3] = (byte)((n0_1 >> 11) | (n0_2 <<  2));
 1406            t0[ 4] = (byte)((n0_2 >>  6) | (n0_3 <<  7));
 1407            t0[ 5] = (byte)((n0_3 >>  1)               );
 1408            t0[ 6] = (byte)((n0_3 >>  9) | (n0_4 <<  4));
 1409            t0[ 7] = (byte)((n0_4 >>  4)               );
 1410            t0[ 8] = (byte)((n0_4 >> 12) | (n0_5 <<  1));
 1411            t0[ 9] = (byte)((n0_5 >>  7) | (n0_6 <<  6));
 1412            t0[10] = (byte)((n0_6 >>  2)               );
 1413            t0[11] = (byte)((n0_6 >> 10) | (n0_7 <<  3));
 1414        #endif
 1415            t0[12] = (byte)((n0_7 >>  5)               );
 1416
 1417            /* 10 bits per number.
 1418             * 8 bytes become 10 bytes. (8 * 10 bits = 10 * 8 bits) */
 1419        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 2)
 1420            tp = (word32*)t1;
 1421            tp[0] =  (n1_0      ) | ((word32)n1_1 << 10) |
 1422                     ((word32)n1_2 << 20) | ((word32)n1_3 << 30);
 1423            tp[1] =  (n1_3 >>  2) | ((word32)n1_4 <<  8) |
 1424                     ((word32)n1_5 << 18) | ((word32)n1_6 << 28);
 1425        #else
 1426            t1[0] = (byte)(               (n1_0 << 0));
 1427            t1[1] = (byte)((n1_0 >> 8) |  (n1_1 << 2));
 1428            t1[2] = (byte)((n1_1 >> 6) |  (n1_2 << 4));
 1429            t1[3] = (byte)((n1_2 >> 4) |  (n1_3 << 6));
 1430            t1[4] = (byte)((n1_3 >> 2)               );
 1431            t1[5] = (byte)(               (n1_4 << 0));
 1432            t1[6] = (byte)((n1_4 >> 8) |  (n1_5 << 2));
 1433            t1[7] = (byte)((n1_5 >> 6) |  (n1_6 << 4));
 1434        #endif
 1435            t1[8] = (byte)((n1_6 >> 4) |  (n1_7 << 6));
 1436            t1[9] = (byte)((n1_7 >> 2)               );
 1437
 1438            /* Move to next place to encode bottom bits to. */
 1439            t0 += DILITHIUM_D;
 1440            /* Move to next place to encode top bits to. */
 1441            t1 += DILITHIUM_U;
 1442        }
 1443        /* Next polynomial. */
 1444        t += DILITHIUM_N;
 1445    }
 1446}
 1447
 1448/* Encode t into t0 and t1.
 1449 *
 1450 * @param [in]  t    Vector of polynomials.
 1451 * @param [in]  d    Dimension of vector.
 1452 * @param [out] t0   Buffer to encode bottom part of value of t into.
 1453 * @param [out] t1   Buffer to encode top part of value of t into.
 1454 */
 1455static void dilithium_vec_encode_t0_t1(const sword32* t, byte d, byte* t0,
 1456    byte* t1)
 1457{
 1458#ifdef USE_INTEL_SPEEDUP
 1459    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1460        wc_mldsa_vec_encode_t0_t1_avx2(t, d, t0, t1);
 1461        RESTORE_VECTOR_REGISTERS();
 1462    }
 1463    else
 1464#endif
 1465    {
 1466        dilithium_vec_encode_t0_t1_c(t, d, t0, t1);
 1467    }
 1468}
 1469#endif /* !WOLFSSL_DILITHIUM_NO_MAKE_KEY */
 1470
 1471#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 1472/* Decode bottom D bits of t as t0.
 1473 *
 1474 * FIPS 204. 7.2: Algorithm 25 skDecode(sk)
 1475 *   ...
 1476 *   12:     t0[i] <- BitUnpack(wi, 2^(d-1) - 1, 2^(d-1)
 1477 *   ...
 1478 *
 1479 * @param [in]  t0  Encoded values of t0.
 1480 * @param [in]  d   Dimensions of vector t0.
 1481 * @param [out] t   Vector of polynomials.
 1482 */
 1483static void dilithium_decode_t0_c(const byte* t0, sword32* t)
 1484{
 1485    unsigned int j;
 1486
 1487    /* Step 12. Get 13 bits and convert to range (2^(d-1)-1)..2^(d-1). */
 1488    for (j = 0; j < DILITHIUM_N; j += 8) {
 1489        /* 13 bits used per number.
 1490         * 8 numbers from 13 bytes. (8 * 13 bits = 13 * 8 bits) */
 1491#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1492        word32 t32_2 = ((const word32*)t0)[2];
 1493    #ifdef WC_64BIT_CPU
 1494        word64 t64 = *(const word64*)t0;
 1495        t[j + 0] = DILITHIUM_D_MAX_HALF - (sword32)( t64        & 0x1fff);
 1496        t[j + 1] = DILITHIUM_D_MAX_HALF - (sword32)((t64 >> 13) & 0x1fff);
 1497        t[j + 2] = DILITHIUM_D_MAX_HALF - (sword32)((t64 >> 26) & 0x1fff);
 1498        t[j + 3] = DILITHIUM_D_MAX_HALF - (sword32)((t64 >> 39) & 0x1fff);
 1499        t[j + 4] = DILITHIUM_D_MAX_HALF - (sword32)
 1500                   ((t64 >> 52) | ((t32_2 & 0x0001) << 12));
 1501    #else
 1502        word32 t32_0 = ((const word32*)t0)[0];
 1503        word32 t32_1 = ((const word32*)t0)[1];
 1504        t[j + 0] = DILITHIUM_D_MAX_HALF - (sword32)
 1505                    ( t32_0        & 0x1fff);
 1506        t[j + 1] = DILITHIUM_D_MAX_HALF - (sword32)
 1507                    ((t32_0 >> 13) & 0x1fff);
 1508        t[j + 2] = DILITHIUM_D_MAX_HALF - (sword32)
 1509                   (( t32_0 >> 26          ) | ((t32_1 & 0x007f) <<  6));
 1510        t[j + 3] = DILITHIUM_D_MAX_HALF - (sword32)
 1511                    ((t32_1 >>  7) & 0x1fff);
 1512        t[j + 4] = DILITHIUM_D_MAX_HALF - (sword32)
 1513                   (( t32_1 >> 20          ) | ((t32_2 & 0x0001) << 12));
 1514    #endif
 1515        t[j + 5] = DILITHIUM_D_MAX_HALF - (sword32)
 1516                    ((t32_2 >>  1) & 0x1fff);
 1517        t[j + 6] = DILITHIUM_D_MAX_HALF - (sword32)
 1518                    ((t32_2 >> 14) & 0x1fff);
 1519        t[j + 7] = DILITHIUM_D_MAX_HALF - (sword32)
 1520                   (( t32_2 >> 27          ) | ((word32)t0[12] ) <<  5 );
 1521#else
 1522        t[j + 0] = DILITHIUM_D_MAX_HALF -
 1523                   ((t0[ 0]     ) | (((word16)(t0[ 1] & 0x1f)) <<  8));
 1524        t[j + 1] = DILITHIUM_D_MAX_HALF -
 1525                   ((t0[ 1] >> 5) | (((word16)(t0[ 2]       )) <<  3) |
 1526                                    (((word16)(t0[ 3] & 0x03)) << 11));
 1527        t[j + 2] = DILITHIUM_D_MAX_HALF -
 1528                   ((t0[ 3] >> 2) | (((word16)(t0[ 4] & 0x7f)) <<  6));
 1529        t[j + 3] = DILITHIUM_D_MAX_HALF -
 1530                   ((t0[ 4] >> 7) | (((word16)(t0[ 5]       )) <<  1) |
 1531                                    (((word16)(t0[ 6] & 0x0f)) <<  9));
 1532        t[j + 4] = DILITHIUM_D_MAX_HALF -
 1533                   ((t0[ 6] >> 4) | (((word16)(t0[ 7]       )) <<  4) |
 1534                                    (((word16)(t0[ 8] & 0x01)) << 12));
 1535        t[j + 5] = DILITHIUM_D_MAX_HALF -
 1536                   ((t0[ 8] >> 1) | (((word16)(t0[ 9] & 0x3f)) <<  7));
 1537        t[j + 6] = DILITHIUM_D_MAX_HALF -
 1538                   ((t0[ 9] >> 6) | (((word16)(t0[10]       )) <<  2) |
 1539                                    (((word16)(t0[11] & 0x07)) << 10));
 1540        t[j + 7] = DILITHIUM_D_MAX_HALF -
 1541                   ((t0[11] >> 3) | (((word16)(t0[12]       )) <<  5));
 1542#endif
 1543        /* Move to next place to decode from. */
 1544        t0 += DILITHIUM_D;
 1545    }
 1546}
 1547
 1548/* Decode bottom D bits of t as t0.
 1549 *
 1550 * @param [in]  t0  Encoded values of t0.
 1551 * @param [in]  d   Dimensions of vector t0.
 1552 * @param [out] t   Vector of polynomials.
 1553 */
 1554static void dilithium_decode_t0(const byte* t0, sword32* t)
 1555{
 1556#ifdef USE_INTEL_SPEEDUP
 1557    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1558        wc_mldsa_decode_t0_avx2(t0, t);
 1559        RESTORE_VECTOR_REGISTERS();
 1560    }
 1561    else
 1562#endif
 1563    {
 1564        dilithium_decode_t0_c(t0, t);
 1565    }
 1566}
 1567
 1568#if defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 1569    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 1570     (defined(WC_DILITHIUM_CACHE_PRIV_VECTORS) || \
 1571      !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)))
 1572/* Decode bottom D bits of t as t0.
 1573 *
 1574 * FIPS 204. 7.2: Algorithm 25 skDecode(sk)
 1575 *   ...
 1576 *   11: for i from 0 to k - 1 do
 1577 *   12:     t0[i] <- BitUnpack(wi, 2^(d-1) - 1, 2^(d-1)
 1578 *   13: end for
 1579 *   ...
 1580 *
 1581 * @param [in]  t0  Encoded values of t0.
 1582 * @param [in]  d   Dimensions of vector t0.
 1583 * @param [out] t   Vector of polynomials.
 1584 */
 1585static void dilithium_vec_decode_t0(const byte* t0, byte d, sword32* t)
 1586{
 1587    unsigned int i;
 1588
 1589    /* Step 11. For each polynomial of vector. */
 1590    for (i = 0; i < d; i++) {
 1591        dilithium_decode_t0(t0, t);
 1592        t0 += DILITHIUM_D * DILITHIUM_N / 8;
 1593        /* Next polynomial. */
 1594        t += DILITHIUM_N;
 1595    }
 1596}
 1597#endif
 1598#endif /* !WOLFSSL_DILITHIUM_NO_SIGN || WOLFSSL_DILITHIUM_CHECK_KEY */
 1599
 1600#if !defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 1601    defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 1602/* Decode top bits of t as t1.
 1603 *
 1604 * FIPS 204. 7.2: Algorithm 23 pkDecode(pk)
 1605 *   ...
 1606 *   4:     t1[i] <- SimpleBitUnpack(zi, 2^(bitlen(q-1)-d) - 1)
 1607 *   ...
 1608 *
 1609 * @param [in]  t1  Encoded values of t1.
 1610 * @param [out] t   Polynomials.
 1611 */
 1612static void dilithium_decode_t1_c(const byte* t1, sword32* t)
 1613{
 1614    unsigned int j;
 1615    /* Step 4. Get 10 bits as a number. */
 1616    for (j = 0; j < DILITHIUM_N; j += 8) {
 1617        /* 10 bits used per number.
 1618         * 8 numbers from 10 bytes. (8 * 10 bits = 10 * 8 bits) */
 1619#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1620    #ifdef WC_64BIT_CPU
 1621        word64 t64 = *(const word64*) t1;
 1622        word16 t16 = *(const word16*)(t1 + 8);
 1623        t[j+0] = (sword32)( ( t64                     & 0x03ff) << DILITHIUM_D);
 1624        t[j+1] = (sword32)( ((t64 >> 10)              & 0x03ff) << DILITHIUM_D);
 1625        t[j+2] = (sword32)( ((t64 >> 20)              & 0x03ff) << DILITHIUM_D);
 1626        t[j+3] = (sword32)( ((t64 >> 30)              & 0x03ff) << DILITHIUM_D);
 1627        t[j+4] = (sword32)( ((t64 >> 40)              & 0x03ff) << DILITHIUM_D);
 1628        t[j+5] = (sword32)( ((t64 >> 50)              & 0x03ff) << DILITHIUM_D);
 1629        t[j+6] = (sword32)((((t64 >> 60) |
 1630                             (word64)(t16 << 4))      & 0x03ff) << DILITHIUM_D);
 1631        t[j+7] = (sword32)( ((t16 >>  6)              & 0x03ff) << DILITHIUM_D);
 1632    #else
 1633        word32 t32 = *((const word32*)t1);
 1634        t[j + 0] = (sword32)(( t32        & 0x03ff                       ) <<
 1635                   DILITHIUM_D);
 1636        t[j + 1] = (sword32)(((t32 >> 10) & 0x03ff                       ) <<
 1637                   DILITHIUM_D);
 1638        t[j + 2] = (sword32)(((t32 >> 20) & 0x03ff                       ) <<
 1639                   DILITHIUM_D);
 1640        t[j + 3] = (sword32)(((t32 >> 30)          | ((word32)t1[4] << 2)) <<
 1641                   DILITHIUM_D);
 1642        t32 = *((const word32*)(t1 + 5));
 1643        t[j + 4] = (sword32)(( t32        & 0x03ff                       ) <<
 1644                   DILITHIUM_D);
 1645        t[j + 5] = (sword32)(((t32 >> 10) & 0x03ff                       ) <<
 1646                   DILITHIUM_D);
 1647        t[j + 6] = (sword32)(((t32 >> 20) & 0x03ff                       ) <<
 1648                   DILITHIUM_D);
 1649        t[j + 7] = (sword32)(((t32 >> 30)          | ((word32)t1[9] << 2)) <<
 1650                   DILITHIUM_D);
 1651    #endif
 1652#else
 1653        t[j + 0] = (sword32)((t1[0] >> 0) | (((word16)(t1[1] & 0x03)) << 8))
 1654                   << DILITHIUM_D;
 1655        t[j + 1] = (sword32)((t1[1] >> 2) | (((word16)(t1[2] & 0x0f)) << 6))
 1656                   << DILITHIUM_D;
 1657        t[j + 2] = (sword32)((t1[2] >> 4) | (((word16)(t1[3] & 0x3f)) << 4))
 1658                   << DILITHIUM_D;
 1659        t[j + 3] = (sword32)((t1[3] >> 6) | (((word16)(t1[4]       )) << 2))
 1660                   << DILITHIUM_D;
 1661        t[j + 4] = (sword32)((t1[5] >> 0) | (((word16)(t1[6] & 0x03)) << 8))
 1662                   << DILITHIUM_D;
 1663        t[j + 5] = (sword32)((t1[6] >> 2) | (((word16)(t1[7] & 0x0f)) << 6))
 1664                   << DILITHIUM_D;
 1665        t[j + 6] = (sword32)((t1[7] >> 4) | (((word16)(t1[8] & 0x3f)) << 4))
 1666                   << DILITHIUM_D;
 1667        t[j + 7] = (sword32)((t1[8] >> 6) | (((word16)(t1[9]       )) << 2))
 1668                   << DILITHIUM_D;
 1669#endif
 1670        /* Move to next place to decode from. */
 1671        t1 += DILITHIUM_U;
 1672    }
 1673}
 1674
 1675/* Decode top bits of t as t1.
 1676 *
 1677 * FIPS 204. 7.2: Algorithm 23 pkDecode(pk)
 1678 *   ...
 1679 *   4:     t1[i] <- SimpleBitUnpack(zi, 2^(bitlen(q-1)-d) - 1)
 1680 *   ...
 1681 *
 1682 * @param [in]  t1  Encoded values of t1.
 1683 * @param [out] t   Polynomials.
 1684 */
 1685static void dilithium_decode_t1(const byte* t1, sword32* t)
 1686{
 1687#ifdef USE_INTEL_SPEEDUP
 1688    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1689        wc_mldsa_decode_t1_avx2(t1, t);
 1690        RESTORE_VECTOR_REGISTERS();
 1691    }
 1692    else
 1693#endif
 1694    {
 1695        dilithium_decode_t1_c(t1, t);
 1696    }
 1697}
 1698#endif
 1699
 1700#if (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 1701     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 1702    defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 1703/* Decode top bits of t as t1.
 1704 *
 1705 * FIPS 204. 7.2: Algorithm 23 pkDecode(pk)
 1706 *   ...
 1707 *   3: for i from 0 to k - 1 do
 1708 *   4:     t1[i] <- SimpleBitUnpack(zi, 2^(bitlen(q-1)-d) - 1)
 1709 *   5: end for
 1710 *   ...
 1711 *
 1712 * @param [in]  t1  Encoded values of t1.
 1713 * @param [in]  d   Dimensions of vector t1.
 1714 * @param [out] t   Vector of polynomials.
 1715 */
 1716static void dilithium_vec_decode_t1(const byte* t1, byte d, sword32* t)
 1717{
 1718    unsigned int i;
 1719
 1720    /* Step 3. For each polynomial of vector. */
 1721    for (i = 0; i < d; i++) {
 1722        dilithium_decode_t1(t1, t);
 1723        /* Next polynomial. */
 1724        t1 += DILITHIUM_U * DILITHIUM_N / 8;
 1725        t += DILITHIUM_N;
 1726    }
 1727}
 1728#endif
 1729
 1730#ifndef WOLFSSL_DILITHIUM_NO_SIGN
 1731
 1732#ifndef WOLFSSL_NO_ML_DSA_44
 1733/* Encode z with range of -(GAMMA1-1)...GAMMA1
 1734 *
 1735 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 1736 *   ...
 1737 *   3:     sigma <- sigma || BitPack(z[i], GAMMA1 - 1, GAMMA1)
 1738 *   ...
 1739 *
 1740 * @param [in]  z     Polynomial to encode.
 1741 * @param [out] s     Buffer to encode into.
 1742 */
 1743static void dilithium_encode_gamma1_17_bits_c(const sword32* z, byte* s)
 1744{
 1745    unsigned int j;
 1746
 1747    /* Step 3. Get 18 bits as a number. */
 1748    for (j = 0; j < DILITHIUM_N; j += 4) {
 1749        word32 z0 = (word32)(DILITHIUM_GAMMA1_17 - z[j + 0]);
 1750        word32 z1 = (word32)(DILITHIUM_GAMMA1_17 - z[j + 1]);
 1751        word32 z2 = (word32)(DILITHIUM_GAMMA1_17 - z[j + 2]);
 1752        word32 z3 = (word32)(DILITHIUM_GAMMA1_17 - z[j + 3]);
 1753
 1754        /* 18 bits per number.
 1755         * 8 numbers become 9 bytes. (8 * 9 bits = 9 * 8 bits) */
 1756#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1757    #ifdef WC_64BIT_CPU
 1758        word64* s64p = (word64*)s;
 1759        s64p[0] =           z0        | ((word64)z1 << 18) |
 1760                   ((word64)z2 << 36) | ((word64)z3 << 54);
 1761    #else
 1762        word32* s32p = (word32*)s;
 1763        s32p[0] =  z0        | (z1 << 18)             ;
 1764        s32p[1] = (z1 >> 14) | (z2 <<  4) | (z3 << 22);
 1765    #endif
 1766#else
 1767        s[0] = (byte)( z0                   );
 1768        s[1] = (byte)( z0 >>  8             );
 1769        s[2] = (byte)((z0 >> 16) | (z1 << 2));
 1770        s[3] = (byte)( z1 >>  6             );
 1771        s[4] = (byte)((z1 >> 14) | (z2 << 4));
 1772        s[5] = (byte)( z2 >>  4             );
 1773        s[6] = (byte)((z2 >> 12) | (z3 << 6));
 1774        s[7] = (byte)( z3 >>  2             );
 1775#endif
 1776        s[8] = (byte)( z3 >> 10             );
 1777        /* Move to next place to encode to. */
 1778        s += DILITHIUM_GAMMA1_17_ENC_BITS / 2;
 1779    }
 1780}
 1781
 1782/* Encode z with range of -(GAMMA1-1)...GAMMA1
 1783 *
 1784 * @param [in]  z     Polynomial to encode.
 1785 * @param [out] s     Buffer to encode into.
 1786 */
 1787static void dilithium_encode_gamma1_17_bits(const sword32* z, byte* s)
 1788{
 1789#ifdef USE_INTEL_SPEEDUP
 1790    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1791        wc_mldsa_encode_gamma1_17_avx2(z, s);
 1792        RESTORE_VECTOR_REGISTERS();
 1793    }
 1794    else
 1795#endif
 1796    {
 1797        dilithium_encode_gamma1_17_bits_c(z, s);
 1798    }
 1799}
 1800#endif
 1801#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 1802/* Encode z with range of -(GAMMA1-1)...GAMMA1
 1803 *
 1804 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 1805 *   ...
 1806 *   3:     sigma <- sigma || BitPack(z[i], GAMMA1 - 1, GAMMA1)
 1807 *   ...
 1808 *
 1809 * @param [in]  z     Polynomial to encode.
 1810 * @param [out] s     Buffer to encode into.
 1811 */
 1812static void dilithium_encode_gamma1_19_bits_c(const sword32* z, byte* s)
 1813{
 1814    unsigned int j;
 1815
 1816    /* Step 3. Get 20 bits as a number. */
 1817    for (j = 0; j < DILITHIUM_N; j += 4) {
 1818        sword32 z0 = DILITHIUM_GAMMA1_19 - z[j + 0];
 1819        sword32 z1 = DILITHIUM_GAMMA1_19 - z[j + 1];
 1820        sword32 z2 = DILITHIUM_GAMMA1_19 - z[j + 2];
 1821        sword32 z3 = DILITHIUM_GAMMA1_19 - z[j + 3];
 1822
 1823        /* 20 bits per number.
 1824         * 4 numbers become 10 bytes. (4 * 20 bits = 10 * 8 bits) */
 1825#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 2)
 1826        word16* s16p = (word16*)s;
 1827    #ifdef WC_64BIT_CPU
 1828        word64* s64p = (word64*)s;
 1829        s64p[0] =   (word64)z0        | ((word64)z1 << 20) |
 1830                   ((word64)z2 << 40) | ((word64)z3 << 60);
 1831    #else
 1832        word32* s32p = (word32*)s;
 1833        s32p[0] = (word32)( z0        | (z1 << 20)             );
 1834        s32p[1] = (word32)((z1 >> 12) | (z2 <<  8) | (z3 << 28));
 1835    #endif
 1836        s16p[4] = (word16)((z3 >>  4)                          );
 1837#else
 1838        s[0] = (byte)  z0                    ;
 1839        s[1] = (byte) (z0 >>  8)             ;
 1840        s[2] = (byte)((z0 >> 16) | (z1 << 4));
 1841        s[3] = (byte) (z1 >>  4)             ;
 1842        s[4] = (byte) (z1 >> 12)             ;
 1843        s[5] = (byte)  z2                    ;
 1844        s[6] = (byte) (z2 >>  8)             ;
 1845        s[7] = (byte)((z2 >> 16) | (z3 << 4));
 1846        s[8] = (byte) (z3 >>  4)             ;
 1847        s[9] = (byte) (z3 >> 12)             ;
 1848#endif
 1849        /* Move to next place to encode to. */
 1850        s += DILITHIUM_GAMMA1_19_ENC_BITS / 2;
 1851    }
 1852}
 1853
 1854/* Encode z with range of -(GAMMA1-1)...GAMMA1
 1855 *
 1856 * @param [in]  z     Polynomial to encode.
 1857 * @param [out] s     Buffer to encode into.
 1858 */
 1859static void dilithium_encode_gamma1_19_bits(const sword32* z, byte* s)
 1860{
 1861#ifdef USE_INTEL_SPEEDUP
 1862    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 1863        wc_mldsa_encode_gamma1_19_avx2(z, s);
 1864        RESTORE_VECTOR_REGISTERS();
 1865    }
 1866    else
 1867#endif
 1868    {
 1869        dilithium_encode_gamma1_19_bits_c(z, s);
 1870    }
 1871}
 1872#endif
 1873
 1874#ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM
 1875/* Encode z with range of -(GAMMA1-1)...GAMMA1
 1876 *
 1877 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 1878 *   ...
 1879 *   2: for i form 0 to l - 1 do
 1880 *   3:     sigma <- sigma || BitPack(z[i], GAMMA1 - 1, GAMMA1)
 1881 *   4: end for
 1882 *   ...
 1883 *
 1884 * @param [in]  z     Vector of polynomials to encode.
 1885 * @param [in]  l     Dimension of vector.
 1886 * @param [in]  bits  Number of bits used in encoding - GAMMA1 bits.
 1887 * @param [out] s     Buffer to encode into.
 1888 */
 1889static void dilithium_vec_encode_gamma1(const sword32* z, byte l, int bits,
 1890    byte* s)
 1891{
 1892    unsigned int i;
 1893
 1894    (void)l;
 1895
 1896#ifndef WOLFSSL_NO_ML_DSA_44
 1897    if (bits == DILITHIUM_GAMMA1_BITS_17) {
 1898        /* Step 2. For each polynomial of vector. */
 1899        for (i = 0; i < PARAMS_ML_DSA_44_L; i++) {
 1900            dilithium_encode_gamma1_17_bits(z, s);
 1901            /* Move to next place to encode to. */
 1902            s += DILITHIUM_GAMMA1_17_ENC_BITS / 2 * DILITHIUM_N / 4;
 1903            /* Next polynomial. */
 1904            z += DILITHIUM_N;
 1905        }
 1906    }
 1907#endif
 1908#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 1909    if (bits == DILITHIUM_GAMMA1_BITS_19) {
 1910        /* Step 2. For each polynomial of vector. */
 1911        for (i = 0; i < l; i++) {
 1912            dilithium_encode_gamma1_19_bits(z, s);
 1913            /* Move to next place to encode to. */
 1914            s += DILITHIUM_GAMMA1_19_ENC_BITS / 2 * DILITHIUM_N / 4;
 1915            /* Next polynomial. */
 1916            z += DILITHIUM_N;
 1917        }
 1918    }
 1919#endif
 1920}
 1921#endif /* WOLFSSL_DILITHIUM_SIGN_SMALL_MEM */
 1922
 1923#endif /* !WOLFSSL_DILITHIUM_NO_SIGN */
 1924
 1925#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
 1926/* Decode polynomial with range -(GAMMA1-1)..GAMMA1.
 1927 *
 1928 * FIPS 204. 8.2: Algorithm 21 sigDecode(sigma)
 1929 *   ...
 1930 *   4:     z[i] <- BitUnpack(xi, GAMMA1 - 1, GAMMA1)
 1931 *   ...
 1932 *
 1933 * @param [in]  s     Encoded values of z.
 1934 * @param [in]  bits  Number of bits used in encoding - GAMMA1 bits.
 1935 * @param [out] z     Polynomial to fill.
 1936 */
 1937static void dilithium_decode_gamma1_c(const byte* s, int bits, sword32* z)
 1938{
 1939    unsigned int i;
 1940
 1941#ifndef WOLFSSL_NO_ML_DSA_44
 1942    if (bits == DILITHIUM_GAMMA1_BITS_17) {
 1943#if defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE) || defined(WOLFSSL_DILITHIUM_SMALL)
 1944        /* Step 4: Get 18 bits as a number. */
 1945        for (i = 0; i < DILITHIUM_N; i += 4) {
 1946            /* 18 bits per number.
 1947             * 4 numbers from 9 bytes. (4 * 18 bits = 9 * 8 bits) */
 1948    #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1949        #ifdef WC_64BIT_CPU
 1950            word64 s64_0 = *(const word64*)(s+0);
 1951            z[i+0] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1952                               ( s64_0        & 0x3ffff                   ));
 1953            z[i+1] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1954                               ((s64_0 >> 18) & 0x3ffff                   ));
 1955            z[i+2] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1956                               ((s64_0 >> 36) & 0x3ffff                   ));
 1957            z[i+3] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1958                               ((s64_0 >> 54) | (((word32)s[8])     << 10)));
 1959        #else
 1960            word32 s32_0 = ((const word32*)(s+0))[0];
 1961            word32 s32_1 = ((const word32*)(s+0))[1];
 1962            z[i+0] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1963                               ( s32_0        & 0x3ffff                    ));
 1964            z[i+1] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1965                               ((s32_0 >> 18) | (((s32_1 & 0x0000f) << 14))));
 1966            z[i+2] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1967                               ((s32_1 >>  4) & 0x3ffff                    ));
 1968            z[i+3] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1969                               ((s32_1 >> 22) | (((word32)s[8])     << 10 )));
 1970        #endif
 1971    #else
 1972            z[i+0] = DILITHIUM_GAMMA1_17 -
 1973                     ( (sword32)s[ 0]       | (((sword32)s[ 1] << 8) |
 1974                      (sword32)(s[ 2] & 0x03) << 16));
 1975            z[i+1] = DILITHIUM_GAMMA1_17 -
 1976                     (((sword32)s[ 2] >> 2) | (((sword32)s[ 3] << 6) |
 1977                      (sword32)(s[ 4] & 0x0f) << 14));
 1978            z[i+2] = DILITHIUM_GAMMA1_17 -
 1979                     (((sword32)s[ 4] >> 4) | (((sword32)s[ 5] << 4) |
 1980                      (sword32)(s[ 6] & 0x3f) << 12));
 1981            z[i+3] = DILITHIUM_GAMMA1_17 -
 1982                     (((sword32)s[ 6] >> 6) | (((sword32)s[ 7] << 2) |
 1983                      ((sword32)s[ 8]       ) << 10));
 1984    #endif
 1985            /* Move to next place to decode from. */
 1986            s += DILITHIUM_GAMMA1_17_ENC_BITS / 2;
 1987        }
 1988#else
 1989        /* Step 4: Get 18 bits as a number. */
 1990        for (i = 0; i < DILITHIUM_N; i += 8) {
 1991            /* 18 bits per number.
 1992             * 8 numbers from 9 bytes. (8 * 18 bits = 18 * 8 bits) */
 1993    #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 1994        #ifdef WC_64BIT_CPU
 1995            word64 s64_0 = *(const word64*)(s+0);
 1996            word64 s64_1 = *(const word64*)(s+9);
 1997            z[i+0] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 1998                               ( s64_0        & 0x3ffff                   ));
 1999            z[i+1] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2000                               ((s64_0 >> 18) & 0x3ffff                   ));
 2001            z[i+2] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2002                               ((s64_0 >> 36) & 0x3ffff                   ));
 2003            z[i+3] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2004                               ((s64_0 >> 54) | (((word32)s[8])     << 10)));
 2005            z[i+4] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2006                               ( s64_1        & 0x3ffff                   ));
 2007            z[i+5] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2008                               ((s64_1 >> 18) & 0x3ffff                   ));
 2009            z[i+6] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2010                               ((s64_1 >> 36) & 0x3ffff                   ));
 2011            z[i+7] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2012                               ((s64_1 >> 54) | (((word32)s[17])    << 10)));
 2013        #else
 2014            word32 s32_0 = ((const word32*)(s+0))[0];
 2015            word32 s32_1 = ((const word32*)(s+0))[1];
 2016            word32 s32_2 = ((const word32*)(s+9))[0];
 2017            word32 s32_3 = ((const word32*)(s+9))[1];
 2018            z[i+0] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2019                               ( s32_0        & 0x3ffff                    ));
 2020            z[i+1] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2021                               ((s32_0 >> 18) | (((s32_1 & 0x0000f) << 14))));
 2022            z[i+2] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2023                               ((s32_1 >>  4) & 0x3ffff                    ));
 2024            z[i+3] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2025                               ((s32_1 >> 22) | (((word32)s[8])     << 10 )));
 2026            z[i+4] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2027                               ( s32_2        & 0x3ffff                    ));
 2028            z[i+5] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2029                               ((s32_2 >> 18) | (((s32_3 & 0x0000f) << 14))));
 2030            z[i+6] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2031                               ((s32_3 >>  4) & 0x3ffff                    ));
 2032            z[i+7] = (sword32)((word32)DILITHIUM_GAMMA1_17 -
 2033                               ((s32_3 >> 22) | (((word32)s[17])    << 10 )));
 2034        #endif
 2035    #else
 2036            z[i+0] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2037                               ( (sword32)s[ 0]       | (((sword32)s[ 1] << 8) |
 2038                                ((sword32)s[ 2] & 0x03) << 16)));
 2039            z[i+1] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2040                               (((sword32)s[ 2] >> 2) | (((sword32)s[ 3] << 6) |
 2041                                ((sword32)s[ 4] & 0x0f) << 14)));
 2042            z[i+2] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2043                               (((sword32)s[ 4] >> 4) | (((sword32)s[ 5] << 4) |
 2044                                ((sword32)s[ 6] & 0x3f) << 12)));
 2045            z[i+3] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2046                               (((sword32)s[ 6] >> 6) | (((sword32)s[ 7] << 2) |
 2047                                ((sword32)s[ 8]       ) << 10)));
 2048            z[i+4] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2049                               ( (sword32)s[ 9]       | (((sword32)s[10] << 8) |
 2050                                ((sword32)s[11] & 0x03) << 16)));
 2051            z[i+5] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2052                               (((sword32)s[11] >> 2) | (((sword32)s[12] << 6) |
 2053                                ((sword32)s[13] & 0x0f) << 14)));
 2054            z[i+6] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2055                               (((sword32)s[13] >> 4) | (((sword32)s[14] << 4) |
 2056                                ((sword32)s[15] & 0x3f) << 12)));
 2057            z[i+7] = (sword32)((word32)DILITHIUM_GAMMA1_17 - (word32)
 2058                               (((sword32)s[15] >> 6) | (((sword32)s[16] << 2) |
 2059                                ((sword32)s[17]       ) << 10)));
 2060    #endif
 2061            /* Move to next place to decode from. */
 2062            s += DILITHIUM_GAMMA1_17_ENC_BITS;
 2063        }
 2064#endif
 2065    }
 2066#endif
 2067#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 2068    if (bits == DILITHIUM_GAMMA1_BITS_19) {
 2069#if defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE) || defined(WOLFSSL_DILITHIUM_SMALL)
 2070        /* Step 4: Get 20 bits as a number. */
 2071        for (i = 0; i < DILITHIUM_N; i += 4) {
 2072            /* 20 bits per number.
 2073             * 4 numbers from 10 bytes. (4 * 20 bits = 10 * 8 bits) */
 2074    #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 2)
 2075            word16 s16_0 = ((const word16*)s)[4];
 2076        #ifdef WC_64BIT_CPU
 2077            word64 s64_0 = *(const word64*)s;
 2078            z[i+0] = DILITHIUM_GAMMA1_19 - ((sword32)( s64_0        & 0xfffff));
 2079            z[i+1] = DILITHIUM_GAMMA1_19 - ((sword32)((s64_0 >> 20) & 0xfffff));
 2080            z[i+2] = DILITHIUM_GAMMA1_19 - ((sword32)((s64_0 >> 40) & 0xfffff));
 2081            z[i+3] = DILITHIUM_GAMMA1_19 - ((sword32)((s64_0 >> 60) & 0xfffff) |
 2082                                            ((sword32) s16_0 <<  4));
 2083        #else
 2084            word32 s32_0 = ((const word32*)s)[0];
 2085            word32 s32_1 = ((const word32*)s)[1];
 2086            z[i+0] = DILITHIUM_GAMMA1_19 - (sword32)(  s32_0       & 0xfffff);
 2087            z[i+1] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_0            >> 20) |
 2088                                                     ((s32_1 & 0x000ff) << 12));
 2089            z[i+2] = DILITHIUM_GAMMA1_19 - (sword32)( (s32_1 >>  8) & 0xfffff);
 2090            z[i+3] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_1            >> 28) |
 2091                                                     ((word32)s16_0     <<  4));
 2092        #endif
 2093    #else
 2094            z[i+0] = DILITHIUM_GAMMA1_19 -
 2095                     ( (sword32)s[0]       | ((sword32)s[1] << 8) |
 2096                     (((sword32)s[2] & 0x0f) << 16));
 2097            z[i+1] = DILITHIUM_GAMMA1_19 -
 2098                     (((sword32)s[2] >> 4) | ((sword32)s[3] << 4) |
 2099                     (((sword32)s[4]       ) << 12));
 2100            z[i+2] = DILITHIUM_GAMMA1_19 -
 2101                     ( (sword32)s[5]       | ((sword32)s[6] << 8) |
 2102                     (((sword32)s[7] & 0x0f) << 16));
 2103            z[i+3] = DILITHIUM_GAMMA1_19 -
 2104                     (((sword32)s[7] >> 4) | ((sword32)s[8] << 4) |
 2105                     (((sword32)s[9]       ) << 12));
 2106    #endif
 2107            /* Move to next place to decode from. */
 2108            s += DILITHIUM_GAMMA1_19_ENC_BITS / 2;
 2109        }
 2110#else
 2111        /* Step 4: Get 20 bits as a number. */
 2112        for (i = 0; i < DILITHIUM_N; i += 8) {
 2113            /* 20 bits per number.
 2114             * 8 numbers from 20 bytes. (8 * 20 bits = 20 * 8 bits) */
 2115    #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 2)
 2116            word16 s16_0 = ((const word16*)s)[4];
 2117            word16 s16_1 = ((const word16*)s)[9];
 2118        #ifdef WC_64BIT_CPU
 2119            word64 s64_0 = *(const word64*)(s+0);
 2120            word64 s64_1 = *(const word64*)(s+10);
 2121            z[i+0] = DILITHIUM_GAMMA1_19 -
 2122                     ((sword32)(  s64_0        & 0xfffff))   ;
 2123            z[i+1] = DILITHIUM_GAMMA1_19 -
 2124                     ((sword32)( (s64_0 >> 20) & 0xfffff))   ;
 2125            z[i+2] = DILITHIUM_GAMMA1_19 -
 2126                     ((sword32)( (s64_0 >> 40) & 0xfffff))   ;
 2127            z[i+3] = DILITHIUM_GAMMA1_19 -
 2128                     ((sword32)(((s64_0 >> 60) & 0xfffff))   |
 2129                      ((sword32)s16_0      <<  4));
 2130            z[i+4] = DILITHIUM_GAMMA1_19 -
 2131                     ((sword32)(  s64_1        & 0xfffff))   ;
 2132            z[i+5] = DILITHIUM_GAMMA1_19 -
 2133                     ((sword32)( (s64_1 >> 20) & 0xfffff))   ;
 2134            z[i+6] = DILITHIUM_GAMMA1_19 -
 2135                     ((sword32)( (s64_1 >> 40) & 0xfffff))   ;
 2136            z[i+7] = DILITHIUM_GAMMA1_19 -
 2137                     ((sword32)(((s64_1 >> 60) & 0xfffff))   |
 2138                      ((sword32)s16_1      <<  4));
 2139        #else
 2140            word32 s32_0 = ((const word32*)(s+ 0))[0];
 2141            word32 s32_1 = ((const word32*)(s+ 0))[1];
 2142            word32 s32_2 = ((const word32*)(s+10))[0];
 2143            word32 s32_3 = ((const word32*)(s+10))[1];
 2144            z[i+0] = DILITHIUM_GAMMA1_19 - (sword32)(  s32_0       & 0xfffff);
 2145            z[i+1] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_0            >> 20) |
 2146                                                     ((s32_1 & 0x000ff) << 12));
 2147            z[i+2] = DILITHIUM_GAMMA1_19 - (sword32)( (s32_1 >>  8) & 0xfffff);
 2148            z[i+3] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_1            >> 28) |
 2149                                                     ((word32)s16_0     <<  4));
 2150            z[i+4] = DILITHIUM_GAMMA1_19 - (sword32)(  s32_2       & 0xfffff);
 2151            z[i+5] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_2            >> 20) |
 2152                                                     ((s32_3 & 0x000ff) << 12));
 2153            z[i+6] = DILITHIUM_GAMMA1_19 - (sword32)( (s32_3 >>  8) & 0xfffff);
 2154            z[i+7] = DILITHIUM_GAMMA1_19 - (sword32)(( s32_3            >> 28) |
 2155                                                     ((word32)s16_1     <<  4));
 2156        #endif
 2157    #else
 2158            z[i+0] = DILITHIUM_GAMMA1_19 -
 2159                     (  (sword32)s[ 0]       |
 2160                      ( (sword32)s[ 1] << 8) |
 2161                      (((sword32)s[ 2] & 0x0f) << 16));
 2162            z[i+1] = DILITHIUM_GAMMA1_19 -
 2163                     (((sword32)s[ 2] >> 4) |
 2164                      ( (sword32)s[ 3] << 4) |
 2165                      (((sword32)s[ 4]       ) << 12));
 2166            z[i+2] = DILITHIUM_GAMMA1_19 -
 2167                     (  (sword32)s[ 5]       |
 2168                      ( (sword32)s[ 6] << 8) |
 2169                      (((sword32)s[ 7] & 0x0f) << 16));
 2170            z[i+3] = DILITHIUM_GAMMA1_19 -
 2171                     ( ((sword32)s[ 7] >> 4) |
 2172                      ( (sword32)s[ 8] << 4) |
 2173                      (((sword32)s[ 9]       ) << 12));
 2174            z[i+4] = DILITHIUM_GAMMA1_19 -
 2175                     (  (sword32)s[10]       |
 2176                      ( (sword32)s[11] << 8) |
 2177                      (((sword32)s[12] & 0x0f) << 16));
 2178            z[i+5] = DILITHIUM_GAMMA1_19 -
 2179                     ( ((sword32)s[12] >> 4) |
 2180                      ( (sword32)s[13] << 4) |
 2181                      (((sword32)s[14]       ) << 12));
 2182            z[i+6] = DILITHIUM_GAMMA1_19 -
 2183                     (  (sword32)s[15]       |
 2184                      ( (sword32)s[16] << 8) |
 2185                      (((sword32)s[17] & 0x0f) << 16));
 2186            z[i+7] = DILITHIUM_GAMMA1_19 -
 2187                     ( ((sword32)s[17] >> 4) |
 2188                      ( (sword32)s[18] << 4) |
 2189                      (((sword32)s[19]       ) << 12));
 2190    #endif
 2191            /* Move to next place to decode from. */
 2192            s += DILITHIUM_GAMMA1_19_ENC_BITS;
 2193        }
 2194#endif
 2195    }
 2196#endif
 2197}
 2198
 2199/* Decode polynomial with range -(GAMMA1-1)..GAMMA1.
 2200 *
 2201 * @param [in]  s     Encoded values of z.
 2202 * @param [in]  bits  Number of bits used in encoding - GAMMA1 bits.
 2203 * @param [out] z     Polynomial to fill.
 2204 */
 2205static void dilithium_decode_gamma1(const byte* s, int bits, sword32* z)
 2206{
 2207#ifdef USE_INTEL_SPEEDUP
 2208    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 2209        if (bits == DILITHIUM_GAMMA1_BITS_17) {
 2210            wc_mldsa_decode_gamma1_17_avx2(s, z);
 2211        }
 2212        else {
 2213            wc_mldsa_decode_gamma1_19_avx2(s, z);
 2214        }
 2215        RESTORE_VECTOR_REGISTERS();
 2216    }
 2217    else
 2218#endif
 2219    {
 2220        dilithium_decode_gamma1_c(s, bits, z);
 2221    }
 2222}
 2223#endif
 2224
 2225#ifndef WOLFSSL_DILITHIUM_NO_VERIFY
 2226/* Decode polynomial with range -(GAMMA1-1)..GAMMA1.
 2227 *
 2228 * FIPS 204. 8.2: Algorithm 21 sigDecode(sigma)
 2229 *   ...
 2230 *   3: for i from 0 to l - 1 do
 2231 *   4:     z[i] <- BitUnpack(xi, GAMMA1 - 1, GAMMA1)
 2232 *   5: end for
 2233 *   ...
 2234 *
 2235 * @param [in]  x  Encoded values of t0.
 2236 * @param [in]  l  Dimensions of vector z.
 2237 * @param [in]  bits  Number of bits used in encoding - GAMMA1 bits.
 2238 * @param [out] z  Vector of polynomials.
 2239 */
 2240static void dilithium_vec_decode_gamma1(const byte* x, byte l, int bits,
 2241    sword32* z)
 2242{
 2243    unsigned int i;
 2244
 2245    /* Step 3: For each polynomial of vector. */
 2246    for (i = 0; i < l; i++) {
 2247        /* Step 4: Unpack a polynomial. */
 2248        dilithium_decode_gamma1(x, bits, z);
 2249        /* Move pointers on to next polynomial. */
 2250        x += DILITHIUM_N / 8 * (bits + 1);
 2251        z += DILITHIUM_N;
 2252    }
 2253}
 2254#endif
 2255
 2256#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
 2257#ifndef WOLFSSL_NO_ML_DSA_44
 2258/* Encode w1 with range of 0..((q-1)/(2*GAMMA2)-1).
 2259 *
 2260 * FIPS 204. 8.2: Algorithm 22 w1Encode(w1)
 2261 *   ...
 2262 *   3:     w1_tilde <- w1_tilde ||
 2263 *                      ByteToBits(SimpleBitPack(w1[i], (q-1)/(2*GAMMA2)-1))
 2264 *   ...
 2265 *
 2266 * @param [in]  w1      Vector of polynomials to encode.
 2267 * @param [in]  gamma2  Maximum value in range.
 2268 * @param [out] w1e     Buffer to encode into.
 2269 */
 2270static void dilithium_encode_w1_88_c(const sword32* w1, byte* w1e)
 2271{
 2272    unsigned int j;
 2273
 2274    /* Step 3: Encode a polynomial values 6 bits at a time. */
 2275    for (j = 0; j < DILITHIUM_N; j += 16) {
 2276        /* 6 bits per number.
 2277         * 16 numbers in 12 bytes. (16 * 6 bits = 12 * 8 bits) */
 2278#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 4)
 2279        word32* w1e32 = (word32*)w1e;
 2280        w1e32[0] = (word32)( (word32)w1[j+ 0]        |
 2281                            ((word32)w1[j+ 1] <<  6) |
 2282                            ((word32)w1[j+ 2] << 12) |
 2283                            ((word32)w1[j+ 3] << 18) |
 2284                            ((word32)w1[j+ 4] << 24) |
 2285                            ((word32)w1[j+ 5] << 30));
 2286        w1e32[1] = (word32)(((word32)w1[j+ 5] >>  2) |
 2287                            ((word32)w1[j+ 6] <<  4) |
 2288                            ((word32)w1[j+ 7] << 10) |
 2289                            ((word32)w1[j+ 8] << 16) |
 2290                            ((word32)w1[j+ 9] << 22) |
 2291                            ((word32)w1[j+10] << 28));
 2292        w1e32[2] = (word32)(((word32)w1[j+10] >>  4) |
 2293                            ((word32)w1[j+11] <<  2) |
 2294                            ((word32)w1[j+12] <<  8) |
 2295                            ((word32)w1[j+13] << 14) |
 2296                            ((word32)w1[j+14] << 20) |
 2297                            ((word32)w1[j+15] << 26));
 2298#else
 2299        w1e[ 0] = (byte)( w1[j+ 0]       | (w1[j+ 1] << 6));
 2300        w1e[ 1] = (byte)((w1[j+ 1] >> 2) | (w1[j+ 2] << 4));
 2301        w1e[ 2] = (byte)((w1[j+ 2] >> 4) | (w1[j+ 3] << 2));
 2302        w1e[ 3] = (byte)( w1[j+ 4]       | (w1[j+ 5] << 6));
 2303        w1e[ 4] = (byte)((w1[j+ 5] >> 2) | (w1[j+ 6] << 4));
 2304        w1e[ 5] = (byte)((w1[j+ 6] >> 4) | (w1[j+ 7] << 2));
 2305        w1e[ 6] = (byte)( w1[j+ 8]       | (w1[j+ 9] << 6));
 2306        w1e[ 7] = (byte)((w1[j+ 9] >> 2) | (w1[j+10] << 4));
 2307        w1e[ 8] = (byte)((w1[j+10] >> 4) | (w1[j+11] << 2));
 2308        w1e[ 9] = (byte)( w1[j+12]       | (w1[j+13] << 6));
 2309        w1e[10] = (byte)((w1[j+13] >> 2) | (w1[j+14] << 4));
 2310        w1e[11] = (byte)((w1[j+14] >> 4) | (w1[j+15] << 2));
 2311#endif
 2312        /* Move to next place to encode to. */
 2313        w1e += DILITHIUM_Q_HI_88_ENC_BITS * 2;
 2314    }
 2315}
 2316
 2317/* Encode w1 with range of 0..((q-1)/(2*GAMMA2)-1).
 2318 *
 2319 * @param [in]  w1      Vector of polynomials to encode.
 2320 * @param [in]  gamma2  Maximum value in range.
 2321 * @param [out] w1e     Buffer to encode into.
 2322 */
 2323static void dilithium_encode_w1_88(const sword32* w1, byte* w1e)
 2324{
 2325#ifdef USE_INTEL_SPEEDUP
 2326    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 2327        wc_mldsa_encode_w1_88_avx2(w1, w1e);
 2328        RESTORE_VECTOR_REGISTERS();
 2329    }
 2330    else
 2331#endif
 2332    {
 2333        dilithium_encode_w1_88_c(w1, w1e);
 2334    }
 2335}
 2336
 2337WOLFSSL_TEST_VIS void wc_dilithium_encode_w1_88(const sword32* w1, byte* w1e)
 2338{
 2339    dilithium_encode_w1_88(w1, w1e);
 2340}
 2341#endif /* !WOLFSSL_NO_ML_DSA_44 */
 2342
 2343#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 2344/* Encode w1 with range of 0..((q-1)/(2*GAMMA2)-1).
 2345 *
 2346 * FIPS 204. 8.2: Algorithm 22 w1Encode(w1)
 2347 *   ...
 2348 *   3:     w1_tilde <- w1_tilde ||
 2349 *                      ByteToBits(SimpleBitPack(w1[i], (q-1)/(2*GAMMA2)-1))
 2350 *   ...
 2351 *
 2352 * @param [in]  w1      Vector of polynomials to encode.
 2353 * @param [in]  gamma2  Maximum value in range.
 2354 * @param [out] w1e     Buffer to encode into.
 2355 */
 2356static void dilithium_encode_w1_32_c(const sword32* w1, byte* w1e)
 2357{
 2358    unsigned int j;
 2359
 2360    /* Step 3: Encode a polynomial values 4 bits at a time. */
 2361    for (j = 0; j < DILITHIUM_N; j += 16) {
 2362        /* 4 bits per number.
 2363         * 16 numbers in 8 bytes. (16 * 4 bits = 8 * 8 bits) */
 2364#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT <= 8)
 2365        word32* w1e32 = (word32*)w1e;
 2366        w1e32[0] = (word32)(((word32)w1[j +  0] <<  0) |
 2367                            ((word32)w1[j +  1] <<  4) |
 2368                            ((word32)w1[j +  2] <<  8) |
 2369                            ((word32)w1[j +  3] << 12) |
 2370                            ((word32)w1[j +  4] << 16) |
 2371                            ((word32)w1[j +  5] << 20) |
 2372                            ((word32)w1[j +  6] << 24) |
 2373                            ((word32)w1[j +  7] << 28));
 2374        w1e32[1] = (word32)(((word32)w1[j +  8] <<  0) |
 2375                            ((word32)w1[j +  9] <<  4) |
 2376                            ((word32)w1[j + 10] <<  8) |
 2377                            ((word32)w1[j + 11] << 12) |
 2378                            ((word32)w1[j + 12] << 16) |
 2379                            ((word32)w1[j + 13] << 20) |
 2380                            ((word32)w1[j + 14] << 24) |
 2381                            ((word32)w1[j + 15] << 28));
 2382#else
 2383        w1e[0] = (byte)(w1[j +  0] | (w1[j +  1] << 4));
 2384        w1e[1] = (byte)(w1[j +  2] | (w1[j +  3] << 4));
 2385        w1e[2] = (byte)(w1[j +  4] | (w1[j +  5] << 4));
 2386        w1e[3] = (byte)(w1[j +  6] | (w1[j +  7] << 4));
 2387        w1e[4] = (byte)(w1[j +  8] | (w1[j +  9] << 4));
 2388        w1e[5] = (byte)(w1[j + 10] | (w1[j + 11] << 4));
 2389        w1e[6] = (byte)(w1[j + 12] | (w1[j + 13] << 4));
 2390        w1e[7] = (byte)(w1[j + 14] | (w1[j + 15] << 4));
 2391#endif
 2392        /* Move to next place to encode to. */
 2393        w1e += DILITHIUM_Q_HI_32_ENC_BITS * 2;
 2394    }
 2395}
 2396
 2397/* Encode w1 with range of 0..((q-1)/(2*GAMMA2)-1).
 2398 *
 2399 * @param [in]  w1      Vector of polynomials to encode.
 2400 * @param [in]  gamma2  Maximum value in range.
 2401 * @param [out] w1e     Buffer to encode into.
 2402 */
 2403static void dilithium_encode_w1_32(const sword32* w1, byte* w1e)
 2404{
 2405#ifdef USE_INTEL_SPEEDUP
 2406    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 2407        wc_mldsa_encode_w1_32_avx2(w1, w1e);
 2408        RESTORE_VECTOR_REGISTERS();
 2409    }
 2410    else
 2411#endif
 2412    {
 2413        dilithium_encode_w1_32_c(w1, w1e);
 2414    }
 2415}
 2416
 2417WOLFSSL_TEST_VIS void wc_dilithium_encode_w1_32(const sword32* w1, byte* w1e)
 2418{
 2419    dilithium_encode_w1_32(w1, w1e);
 2420}
 2421#endif
 2422#endif
 2423
 2424#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 2425     (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 2426      !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))
 2427/* Encode w1 with range of 0..((q-1)/(2*GAMMA2)-1).
 2428 *
 2429 * FIPS 204. 8.2: Algorithm 22 w1Encode(w1)
 2430 *   1: w1_tilde = ()
 2431 *   2: for i form 0 to k - 1 do
 2432 *   3:     w1_tilde <- w1_tilde ||
 2433 *                      ByteToBits(SimpleBitPack(w1[i], (q-1)/(2*GAMMA2)-1))
 2434 *   4: end for
 2435 *   5: return w1_tilde
 2436 *
 2437 * @param [in]  w1      Vector of polynomials to encode.
 2438 * @param [in]  k       Dimension of vector.
 2439 * @param [in]  gamma2  Maximum value in range.
 2440 * @param [out] w1e     Buffer to encode into.
 2441 */
 2442static void dilithium_vec_encode_w1(const sword32* w1, byte k, sword32 gamma2,
 2443    byte* w1e)
 2444{
 2445    unsigned int i;
 2446
 2447    (void)k;
 2448
 2449#ifndef WOLFSSL_NO_ML_DSA_44
 2450    if (gamma2 == DILITHIUM_Q_LOW_88) {
 2451        /* Step 2. For each polynomial of vector. */
 2452        for (i = 0; i < PARAMS_ML_DSA_44_K; i++) {
 2453            dilithium_encode_w1_88(w1, w1e);
 2454            /* Next polynomial. */
 2455            w1 += DILITHIUM_N;
 2456            w1e += DILITHIUM_Q_HI_88_ENC_BITS * 2 * DILITHIUM_N / 16;
 2457        }
 2458    }
 2459    else
 2460#endif
 2461#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 2462    if (gamma2 == DILITHIUM_Q_LOW_32) {
 2463        /* Step 2. For each polynomial of vector. */
 2464        for (i = 0; i < k; i++) {
 2465            dilithium_encode_w1_32(w1, w1e);
 2466            /* Next polynomial. */
 2467            w1 += DILITHIUM_N;
 2468            w1e += DILITHIUM_Q_HI_32_ENC_BITS * 2 * DILITHIUM_N / 16;
 2469        }
 2470    }
 2471    else
 2472#endif
 2473    {
 2474    }
 2475}
 2476#endif
 2477
 2478/******************************************************************************
 2479 * Expand operations
 2480 ******************************************************************************/
 2481
 2482/* Generate a random polynomial by rejection.
 2483 *
 2484 * FIPS 204. 8.3: Algorithm 24 RejNTTPoly(rho)
 2485 *   1: j <- 0
 2486 *   2: c <- 0
 2487 *   3: while j < 256 do
 2488 *   4:    a_hat[j] <- CoeffFromThreeBytes(H128(rho)[[c]], H128(rho)[[c+1]],
 2489 *                                         H128(rho)[[c+2]])
 2490 *   5:    c <- c + 3
 2491 *   6:    if a_hat[j] != falsam then
 2492 *   7:       j <- j + 1
 2493 *   8:    end if
 2494 *   9: end while
 2495 *  10: return a_hat
 2496 *
 2497 * FIPS 204. 8.1: Algorithm 8 CoeffFromThreeBytes(b0,b1,b2)
 2498 *   1: if b2 > 127 then
 2499 *   2:    b2 <- b2 - 128
 2500 *   3. end if
 2501 *   4. z <- 2^16.b2 + s^8.b1 + b0
 2502 *   5. if z < q then return z
 2503 *   6. else return falsam
 2504 *   7. end if
 2505 *
 2506 * @param [in, out] shake128  SHAKE-128 object.
 2507 * @param [in]      seed      Seed to hash to generate values.
 2508 * @param [out]     a         Polynomial.
 2509 * @param [in]      h         Buffer to hold hashes.
 2510 * @return  0 on success.
 2511 * @return  Negative on hash error.
 2512 */
 2513static int dilithium_rej_ntt_poly_ex(wc_Shake* shake128, byte* seed, sword32* a,
 2514    byte* h)
 2515{
 2516    int ret = 0;
 2517#ifdef WOLFSSL_DILITHIUM_SMALL
 2518    int j = 0;
 2519
 2520#if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2521    /* Reading 4 bytes for 3 so need to set 1 past for last read. */
 2522    h[DILITHIUM_GEN_A_BLOCK_BYTES] = 0;
 2523#endif
 2524
 2525    /* Initialize SHAKE-128 object for new hash. */
 2526    ret = wc_InitShake128(shake128, NULL, INVALID_DEVID);
 2527    if (ret == 0) {
 2528        /* Absorb the seed. */
 2529        ret = wc_Shake128_Absorb(shake128, seed, DILITHIUM_GEN_A_SEED_SZ);
 2530    }
 2531    /* Keep generating more blocks and using triplets until we have enough.
 2532     */
 2533    while ((ret == 0) && (j < DILITHIUM_N)) {
 2534         /* Squeeze out a block - 168 bytes = 56 values. */
 2535        ret = wc_Shake128_SqueezeBlocks(shake128, h, 1);
 2536        if (ret == 0) {
 2537            int c;
 2538            /* Use triplets until run out or have enough for polynomial. */
 2539            for (c = 0; c < DILITHIUM_GEN_A_BLOCK_BYTES; c += 3) {
 2540            #if defined(LITTLE_ENDIAN_ORDER) && \
 2541                (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2542                /* Load 32-bit value and mask out 23 bits. */
 2543                sword32 t = *((sword32*)(h + c)) & 0x7fffff;
 2544            #else
 2545                /* Load 24-bit value and mask out 23 bits. */
 2546                sword32 t = (h[c] + ((sword32)h[c+1] << 8) +
 2547                             ((sword32)h[c+2] << 16)) & 0x7fffff;
 2548            #endif
 2549                /* Check if value is in valid range. */
 2550                if (t < DILITHIUM_Q) {
 2551                    /* Store value in polynomial and increment count of values.
 2552                     */
 2553                    a[j++] = t;
 2554                    /* Check we whether we have enough yet. */
 2555                    if (j == DILITHIUM_N) {
 2556                        break;
 2557                    }
 2558                }
 2559            }
 2560        }
 2561    }
 2562#else
 2563    unsigned int j = 0;
 2564    unsigned int c;
 2565
 2566    /* Generate enough SHAKE-128 output blocks to give high probability of
 2567     * being able to get 256 valid 3-byte, 23-bit values from it. */
 2568    ret = dilithium_squeeze128(shake128, seed, DILITHIUM_GEN_A_SEED_SZ, h,
 2569        DILITHIUM_GEN_A_NBLOCKS);
 2570    if (ret == 0) {
 2571    #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2572        /* Reading 4 bytes for 3 so need to set 1 past for last read. */
 2573        h[DILITHIUM_GEN_A_BYTES] = 0;
 2574    #endif
 2575
 2576        /* Use the first 256 triplets and know we won't exceed required. */
 2577    #ifdef WOLFSSL_DILITHIUM_NO_LARGE_CODE
 2578        for (c = 0; c < (DILITHIUM_N - 1) * 3; c += 3) {
 2579        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2580            /* Load 32-bit value and mask out 23 bits. */
 2581            sword32 t = *((sword32*)(h + c)) & 0x7fffff;
 2582        #else
 2583            /* Load 24-bit value and mask out 23 bits. */
 2584            sword32 t = (h[c] + ((sword32)h[c+1] << 8) +
 2585                         ((sword32)h[c+2] << 16)) & 0x7fffff;
 2586        #endif
 2587            /* Check if value is in valid range. */
 2588            if (t < DILITHIUM_Q) {
 2589                /* Store value in polynomial and increment count of values. */
 2590                a[j++] = t;
 2591            }
 2592        }
 2593        /* Use the remaining triplets, checking we have enough. */
 2594        for (; c < DILITHIUM_GEN_A_BYTES; c += 3) {
 2595        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2596            /* Load 32-bit value and mask out 23 bits. */
 2597            sword32 t = *((sword32*)(h + c)) & 0x7fffff;
 2598        #else
 2599            /* Load 24-bit value and mask out 23 bits. */
 2600            sword32 t = (h[c] + ((sword32)h[c+1] << 8) +
 2601                         ((sword32)h[c+2] << 16)) & 0x7fffff;
 2602        #endif
 2603            /* Check if value is in valid range. */
 2604            if (t < DILITHIUM_Q) {
 2605                /* Store value in polynomial and increment count of values. */
 2606                a[j++] = t;
 2607                /* Check we whether we have enough yet. */
 2608                if (j == DILITHIUM_N) {
 2609                    break;
 2610                }
 2611            }
 2612        }
 2613    #else
 2614        /* Do 24 bytes at a time: 256 * 3 / 24 = 32 */
 2615        for (c = 0; c < DILITHIUM_N * 3; c += 24) {
 2616        #if defined(LITTLE_ENDIAN_ORDER) && (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2617            /* Load 32-bit value and mask out 23 bits. */
 2618            sword32 t0 = *((sword32*)(h + c +  0)) & 0x7fffff;
 2619            sword32 t1 = *((sword32*)(h + c +  3)) & 0x7fffff;
 2620            sword32 t2 = *((sword32*)(h + c +  6)) & 0x7fffff;
 2621            sword32 t3 = *((sword32*)(h + c +  9)) & 0x7fffff;
 2622            sword32 t4 = *((sword32*)(h + c + 12)) & 0x7fffff;
 2623            sword32 t5 = *((sword32*)(h + c + 15)) & 0x7fffff;
 2624            sword32 t6 = *((sword32*)(h + c + 18)) & 0x7fffff;
 2625            sword32 t7 = *((sword32*)(h + c + 21)) & 0x7fffff;
 2626        #else
 2627            /* Load 24-bit value and mask out 23 bits. */
 2628            sword32 t0 = (h[c +  0] + ((sword32)h[c +  1] << 8) +
 2629                          ((sword32)h[c +  2] << 16)) & 0x7fffff;
 2630            sword32 t1 = (h[c +  3] + ((sword32)h[c +  4] << 8) +
 2631                          ((sword32)h[c +  5] << 16)) & 0x7fffff;
 2632            sword32 t2 = (h[c +  6] + ((sword32)h[c +  7] << 8) +
 2633                          ((sword32)h[c +  8] << 16)) & 0x7fffff;
 2634            sword32 t3 = (h[c +  9] + ((sword32)h[c + 10] << 8) +
 2635                          ((sword32)h[c + 11] << 16)) & 0x7fffff;
 2636            sword32 t4 = (h[c + 12] + ((sword32)h[c + 13] << 8) +
 2637                          ((sword32)h[c + 14] << 16)) & 0x7fffff;
 2638            sword32 t5 = (h[c + 15] + ((sword32)h[c + 16] << 8) +
 2639                          ((sword32)h[c + 17] << 16)) & 0x7fffff;
 2640            sword32 t6 = (h[c + 18] + ((sword32)h[c + 19] << 8) +
 2641                          ((sword32)h[c + 20] << 16)) & 0x7fffff;
 2642            sword32 t7 = (h[c + 21] + ((sword32)h[c + 22] << 8) +
 2643                          ((sword32)h[c + 23] << 16)) & 0x7fffff;
 2644        #endif
 2645            a[j] = t0;
 2646            j += (t0 < DILITHIUM_Q);
 2647            a[j] = t1;
 2648            j += (t1 < DILITHIUM_Q);
 2649            a[j] = t2;
 2650            j += (t2 < DILITHIUM_Q);
 2651            a[j] = t3;
 2652            j += (t3 < DILITHIUM_Q);
 2653            a[j] = t4;
 2654            j += (t4 < DILITHIUM_Q);
 2655            a[j] = t5;
 2656            j += (t5 < DILITHIUM_Q);
 2657            a[j] = t6;
 2658            j += (t6 < DILITHIUM_Q);
 2659            a[j] = t7;
 2660            j += (t7 < DILITHIUM_Q);
 2661        }
 2662        if (j < DILITHIUM_N) {
 2663            /* Use the remaining triplets, checking we have enough. */
 2664            for (; c < DILITHIUM_GEN_A_BYTES; c += 3) {
 2665            #if defined(LITTLE_ENDIAN_ORDER) && \
 2666                (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2667                /* Load 32-bit value and mask out 23 bits. */
 2668                sword32 t = *((sword32*)(h + c)) & 0x7fffff;
 2669            #else
 2670                /* Load 24-bit value and mask out 23 bits. */
 2671                sword32 t = (h[c] + ((sword32)h[c+1] << 8) +
 2672                             ((sword32)h[c+2] << 16)) & 0x7fffff;
 2673            #endif
 2674                /* Check if value is in valid range. */
 2675                if (t < DILITHIUM_Q) {
 2676                    /* Store value in polynomial and increment count of values.
 2677                     */
 2678                    a[j++] = t;
 2679                    /* Check we whether we have enough yet. */
 2680                    if (j == DILITHIUM_N) {
 2681                        break;
 2682                    }
 2683                }
 2684            }
 2685        }
 2686    #endif
 2687        /* Keep generating more blocks and using triplets until we have enough.
 2688         */
 2689        while (j < DILITHIUM_N) {
 2690            /* Squeeze out a block - 168 bytes = 56 values. */
 2691            ret = wc_Shake128_SqueezeBlocks(shake128, h, 1);
 2692            if (ret != 0) {
 2693                break;
 2694            }
 2695            /* Use triplets until run out or have enough for polynomial. */
 2696            for (c = 0; c < DILITHIUM_GEN_A_BLOCK_BYTES; c += 3) {
 2697            #if defined(LITTLE_ENDIAN_ORDER) && \
 2698                (WOLFSSL_DILITHIUM_ALIGNMENT == 0)
 2699                /* Load 32-bit value and mask out 23 bits. */
 2700                sword32 t = *((sword32*)(h + c)) & 0x7fffff;
 2701            #else
 2702                /* Load 24-bit value and mask out 23 bits. */
 2703                sword32 t = (h[c] + ((sword32)h[c+1] << 8) +
 2704                             ((sword32)h[c+2] << 16)) & 0x7fffff;
 2705            #endif
 2706                /* Check if value is in valid range. */
 2707                if (t < DILITHIUM_Q) {
 2708                    /* Store value in polynomial and increment count of values.
 2709                     */
 2710                    a[j++] = t;
 2711                    /* Check we whether we have enough yet. */
 2712                    if (j == DILITHIUM_N) {
 2713                        break;
 2714                    }
 2715                }
 2716            }
 2717        }
 2718    }
 2719#endif
 2720
 2721    return ret;
 2722}
 2723
 2724#if (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) && \
 2725     !defined(WOLFSSL_DILITHIUM_MAKE_KEY_SMALL_MEM)) || \
 2726    defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 2727    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 2728     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)) || \
 2729    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 2730     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))
 2731/* Generate a random polynomial by rejection.
 2732 *
 2733 * @param [in, out] shake128  SHAKE-128 object.
 2734 * @param [in]      seed      Seed to hash to generate values.
 2735 * @param [out]     a         Polynomial.
 2736 * @param [in]      heap      Dynamic memory hint.
 2737 * @return  0 on success.
 2738 * @return  MEMORY_E when dynamic memory allocation fails.
 2739 * @return  Negative on hash error.
 2740 */
 2741static int dilithium_rej_ntt_poly(wc_Shake* shake128, byte* seed, sword32* a,
 2742    void* heap)
 2743{
 2744    int ret = 0;
 2745#if defined(WOLFSSL_SMALL_STACK)
 2746    byte* h = NULL;
 2747#else
 2748    byte h[DILITHIUM_REJ_NTT_POLY_H_SIZE];
 2749#endif
 2750
 2751    (void)heap;
 2752
 2753#if defined(WOLFSSL_SMALL_STACK)
 2754    h = (byte*)XMALLOC(DILITHIUM_REJ_NTT_POLY_H_SIZE, heap,
 2755        DYNAMIC_TYPE_DILITHIUM);
 2756    if (h == NULL) {
 2757        ret = MEMORY_E;
 2758    }
 2759#endif
 2760
 2761    if (ret == 0)
 2762        ret = dilithium_rej_ntt_poly_ex(shake128, seed, a, h);
 2763
 2764#if defined(WOLFSSL_SMALL_STACK)
 2765    XFREE(h, heap, DYNAMIC_TYPE_DILITHIUM);
 2766#endif
 2767
 2768    return ret;
 2769}
 2770#endif
 2771
 2772#if (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) && \
 2773     !defined(WOLFSSL_DILITHIUM_MAKE_KEY_SMALL_MEM)) || \
 2774    defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 2775    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 2776     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 2777    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 2778     (!defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 2779      defined(WC_DILITHIUM_CACHE_MATRIX_A)))
 2780#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 2781
 2782#define SHA3_128_BYTES   (WC_SHA3_128_COUNT * 8)
 2783/* Number of blocks to generate for matrix. */
 2784#define GEN_MATRIX_NBLOCKS  5
 2785/* Number of bytes to generate for matrix. */
 2786#define GEN_MATRIX_SIZE     GEN_MATRIX_NBLOCKS * SHA3_128_BYTES
 2787
 2788#ifndef WOLFSSL_NO_ML_DSA_44
 2789/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 2790 *
 2791 * Seed used with XOF to generate random bytes.
 2792 *
 2793 * @param  [out]  a           Matrix of uniform integers.
 2794 * @param  [in]   seed        Bytes to seed XOF generation.
 2795 * @return  0 on success.
 2796 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 2797 * WOLFSSL_SMALL_STACK is defined.
 2798 */
 2799static int wc_mldsa_gen_matrix_4x4_avx2(sword32* a, byte* seed)
 2800{
 2801    int i;
 2802    int k;
 2803    int l;
 2804#ifdef WOLFSSL_SMALL_STACK
 2805    byte *rand = NULL;
 2806    word64 *state = NULL;
 2807#else
 2808    byte rand[4 * GEN_MATRIX_SIZE + 2];
 2809    word64 state[25 * 4];
 2810#endif
 2811    unsigned int ctr0;
 2812    unsigned int ctr1;
 2813    unsigned int ctr2;
 2814    unsigned int ctr3;
 2815    byte* p;
 2816
 2817#ifdef WOLFSSL_SMALL_STACK
 2818    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 2, NULL,
 2819                          DYNAMIC_TYPE_TMP_BUFFER);
 2820    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 2821                          DYNAMIC_TYPE_TMP_BUFFER);
 2822    if ((rand == NULL) || (state == NULL)) {
 2823        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2824        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2825        return MEMORY_E;
 2826    }
 2827#endif
 2828
 2829    /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
 2830    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
 2831    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
 2832
 2833    for (k = 0; k < 4; k++) {
 2834        for (l = 0; l < 4; l++) {
 2835            state[4*4 + l] = 0x1f0000U + ((word32)k << 8) + (word32)l;
 2836        }
 2837
 2838        sha3_128_blocksx4_seed_avx2(state, seed);
 2839        wc_mldsa_redistribute_21_rand_avx2(state,
 2840            rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 2841            rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 2842        for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
 2843            sha3_blocksx4_avx2(state);
 2844            wc_mldsa_redistribute_21_rand_avx2(state,
 2845                rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
 2846                rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
 2847        }
 2848
 2849        /* Sample random bytes to create a polynomial. */
 2850        p = rand;
 2851        ctr0 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 0 * MLDSA_N, MLDSA_N, p,
 2852            GEN_MATRIX_SIZE);
 2853        p += GEN_MATRIX_SIZE;
 2854        ctr1 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 1 * MLDSA_N, MLDSA_N, p,
 2855            GEN_MATRIX_SIZE);
 2856        p += GEN_MATRIX_SIZE;
 2857        ctr2 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 2 * MLDSA_N, MLDSA_N, p,
 2858            GEN_MATRIX_SIZE);
 2859        p += GEN_MATRIX_SIZE;
 2860        ctr3 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 3 * MLDSA_N, MLDSA_N, p,
 2861            GEN_MATRIX_SIZE);
 2862
 2863        /* Create more blocks if too many rejected. */
 2864        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 2865               (ctr3 < MLDSA_N)) {
 2866            sha3_blocksx4_avx2(state);
 2867            wc_mldsa_redistribute_21_rand_avx2(state,
 2868                rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 2869                rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 2870
 2871            p = rand;
 2872            ctr0 += (word32)wc_mldsa_rej_uniform_avx2(a + 0 * MLDSA_N + ctr0,
 2873                MLDSA_N - ctr0, p, SHA3_128_BYTES);
 2874            p += GEN_MATRIX_SIZE;
 2875            ctr1 += (word32)wc_mldsa_rej_uniform_avx2(a + 1 * MLDSA_N + ctr1,
 2876                MLDSA_N - ctr1, p, SHA3_128_BYTES);
 2877            p += GEN_MATRIX_SIZE;
 2878            ctr2 += (word32)wc_mldsa_rej_uniform_avx2(a + 2 * MLDSA_N + ctr2,
 2879                MLDSA_N - ctr2, p, SHA3_128_BYTES);
 2880            p += GEN_MATRIX_SIZE;
 2881            ctr3 += (word32)wc_mldsa_rej_uniform_avx2(a + 3 * MLDSA_N + ctr3,
 2882                MLDSA_N - ctr3, p, SHA3_128_BYTES);
 2883        }
 2884
 2885        a += 4 * MLDSA_N;
 2886    }
 2887
 2888    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2889    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2890
 2891    return 0;
 2892}
 2893#endif
 2894
 2895#ifndef WOLFSSL_NO_ML_DSA_65
 2896/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 2897 *
 2898 * Seed used with XOF to generate random bytes.
 2899 *
 2900 * @param  [out]  a           Matrix of uniform integers.
 2901 * @param  [in]   seed        Bytes to seed XOF generation.
 2902 * @return  0 on success.
 2903 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 2904 * WOLFSSL_SMALL_STACK is defined.
 2905 */
 2906static int wc_mldsa_gen_matrix_6x5_avx2(sword32* a, byte* seed)
 2907{
 2908    int i;
 2909    int k;
 2910    int l;
 2911#ifdef WOLFSSL_SMALL_STACK
 2912    byte *rand = NULL;
 2913    word64 *state = NULL;
 2914#else
 2915    byte rand[4 * GEN_MATRIX_SIZE + 2];
 2916    word64 state[25 * 4];
 2917#endif
 2918    unsigned int ctr0;
 2919    unsigned int ctr1;
 2920    unsigned int ctr2;
 2921    unsigned int ctr3;
 2922    byte* p;
 2923
 2924#ifdef WOLFSSL_SMALL_STACK
 2925    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 2, NULL,
 2926                          DYNAMIC_TYPE_TMP_BUFFER);
 2927    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 2928                          DYNAMIC_TYPE_TMP_BUFFER);
 2929    if ((rand == NULL) || (state == NULL)) {
 2930        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2931        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 2932        return MEMORY_E;
 2933    }
 2934#endif
 2935
 2936    /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
 2937    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
 2938    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
 2939
 2940    for (k = 0; k < 6 * 5 - 2; k += 4) {
 2941        for (l = 0; l < 4; l++) {
 2942            state[4*4 + l] = 0x1f0000U + ((word32)((k + l) / 5) << 8) +
 2943                             (word32)((k + l) % 5);
 2944        }
 2945
 2946        sha3_128_blocksx4_seed_avx2(state, seed);
 2947        wc_mldsa_redistribute_21_rand_avx2(state,
 2948            rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 2949            rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 2950        for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
 2951            sha3_blocksx4_avx2(state);
 2952            wc_mldsa_redistribute_21_rand_avx2(state,
 2953                rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
 2954                rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
 2955        }
 2956
 2957        /* Sample random bytes to create a polynomial. */
 2958        p = rand;
 2959        ctr0 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 0 * MLDSA_N, MLDSA_N, p,
 2960            GEN_MATRIX_SIZE);
 2961        p += GEN_MATRIX_SIZE;
 2962        ctr1 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 1 * MLDSA_N, MLDSA_N, p,
 2963            GEN_MATRIX_SIZE);
 2964        p += GEN_MATRIX_SIZE;
 2965        ctr2 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 2 * MLDSA_N, MLDSA_N, p,
 2966            GEN_MATRIX_SIZE);
 2967        p += GEN_MATRIX_SIZE;
 2968        ctr3 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 3 * MLDSA_N, MLDSA_N, p,
 2969            GEN_MATRIX_SIZE);
 2970
 2971        /* Create more blocks if too many rejected. */
 2972        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 2973               (ctr3 < MLDSA_N)) {
 2974            sha3_blocksx4_avx2(state);
 2975            wc_mldsa_redistribute_21_rand_avx2(state,
 2976                rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 2977                rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 2978
 2979            p = rand;
 2980            ctr0 += (word32)wc_mldsa_rej_uniform_avx2(a + 0 * MLDSA_N + ctr0,
 2981                MLDSA_N - ctr0, p, SHA3_128_BYTES);
 2982            p += GEN_MATRIX_SIZE;
 2983            ctr1 += (word32)wc_mldsa_rej_uniform_avx2(a + 1 * MLDSA_N + ctr1,
 2984                MLDSA_N - ctr1, p, SHA3_128_BYTES);
 2985            p += GEN_MATRIX_SIZE;
 2986            ctr2 += (word32)wc_mldsa_rej_uniform_avx2(a + 2 * MLDSA_N + ctr2,
 2987                MLDSA_N - ctr2, p, SHA3_128_BYTES);
 2988            p += GEN_MATRIX_SIZE;
 2989            ctr3 += (word32)wc_mldsa_rej_uniform_avx2(a + 3 * MLDSA_N + ctr3,
 2990                MLDSA_N - ctr3, p, SHA3_128_BYTES);
 2991        }
 2992
 2993        a += 4 * MLDSA_N;
 2994    }
 2995
 2996    for (l = 0; l < 2; l++) {
 2997        state[4*4 + l] = 0x1f0000U + (5U << 8) + (word32)(l + 3);
 2998    }
 2999
 3000    sha3_128_blocksx4_seed_avx2(state, seed);
 3001    wc_mldsa_redistribute_21_rand_avx2(state,
 3002        rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 3003        rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 3004    for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
 3005        sha3_blocksx4_avx2(state);
 3006        wc_mldsa_redistribute_21_rand_avx2(state,
 3007            rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
 3008            rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
 3009    }
 3010
 3011    /* Sample random bytes to create a polynomial. */
 3012    p = rand;
 3013    ctr0 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 0 * MLDSA_N, MLDSA_N, p,
 3014        GEN_MATRIX_SIZE);
 3015    p += GEN_MATRIX_SIZE;
 3016    ctr1 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 1 * MLDSA_N, MLDSA_N, p,
 3017        GEN_MATRIX_SIZE);
 3018
 3019    /* Create more blocks if too many rejected. */
 3020    while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N)) {
 3021        sha3_blocksx4_avx2(state);
 3022        wc_mldsa_redistribute_21_rand_avx2(state, rand + 0 * GEN_MATRIX_SIZE,
 3023            rand + 1 * GEN_MATRIX_SIZE, rand + 2 * GEN_MATRIX_SIZE,
 3024            rand + 3 * GEN_MATRIX_SIZE);
 3025
 3026        p = rand;
 3027        ctr0 += (word32)wc_mldsa_rej_uniform_avx2(a + 0 * MLDSA_N + ctr0,
 3028            MLDSA_N - ctr0, p, SHA3_128_BYTES);
 3029        p += GEN_MATRIX_SIZE;
 3030        ctr1 += (word32)wc_mldsa_rej_uniform_avx2(a + 1 * MLDSA_N + ctr1,
 3031            MLDSA_N - ctr1, p, SHA3_128_BYTES);
 3032    }
 3033
 3034    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3035    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3036
 3037    return 0;
 3038}
 3039#endif
 3040
 3041#ifndef WOLFSSL_NO_ML_DSA_87
 3042/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 3043 *
 3044 * Seed used with XOF to generate random bytes.
 3045 *
 3046 * @param  [out]  a           Matrix of uniform integers.
 3047 * @param  [in]   seed        Bytes to seed XOF generation.
 3048 * @return  0 on success.
 3049 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 3050 * WOLFSSL_SMALL_STACK is defined.
 3051 */
 3052static int wc_mldsa_gen_matrix_8x7_avx2(sword32* a, byte* seed)
 3053{
 3054    int i;
 3055    int k;
 3056    int l;
 3057#ifdef WOLFSSL_SMALL_STACK
 3058    byte *rand = NULL;
 3059    word64 *state = NULL;
 3060#else
 3061    byte rand[4 * GEN_MATRIX_SIZE + 2];
 3062    word64 state[25 * 4];
 3063#endif
 3064    unsigned int ctr0;
 3065    unsigned int ctr1;
 3066    unsigned int ctr2;
 3067    unsigned int ctr3;
 3068    byte* p;
 3069
 3070#ifdef WOLFSSL_SMALL_STACK
 3071    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 2, NULL,
 3072                          DYNAMIC_TYPE_TMP_BUFFER);
 3073    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 3074                          DYNAMIC_TYPE_TMP_BUFFER);
 3075    if ((rand == NULL) || (state == NULL)) {
 3076        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3077        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3078        return MEMORY_E;
 3079    }
 3080#endif
 3081
 3082    /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
 3083    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
 3084    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
 3085
 3086    for (k = 0; k < 8 * 7; k += 4) {
 3087        for (l = 0; l < 4; l++) {
 3088            state[4*4 + l] = 0x1f0000U + ((word32)((k + l) / 7) << 8) +
 3089                             (word32)((k + l) % 7);
 3090        }
 3091
 3092        sha3_128_blocksx4_seed_avx2(state, seed);
 3093        wc_mldsa_redistribute_21_rand_avx2(state,
 3094            rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 3095            rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 3096        for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
 3097            sha3_blocksx4_avx2(state);
 3098            wc_mldsa_redistribute_21_rand_avx2(state,
 3099                rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
 3100                rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
 3101        }
 3102
 3103        /* Sample random bytes to create a polynomial. */
 3104        p = rand;
 3105        ctr0 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 0 * MLDSA_N, MLDSA_N, p,
 3106            GEN_MATRIX_SIZE);
 3107        p += GEN_MATRIX_SIZE;
 3108        ctr1 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 1 * MLDSA_N, MLDSA_N, p,
 3109            GEN_MATRIX_SIZE);
 3110        p += GEN_MATRIX_SIZE;
 3111        ctr2 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 2 * MLDSA_N, MLDSA_N, p,
 3112            GEN_MATRIX_SIZE);
 3113        p += GEN_MATRIX_SIZE;
 3114        ctr3 = (word32)wc_mldsa_rej_uniform_n_avx2(a + 3 * MLDSA_N, MLDSA_N, p,
 3115            GEN_MATRIX_SIZE);
 3116
 3117        /* Create more blocks if too many rejected. */
 3118        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 3119               (ctr3 < MLDSA_N)) {
 3120            sha3_blocksx4_avx2(state);
 3121            wc_mldsa_redistribute_21_rand_avx2(state,
 3122                rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
 3123                rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
 3124
 3125            p = rand;
 3126            ctr0 += (word32)wc_mldsa_rej_uniform_avx2(a + 0 * MLDSA_N + ctr0,
 3127                MLDSA_N - ctr0, p, SHA3_128_BYTES);
 3128            p += GEN_MATRIX_SIZE;
 3129            ctr1 += (word32)wc_mldsa_rej_uniform_avx2(a + 1 * MLDSA_N + ctr1,
 3130                MLDSA_N - ctr1, p, SHA3_128_BYTES);
 3131            p += GEN_MATRIX_SIZE;
 3132            ctr2 += (word32)wc_mldsa_rej_uniform_avx2(a + 2 * MLDSA_N + ctr2,
 3133                MLDSA_N - ctr2, p, SHA3_128_BYTES);
 3134            p += GEN_MATRIX_SIZE;
 3135            ctr3 += (word32)wc_mldsa_rej_uniform_avx2(a + 3 * MLDSA_N + ctr3,
 3136                MLDSA_N - ctr3, p, SHA3_128_BYTES);
 3137        }
 3138
 3139        a += 4 * MLDSA_N;
 3140    }
 3141
 3142    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3143    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3144
 3145    return 0;
 3146}
 3147#endif
 3148
 3149#endif
 3150
 3151/* Expand the seed to create matrix a.
 3152 *
 3153 * FIPS 204. 8.3: Algorithm 26 ExpandA(rho)
 3154 *   1: for r from 0 to k - 1 do
 3155 *   2:     for s from 0 to l - 1 do
 3156 *   3:         A_hat[r,s] <- RejNTTPoly(rho||IntegerToBits(s,8)||
 3157 *                                       IntegerToBits(r,8))
 3158 *   4:     end for
 3159 *   5: end for
 3160 *   6: return A_hat
 3161 *
 3162 * @param [in, out] shake128  SHAKE-128 object.
 3163 * @param [in]      pub_seed  Seed to generate stream of data.
 3164 * @param [in]      k         First dimension of matrix a.
 3165 * @param [in]      l         Second dimension of matrix a.
 3166 * @param [out]     a         Matrix of polynomials.
 3167 * @param [in]      heap      Dynamic memory hint.
 3168 * @return  0 on success.
 3169 * @return  Negative on hash error.
 3170 */
 3171static int dilithium_expand_a_c(wc_Shake* shake128, const byte* pub_seed,
 3172    byte k, byte l, sword32* a, void* heap)
 3173{
 3174    int ret = 0;
 3175    byte r;
 3176    byte s;
 3177    byte seed[DILITHIUM_GEN_A_SEED_SZ];
 3178
 3179    /* Copy the seed into a buffer that has space for s and r. */
 3180    XMEMCPY(seed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 3181    /* Step 1: Loop over first dimension of matrix. */
 3182    for (r = 0; (ret == 0) && (r < k); r++) {
 3183        /* Put r into buffer to be hashed. */
 3184        seed[DILITHIUM_PUB_SEED_SZ + 1] = r;
 3185        /* Step 2: Loop over second dimension of matrix. */
 3186        for (s = 0; (ret == 0) && (s < l); s++) {
 3187            /* Put s into buffer to be hashed. */
 3188            seed[DILITHIUM_PUB_SEED_SZ + 0] = s;
 3189            /* Step 3: Create polynomial from hashing seed. */
 3190            ret = dilithium_rej_ntt_poly(shake128, seed, a, heap);
 3191            /* Next polynomial. */
 3192            a += DILITHIUM_N;
 3193        }
 3194    }
 3195
 3196    return ret;
 3197}
 3198
 3199/* Expand the seed to create matrix a.
 3200 *
 3201 * FIPS 204. 8.3: Algorithm 26 ExpandA(rho)
 3202 *   1: for r from 0 to k - 1 do
 3203 *   2:     for s from 0 to l - 1 do
 3204 *   3:         A_hat[r,s] <- RejNTTPoly(rho||IntegerToBits(s,8)||
 3205 *                                       IntegerToBits(r,8))
 3206 *   4:     end for
 3207 *   5: end for
 3208 *   6: return A_hat
 3209 *
 3210 * @param [in, out] shake128  SHAKE-128 object.
 3211 * @param [in]      pub_seed  Seed to generate stream of data.
 3212 * @param [in]      k         First dimension of matrix a.
 3213 * @param [in]      l         Second dimension of matrix a.
 3214 * @param [out]     a         Matrix of polynomials.
 3215 * @param [in]      heap      Dynamic memory hint.
 3216 * @return  0 on success.
 3217 * @return  Negative on hash error.
 3218 */
 3219static int dilithium_expand_a(wc_Shake* shake128, const byte* pub_seed,
 3220    byte k, byte l, sword32* a, void* heap)
 3221{
 3222    int ret;
 3223#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 3224    byte seed[DILITHIUM_GEN_A_SEED_SZ];
 3225#endif
 3226
 3227#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 3228#ifndef WOLFSSL_NO_ML_DSA_44
 3229    if ((k == 4) && (l == 4) && IS_INTEL_AVX2(cpuid_flags) &&
 3230            IS_INTEL_BMI2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 3231        XMEMCPY(seed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 3232        ret = wc_mldsa_gen_matrix_4x4_avx2(a, seed);
 3233        RESTORE_VECTOR_REGISTERS();
 3234    }
 3235    else
 3236#endif
 3237#ifndef WOLFSSL_NO_ML_DSA_65
 3238    if ((k == 6) && (l == 5) && IS_INTEL_AVX2(cpuid_flags) &&
 3239            IS_INTEL_BMI2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 3240        XMEMCPY(seed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 3241        ret = wc_mldsa_gen_matrix_6x5_avx2(a, seed);
 3242        RESTORE_VECTOR_REGISTERS();
 3243    }
 3244    else
 3245#endif
 3246#ifndef WOLFSSL_NO_ML_DSA_87
 3247    if ((k == 8) && (l == 7) && IS_INTEL_AVX2(cpuid_flags) &&
 3248            IS_INTEL_BMI2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 3249        XMEMCPY(seed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 3250        ret = wc_mldsa_gen_matrix_8x7_avx2(a, seed);
 3251        RESTORE_VECTOR_REGISTERS();
 3252    }
 3253    else
 3254#endif
 3255#endif
 3256    {
 3257        ret = dilithium_expand_a_c(shake128, pub_seed, k, l, a, heap);
 3258    }
 3259
 3260    return ret;
 3261}
 3262#endif
 3263
 3264#ifndef WOLFSSL_DILITHIUM_NO_MAKE_KEY
 3265
 3266#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 3267/* Check random value is in valid range.
 3268 *
 3269 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3270 *   1: if             b < 15
 3271 *
 3272 * @param [in] b    Random half-byte (nibble) value.
 3273 * @param [in] eta  Range specifier of result. Will always be 2 - unused.
 3274 * @return  1 when value less than 9.
 3275 * @return  0 when value greater than or equal to 9.
 3276 */
 3277#define DILITHIUM_COEFF_S_VALID_ETA2(b) \
 3278    ((b) < DILITHIUM_ETA_2_MOD)
 3279
 3280static const signed char dilithium_coeff_eta2[] = {
 3281    2, 1, 0, -1, -2,
 3282    2, 1, 0, -1, -2,
 3283    2, 1, 0, -1, -2
 3284};
 3285/* Convert random value 0..15 to a value in range of -2..2.
 3286 *
 3287 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3288 *   1:                            return 2 - (b mod 5)
 3289 *
 3290 * @param [in] b    Random half-byte (nibble) value.
 3291 * @return  Value in range of -2..2 on success.
 3292 */
 3293#define DILITHIUM_COEFF_S_ETA2(b)       \
 3294    (dilithium_coeff_eta2[b])
 3295#endif
 3296
 3297#ifndef WOLFSSL_NO_ML_DSA_65
 3298/* Check random value is in valid range.
 3299 *
 3300 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3301 *   3:     if             b < 9
 3302 *
 3303 * @param [in] b    Random half-byte (nibble) value.
 3304 * @param [in] eta  Range specifier of result. Will always be 4 - unused.
 3305 * @return  1 when value less than 9.
 3306 * @return  0 when value greater than or equal to 9.
 3307 */
 3308#define DILITHIUM_COEFF_S_VALID_ETA4(b) \
 3309    ((b) < DILITHIUM_ETA_4_MOD)
 3310
 3311/* Convert random value 0..15 to a value in range of -4..4.
 3312 *
 3313 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3314 *   3:                               return 4 - b
 3315 *
 3316 * @param [in] b    Random half-byte (nibble) value.
 3317 * @param [in] eta  Range specifier of result. Will always be 4 - unused.
 3318 * @return  Value in range of -4..4 on success.
 3319 */
 3320#define DILITHIUM_COEFF_S_ETA4(b)       \
 3321    (4 - (b))
 3322#endif
 3323
 3324#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_87)
 3325#ifndef WOLFSSL_NO_ML_DSA_65
 3326
 3327/* Check random value is in valid range.
 3328 *
 3329 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3330 *   1: if eta = 2 and b < 15
 3331 *   2: else
 3332 *   3:     if eta = 4 and b < 9
 3333 *
 3334 * @param [in] b    Random half-byte (nibble) value.
 3335 * @param [in] eta  Range specifier of result.
 3336 * @return  Value in range of -ETA..ETA on success.
 3337 */
 3338#define DILITHIUM_COEFF_S_VALID(b, eta)                             \
 3339    (((eta) == DILITHIUM_ETA_2) ? DILITHIUM_COEFF_S_VALID_ETA2(b)   \
 3340                                : DILITHIUM_COEFF_S_VALID_ETA4(b))
 3341
 3342/* Convert random value 0..15 to a value in range of -ETA..ETA.
 3343 *
 3344 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3345 *   1: if eta = 2            then return 2 - (b mod 5)
 3346 *   2: else
 3347 *   3:     if eta = 4           then return 4 - b
 3348 *   ...
 3349 *   6: end if
 3350 *
 3351 * @param [in] b    Random half-byte (nibble) value.
 3352 * @param [in] eta  Range specifier of result.
 3353 * @return  Value in range of -ETA..ETA on success.
 3354 */
 3355#define DILITHIUM_COEFF_S(b, eta)                               \
 3356    (((eta) == DILITHIUM_ETA_2) ? DILITHIUM_COEFF_S_ETA2(b)     \
 3357                                : DILITHIUM_COEFF_S_ETA4(b))
 3358
 3359#else
 3360
 3361/* Check random value is in valid range.
 3362 *
 3363 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3364 *   1: if             b < 15
 3365 *
 3366 * @param [in] b    Random half-byte (nibble) value.
 3367 * @param [in] eta  Range specifier of result. Will always be 2 - unused.
 3368 * @return  1 when value less than 9.
 3369 * @return  0 when value greater than or equal to 9.
 3370 */
 3371#define DILITHIUM_COEFF_S_VALID(b, eta) \
 3372    DILITHIUM_COEFF_S_VALID_ETA2(b)
 3373
 3374/* Convert random value 0..15 to a value in range of -2..2.
 3375 *
 3376 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3377 *   1:                            return 2 - (b mod 5)
 3378 *
 3379 * @param [in] b    Random half-byte (nibble) value.
 3380 * @param [in] eta  Range specifier of result. Will always be 2 - unused.
 3381 * @return  Value in range of -2..2 on success.
 3382 */
 3383#define DILITHIUM_COEFF_S(b, eta)   \
 3384    DILITHIUM_COEFF_S_ETA2(b)
 3385
 3386#endif /*  WOLFSSL_NO_ML_DSA_65 */
 3387
 3388#else
 3389
 3390/* Check random value is in valid range.
 3391 *
 3392 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3393 *   3:     if             b < 9
 3394 *
 3395 * @param [in] b    Random half-byte (nibble) value.
 3396 * @param [in] eta  Range specifier of result. Will always be 4 - unused.
 3397 * @return  1 when value less than 9.
 3398 * @return  0 when value greater than or equal to 9.
 3399 */
 3400#define DILITHIUM_COEFF_S_VALID(b, eta) \
 3401    DILITHIUM_COEFF_S_VALID_ETA4(b)
 3402
 3403/* Convert random value 0..15 to a value in range of -4..4.
 3404 *
 3405 * FIPS 204. 8.1: Algorithm 9 CoeffFromHalfByte(b)
 3406 *   3:                               return 4 - b
 3407 *
 3408 * @param [in] b    Random half-byte (nibble) value.
 3409 * @param [in] eta  Range specifier of result. Will always be 4 - unused.
 3410 * @return  Value in range of -4..4 on success.
 3411 */
 3412#define DILITHIUM_COEFF_S(b, eta)   \
 3413    DILITHIUM_COEFF_S_ETA4(b)
 3414
 3415#endif /* !WOLFSSL_NO_ML_DSA_44 || !WOLFSSL_NO_ML_DSA_87 */
 3416
 3417/* Extract a coefficient from a nibble of z.
 3418 *
 3419 * Breaks out of loop when we have enough coefficients.
 3420 *
 3421 * @param [in] z    A random value.
 3422 * @param [in] rs   Amount to shift right.
 3423 * @param [in] t    Temporary result.
 3424 * @param [in] eta  ETA value from parameters.
 3425 * @return  Value in range -eta..eta on success.
 3426 * @return  Falsam (0x10) when random value out of range.
 3427 */
 3428#define EXTRACT_COEFF_NIBBLE_CHECK_J(z, rs, t, eta)                     \
 3429        (t) = (sword8)(((z) >> (rs)) & 0xf);                            \
 3430        /* Step 7: Check we have a valid coefficient. */                \
 3431        if (DILITHIUM_COEFF_S_VALID(t, eta)) {                          \
 3432            (t) = (sword8)DILITHIUM_COEFF_S(t, eta);                    \
 3433            /* Step 8: Store coefficient as next polynomial value.      \
 3434             * Step 9: Increment count of polynomial values set. */     \
 3435            s[j++] = (sword32)(t);                                      \
 3436            if (j == DILITHIUM_N) {                                     \
 3437                break;                                                  \
 3438            }                                                           \
 3439        }
 3440
 3441/* Extract a coefficient from a nibble of z.
 3442 *
 3443 * @param [in] z    A random value.
 3444 * @param [in] rs   Amount to shift right.
 3445 * @param [in] t    Temporary result.
 3446 * @param [in] eta  ETA value from parameters.
 3447 * @return  Value in range -eta..eta on success.
 3448 * @return  Falsam (0x10) when random value out of range.
 3449 */
 3450#define EXTRACT_COEFF_NIBBLE(z, rs, t, eta)                             \
 3451        (t) = (sword8)(((z) >> (rs)) & 0xf);                            \
 3452        /* Step 7: Check we have a valid coefficient. */                \
 3453        if (DILITHIUM_COEFF_S_VALID(t, eta)) {                          \
 3454            (t) = (sword8)DILITHIUM_COEFF_S(t, eta);                    \
 3455            /* Step 8: Store coefficient as next polynomial value.      \
 3456             * Step 9: Increment count of polynomial values set. */     \
 3457            s[j++] = (sword32)(t);                                      \
 3458        }
 3459
 3460
 3461/* Extract coefficients from hash - z.
 3462 *
 3463 * FIPS 204. 8.3: Algorithm 25 RejBoundedPoly(rho)
 3464 *   2: c <- 0
 3465 *   5:    z0 <- CoeffFromHalfByte(z mod 16, eta)
 3466 *   6:    z1 <- CoeffFromHalfByte(lower(z / 16), eta)
 3467 *   7:    if z0 != falsam then
 3468 *   8:         aj <- z0
 3469 *   9:         j <- j + 1
 3470 *  10:    end if
 3471 *  11:    if z1 != falsam then
 3472 *  12:         aj <- z1
 3473 *  13:         j <- j + 1
 3474 *  14:    end if
 3475 *  15:    c <- c + 1
 3476 *
 3477 * @param [in]      z     Hash data to extract coefficients from.
 3478 * @param [in]      zLen  Length of z in bytes.
 3479 * @param [in]      eta   Range specifier of each value.
 3480 * @param [out]     s     Polynomial to fill with coefficients.
 3481 * @param [in, out] cnt   Current count of coefficients in polynomial.
 3482 */
 3483static void dilithium_extract_coeffs(const byte* z, unsigned int zLen,
 3484    byte eta, sword32* s, unsigned int* cnt)
 3485{
 3486#ifdef WOLFSSL_DILITHIUM_NO_LARGE_CODE
 3487    unsigned int j = *cnt;
 3488    unsigned int c;
 3489
 3490    (void)eta;
 3491
 3492    /* Extract values from the squeezed data. */
 3493    for (c = 0; c < zLen; c++) {
 3494        sword8 t;
 3495
 3496        /* Step 5: Get coefficient from bottom nibble. */
 3497        EXTRACT_COEFF_NIBBLE_CHECK_J(z[c], 0, t, eta);
 3498        /* Step 6: Get coefficient from top nibble. */
 3499        EXTRACT_COEFF_NIBBLE_CHECK_J(z[c], 4, t, eta);
 3500    }
 3501
 3502    *cnt = j;
 3503#else
 3504    unsigned int j = *cnt;
 3505    unsigned int c;
 3506    unsigned int min = (DILITHIUM_N - j) / 2;
 3507
 3508    (void)eta;
 3509
 3510#if defined(LITTLE_ENDIAN_ORDER)
 3511#ifdef WC_64BIT_CPU
 3512    min &= ~(unsigned int)7;
 3513    /* Extract values from the squeezed data. */
 3514    for (c = 0; c < min; c += 8) {
 3515        word64 z64 = *(const word64*)(z + c);
 3516        sword8 t;
 3517
 3518        /* Do each nibble from lowest to highest 16 at a time. */
 3519        EXTRACT_COEFF_NIBBLE(z64,  0, t, eta);
 3520        EXTRACT_COEFF_NIBBLE(z64,  4, t, eta);
 3521        EXTRACT_COEFF_NIBBLE(z64,  8, t, eta);
 3522        EXTRACT_COEFF_NIBBLE(z64, 12, t, eta);
 3523        EXTRACT_COEFF_NIBBLE(z64, 16, t, eta);
 3524        EXTRACT_COEFF_NIBBLE(z64, 20, t, eta);
 3525        EXTRACT_COEFF_NIBBLE(z64, 24, t, eta);
 3526        EXTRACT_COEFF_NIBBLE(z64, 28, t, eta);
 3527        EXTRACT_COEFF_NIBBLE(z64, 32, t, eta);
 3528        EXTRACT_COEFF_NIBBLE(z64, 36, t, eta);
 3529        EXTRACT_COEFF_NIBBLE(z64, 40, t, eta);
 3530        EXTRACT_COEFF_NIBBLE(z64, 44, t, eta);
 3531        EXTRACT_COEFF_NIBBLE(z64, 48, t, eta);
 3532        EXTRACT_COEFF_NIBBLE(z64, 52, t, eta);
 3533        EXTRACT_COEFF_NIBBLE(z64, 56, t, eta);
 3534        EXTRACT_COEFF_NIBBLE(z64, 60, t, eta);
 3535    }
 3536#else
 3537    min &= ~(unsigned int)3;
 3538    /* Extract values from the squeezed data. */
 3539    for (c = 0; c < min; c += 4) {
 3540        word32 z32 = *(const word32*)(z + c);
 3541        sword8 t;
 3542
 3543        /* Do each nibble from lowest to highest 8 at a time. */
 3544        EXTRACT_COEFF_NIBBLE(z32,  0, t, eta);
 3545        EXTRACT_COEFF_NIBBLE(z32,  4, t, eta);
 3546        EXTRACT_COEFF_NIBBLE(z32,  8, t, eta);
 3547        EXTRACT_COEFF_NIBBLE(z32, 12, t, eta);
 3548        EXTRACT_COEFF_NIBBLE(z32, 16, t, eta);
 3549        EXTRACT_COEFF_NIBBLE(z32, 20, t, eta);
 3550        EXTRACT_COEFF_NIBBLE(z32, 24, t, eta);
 3551        EXTRACT_COEFF_NIBBLE(z32, 28, t, eta);
 3552    }
 3553#endif
 3554#else
 3555    /* Extract values from the squeezed data. */
 3556    for (c = 0; c < min; c++) {
 3557        sword8 t;
 3558
 3559        /* Step 5: Get coefficient from bottom nibble. */
 3560        EXTRACT_COEFF_NIBBLE(z[c], 0, t, eta);
 3561        EXTRACT_COEFF_NIBBLE(z[c], 4, t, eta);
 3562    }
 3563#endif
 3564    if (j != DILITHIUM_N) {
 3565        /* Extract values from the squeezed data. */
 3566        for (; c < zLen; c++) {
 3567            sword8 t;
 3568
 3569            EXTRACT_COEFF_NIBBLE_CHECK_J(z[c], 0, t, eta);
 3570            EXTRACT_COEFF_NIBBLE_CHECK_J(z[c], 4, t, eta);
 3571        }
 3572    }
 3573
 3574    *cnt = j;
 3575#endif
 3576}
 3577
 3578/* Create polynomial from hashing the seed with bounded values.
 3579 *
 3580 * FIPS 204. 8.3: Algorithm 25 RejBoundedPoly(rho)
 3581 *   1: j <- 0
 3582 *      ...
 3583 *   3: while j < 256 do
 3584 *   4:    z <- H(rho)[[c]]
 3585 *         ... [Extract coefficients into polynomial from z]
 3586 *  16: end while
 3587 *  17: return a
 3588 *
 3589 * @param [in, out] shake256  SHAKE-256 object.
 3590 * @param [in]      seed      Seed, rho, to hash to generate values.
 3591 * @param [in]      eta       Range specifier of each value.
 3592 * @return  0 on success.
 3593 * @return  Negative on hash error.
 3594 */
 3595static int dilithium_rej_bound_poly(wc_Shake* shake256, byte* seed, sword32* s,
 3596    byte eta)
 3597{
 3598#ifdef WOLFSSL_DILITHIUM_SMALL
 3599    int ret;
 3600    unsigned int j = 0;
 3601    byte z[DILITHIUM_GEN_S_BLOCK_BYTES];
 3602
 3603    /* Initialize SHAKE-256 object for new hash. */
 3604    ret = wc_InitShake256(shake256, NULL, INVALID_DEVID);
 3605    if (ret == 0) {
 3606        /* Absorb the seed. */
 3607        ret = wc_Shake256_Absorb(shake256, seed, DILITHIUM_GEN_S_SEED_SZ);
 3608    }
 3609    if (ret == 0) {
 3610        do {
 3611            /* Squeeze out another block. */
 3612            ret = wc_Shake256_SqueezeBlocks(shake256, z, 1);
 3613            if (ret != 0) {
 3614                break;
 3615            }
 3616            /* Extract up to the 256 valid coefficients for polynomial. */
 3617            dilithium_extract_coeffs(z, DILITHIUM_GEN_S_BLOCK_BYTES, eta, s,
 3618                &j);
 3619        }
 3620        /* Check we got enough values to fill polynomial. */
 3621        while (j < DILITHIUM_N);
 3622    }
 3623
 3624    return ret;
 3625#else
 3626    int ret;
 3627    unsigned int j = 0;
 3628    WC_DECLARE_VAR(z, byte, DILITHIUM_GEN_S_BYTES, NULL);
 3629
 3630    WC_ALLOC_VAR_EX(z, byte, DILITHIUM_GEN_S_BYTES, NULL, DYNAMIC_TYPE_DILITHIUM,
 3631                    return MEMORY_E);
 3632
 3633    /* Absorb seed and squeeze out some blocks. */
 3634    ret = dilithium_squeeze256(shake256, seed, DILITHIUM_GEN_S_SEED_SZ, z,
 3635        DILITHIUM_GEN_S_NBLOCKS);
 3636    if (ret == 0) {
 3637        /* Extract up to 256 valid coefficients for polynomial. */
 3638        dilithium_extract_coeffs(z, DILITHIUM_GEN_S_BYTES, eta, s, &j);
 3639        /* Check we got enough values to fill polynomial. */
 3640        while (j < DILITHIUM_N) {
 3641            /* Squeeze out another block. */
 3642            ret = wc_Shake256_SqueezeBlocks(shake256, z, 1);
 3643            if (ret != 0) {
 3644                break;
 3645            }
 3646            /* Extract up to the 256 valid coefficients for polynomial. */
 3647            dilithium_extract_coeffs(z, DILITHIUM_GEN_S_BLOCK_BYTES, eta, s,
 3648                &j);
 3649        }
 3650    }
 3651
 3652    WC_FREE_VAR_EX(z, NULL, DYNAMIC_TYPE_DILITHIUM);
 3653    return ret;
 3654#endif
 3655}
 3656
 3657#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 3658#ifndef WOLFSSL_NO_ML_DSA_44
 3659/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 3660 *
 3661 * Seed used with XOF to generate random bytes.
 3662 *
 3663 * @param  [out]  a           Vectos of uniform integers.
 3664 * @param  [in]   seed        Bytes to seed XOF generation.
 3665 * @return  0 on success.
 3666 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 3667 * WOLFSSL_SMALL_STACK is defined.
 3668 */
 3669static int wc_mldsa_gen_s_4_4_avx2(sword32* s[2], byte* seed)
 3670{
 3671    int k;
 3672    int l;
 3673#ifdef WOLFSSL_SMALL_STACK
 3674    byte *rand = NULL;
 3675    word64 *state = NULL;
 3676#else
 3677    byte rand[4 * DILITHIUM_GEN_S_BLOCK_BYTES];
 3678    word64 state[25 * 4];
 3679#endif
 3680    unsigned int ctr0;
 3681    unsigned int ctr1;
 3682    unsigned int ctr2;
 3683    unsigned int ctr3;
 3684    byte* p;
 3685
 3686#ifdef WOLFSSL_SMALL_STACK
 3687    rand = (byte*)XMALLOC(4 * DILITHIUM_GEN_S_BLOCK_BYTES, NULL,
 3688                          DYNAMIC_TYPE_TMP_BUFFER);
 3689    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 3690                          DYNAMIC_TYPE_TMP_BUFFER);
 3691    if ((rand == NULL) || (state == NULL)) {
 3692        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3693        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3694        return MEMORY_E;
 3695    }
 3696#endif
 3697
 3698    for (k = 0; k < 2; k++) {
 3699        for (l = 0; l < 4; l++) {
 3700            state[8*4 + l] = 0x1f0000U + ((word32)k * 4U + (word32)l);
 3701        }
 3702
 3703        ctr0 = 0;
 3704        ctr1 = 0;
 3705        ctr2 = 0;
 3706        ctr3 = 0;
 3707
 3708        sha3_256_blocksx4_seed_64_avx2(state, seed);
 3709        wc_mldsa_redistribute_17_rand_avx2(state,
 3710            rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3711            rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3712            rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3713            rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3714
 3715        do {
 3716            p = rand;
 3717            if (ctr0 < MLDSA_N) {
 3718                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3719                    DILITHIUM_GEN_S_BLOCK_BYTES, s[k] + 0 * MLDSA_N + ctr0,
 3720                    &ctr0);
 3721            }
 3722            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3723            if (ctr1 < MLDSA_N) {
 3724                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3725                    DILITHIUM_GEN_S_BLOCK_BYTES, s[k] + 1 * MLDSA_N + ctr1,
 3726                    &ctr1);
 3727            }
 3728            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3729            if (ctr2 < MLDSA_N) {
 3730                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3731                    DILITHIUM_GEN_S_BLOCK_BYTES, s[k] + 2 * MLDSA_N + ctr2,
 3732                    &ctr2);
 3733            }
 3734            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3735            if (ctr3 < MLDSA_N) {
 3736                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3737                    DILITHIUM_GEN_S_BLOCK_BYTES, s[k] + 3 * MLDSA_N + ctr3,
 3738                    &ctr3);
 3739            }
 3740
 3741            if ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 3742                   (ctr3 < MLDSA_N)) {
 3743                sha3_blocksx4_avx2(state);
 3744                wc_mldsa_redistribute_17_rand_avx2(state,
 3745                    rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3746                    rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3747                    rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3748                    rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3749            }
 3750        }
 3751        /* Create more blocks if too many rejected. */
 3752        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 3753               (ctr3 < MLDSA_N));
 3754    }
 3755
 3756    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3757    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3758
 3759    return 0;
 3760}
 3761#endif
 3762
 3763#ifndef WOLFSSL_NO_ML_DSA_65
 3764/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 3765 *
 3766 * Seed used with XOF to generate random bytes.
 3767 *
 3768 * @param  [out]  a           Vectos of uniform integers.
 3769 * @param  [in]   seed        Bytes to seed XOF generation.
 3770 * @return  0 on success.
 3771 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 3772 * WOLFSSL_SMALL_STACK is defined.
 3773 */
 3774static int wc_mldsa_gen_s_5_6_avx2(sword32* s[2], byte* seed)
 3775{
 3776    int k;
 3777    int l;
 3778#ifdef WOLFSSL_SMALL_STACK
 3779    byte *rand = NULL;
 3780    word64 *state = NULL;
 3781#else
 3782    byte rand[4 * DILITHIUM_GEN_S_BLOCK_BYTES];
 3783    word64 state[25 * 4];
 3784#endif
 3785    unsigned int ctr0;
 3786    unsigned int ctr1;
 3787    unsigned int ctr2;
 3788    unsigned int ctr3;
 3789    byte* p;
 3790    sword32* sa[3][4] = {
 3791        { &s[0][0 * MLDSA_N], &s[0][1 * MLDSA_N],
 3792          &s[0][2 * MLDSA_N], &s[0][3 * MLDSA_N] },
 3793        { &s[0][4 * MLDSA_N], &s[1][0 * MLDSA_N],
 3794          &s[1][1 * MLDSA_N], &s[1][2 * MLDSA_N] },
 3795        { &s[1][3 * MLDSA_N], &s[1][4 * MLDSA_N],
 3796          &s[1][5 * MLDSA_N], NULL               }
 3797    };
 3798
 3799#ifdef WOLFSSL_SMALL_STACK
 3800    rand = (byte*)XMALLOC(4 * DILITHIUM_GEN_S_BLOCK_BYTES, NULL,
 3801                          DYNAMIC_TYPE_TMP_BUFFER);
 3802    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 3803                          DYNAMIC_TYPE_TMP_BUFFER);
 3804    if ((rand == NULL) || (state == NULL)) {
 3805        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3806        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3807        return MEMORY_E;
 3808    }
 3809#endif
 3810
 3811    for (k = 0; k < 2; k++) {
 3812        for (l = 0; l < 4; l++) {
 3813            state[8*4 + l] = 0x1f0000U + ((word32)k * 4U + (word32)l);
 3814        }
 3815
 3816        ctr0 = 0;
 3817        ctr1 = 0;
 3818        ctr2 = 0;
 3819        ctr3 = 0;
 3820
 3821        sha3_256_blocksx4_seed_64_avx2(state, seed);
 3822        wc_mldsa_redistribute_17_rand_avx2(state,
 3823            rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3824            rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3825            rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3826            rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3827
 3828        do {
 3829            p = rand;
 3830            if (ctr0 < MLDSA_N) {
 3831                wc_mldsa_extract_coeffs_eta4_avx2(p,
 3832                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][0] + ctr0, &ctr0);
 3833            }
 3834            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3835            if (ctr1 < MLDSA_N) {
 3836                wc_mldsa_extract_coeffs_eta4_avx2(p,
 3837                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][1] + ctr1, &ctr1);
 3838            }
 3839            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3840            if (ctr2 < MLDSA_N) {
 3841                wc_mldsa_extract_coeffs_eta4_avx2(p,
 3842                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][2] + ctr2, &ctr2);
 3843            }
 3844            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3845            if (ctr3 < MLDSA_N) {
 3846                wc_mldsa_extract_coeffs_eta4_avx2(p,
 3847                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][3] + ctr3, &ctr3);
 3848            }
 3849
 3850            if ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 3851                   (ctr3 < MLDSA_N)) {
 3852                sha3_blocksx4_avx2(state);
 3853                wc_mldsa_redistribute_17_rand_avx2(state,
 3854                    rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3855                    rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3856                    rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3857                    rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3858            }
 3859        }
 3860        /* Create more blocks if too many rejected. */
 3861        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 3862               (ctr3 < MLDSA_N));
 3863    }
 3864
 3865    for (l = 0; l < 4; l++) {
 3866        state[8*4 + l] = 0x1f0000U + 8U + (word32)l;
 3867    }
 3868
 3869    ctr0 = 0;
 3870    ctr1 = 0;
 3871    ctr2 = 0;
 3872
 3873    sha3_256_blocksx4_seed_64_avx2(state, seed);
 3874    wc_mldsa_redistribute_17_rand_avx2(state,
 3875        rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3876        rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3877        rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3878        rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3879
 3880    do {
 3881        p = rand;
 3882        if (ctr0 < MLDSA_N) {
 3883            wc_mldsa_extract_coeffs_eta4_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 3884                sa[k][0] + ctr0, &ctr0);
 3885        }
 3886        p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3887        if (ctr1 < MLDSA_N) {
 3888            wc_mldsa_extract_coeffs_eta4_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 3889                sa[k][1] + ctr1, &ctr1);
 3890        }
 3891        p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3892        if (ctr2 < MLDSA_N) {
 3893            wc_mldsa_extract_coeffs_eta4_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 3894                sa[k][2] + ctr2, &ctr2);
 3895        }
 3896
 3897        if ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N)) {
 3898            sha3_blocksx4_avx2(state);
 3899            wc_mldsa_redistribute_17_rand_avx2(state,
 3900                rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3901                rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3902                rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3903                rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3904        }
 3905    }
 3906    /* Create more blocks if too many rejected. */
 3907    while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N));
 3908
 3909    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3910    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3911
 3912    return 0;
 3913}
 3914#endif
 3915
 3916#ifndef WOLFSSL_NO_ML_DSA_87
 3917/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 3918 *
 3919 * Seed used with XOF to generate random bytes.
 3920 *
 3921 * @param  [out]  a           Vectos of uniform integers.
 3922 * @param  [in]   seed        Bytes to seed XOF generation.
 3923 * @return  0 on success.
 3924 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 3925 * WOLFSSL_SMALL_STACK is defined.
 3926 */
 3927static int wc_mldsa_gen_s_7_8_avx2(sword32* s[2], byte* seed)
 3928{
 3929    int k;
 3930    int l;
 3931#ifdef WOLFSSL_SMALL_STACK
 3932    byte *rand = NULL;
 3933    word64 *state = NULL;
 3934#else
 3935    byte rand[4 * DILITHIUM_GEN_S_BLOCK_BYTES];
 3936    word64 state[25 * 4];
 3937#endif
 3938    unsigned int ctr0;
 3939    unsigned int ctr1;
 3940    unsigned int ctr2;
 3941    unsigned int ctr3;
 3942    byte* p;
 3943    sword32* sa[4][4] = {
 3944        { &s[0][0 * MLDSA_N], &s[0][1 * MLDSA_N],
 3945          &s[0][2 * MLDSA_N], &s[0][3 * MLDSA_N] },
 3946        { &s[0][4 * MLDSA_N], &s[0][5 * MLDSA_N],
 3947          &s[0][6 * MLDSA_N], &s[1][0 * MLDSA_N] },
 3948        { &s[1][1 * MLDSA_N], &s[1][2 * MLDSA_N],
 3949          &s[1][3 * MLDSA_N], &s[1][4 * MLDSA_N] },
 3950        { &s[1][5 * MLDSA_N], &s[1][6 * MLDSA_N],
 3951          &s[1][7 * MLDSA_N], NULL               }
 3952    };
 3953
 3954#ifdef WOLFSSL_SMALL_STACK
 3955    rand = (byte*)XMALLOC(4 * DILITHIUM_GEN_S_BLOCK_BYTES, NULL,
 3956                          DYNAMIC_TYPE_TMP_BUFFER);
 3957    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 3958                          DYNAMIC_TYPE_TMP_BUFFER);
 3959    if ((rand == NULL) || (state == NULL)) {
 3960        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3961        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 3962        return MEMORY_E;
 3963    }
 3964#endif
 3965
 3966    for (k = 0; k < 3; k++) {
 3967        for (l = 0; l < 4; l++) {
 3968            state[8*4 + l] = 0x1f0000U + ((word32)k * 4U + (word32)l);
 3969        }
 3970
 3971        ctr0 = 0;
 3972        ctr1 = 0;
 3973        ctr2 = 0;
 3974        ctr3 = 0;
 3975
 3976        sha3_256_blocksx4_seed_64_avx2(state, seed);
 3977        wc_mldsa_redistribute_17_rand_avx2(state,
 3978            rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3979            rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3980            rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 3981            rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 3982
 3983        do {
 3984            p = rand;
 3985            if (ctr0 < MLDSA_N) {
 3986                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3987                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][0] + ctr0, &ctr0);
 3988            }
 3989            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3990            if (ctr1 < MLDSA_N) {
 3991                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3992                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][1] + ctr1, &ctr1);
 3993            }
 3994            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 3995            if (ctr2 < MLDSA_N) {
 3996                wc_mldsa_extract_coeffs_eta2_avx2(p,
 3997                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][2] + ctr2, &ctr2);
 3998            }
 3999            p += DILITHIUM_GEN_S_BLOCK_BYTES;
 4000            if (ctr3 < MLDSA_N) {
 4001                wc_mldsa_extract_coeffs_eta2_avx2(p,
 4002                    DILITHIUM_GEN_S_BLOCK_BYTES, sa[k][3] + ctr3, &ctr3);
 4003            }
 4004
 4005            if ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 4006                   (ctr3 < MLDSA_N)) {
 4007                sha3_blocksx4_avx2(state);
 4008                wc_mldsa_redistribute_17_rand_avx2(state,
 4009                    rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4010                    rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4011                    rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4012                    rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 4013            }
 4014        }
 4015        /* Create more blocks if too many rejected. */
 4016        while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N) ||
 4017               (ctr3 < MLDSA_N));
 4018    }
 4019
 4020    for (l = 0; l < 4; l++) {
 4021        state[8*4 + l] = 0x1f0000U + 12U + (word32)l;
 4022    }
 4023
 4024    ctr0 = 0;
 4025    ctr1 = 0;
 4026    ctr2 = 0;
 4027
 4028    sha3_256_blocksx4_seed_64_avx2(state, seed);
 4029    wc_mldsa_redistribute_17_rand_avx2(state,
 4030        rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4031        rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4032        rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4033        rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 4034
 4035    do {
 4036        p = rand;
 4037        if (ctr0 < MLDSA_N) {
 4038            wc_mldsa_extract_coeffs_eta2_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 4039                sa[k][0] + ctr0, &ctr0);
 4040        }
 4041        p += DILITHIUM_GEN_S_BLOCK_BYTES;
 4042        if (ctr1 < MLDSA_N) {
 4043            wc_mldsa_extract_coeffs_eta2_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 4044                sa[k][1] + ctr1, &ctr1);
 4045        }
 4046        p += DILITHIUM_GEN_S_BLOCK_BYTES;
 4047        if (ctr2 < MLDSA_N) {
 4048            wc_mldsa_extract_coeffs_eta2_avx2(p, DILITHIUM_GEN_S_BLOCK_BYTES,
 4049                sa[k][2] + ctr2, &ctr2);
 4050        }
 4051
 4052        if ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N)) {
 4053            sha3_blocksx4_avx2(state);
 4054            wc_mldsa_redistribute_17_rand_avx2(state,
 4055                rand + 0 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4056                rand + 1 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4057                rand + 2 * DILITHIUM_GEN_S_BLOCK_BYTES,
 4058                rand + 3 * DILITHIUM_GEN_S_BLOCK_BYTES);
 4059        }
 4060    }
 4061    /* Create more blocks if too many rejected. */
 4062    while ((ctr0 < MLDSA_N) || (ctr1 < MLDSA_N) || (ctr2 < MLDSA_N));
 4063
 4064    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4065    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4066
 4067    return 0;
 4068}
 4069#endif
 4070#endif
 4071
 4072/* Expand private seed into vectors s1 and s2.
 4073 *
 4074 * FIPS 204. 8.3: Algorithm 27 ExpandS(rho)
 4075 *   1: for r from 0 to l - 1 do
 4076 *   2:     s1[r] <- RejBoundedPoly(rho||IntegerToBits(r,16))
 4077 *   3: end for
 4078 *   4: for r from 0 to k - 1 do
 4079 *   5:     s2[r] <- RejBoundedPoly(rho||IntegerToBits(r + l,16))
 4080 *   6: end for
 4081 *   7: return (s1,s2)
 4082 *
 4083 * @param [in, out] shake256   SHAKE-256 object.
 4084 * @param [in]      priv_seed  Private seed, rho, to expand.
 4085 * @param [in]      eta        Range specifier of each value.
 4086 * @param [out]     s1         First vector of polynomials.
 4087 * @param [in]      s1Len      Dimension of first vector.
 4088 * @param [out]     s2         Second vector of polynomials.
 4089 * @param [in]      s2Len      Dimension of second vector.
 4090 * @return  0 on success.
 4091 * @return  Negative on hash error.
 4092 */
 4093static int dilithium_expand_s_c(wc_Shake* shake256, byte* priv_seed, byte eta,
 4094    sword32* s1, byte s1Len, sword32* s2, byte s2Len)
 4095{
 4096    int ret = 0;
 4097    byte r;
 4098    byte seed[DILITHIUM_GEN_S_SEED_SZ];
 4099
 4100    /* Copy the seed into a buffer that has space for r. */
 4101    XMEMCPY(seed, priv_seed, DILITHIUM_PRIV_SEED_SZ);
 4102    /* Set top 8-bits of r in buffer to 0. */
 4103    seed[DILITHIUM_PRIV_SEED_SZ + 1] = 0;
 4104    /* Step 1: Each polynomial in s1. */
 4105    for (r = 0; (ret == 0) && (r < s1Len); r++) {
 4106        /* Set bottom 8-bits of r into buffer - little endian. */
 4107        seed[DILITHIUM_PRIV_SEED_SZ] = r;
 4108
 4109        /* Step 2: Generate polynomial for s1. */
 4110        ret = dilithium_rej_bound_poly(shake256, seed, s1, eta);
 4111        /* Next polynomial in s1. */
 4112        s1 += DILITHIUM_N;
 4113    }
 4114    /* Step 4: Each polynomial in s2. */
 4115    for (r = 0; (ret == 0) && (r < s2Len); r++) {
 4116        /* Set bottom 8-bits of r + l into buffer - little endian. */
 4117        seed[DILITHIUM_PRIV_SEED_SZ] = (byte)(r + s1Len);
 4118        /* Step 5: Generate polynomial for s1. */
 4119        ret = dilithium_rej_bound_poly(shake256, seed, s2, eta);
 4120        /* Next polynomial in s2. */
 4121        s2 += DILITHIUM_N;
 4122    }
 4123
 4124    return ret;
 4125}
 4126
 4127/* Expand private seed into vectors s1 and s2.
 4128 *
 4129 * @param [in, out] shake256   SHAKE-256 object.
 4130 * @param [in]      priv_seed  Private seed, rho, to expand.
 4131 * @param [in]      eta        Range specifier of each value.
 4132 * @param [out]     s1         First vector of polynomials.
 4133 * @param [in]      s1Len      Dimension of first vector.
 4134 * @param [out]     s2         Second vector of polynomials.
 4135 * @param [in]      s2Len      Dimension of second vector.
 4136 * @return  0 on success.
 4137 * @return  Negative on hash error.
 4138 */
 4139static int dilithium_expand_s(wc_Shake* shake256, byte* priv_seed, byte eta,
 4140    sword32* s1, byte s1Len, sword32* s2, byte s2Len)
 4141{
 4142    int ret = 0;
 4143
 4144#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 4145    #ifndef WOLFSSL_NO_ML_DSA_44
 4146    if ((s1Len == 4) && IS_INTEL_AVX2(cpuid_flags) &&
 4147        (SAVE_VECTOR_REGISTERS2() == 0))
 4148    {
 4149        sword32* s[2] = { s1, s2 };
 4150        ret = wc_mldsa_gen_s_4_4_avx2(s, priv_seed);
 4151        RESTORE_VECTOR_REGISTERS();
 4152    }
 4153    else
 4154    #endif
 4155    #ifndef WOLFSSL_NO_ML_DSA_65
 4156    if ((s1Len == 5) && IS_INTEL_AVX2(cpuid_flags) &&
 4157        (SAVE_VECTOR_REGISTERS2() == 0))
 4158    {
 4159        sword32* s[2] = { s1, s2 };
 4160        ret = wc_mldsa_gen_s_5_6_avx2(s, priv_seed);
 4161        RESTORE_VECTOR_REGISTERS();
 4162    }
 4163    else
 4164    #endif
 4165    #ifndef WOLFSSL_NO_ML_DSA_87
 4166    if ((s1Len == 7) && IS_INTEL_AVX2(cpuid_flags) &&
 4167        (SAVE_VECTOR_REGISTERS2() == 0))
 4168    {
 4169        sword32* s[2] = { s1, s2 };
 4170        ret = wc_mldsa_gen_s_7_8_avx2(s, priv_seed);
 4171        RESTORE_VECTOR_REGISTERS();
 4172    }
 4173    else
 4174    #endif
 4175#endif /* USE_INTEL_SPEEDUP && !WC_SHA3_NO_ASM */
 4176    {
 4177        ret = dilithium_expand_s_c(shake256, priv_seed, eta, s1, s1Len, s2,
 4178            s2Len);
 4179    }
 4180
 4181    return ret;
 4182}
 4183#endif /* !WOLFSSL_DILITHIUM_NO_MAKE_KEY */
 4184
 4185#ifndef WOLFSSL_DILITHIUM_NO_SIGN
 4186#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 4187#define SHA3_256_BYTES   (WC_SHA3_256_COUNT * 8)
 4188
 4189#ifndef WOLFSSL_NO_ML_DSA_44
 4190/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 4191 *
 4192 * Seed used with XOF to generate random bytes.
 4193 *
 4194 * @param  [out]  a           Vectos of uniform integers.
 4195 * @param  [in]   seed        Bytes to seed XOF generation.
 4196 * @return  0 on success.
 4197 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 4198 * WOLFSSL_SMALL_STACK is defined.
 4199 */
 4200static int wc_mldsa_gen_y_4_avx2(sword32* y, byte* seed, word16 kappa)
 4201{
 4202    int l;
 4203#ifdef WOLFSSL_SMALL_STACK
 4204    byte *rand = NULL;
 4205    word64 *state = NULL;
 4206#else
 4207    byte rand[4 * DILITHIUM_MAX_V];
 4208    word64 state[25 * 4];
 4209#endif
 4210
 4211#ifdef WOLFSSL_SMALL_STACK
 4212    rand = (byte*)XMALLOC(4 * DILITHIUM_MAX_V, NULL,
 4213                          DYNAMIC_TYPE_TMP_BUFFER);
 4214    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 4215                          DYNAMIC_TYPE_TMP_BUFFER);
 4216    if ((rand == NULL) || (state == NULL)) {
 4217        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4218        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4219        return MEMORY_E;
 4220    }
 4221#endif
 4222
 4223    for (l = 0; l < 4; l++) {
 4224        state[8*4 + l] = 0x1f0000U + (word32)kappa + (word32)l;
 4225    }
 4226    sha3_256_blocksx4_seed_64_avx2(state, seed);
 4227    wc_mldsa_redistribute_17_rand_avx2(state,
 4228        rand + 0 * DILITHIUM_MAX_V,
 4229        rand + 1 * DILITHIUM_MAX_V,
 4230        rand + 2 * DILITHIUM_MAX_V,
 4231        rand + 3 * DILITHIUM_MAX_V);
 4232    for (l = 1; l < DILITHIUM_MAX_V_BLOCKS; l++) {
 4233        sha3_blocksx4_avx2(state);
 4234        wc_mldsa_redistribute_17_rand_avx2(state,
 4235            rand + 0 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4236            rand + 1 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4237            rand + 2 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4238            rand + 3 * DILITHIUM_MAX_V + l * SHA3_256_BYTES);
 4239    }
 4240    wc_mldsa_decode_gamma1_17_avx2(rand + 0 * DILITHIUM_MAX_V,
 4241        y + 0 * DILITHIUM_N);
 4242    wc_mldsa_decode_gamma1_17_avx2(rand + 1 * DILITHIUM_MAX_V,
 4243        y + 1 * DILITHIUM_N);
 4244    wc_mldsa_decode_gamma1_17_avx2(rand + 2 * DILITHIUM_MAX_V,
 4245        y + 2 * DILITHIUM_N);
 4246    wc_mldsa_decode_gamma1_17_avx2(rand + 3 * DILITHIUM_MAX_V,
 4247        y + 3 * DILITHIUM_N);
 4248
 4249    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4250    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4251
 4252    return 0;
 4253}
 4254#endif
 4255
 4256#ifndef WOLFSSL_NO_ML_DSA_65
 4257/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 4258 *
 4259 * Seed used with XOF to generate random bytes.
 4260 *
 4261 * @param  [out]  a           Vectos of uniform integers.
 4262 * @param  [in]   seed        Bytes to seed XOF generation.
 4263 * @return  0 on success.
 4264 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 4265 * WOLFSSL_SMALL_STACK is defined.
 4266 */
 4267static int wc_mldsa_gen_y_5_avx2(sword32* y, byte* seed, word16 kappa,
 4268    wc_Shake* shake256)
 4269{
 4270    int ret;
 4271    int l;
 4272#ifdef WOLFSSL_SMALL_STACK
 4273    byte *rand = NULL;
 4274    word64 *state = NULL;
 4275#else
 4276    byte rand[4 * DILITHIUM_MAX_V];
 4277    word64 state[25 * 4];
 4278#endif
 4279
 4280#ifdef WOLFSSL_SMALL_STACK
 4281    rand = (byte*)XMALLOC(4 * DILITHIUM_MAX_V, NULL,
 4282                          DYNAMIC_TYPE_TMP_BUFFER);
 4283    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 4284                          DYNAMIC_TYPE_TMP_BUFFER);
 4285    if ((rand == NULL) || (state == NULL)) {
 4286        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4287        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4288        return MEMORY_E;
 4289    }
 4290#endif
 4291
 4292    /* Polynomials: 0-3 */
 4293    for (l = 0; l < 4; l++) {
 4294        state[8*4 + l] = 0x1f0000U + (word32)kappa + (word32)l;
 4295    }
 4296    sha3_256_blocksx4_seed_64_avx2(state, seed);
 4297    wc_mldsa_redistribute_17_rand_avx2(state,
 4298        rand + 0 * DILITHIUM_MAX_V,
 4299        rand + 1 * DILITHIUM_MAX_V,
 4300        rand + 2 * DILITHIUM_MAX_V,
 4301        rand + 3 * DILITHIUM_MAX_V);
 4302    for (l = 1; l < DILITHIUM_MAX_V_BLOCKS; l++) {
 4303        sha3_blocksx4_avx2(state);
 4304        wc_mldsa_redistribute_17_rand_avx2(state,
 4305            rand + 0 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4306            rand + 1 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4307            rand + 2 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4308            rand + 3 * DILITHIUM_MAX_V + l * SHA3_256_BYTES);
 4309    }
 4310    wc_mldsa_decode_gamma1_19_avx2(rand + 0 * DILITHIUM_MAX_V,
 4311        y + 0 * DILITHIUM_N);
 4312    wc_mldsa_decode_gamma1_19_avx2(rand + 1 * DILITHIUM_MAX_V,
 4313        y + 1 * DILITHIUM_N);
 4314    wc_mldsa_decode_gamma1_19_avx2(rand + 2 * DILITHIUM_MAX_V,
 4315        y + 2 * DILITHIUM_N);
 4316    wc_mldsa_decode_gamma1_19_avx2(rand + 3 * DILITHIUM_MAX_V,
 4317        y + 3 * DILITHIUM_N);
 4318
 4319    kappa = (word16)(kappa + 4);
 4320
 4321    seed[DILITHIUM_PRIV_RAND_SEED_SZ + 0] = (byte)kappa;
 4322    seed[DILITHIUM_PRIV_RAND_SEED_SZ + 1] = (byte)(kappa >> 8);
 4323    ret = dilithium_squeeze256(shake256, seed, DILITHIUM_Y_SEED_SZ, rand,
 4324        DILITHIUM_MAX_V_BLOCKS);
 4325    if (ret == 0) {
 4326        wc_mldsa_decode_gamma1_19_avx2(rand, y + 4 * DILITHIUM_N);
 4327    }
 4328
 4329    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4330    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4331
 4332    return ret;
 4333}
 4334#endif
 4335
 4336#ifndef WOLFSSL_NO_ML_DSA_87
 4337/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
 4338 *
 4339 * Seed used with XOF to generate random bytes.
 4340 *
 4341 * @param  [out]  a           Vectos of uniform integers.
 4342 * @param  [in]   seed        Bytes to seed XOF generation.
 4343 * @return  0 on success.
 4344 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
 4345 * WOLFSSL_SMALL_STACK is defined.
 4346 */
 4347static int wc_mldsa_gen_y_7_avx2(sword32* y, byte* seed, word16 kappa)
 4348{
 4349    int l;
 4350#ifdef WOLFSSL_SMALL_STACK
 4351    byte *rand = NULL;
 4352    word64 *state = NULL;
 4353#else
 4354    byte rand[4 * DILITHIUM_MAX_V];
 4355    word64 state[25 * 4];
 4356#endif
 4357
 4358#ifdef WOLFSSL_SMALL_STACK
 4359    rand = (byte*)XMALLOC(4 * DILITHIUM_MAX_V, NULL,
 4360                          DYNAMIC_TYPE_TMP_BUFFER);
 4361    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
 4362                          DYNAMIC_TYPE_TMP_BUFFER);
 4363    if ((rand == NULL) || (state == NULL)) {
 4364        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4365        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4366        return MEMORY_E;
 4367    }
 4368#endif
 4369
 4370    /* Polynomials: 0-3 */
 4371    for (l = 0; l < 4; l++) {
 4372        state[8*4 + l] = 0x1f0000U + (word32)kappa + (word32)l;
 4373    }
 4374    sha3_256_blocksx4_seed_64_avx2(state, seed);
 4375    wc_mldsa_redistribute_17_rand_avx2(state,
 4376        rand + 0 * DILITHIUM_MAX_V,
 4377        rand + 1 * DILITHIUM_MAX_V,
 4378        rand + 2 * DILITHIUM_MAX_V,
 4379        rand + 3 * DILITHIUM_MAX_V);
 4380    for (l = 1; l < DILITHIUM_MAX_V_BLOCKS; l++) {
 4381        sha3_blocksx4_avx2(state);
 4382        wc_mldsa_redistribute_17_rand_avx2(state,
 4383            rand + 0 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4384            rand + 1 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4385            rand + 2 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4386            rand + 3 * DILITHIUM_MAX_V + l * SHA3_256_BYTES);
 4387    }
 4388    wc_mldsa_decode_gamma1_19_avx2(rand + 0 * DILITHIUM_MAX_V,
 4389        y + 0 * DILITHIUM_N);
 4390    wc_mldsa_decode_gamma1_19_avx2(rand + 1 * DILITHIUM_MAX_V,
 4391        y + 1 * DILITHIUM_N);
 4392    wc_mldsa_decode_gamma1_19_avx2(rand + 2 * DILITHIUM_MAX_V,
 4393        y + 2 * DILITHIUM_N);
 4394    wc_mldsa_decode_gamma1_19_avx2(rand + 3 * DILITHIUM_MAX_V,
 4395        y + 3 * DILITHIUM_N);
 4396
 4397    kappa = (word16)(kappa + 4);
 4398
 4399    /* Polynomials: 4-7 */
 4400    for (l = 0; l < 3; l++) {
 4401        state[8*4 + l] = 0x1f0000U + (word32)kappa + (word32)l;
 4402    }
 4403    sha3_256_blocksx4_seed_64_avx2(state, seed);
 4404    wc_mldsa_redistribute_17_rand_avx2(state,
 4405        rand + 0 * DILITHIUM_MAX_V,
 4406        rand + 1 * DILITHIUM_MAX_V,
 4407        rand + 2 * DILITHIUM_MAX_V,
 4408        rand + 3 * DILITHIUM_MAX_V);
 4409    for (l = 1; l < DILITHIUM_MAX_V_BLOCKS; l++) {
 4410        sha3_blocksx4_avx2(state);
 4411        wc_mldsa_redistribute_17_rand_avx2(state,
 4412            rand + 0 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4413            rand + 1 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4414            rand + 2 * DILITHIUM_MAX_V + l * SHA3_256_BYTES,
 4415            rand + 3 * DILITHIUM_MAX_V + l * SHA3_256_BYTES);
 4416    }
 4417    wc_mldsa_decode_gamma1_19_avx2(rand + 0 * DILITHIUM_MAX_V,
 4418        y + 4 * DILITHIUM_N);
 4419    wc_mldsa_decode_gamma1_19_avx2(rand + 1 * DILITHIUM_MAX_V,
 4420        y + 5 * DILITHIUM_N);
 4421    wc_mldsa_decode_gamma1_19_avx2(rand + 2 * DILITHIUM_MAX_V,
 4422        y + 6 * DILITHIUM_N);
 4423
 4424    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4425    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
 4426
 4427    return 0;
 4428}
 4429#endif
 4430#endif
 4431
 4432/* Expand the private random seed into vector y.
 4433 *
 4434 * FIPS 204. 8.3: Algorithm 28 ExpandMask(rho, mu)
 4435 *   1: c <- 1 + bitlen(GAMMA1 - 1)
 4436 *   2: for r from 0 to l - 1 do
 4437 *   3:     n <- IntegerToBits(mu + r, 16)
 4438 *   4:     v <- (H(rho||n)[[32rc]], H(rho||n)[[32rc + 1]], ...,
 4439 *                H(rho||n)[[32rc + 32c - 1]])
 4440 *   5:     s[r] <- BitUnpack(v, GAMMA-1, GAMMA1)
 4441 *   6: end for
 4442 *   7: return s
 4443 *
 4444 * @param [in, out] shake256     SHAKE-256 object.
 4445 * @param [in, out] seed         Buffer containing seed to expand.
 4446 *                               Has space for two bytes to be appended.
 4447 * @param [in]      kappa        Base value to append to seed.
 4448 * @param [in]      gamma1_bits  Number of bits per value.
 4449 * @param [out]     y            Vector of polynomials.
 4450 * @param [in]      l            Dimension of vector.
 4451 * @return  0 on success.
 4452 * @return  Negative on hash error.
 4453 */
 4454static int dilithium_vec_expand_mask_c(wc_Shake* shake256, byte* seed,
 4455    word16 kappa, byte gamma1_bits, sword32* y, byte l)
 4456{
 4457    int ret = 0;
 4458    byte r;
 4459    WC_DECLARE_VAR(v, byte, DILITHIUM_MAX_V, NULL);
 4460
 4461    WC_ALLOC_VAR_EX(v, byte, DILITHIUM_MAX_V, NULL, DYNAMIC_TYPE_DILITHIUM,
 4462                    return MEMORY_E);
 4463
 4464    /* Step 2: For each polynomial of vector. */
 4465    for (r = 0; (ret == 0) && (r < l); r++) {
 4466        /* Step 3: Calculate value to append to seed. */
 4467        word16 n = (word16)(kappa + r);
 4468
 4469        /* Step 4: Append to seed and squeeze out data. */
 4470        seed[DILITHIUM_PRIV_RAND_SEED_SZ + 0] = (byte)n;
 4471        seed[DILITHIUM_PRIV_RAND_SEED_SZ + 1] = (byte)(n >> 8);
 4472        ret = dilithium_squeeze256(shake256, seed, DILITHIUM_Y_SEED_SZ, v,
 4473            DILITHIUM_MAX_V_BLOCKS);
 4474        if (ret == 0) {
 4475            /* Decode v into polynomial. */
 4476            dilithium_decode_gamma1(v, gamma1_bits, y);
 4477            /* Next polynomial. */
 4478            y += DILITHIUM_N;
 4479        }
 4480    }
 4481
 4482    WC_FREE_VAR_EX(v, NULL, DYNAMIC_TYPE_DILITHIUM);
 4483    return ret;
 4484}
 4485
 4486/* Expand the private random seed into vector y.
 4487 *
 4488 * @param [in, out] shake256     SHAKE-256 object.
 4489 * @param [in, out] seed         Buffer containing seed to expand.
 4490 *                               Has space for two bytes to be appended.
 4491 * @param [in]      kappa        Base value to append to seed.
 4492 * @param [in]      gamma1_bits  Number of bits per value.
 4493 * @param [out]     y            Vector of polynomials.
 4494 * @param [in]      l            Dimension of vector.
 4495 * @return  0 on success.
 4496 * @return  Negative on hash error.
 4497 */
 4498static int dilithium_vec_expand_mask(wc_Shake* shake256, byte* seed,
 4499    word16 kappa, byte gamma1_bits, sword32* y, byte l)
 4500{
 4501    int ret = 0;
 4502
 4503#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
 4504    if (IS_INTEL_AVX2(cpuid_flags) && IS_INTEL_BMI2(cpuid_flags) &&
 4505            (SAVE_VECTOR_REGISTERS2() == 0)) {
 4506    #ifndef WOLFSSL_NO_ML_DSA_44
 4507        if (l == 4) {
 4508            ret = wc_mldsa_gen_y_4_avx2(y, seed, kappa);
 4509        }
 4510    #endif
 4511    #ifndef WOLFSSL_NO_ML_DSA_65
 4512        if (l == 5) {
 4513            ret = wc_mldsa_gen_y_5_avx2(y, seed, kappa, shake256);
 4514        }
 4515    #endif
 4516    #ifndef WOLFSSL_NO_ML_DSA_87
 4517        if (l == 7) {
 4518            ret = wc_mldsa_gen_y_7_avx2(y, seed, kappa);
 4519        }
 4520    #endif
 4521        RESTORE_VECTOR_REGISTERS();
 4522    }
 4523    else
 4524#endif
 4525    {
 4526        ret = dilithium_vec_expand_mask_c(shake256, seed, kappa, gamma1_bits, y,
 4527            l);
 4528    }
 4529
 4530    return ret;
 4531}
 4532#endif
 4533
 4534#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
 4535/* Expand commit to a polynomial.
 4536 *
 4537 * FIPS 204. 8.3: Algorithm 23 SampleInBall(rho)
 4538 *   1: c <- 0
 4539 *   2: k <- 8
 4540 *   3: for i from 256 - TAU to 255 do
 4541 *   4:     while H(rho)[[k]] > i do
 4542 *   5:        k <- k + 1
 4543 *   6:     end while
 4544 *   7:     j <- H(rho)[[k]]
 4545 *   8:     c[i] <- c[j]
 4546 *   9:     c[j] <- (-1)^H(rho)[i+TAU-256]
 4547 *  10:     k <- k + 1
 4548 *  11: end for
 4549 *  12: return c
 4550 *
 4551 * @param [in]  shake256  SHAKE-256 object.
 4552 * @param [in]  seed      Buffer containing seed to expand.
 4553 * @param [in]  seedLen   Length of seed in bytes.
 4554 * @param [in]  tau       Number of +/- 1s in polynomial.
 4555 * @param [out] c         Commit polynomial.
 4556 * @param [in]  block     Memory to use for block from key.
 4557 * @return  0 on success.
 4558 * @return  Negative on hash error.
 4559 */
 4560static int dilithium_sample_in_ball_ex(int level, wc_Shake* shake256,
 4561   const byte* seed, word32 seedLen, byte tau, sword32* c, byte* block)
 4562{
 4563#ifndef USE_INTEL_SPEEDUP
 4564    int ret = 0;
 4565    byte signs[DILITHIUM_SIGN_BYTES];
 4566    unsigned int i;
 4567    /* Step 1: Initialize sign bit index. */
 4568    unsigned int s = 0;
 4569    /* Step 2: First 8 bytes are used for sign. */
 4570    unsigned int k = DILITHIUM_SIGN_BYTES;
 4571
 4572    /* Set polynomial to all zeros. */
 4573    XMEMSET(c, 0, DILITHIUM_POLY_SIZE);
 4574
 4575    /* Generate a block of data from seed. */
 4576#ifdef WOLFSSL_DILITHIUM_FIPS204_DRAFT
 4577    if (level >= WC_ML_DSA_DRAFT) {
 4578        ret = dilithium_shake256(shake256, seed, DILITHIUM_SEED_SZ, block,
 4579            DILITHIUM_GEN_C_BLOCK_BYTES);
 4580    }
 4581    else
 4582#endif
 4583    {
 4584        (void)level;
 4585        ret = dilithium_shake256(shake256, seed, seedLen, block,
 4586            DILITHIUM_GEN_C_BLOCK_BYTES);
 4587    }
 4588    if (ret == 0) {
 4589        /* Copy first 8 bytes of first hash block as random sign bits. */
 4590        XMEMCPY(signs, block, DILITHIUM_SIGN_BYTES);
 4591    }
 4592
 4593    /* Step 3: Put in TAU +/- 1s. */
 4594    for (i = (unsigned int)DILITHIUM_N - tau;
 4595                                         (ret == 0) && (i < DILITHIUM_N); i++) {
 4596        unsigned int j;
 4597        do {
 4598            /* Check whether block is exhausted. */
 4599            if (k == DILITHIUM_GEN_C_BLOCK_BYTES) {
 4600                /* Generate a new block. */
 4601                ret = wc_Shake256_SqueezeBlocks(shake256, block, 1);
 4602
 4603                /* Restart hash block index. */
 4604                k = 0;
 4605            }
 4606            /* Step 7: Get random byte from block as index.
 4607             * Step 5 and 10: Increment hash block index.
 4608             */
 4609            j = block[k++];
 4610        }
 4611        /* Step 4: Get another random if random index is a future swap index. */
 4612        while ((ret == 0) && (j > i));
 4613
 4614        /* Step 8: Move value from random index to current index. */
 4615        c[i] = c[j];
 4616        /* Step 9: Set value at random index to +/- 1. */
 4617        c[j] = 1 - ((((signs[s >> 3]) >> (s & 0x7)) & 0x1) << 1);
 4618        /* Next sign bit index. */
 4619        s++;
 4620    }
 4621
 4622    return ret;
 4623#else
 4624    int ret = 0;
 4625    /* SHA-3 state. */
 4626    word64* state = shake256->s;
 4627    word64 signs;
 4628    unsigned int i;
 4629    /* Step 2: First 8 bytes are used for sign. */
 4630    unsigned int k = DILITHIUM_SIGN_BYTES;
 4631
 4632    block = (byte*)state;
 4633
 4634    /* Set polynomial to all zeros. */
 4635    XMEMSET(c, 0, DILITHIUM_POLY_SIZE);
 4636
 4637    /* Generate a block of data from seed. */
 4638#ifdef WOLFSSL_DILITHIUM_FIPS204_DRAFT
 4639    if (level >= WC_ML_DSA_DRAFT) {
 4640        ret = dilithium_shake256(shake256, seed, DILITHIUM_SEED_SZ, block,
 4641            DILITHIUM_GEN_C_BLOCK_BYTES);
 4642    }
 4643    else
 4644#endif
 4645    {
 4646        (void)level;
 4647        ret = dilithium_shake256(shake256, seed, seedLen, block,
 4648            DILITHIUM_GEN_C_BLOCK_BYTES);
 4649    }
 4650    if (ret == 0) {
 4651        /* Step 1: Initialize sign bit index. */
 4652        /* Copy first 8 bytes of first hash block as random sign bits. */
 4653        signs = *(word64*)block;
 4654
 4655        /* Step 3: Put in TAU +/- 1s. */
 4656        for (i = (unsigned int)DILITHIUM_N - tau; i < DILITHIUM_N; i++) {
 4657            unsigned int j;
 4658            do {
 4659                /* Check whether block is exhausted. */
 4660                if (k == DILITHIUM_GEN_C_BLOCK_BYTES) {
 4661                    /* Generate a new block. */
 4662#ifndef WC_SHA3_NO_ASM
 4663                    if (IS_INTEL_AVX2(cpuid_flags) &&
 4664                             (SAVE_VECTOR_REGISTERS2() == 0)) {
 4665                        sha3_block_avx2(state);
 4666                        RESTORE_VECTOR_REGISTERS();
 4667                    }
 4668                    else if (IS_INTEL_BMI2(cpuid_flags)) {
 4669                        sha3_block_bmi2(state);
 4670                    }
 4671                    else
 4672#endif
 4673                    {
 4674                        BlockSha3(state);
 4675                    }
 4676
 4677                    /* Restart hash block index. */
 4678                    k = 0;
 4679                }
 4680                /* Step 7: Get random byte from block as index.
 4681                 * Step 5 and 10: Increment hash block index.
 4682                 */
 4683                j = block[k++];
 4684            }
 4685            /* Step 4: Get another random if random index is a future swap
 4686             * index. */
 4687            while (j > i);
 4688
 4689            /* Step 8: Move value from random index to current index. */
 4690            c[i] = c[j];
 4691            /* Step 9: Set value at random index to +/- 1. */
 4692            c[j] = (sword32)1 - (sword32)((sword32)(signs & 0x1) << 1);
 4693            /* Next sign bit index. */
 4694            signs >>= 1;
 4695        }
 4696    }
 4697
 4698    return ret;
 4699#endif
 4700}
 4701
 4702#if (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 4703     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM)) || \
 4704    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 4705     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))
 4706/* Expand commit to a polynomial.
 4707 *
 4708 * @param [in]  shake256  SHAKE-256 object.
 4709 * @param [in]  seed      Buffer containing seed to expand.
 4710 * @param [in]  seedLen   Length of seed in bytes.
 4711 * @param [in]  tau       Number of +/- 1s in polynomial.
 4712 * @param [out] c         Commit polynomial.
 4713 * @param [in]  heap      Dynamic memory hint.
 4714 * @return  0 on success.
 4715 * @return  MEMORY_E when dynamic memory allocation fails.
 4716 * @return  Negative on hash error.
 4717 */
 4718static int dilithium_sample_in_ball(int level, wc_Shake* shake256,
 4719   const byte* seed, word32 seedLen, byte tau, sword32* c, void* heap)
 4720{
 4721    int ret = 0;
 4722#if defined(WOLFSSL_SMALL_STACK)
 4723    byte* block = NULL;
 4724#else
 4725    byte block[DILITHIUM_GEN_C_BLOCK_BYTES];
 4726#endif
 4727
 4728    (void)heap;
 4729
 4730#if defined(WOLFSSL_SMALL_STACK)
 4731    block = (byte*)XMALLOC(DILITHIUM_GEN_C_BLOCK_BYTES, heap,
 4732        DYNAMIC_TYPE_DILITHIUM);
 4733    if (block == NULL) {
 4734        ret = MEMORY_E;
 4735    }
 4736#endif
 4737
 4738    if (ret == 0) {
 4739        ret = dilithium_sample_in_ball_ex(level, shake256, seed, seedLen, tau,
 4740            c, block);
 4741    }
 4742
 4743#if defined(WOLFSSL_SMALL_STACK)
 4744    XFREE(block, heap, DYNAMIC_TYPE_DILITHIUM);
 4745#endif
 4746    return ret;
 4747}
 4748#endif
 4749
 4750#endif
 4751
 4752/******************************************************************************
 4753 * Decompose operations
 4754 ******************************************************************************/
 4755
 4756#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
 4757#ifndef WOLFSSL_NO_ML_DSA_44
 4758/* Decompose value into high and low based on GAMMA2 being ((q-1) / 88).
 4759 *
 4760 * FIPS 204. 8.4: Algorithm 30 Decompose(r)
 4761 *   1: r+ <- r mod q
 4762 *   2: r0 <- r+ mod+/- (2 * GAMMA2)
 4763 *   3: if r+ - r0 = q - 1 then
 4764 *   4:     r1 <- 0
 4765 *   5:     r0 <- r0 - 1
 4766 *   6: else r1 <- (r+ - r0) / (2 * GAMMA2)
 4767 *   7: end if
 4768 *   8: return (r1, r0)
 4769 *
 4770 * DILITHIUM_Q_LOW_88_2 = 0x2e800 = 0b101110100000000000
 4771 * t1 * DILITHIUM_Q_LOW_88_2 = (t1 << 18) - (t1 << 16) - (t1 << 12) - (t1 << 11)
 4772 *                           = ((93 * t1) << 11)
 4773 * Nothing faster than straight multiply.
 4774 *
 4775 * Implementation using Barrett Reduction.
 4776 *
 4777 * @param [in]  r   Value to decompose.
 4778 * @param [out] r0  Low bits.
 4779 * @param [out] r1  High bits.
 4780 */
 4781static void dilithium_decompose_q88(sword32 r, sword32* r0, sword32* r1)
 4782{
 4783    sword32 t0;
 4784    sword32 t1;
 4785#ifdef DILITHIUM_MUL_SLOW
 4786    sword32 t2;
 4787#endif
 4788
 4789    /* Roundup r and calculate approx high value. */
 4790#if !defined(DILITHIUM_MUL_44_SLOW)
 4791    t1 = ((r * 44) + ((DILITHIUM_Q_LOW_88 - 1) * 44)) >> 23;
 4792#elif !defined(DILITHIUM_MUL_11_SLOW)
 4793    t1 = ((r * 11) + ((DILITHIUM_Q_LOW_88 - 1) * 11)) >> 21;
 4794#else
 4795    t0 = r + DILITHIUM_Q_LOW_88 - 1;
 4796    t1 = ((t0 << 3) + (t0 << 1) + t0) >> 21;
 4797#endif
 4798    /* Calculate approx low value. */
 4799    t0 = r - (t1 * DILITHIUM_Q_LOW_88_2);
 4800#ifndef DILITHIUM_MUL_SLOW
 4801    /* Calculate real high value, When t0 > modulus, +1 to approx high value.
 4802     * Sign-extraction trick: each ((word32)A - (word32)B) >> 31 evaluates
 4803     * to 1 if (signed) A - B was negative thanks to two's-complement
 4804     * wrap-around. The outer (sword32) cast lets us add the 0/1 flag back into
 4805     * a sword32 accumulator without -Wsign- conversion firing on the implicit
 4806     * narrowing. */
 4807    t1 += (sword32)(((word32)DILITHIUM_Q_LOW_88 - (word32)t0) >> 31);
 4808    /* Calculate real low value. */
 4809    t0 = r - (t1 * DILITHIUM_Q_LOW_88_2);
 4810#else
 4811    /* Calculate real high value, When t0 > modulus, +1 to approx high value. */
 4812    t2 = (sword32)(((word32)DILITHIUM_Q_LOW_88 - (word32)t0) >> 31);
 4813    t1 += t2;
 4814    /* Calculate real low value. */
 4815    t0 -= (sword32)((word32)(0 - t2) & DILITHIUM_Q_LOW_88_2);
 4816#endif
 4817    /* -1 from low value if high value is 44. Was 43 but low is negative. */
 4818    t0 -= (sword32)(((word32)43 - (word32)t1) >> 31);
 4819    /* When high value is 44, too large, set to 0. */
 4820    t1 &= (sword32)(0U - (((word32)t1 - 44U) >> 31));
 4821
 4822    *r0 = t0;
 4823    *r1 = t1;
 4824}
 4825#endif
 4826
 4827#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 4828/* Decompose value into high and low based on GAMMA2 being ((q-1) / 32).
 4829 *
 4830 * FIPS 204. 8.4: Algorithm 30 Decompose(r)
 4831 *   1: r+ <- r mod q
 4832 *   2: r0 <- r+ mod+/- (2 * GAMMA2)
 4833 *   3: if r+ - r0 = q - 1 then
 4834 *   4:     r1 <- 0
 4835 *   5:     r0 <- r0 - 1
 4836 *   6: else r1 <- (r+ - r0) / (2 * GAMMA2)
 4837 *   7: end if
 4838 *   8: return (r1, r0)
 4839 *
 4840 * DILITHIUM_Q_LOW_32_2 = 0x7fe00 = 0b1111111111000000000
 4841 * t1 * DILITHIUM_Q_LOW_32_2 = (t1 << 19) - (t1 << 9)
 4842 *
 4843 * Implementation using Barrett Reduction.
 4844 *
 4845 * @param [in]  r   Value to decompose.
 4846 * @param [out] r0  Low bits.
 4847 * @param [out] r1  High bits.
 4848 */
 4849static void dilithium_decompose_q32(sword32 r, sword32* r0, sword32* r1)
 4850{
 4851    sword32 t0;
 4852    sword32 t1;
 4853
 4854    /* Roundup r and calculate approx high value. */
 4855    t1 = (r + DILITHIUM_Q_LOW_32 - 1) >> 19;
 4856    /* Calculate approx low value. */
 4857    t0 = r - (t1 << 19) + (t1 << 9);
 4858    /* Calculate real high value, When t0 > modulus, +1 to approx high value. */
 4859    t1 += (sword32)(((word32)DILITHIUM_Q_LOW_32 - (word32)t0) >> 31);
 4860    /* Calculate real low value. */
 4861    t0 = r - (t1 << 19) + (t1 << 9);
 4862    /* -1 from low value if high value is 16. Was 15 but low is negative. */
 4863    t0 -= t1 >> 4;
 4864    /* When high value is 16, too large, set to 0. */
 4865    t1 &= 0xf;
 4866
 4867    *r0 = t0;
 4868    *r1 = t1;
 4869}
 4870#endif
 4871#endif
 4872
 4873#ifndef WOLFSSL_DILITHIUM_NO_SIGN
 4874
 4875#if !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 4876    defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A)
 4877/* Decompose vector of polynomials into high and low based on GAMMA2.
 4878 *
 4879 * @param [in]  r       Vector of polynomials to decompose.
 4880 * @param [in]  k       Dimension of vector.
 4881 * @param [in]  gamma2  Low-order rounding range, GAMMA2.
 4882 * @param [out] r0      Low parts in vector of polynomials.
 4883 * @param [out] r1      High parts in vector of polynomials.
 4884 */
 4885static void dilithium_vec_decompose_c(const sword32* r, byte k, sword32 gamma2,
 4886    sword32* r0, sword32* r1)
 4887{
 4888    unsigned int i;
 4889    unsigned int j;
 4890
 4891#ifndef WOLFSSL_NO_ML_DSA_44
 4892    if (gamma2 == DILITHIUM_Q_LOW_88) {
 4893        /* For each polynomial of vector. */
 4894        for (i = 0; i < k; i++) {
 4895            /* For each value of polynomial. */
 4896            for (j = 0; j < DILITHIUM_N; j++) {
 4897                /* Decompose value into two vectors. */
 4898                dilithium_decompose_q88(r[j], &r0[j], &r1[j]);
 4899            }
 4900            /* Next polynomial of vectors. */
 4901            r += DILITHIUM_N;
 4902            r0 += DILITHIUM_N;
 4903            r1 += DILITHIUM_N;
 4904        }
 4905    }
 4906#endif
 4907#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 4908    if (gamma2 == DILITHIUM_Q_LOW_32) {
 4909        /* For each polynomial of vector. */
 4910        for (i = 0; i < k; i++) {
 4911            /* For each value of polynomial. */
 4912            for (j = 0; j < DILITHIUM_N; j++) {
 4913                /* Decompose value into two vectors. */
 4914                dilithium_decompose_q32(r[j], &r0[j], &r1[j]);
 4915            }
 4916            /* Next polynomial of vectors. */
 4917            r += DILITHIUM_N;
 4918            r0 += DILITHIUM_N;
 4919            r1 += DILITHIUM_N;
 4920        }
 4921    }
 4922#endif
 4923}
 4924
 4925/* Decompose vector of polynomials into high and low based on GAMMA2.
 4926 *
 4927 * @param [in]  r       Vector of polynomials to decompose.
 4928 * @param [in]  k       Dimension of vector.
 4929 * @param [in]  gamma2  Low-order rounding range, GAMMA2.
 4930 * @param [out] r0      Low parts in vector of polynomials.
 4931 * @param [out] r1      High parts in vector of polynomials.
 4932 */
 4933static void dilithium_vec_decompose(const sword32* r, byte k, sword32 gamma2,
 4934    sword32* r0, sword32* r1)
 4935{
 4936#ifdef USE_INTEL_SPEEDUP
 4937    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 4938    #ifndef WOLFSSL_NO_ML_DSA_44
 4939        if (gamma2 == DILITHIUM_Q_LOW_88) {
 4940            wc_mldsa_decompose_q88_avx2(r, r0, r1);
 4941        }
 4942    #endif
 4943    #if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 4944        if (gamma2 == DILITHIUM_Q_LOW_32) {
 4945            wc_mldsa_decompose_q32_avx2(r, k, r0, r1);
 4946        }
 4947    #endif
 4948        RESTORE_VECTOR_REGISTERS();
 4949    }
 4950    else
 4951#endif
 4952    {
 4953        dilithium_vec_decompose_c(r, k, gamma2, r0, r1);
 4954    }
 4955}
 4956#endif
 4957
 4958#endif /* !WOLFSSL_DILITHIUM_NO_SIGN */
 4959
 4960/******************************************************************************
 4961 * Range check operation
 4962 ******************************************************************************/
 4963
 4964#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
 4965/* Check that the values of the polynomial are in range.
 4966 *
 4967 * Many places in FIPS 204. One example from Algorithm 2:
 4968 *   23:    if ||z||inf >= GAMMA1 - BETA or ..., then (z, h) = falsam
 4969 *
 4970 * @param [in] a   Polynomial.
 4971 * @param [in] hi  Largest value in range.
 4972 */
 4973static int dilithium_check_low(const sword32* a, sword32 hi)
 4974{
 4975    int ret = 1;
 4976    unsigned int j;
 4977    /* Calculate lowest range value. */
 4978    sword32 nhi = -hi;
 4979
 4980    /* For each value of polynomial. */
 4981    for (j = 0; j < DILITHIUM_N; j++) {
 4982        /* Check range is -(hi-1)..(hi-1). */
 4983        if ((a[j] <= nhi) || (a[j] >= hi)) {
 4984            /* Check failed. */
 4985            ret = 0;
 4986            break;
 4987        }
 4988    }
 4989
 4990    return ret;
 4991}
 4992
 4993#if !defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 4994    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 4995     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 4996/* Check that the values of the vector are in range.
 4997 *
 4998 * Many places in FIPS 204. One example from Algorithm 2:
 4999 *   23:    if ||z||inf >= GAMMA1 - BETA or ..., then (z, h) = falsam
 5000 *
 5001 * @param [in] a   Vector of polynomials.
 5002 * @param [in] l   Dimension of vector.
 5003 * @param [in] hi  Largest value in range.
 5004 */
 5005static int dilithium_vec_check_low_c(const sword32* a, byte l, sword32 hi)
 5006{
 5007    int ret = 1;
 5008    unsigned int i;
 5009
 5010    /* For each polynomial of vector. */
 5011    for (i = 0; (ret == 1) && (i < l); i++) {
 5012        ret = dilithium_check_low(a, hi);
 5013        if (ret == 0) {
 5014            break;
 5015        }
 5016        /* Next polynomial. */
 5017        a += DILITHIUM_N;
 5018    }
 5019
 5020    return ret;
 5021}
 5022#endif
 5023
 5024/* Check that the values of the vector are in range.
 5025 *
 5026 * @param [in] a   Vector of polynomials.
 5027 * @param [in] l   Dimension of vector.
 5028 * @param [in] hi  Largest value in range.
 5029 */
 5030static int dilithium_vec_check_low(const sword32* a, byte l, sword32 hi)
 5031{
 5032    int ret;
 5033#ifdef USE_INTEL_SPEEDUP
 5034    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 5035        ret = wc_mldsa_vec_check_low_avx2(a, l, hi);
 5036        RESTORE_VECTOR_REGISTERS();
 5037    }
 5038    else
 5039#endif
 5040    {
 5041        ret = dilithium_vec_check_low_c(a, l, hi);
 5042    }
 5043
 5044    return ret;
 5045}
 5046#endif
 5047
 5048/******************************************************************************
 5049 * Hint operations
 5050 ******************************************************************************/
 5051
 5052#ifndef WOLFSSL_DILITHIUM_NO_SIGN
 5053
 5054#ifndef WOLFSSL_NO_ML_DSA_44
 5055/* Compute hints indicating whether adding ct0 to w alters high bits of w.
 5056 *
 5057 * FIPS 204. 6: Algorithm 2 ML-DSA.Sign(sk, M)
 5058 *   ...
 5059 *  26: h <- MakeHint(-<<ct0>>, w - <<sc2>> + <<ct0>>)
 5060 *  27: if ... or the number of 1's in h is greater than OMEGA, then
 5061 *      (z, h) <- falsam
 5062 *   ...
 5063 *  32: sigma <- sigEncode(c_tilda, z mod+/- q, h)
 5064 *   ...
 5065 *
 5066 * FIPS 204. 8.4: Algorithm 33 MakeHint(z, r)
 5067 *   1: r1 <- HighBits(r)
 5068 *   2: v1 <- HightBits(r+z)
 5069 *   3: return [[r1 != v1]]
 5070 *
 5071 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 5072 *   ...
 5073 *   5: sigma <- sigma || HintBitPack(h)
 5074 *   ...
 5075 *
 5076 * FIPS 204. 8.1: Algorithm 14 HintBitPack(h)
 5077 *   ...
 5078 *   4:     for j from 0 to 255 do
 5079 *   5:         if h[i]j != 0 then
 5080 *   6:             y[Index] <- j
 5081 *   7:             Index <- Index + 1
 5082 *   8:         end if
 5083 *   9:     end for
 5084 *   ...
 5085 *
 5086 * @param [in]      s       Vector of polynomials that is sum of ct0 and w0.
 5087 * @param [in]      w1      Vector of polynomials that is high part of w.
 5088 * @param [out]     h       Encoded hints.
 5089 * @param [in, out] idxp    Index to write next hint into.
 5090 * return  Number of hints on success.
 5091 * return  Falsam of -1 when too many hints.
 5092 */
 5093static int dilithium_make_hint_88(const sword32* s, const sword32* w1, byte* h,
 5094    byte *idxp)
 5095{
 5096    unsigned int j;
 5097    byte idx = *idxp;
 5098
 5099    /* Alg 14, Step 3: For each value of polynomial. */
 5100    for (j = 0; j < DILITHIUM_N; j++) {
 5101        /* Alg 14, Step 4: Check whether hint is required.
 5102         * Did sum end up greater than low modulus or
 5103         * sum end up less than the negative of low modulus or
 5104         * sum is the negative of the low modulus and w1 is not zero,
 5105         * then w1 will be modified.
 5106         */
 5107        if ((s[j] > (sword32)DILITHIUM_Q_LOW_88) ||
 5108                (s[j] < -(sword32)DILITHIUM_Q_LOW_88) ||
 5109                ((s[j] == -(sword32)DILITHIUM_Q_LOW_88) &&
 5110                 (w1[j] != 0))) {
 5111            /* Alg 14, Step 6, 7: Put index as hint modifier. */
 5112            h[idx++] = (byte)j;
 5113            /* Alg 2, Step 27: If there are too many hints, return
 5114             *                 falsam of -1. */
 5115            if (idx > PARAMS_ML_DSA_44_OMEGA) {
 5116                return -1;
 5117            }
 5118        }
 5119    }
 5120
 5121    *idxp = idx;
 5122    return 0;
 5123}
 5124#endif
 5125#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 5126/* Compute hints indicating whether adding ct0 to w alters high bits of w.
 5127 *
 5128 * FIPS 204. 6: Algorithm 2 ML-DSA.Sign(sk, M)
 5129 *   ...
 5130 *  26: h <- MakeHint(-<<ct0>>, w - <<sc2>> + <<ct0>>)
 5131 *  27: if ... or the number of 1's in h is greater than OMEGA, then
 5132 *      (z, h) <- falsam
 5133 *   ...
 5134 *  32: sigma <- sigEncode(c_tilda, z mod+/- q, h)
 5135 *   ...
 5136 *
 5137 * FIPS 204. 8.4: Algorithm 33 MakeHint(z, r)
 5138 *   1: r1 <- HighBits(r)
 5139 *   2: v1 <- HightBits(r+z)
 5140 *   3: return [[r1 != v1]]
 5141 *
 5142 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 5143 *   ...
 5144 *   5: sigma <- sigma || HintBitPack(h)
 5145 *   ...
 5146 *
 5147 * FIPS 204. 8.1: Algorithm 14 HintBitPack(h)
 5148 *   ...
 5149 *   4:     for j from 0 to 255 do
 5150 *   5:         if h[i]j != 0 then
 5151 *   6:             y[Index] <- j
 5152 *   7:             Index <- Index + 1
 5153 *   8:         end if
 5154 *   9:     end for
 5155 *   ...
 5156 *
 5157 * @param [in]      s       Vector of polynomials that is sum of ct0 and w0.
 5158 * @param [in]      w1      Vector of polynomials that is high part of w.
 5159 * @param [in]      omega   Maximum number of hints allowed.
 5160 * @param [out]     h       Encoded hints.
 5161 * @param [in, out] idxp    Index to write next hint into.
 5162 * return  Number of hints on success.
 5163 * return  Falsam of -1 when too many hints.
 5164 */
 5165static int dilithium_make_hint_32(const sword32* s, const sword32* w1,
 5166    byte omega, byte* h, byte *idxp)
 5167{
 5168    unsigned int j;
 5169    byte idx = *idxp;
 5170
 5171    (void)omega;
 5172
 5173    /* Alg 14, Step 3: For each value of polynomial. */
 5174    for (j = 0; j < DILITHIUM_N; j++) {
 5175        /* Alg 14, Step 4: Check whether hint is required.
 5176         * Did sum end up greater than low modulus or
 5177         * sum end up less than the negative of low modulus or
 5178         * sum is the negative of the low modulus and w1 is not zero,
 5179         * then w1 will be modified.
 5180         */
 5181        if ((s[j] > (sword32)DILITHIUM_Q_LOW_32) ||
 5182                (s[j] < -(sword32)DILITHIUM_Q_LOW_32) ||
 5183                ((s[j] == -(sword32)DILITHIUM_Q_LOW_32) &&
 5184                 (w1[j] != 0))) {
 5185            /* Alg 14, Step 6, 7: Put index as hint modifier. */
 5186            h[idx++] = (byte)j;
 5187            /* Alg 2, Step 27: If there are too many hints, return
 5188             *                 falsam of -1. */
 5189            if (idx > omega) {
 5190                return -1;
 5191            }
 5192        }
 5193    }
 5194
 5195    *idxp = idx;
 5196    return 0;
 5197}
 5198#endif
 5199
 5200#ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM
 5201/* Compute hints indicating whether adding ct0 to w alters high bits of w.
 5202 *
 5203 * FIPS 204. 6: Algorithm 2 ML-DSA.Sign(sk, M)
 5204 *   ...
 5205 *  26: h <- MakeHint(-<<ct0>>, w - <<sc2>> + <<ct0>>)
 5206 *  27: if ... or the number of 1's in h is greater than OMEGA, then
 5207 *      (z, h) <- falsam
 5208 *   ...
 5209 *  32: sigma <- sigEncode(c_tilda, z mod+/- q, h)
 5210 *   ...
 5211 *
 5212 * FIPS 204. 8.4: Algorithm 33 MakeHint(z, r)
 5213 *   1: r1 <- HighBits(r)
 5214 *   2: v1 <- HightBits(r+z)
 5215 *   3: return [[r1 != v1]]
 5216 *
 5217 * FIPS 204. 8.2: Algorithm 20 sigEncode(c_tilde, z, h)
 5218 *   ...
 5219 *   5: sigma <- sigma || HintBitPack(h)
 5220 *   ...
 5221 *
 5222 * FIPS 204. 8.1: Algorithm 14 HintBitPack(h)
 5223 *   ...
 5224 *   2: Index <- 0
 5225 *   3. for i from 0 to k - 1 do
 5226 *   4:     for j from 0 to 255 do
 5227 *   5:         if h[i]j != 0 then
 5228 *   6:             y[Index] <- j
 5229 *   7:             Index <- Index + 1
 5230 *   8:         end if
 5231 *   9:     end for
 5232 *  10:     y[OMEGA + i] <- Index
 5233 *  11: end for
 5234 *  12: return y
 5235 *
 5236 * @param [in]  s       Vector of polynomials that is sum of ct0 and w0.
 5237 * @param [in]  w1      Vector of polynomials that is high part of w.
 5238 * @param [in]  k       Dimension of vectors.
 5239 * @param [in]  gamma2  Low-order rounding range, GAMMA2.
 5240 * @param [in]  omega   Maximum number of hints allowed.
 5241 * @param [out] h       Encoded hints.
 5242 * return  Number of hints on success.
 5243 * return  Falsam of -1 when too many hints.
 5244 */
 5245static int dilithium_make_hint(const sword32* s, const sword32* w1, byte k,
 5246    sword32 gamma2, byte omega, byte* h)
 5247{
 5248    unsigned int i;
 5249    byte idx = 0;
 5250
 5251    (void)k;
 5252    (void)omega;
 5253
 5254#ifndef WOLFSSL_NO_ML_DSA_44
 5255    if (gamma2 == DILITHIUM_Q_LOW_88) {
 5256        /* Alg 14, Step 2: For each polynomial of vector. */
 5257        for (i = 0; i < PARAMS_ML_DSA_44_K; i++) {
 5258            if (dilithium_make_hint_88(s, w1, h, &idx) == -1) {
 5259                return -1;
 5260            }
 5261            /* Alg 14, Step 10: Store count of hints for polynomial at end of
 5262             *                  list. */
 5263            h[PARAMS_ML_DSA_44_OMEGA + i] = idx;
 5264            /* Next polynomial. */
 5265            s  += DILITHIUM_N;
 5266            w1 += DILITHIUM_N;
 5267        }
 5268    }
 5269    else
 5270#endif
 5271#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 5272    if (gamma2 == DILITHIUM_Q_LOW_32) {
 5273        /* Alg 14, Step 2: For each polynomial of vector. */
 5274        for (i = 0; i < k; i++) {
 5275            if (dilithium_make_hint_32(s, w1, omega, h, &idx) == -1) {
 5276                return -1;
 5277            }
 5278            /* Alg 14, Step 10: Store count of hints for polynomial at end of
 5279             *                  list. */
 5280            h[omega + i] = idx;
 5281            /* Next polynomial. */
 5282            s  += DILITHIUM_N;
 5283            w1 += DILITHIUM_N;
 5284        }
 5285    }
 5286    else
 5287#endif
 5288    {
 5289    }
 5290
 5291    /* Set remaining hints to zero. */
 5292    XMEMSET(h + idx, 0, (size_t)(omega - idx));
 5293    return idx;
 5294}
 5295#endif /* !WOLFSSL_DILITHIUM_SIGN_SMALL_MEM */
 5296
 5297#endif /* !WOLFSSL_DILITHIUM_NO_SIGN */
 5298
 5299#ifndef WOLFSSL_DILITHIUM_NO_VERIFY
 5300/* Check that the hints are valid.
 5301 *
 5302 * @param [in] h      Hints to check
 5303 * @param [in] k      Dimension of vector.
 5304 * @param [in] omega  Maximum number of hints. Hint counts after this index.
 5305 * @return  0 when hints valid.
 5306 * @return  SIG_VERIFY_E when hints invalid.
 5307 */
 5308static int dilithium_check_hint(const byte* h, byte k, byte omega)
 5309{
 5310    int ret = 0;
 5311    unsigned int o = 0;
 5312    unsigned int i;
 5313
 5314    /* Skip polynomial index while count is 0. */
 5315    while ((o < k) && (h[omega + o] == 0)) {
 5316        o++;
 5317    }
 5318    /* Check all possible hints. */
 5319    for (i = 1; (o < k) && (i < omega); i++) {
 5320        /* Done with polynomial if index equals count of hints. */
 5321        if (i == h[omega + o]) {
 5322            /* Next polynomial index while count is index. */
 5323            do {
 5324                o++;
 5325            }
 5326            while ((o < k) && (i == h[omega + o]));
 5327            /* Stop if hints for all polynomials checked. */
 5328            if (o == k) {
 5329                break;
 5330            }
 5331        }
 5332        /* Ensure the last hint is less than the current hint. */
 5333        else if (h[i - 1] >= h[i]) {
 5334            ret = SIG_VERIFY_E;
 5335            break;
 5336        }
 5337    }
 5338    if (ret == 0) {
 5339        /* Use up any sizes that are the last element. */
 5340        while ((o < k) && (i == h[omega + o])) {
 5341            o++;
 5342        }
 5343        /* Ensure all sizes were used. */
 5344        if (o != k) {
 5345            ret = SIG_VERIFY_E;
 5346        }
 5347    }
 5348    /* Check remaining hints are 0. */
 5349    for (; (ret == 0) && (i < omega); i++) {
 5350        if (h[i] != 0) {
 5351           ret = SIG_VERIFY_E;
 5352        }
 5353    }
 5354
 5355    return ret;
 5356}
 5357
 5358#ifndef WOLFSSL_NO_ML_DSA_44
 5359/* Use hints to modify w1.
 5360 *
 5361 * FIPS 204. 8.4: Algorithm 34 UseHint(h, r)
 5362 *   1: m <- (q - 1) / (2 * GAMMA2)
 5363 *   2: (r1, r0) <- Decompose(r)
 5364 *   3: if h == 1 and r0 > 0 return (r1 + 1) mod m
 5365 *   4: if h == 1 and r0 <= 0 return (r1 - 1) mod m
 5366 *   5: return r1
 5367 *
 5368 * @param [in, out] w1  Vector of polynomials needing hints applied to.
 5369 * @param [in]      h   Hints to apply. In signature encoding.
 5370 * @param [in]      i   Dimension index.
 5371 * @param [in, out] op  Pointer to current offset into hints.
 5372 */
 5373static void dilithium_use_hint_88(sword32* w1, const byte* h, unsigned int i,
 5374    byte* op)
 5375{
 5376    byte o = *op;
 5377    unsigned int j;
 5378
 5379    /* For each value of polynomial. */
 5380    for (j = 0; j < DILITHIUM_N; j++) {
 5381        sword32 r;
 5382        sword32 r0;
 5383        sword32 r1;
 5384#ifdef DILITHIUM_USE_HINT_CT
 5385        /* Hint is 1 when index is next in hint list. */
 5386        sword32 hint = ((o < h[PARAMS_ML_DSA_44_OMEGA + i]) &
 5387                        (h[o] == (byte)j));
 5388
 5389        /* Increment hint offset if this index has hint. */
 5390        o += hint;
 5391        /* Convert value to positive only range. */
 5392        r = w1[j] + ((0 - (((word32)w1[j]) >> 31)) & DILITHIUM_Q);
 5393        /* Decompose value into low and high parts. */
 5394        dilithium_decompose_q88(r, &r0, &r1);
 5395        /* Make hint positive or negative based on sign of r0. */
 5396        hint = (1 - (2 * (((word32)r0) >> 31))) & (0 - hint);
 5397        /* Make w1 only the top part plus the hint. */
 5398        w1[j] = r1 + hint;
 5399
 5400        /* Fix up w1 to not be 44 but 0. */
 5401        w1[j] &= (sword32)(0 - (((word32)(w1[j] - 44)) >> 31));
 5402        /* Hint may have reduced 0 to -1 which is actually 43. */
 5403        w1[j] += (sword32)((0 - (((word32)w1[j]) >> 31)) & 44);
 5404#else
 5405        /* Convert value to positive only range. */
 5406        r = w1[j] + (sword32)((0 - (((word32)w1[j]) >> 31)) & DILITHIUM_Q);
 5407        /* Decompose value into low and high parts. */
 5408        dilithium_decompose_q88(r, &r0, &r1);
 5409        /* Check for hint. */
 5410        if ((o < h[PARAMS_ML_DSA_44_OMEGA + i]) && (h[o] == (byte)j)) {
 5411            /* Add or subtract hint based on sign of r0. */
 5412            r1 += (sword32)(1U - (2U * (((word32)r0) >> 31)));
 5413            /* Go to next hint offset. */
 5414            o++;
 5415        }
 5416        /* Fix up w1 to not be 44 but 0. */
 5417        r1 &= (sword32)(0U - (((word32)r1 - 44U) >> 31));
 5418        /* Hint may have reduced 0 to -1 which is actually 43. */
 5419        r1 += (sword32)((0U - (((word32)r1) >> 31)) & 44U);
 5420        /* Make w1 only the top part plus any hint. */
 5421        w1[j] = r1;
 5422#endif
 5423    }
 5424    *op = o;
 5425}
 5426#endif /* !WOLFSSL_NO_ML_DSA_44 */
 5427
 5428#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 5429/* Use hints to modify w1.
 5430 *
 5431 * FIPS 204. 8.4: Algorithm 34 UseHint(h, r)
 5432 *   1: m <- (q - 1) / (2 * GAMMA2)
 5433 *   2: (r1, r0) <- Decompose(r)
 5434 *   3: if h == 1 and r0 > 0 return (r1 + 1) mod m
 5435 *   4: if h == 1 and r0 <= 0 return (r1 - 1) mod m
 5436 *   5: return r1
 5437 *
 5438 * @param [in, out] w1     Vector of polynomials needing hints applied to.
 5439 * @param [in]      h      Hints to apply. In signature encoding.
 5440 * @param [in]      omega  Max number of hints. Hint counts after this index.
 5441 * @param [in]      i      Dimension index.
 5442 * @param [in, out] op     Pointer to current offset into hints.
 5443 */
 5444static void dilithium_use_hint_32(sword32* w1, const byte* h, byte omega,
 5445    unsigned int i, byte* op)
 5446{
 5447    byte o = *op;
 5448    unsigned int j;
 5449
 5450    /* For each value of polynomial. */
 5451    for (j = 0; j < DILITHIUM_N; j++) {
 5452        sword32 r;
 5453        sword32 r0;
 5454        sword32 r1;
 5455#ifdef DILITHIUM_USE_HINT_CT
 5456        /* Hint is 1 when index is next in hint list. */
 5457        sword32 hint = ((o < h[omega + i]) & (h[o] == (byte)j));
 5458
 5459        /* Increment hint offset if this index has hint. */
 5460        o += hint;
 5461        /* Convert value to positive only range. */
 5462        r = w1[j] + (sword32)((0 - (((word32)w1[j]) >> 31)) & DILITHIUM_Q);
 5463        /* Decompose value into low and high parts. */
 5464        dilithium_decompose_q32(r, &r0, &r1);
 5465        /* Make hint positive or negative based on sign of r0. */
 5466        hint = (sword32)((1 - (2 * (((word32)r0) >> 31))) & (0 - hint));
 5467        /* Make w1 only the top part plus the hint. */
 5468        w1[j] = r1 + hint;
 5469
 5470        /* Fix up w1 not be 16 (-> 0) or -1 (-> 15). */
 5471        w1[j] &= 0xf;
 5472#else
 5473        /* Convert value to positive only range. */
 5474        r = w1[j] + (sword32)((0 - (((word32)w1[j]) >> 31)) & DILITHIUM_Q);
 5475        /* Decompose value into low and high parts. */
 5476        dilithium_decompose_q32(r, &r0, &r1);
 5477        /* Check for hint. */
 5478        if ((o < h[omega + i]) && (h[o] == (byte)j)) {
 5479            /* Add or subtract hint based on sign of r0. */
 5480            r1 += (sword32)(1 - (2 * (((word32)r0) >> 31)));
 5481            /* Go to next hint offset. */
 5482            o++;
 5483        }
 5484        /* Fix up w1 not be 16 (-> 0) or -1 (-> 15). */
 5485        w1[j] = r1 & 0xf;
 5486#endif
 5487    }
 5488    *op = o;
 5489}
 5490#endif
 5491
 5492#ifndef WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM
 5493/* Use hints to modify w1.
 5494 *
 5495 * FIPS 204. 8.4: Algorithm 34 UseHint(h, r)
 5496 *   1: m <- (q - 1) / (2 * GAMMA2)
 5497 *   2: (r1, r0) <- Decompose(r)
 5498 *   3: if h == 1 and r0 > 0 return (r1 + 1) mod m
 5499 *   4: if h == 1 and r0 <= 0 return (r1 - 1) mod m
 5500 *   5: return r1
 5501 *
 5502 * @param [in, out] w1      Vector of polynomials needing hints applied to.
 5503 * @param [in]      k       Dimension of vector.
 5504 * @param [in]      gamma2  Low-order rounding range, GAMMA2.
 5505 * @param [in]      omega   Max number of hints. Hint counts after this index.
 5506 * @param [in]      h       Hints to apply. In signature encoding.
 5507 */
 5508static void dilithium_vec_use_hint(sword32* w1, byte k, sword32 gamma2,
 5509    byte omega, const byte* h)
 5510{
 5511    unsigned int i;
 5512    byte o = 0;
 5513
 5514    (void)k;
 5515    (void)omega;
 5516
 5517#ifndef WOLFSSL_NO_ML_DSA_44
 5518    if (gamma2 == DILITHIUM_Q_LOW_88) {
 5519    #ifdef USE_INTEL_SPEEDUP
 5520        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 5521            wc_mldsa_use_hint_88_avx2(w1, h);
 5522            RESTORE_VECTOR_REGISTERS();
 5523        }
 5524        else
 5525    #endif
 5526        {
 5527            /* For each polynomial of vector. */
 5528            for (i = 0; i < PARAMS_ML_DSA_44_K; i++) {
 5529                dilithium_use_hint_88(w1, h, i, &o);
 5530                w1 += DILITHIUM_N;
 5531            }
 5532        }
 5533    }
 5534#endif
 5535#if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 5536    if (gamma2 == DILITHIUM_Q_LOW_32) {
 5537    #ifdef USE_INTEL_SPEEDUP
 5538        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 5539            wc_mldsa_use_hint_32_avx2(w1, k, h);
 5540            RESTORE_VECTOR_REGISTERS();
 5541        }
 5542        else
 5543    #endif
 5544        {
 5545            /* For each polynomial of vector. */
 5546            for (i = 0; i < k; i++) {
 5547                dilithium_use_hint_32(w1, h, omega, i, &o);
 5548                w1 += DILITHIUM_N;
 5549            }
 5550        }
 5551    }
 5552#endif
 5553}
 5554#endif
 5555#endif /* !WOLFSSL_DILITHIUM_NO_VERIFY */
 5556
 5557/******************************************************************************
 5558 * Maths operations
 5559 ******************************************************************************/
 5560
 5561/* q^-1 mod 2^32 (inverse of 8380417 mod 2^32 = 58728449 = 0x3802001) */
 5562#define DILITHIUM_QINV          58728449
 5563
 5564/* Montgomery reduce a.
 5565 *
 5566 * @param  [in]  a  64-bit value to be reduced.
 5567 * @return  Montgomery reduction result.
 5568 */
 5569static sword32 dilithium_mont_red(sword64 a)
 5570{
 5571#ifndef DILITHIUM_MUL_QINV_SLOW
 5572    sword64 t = (sword32)((sword32)a * (sword32)DILITHIUM_QINV);
 5573#else
 5574    sword64 t = (sword32)((sword32)a + (sword32)((sword32)a << 13) -
 5575        (sword32)((sword32)a << 23) + (sword32)((sword32)a << 26));
 5576#endif
 5577#ifndef DILITHIUM_MUL_Q_SLOW
 5578    return (sword32)((a - ((sword32)t * (sword64)DILITHIUM_Q)) >> 32);
 5579#else
 5580    return (sword32)((a - (t << 23) + (t << 13) - t) >> 32);
 5581#endif
 5582}
 5583
 5584#if !defined(WOLFSSL_DILITHIUM_SMALL) || \
 5585    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 5586     (defined(WOLFSSL_DILITHIUM_SMALL) && \
 5587      (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 5588       (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 5589        !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 5590       defined(WOLFSSL_DILITHIUM_CHECK_KEY))))
 5591/* Reduce 32-bit a modulo q. r = a mod q.
 5592 *
 5593 * Barrett reduction.
 5594 *
 5595 * @param  [in]  a  32-bit value to be reduced to range of q.
 5596 * @return  Modulo result.
 5597 */
 5598static sword32 dilithium_red(sword32 a)
 5599{
 5600    sword32 t = (sword32)((a + ((sword32)1 << 22)) >> 23);
 5601#ifndef DILITHIUM_MUL_Q_SLOW
 5602    return (sword32)(a - (t * DILITHIUM_Q));
 5603#else
 5604    return (sword32)(a - (t << 23) + (t << 13) - t);
 5605#endif
 5606}
 5607#endif
 5608
 5609/* Zetas for NTT. */
 5610static const sword32 zetas[DILITHIUM_N] = {
 5611   -41978,    25847, -2608894,  -518909,   237124,  -777960,  -876248,   466468,
 5612  1826347,  2353451,  -359251, -2091905,  3119733, -2884855,  3111497,  2680103,
 5613  2725464,  1024112, -1079900,  3585928,  -549488, -1119584,  2619752, -2108549,
 5614 -2118186, -3859737, -1399561, -3277672,  1757237,   -19422,  4010497,   280005,
 5615  2706023,    95776,  3077325,  3530437, -1661693, -3592148, -2537516,  3915439,
 5616 -3861115, -3043716,  3574422, -2867647,  3539968,  -300467,  2348700,  -539299,
 5617 -1699267, -1643818,  3505694, -3821735,  3507263, -2140649, -1600420,  3699596,
 5618   811944,   531354,   954230,  3881043,  3900724, -2556880,  2071892, -2797779,
 5619 -3930395, -1528703, -3677745, -3041255, -1452451,  3475950,  2176455, -1585221,
 5620 -1257611,  1939314, -4083598, -1000202, -3190144, -3157330, -3632928,   126922,
 5621  3412210,  -983419,  2147896,  2715295, -2967645, -3693493,  -411027, -2477047,
 5622  -671102, -1228525,   -22981, -1308169,  -381987,  1349076,  1852771, -1430430,
 5623 -3343383,   264944,   508951,  3097992,    44288, -1100098,   904516,  3958618,
 5624 -3724342,    -8578,  1653064, -3249728,  2389356,  -210977,   759969, -1316856,
 5625   189548, -3553272,  3159746, -1851402, -2409325,  -177440,  1315589,  1341330,
 5626  1285669, -1584928,  -812732, -1439742, -3019102, -3881060, -3628969,  3839961,
 5627  2091667,  3407706,  2316500,  3817976, -3342478,  2244091, -2446433, -3562462,
 5628   266997,  2434439, -1235728,  3513181, -3520352, -3759364, -1197226, -3193378,
 5629   900702,  1859098,   909542,   819034,   495491, -1613174,   -43260,  -522500,
 5630  -655327, -3122442,  2031748,  3207046, -3556995,  -525098,  -768622, -3595838,
 5631   342297,   286988, -2437823,  4108315,  3437287, -3342277,  1735879,   203044,
 5632  2842341,  2691481, -2590150,  1265009,  4055324,  1247620,  2486353,  1595974,
 5633 -3767016,  1250494,  2635921, -3548272, -2994039,  1869119,  1903435, -1050970,
 5634 -1333058,  1237275, -3318210, -1430225,  -451100,  1312455,  3306115, -1962642,
 5635 -1279661,  1917081, -2546312, -1374803,  1500165,   777191,  2235880,  3406031,
 5636  -542412, -2831860, -1671176, -1846953, -2584293, -3724270,   594136, -3776993,
 5637 -2013608,  2432395,  2454455,  -164721,  1957272,  3369112,   185531, -1207385,
 5638 -3183426,   162844,  1616392,  3014001,   810149,  1652634, -3694233, -1799107,
 5639 -3038916,  3523897,  3866901,   269760,  2213111,  -975884,  1717735,   472078,
 5640  -426683,  1723600, -1803090,  1910376, -1667432, -1104333,  -260646, -3833893,
 5641 -2939036, -2235985,  -420899, -2286327,   183443,  -976891,  1612842, -3545687,
 5642  -554416,  3919660,   -48306, -1362209,  3937738,  1400424,  -846154,  1976782
 5643};
 5644
 5645#ifndef WOLFSSL_DILITHIUM_SMALL
 5646/* Zetas for inverse NTT. */
 5647static const sword32 zetas_inv[DILITHIUM_N] = {
 5648 -1976782,   846154, -1400424, -3937738,  1362209,    48306, -3919660,   554416,
 5649  3545687, -1612842,   976891,  -183443,  2286327,   420899,  2235985,  2939036,
 5650  3833893,   260646,  1104333,  1667432, -1910376,  1803090, -1723600,   426683,
 5651  -472078, -1717735,   975884, -2213111,  -269760, -3866901, -3523897,  3038916,
 5652  1799107,  3694233, -1652634,  -810149, -3014001, -1616392,  -162844,  3183426,
 5653  1207385,  -185531, -3369112, -1957272,   164721, -2454455, -2432395,  2013608,
 5654  3776993,  -594136,  3724270,  2584293,  1846953,  1671176,  2831860,   542412,
 5655 -3406031, -2235880,  -777191, -1500165,  1374803,  2546312, -1917081,  1279661,
 5656  1962642, -3306115, -1312455,   451100,  1430225,  3318210, -1237275,  1333058,
 5657  1050970, -1903435, -1869119,  2994039,  3548272, -2635921, -1250494,  3767016,
 5658 -1595974, -2486353, -1247620, -4055324, -1265009,  2590150, -2691481, -2842341,
 5659  -203044, -1735879,  3342277, -3437287, -4108315,  2437823,  -286988,  -342297,
 5660  3595838,   768622,   525098,  3556995, -3207046, -2031748,  3122442,   655327,
 5661   522500,    43260,  1613174,  -495491,  -819034,  -909542, -1859098,  -900702,
 5662  3193378,  1197226,  3759364,  3520352, -3513181,  1235728, -2434439,  -266997,
 5663  3562462,  2446433, -2244091,  3342478, -3817976, -2316500, -3407706, -2091667,
 5664 -3839961,  3628969,  3881060,  3019102,  1439742,   812732,  1584928, -1285669,
 5665 -1341330, -1315589,   177440,  2409325,  1851402, -3159746,  3553272,  -189548,
 5666  1316856,  -759969,   210977, -2389356,  3249728, -1653064,     8578,  3724342,
 5667 -3958618,  -904516,  1100098,   -44288, -3097992,  -508951,  -264944,  3343383,
 5668  1430430, -1852771, -1349076,   381987,  1308169,    22981,  1228525,   671102,
 5669  2477047,   411027,  3693493,  2967645, -2715295, -2147896,   983419, -3412210,
 5670  -126922,  3632928,  3157330,  3190144,  1000202,  4083598, -1939314,  1257611,
 5671  1585221, -2176455, -3475950,  1452451,  3041255,  3677745,  1528703,  3930395,
 5672  2797779, -2071892,  2556880, -3900724, -3881043,  -954230,  -531354,  -811944,
 5673 -3699596,  1600420,  2140649, -3507263,  3821735, -3505694,  1643818,  1699267,
 5674   539299, -2348700,   300467, -3539968,  2867647, -3574422,  3043716,  3861115,
 5675 -3915439,  2537516,  3592148,  1661693, -3530437, -3077325,   -95776, -2706023,
 5676  -280005, -4010497,    19422, -1757237,  3277672,  1399561,  3859737,  2118186,
 5677  2108549, -2619752,  1119584,   549488, -3585928,  1079900, -1024112, -2725464,
 5678 -2680103, -3111497,  2884855, -3119733,  2091905,   359251, -2353451, -1826347,
 5679  -466468,   876248,   777960,  -237124,   518909,  2608894,   -25847,    41978
 5680};
 5681#endif
 5682
 5683#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 5684    !defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 5685    (!defined(WOLFSSL_DILITHIUM_NO_MAKE) && defined(WOLFSSL_DILITHIUM_SMALL))
 5686
 5687/* One iteration of Number-Theoretic Transform.
 5688 *
 5689 * @param [in] len  Length of sequence.
 5690 */
 5691#define NTT(len)                                                            \
 5692do {                                                                        \
 5693    for (start = 0; start < DILITHIUM_N; start += 2 * (len)) {              \
 5694        zeta = zetas[++k];                                                  \
 5695        for (j = 0; j < (len); ++j) {                                       \
 5696            sword32 t =                                                     \
 5697                dilithium_mont_red((sword64)zeta * r[start + j + (len)]);   \
 5698            sword32 rj = r[start + j];                                      \
 5699            r[start + j + (len)] = rj - t;                                  \
 5700            r[start + j] = rj + t;                                          \
 5701        }                                                                   \
 5702    }                                                                       \
 5703}                                                                           \
 5704while (0)
 5705
 5706/* Number-Theoretic Transform.
 5707 *
 5708 * @param [in, out] r  Polynomial to transform.
 5709 */
 5710static void dilithium_ntt_c(sword32* r)
 5711{
 5712#ifdef WOLFSSL_DILITHIUM_SMALL
 5713    unsigned int len;
 5714    unsigned int k;
 5715    unsigned int j;
 5716
 5717    k = 0;
 5718    for (len = DILITHIUM_N / 2; len >= 1; len >>= 1) {
 5719        unsigned int start;
 5720        for (start = 0; start < DILITHIUM_N; start = j + len) {
 5721            sword32 zeta = zetas[++k];
 5722            for (j = start; j < start + len; ++j) {
 5723                sword32 t = dilithium_mont_red((sword64)zeta * r[j + len]);
 5724                sword32 rj = r[j];
 5725                r[j + len] = rj - t;
 5726                r[j] = rj + t;
 5727            }
 5728        }
 5729    }
 5730#elif defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE)
 5731    unsigned int j;
 5732    unsigned int k;
 5733    unsigned int start;
 5734    sword32 zeta;
 5735
 5736    zeta = zetas[1];
 5737    for (j = 0; j < DILITHIUM_N / 2; j++) {
 5738        sword32 t =
 5739            dilithium_mont_red((sword64)zeta * r[j + DILITHIUM_N / 2]);
 5740        sword32 rj = r[j];
 5741        r[j + DILITHIUM_N / 2] = rj - t;
 5742        r[j] = rj + t;
 5743    }
 5744
 5745    k = 1;
 5746    NTT(64);
 5747    NTT(32);
 5748    NTT(16);
 5749    NTT(8);
 5750    NTT(4);
 5751    NTT(2);
 5752
 5753    for (j = 0; j < DILITHIUM_N; j += 2) {
 5754        sword32 t = dilithium_mont_red((sword64)zetas[++k] * r[j + 1]);
 5755        sword32 rj = r[j];
 5756        r[j + 1] = rj - t;
 5757        r[j] = rj + t;
 5758    }
 5759#elif defined(WC_32BIT_CPU)
 5760    unsigned int j;
 5761    unsigned int k;
 5762    sword32 t0;
 5763    sword32 t2;
 5764
 5765    sword32 zeta128 = zetas[1];
 5766    sword32 zeta640 = zetas[2];
 5767    sword32 zeta641 = zetas[3];
 5768    for (j = 0; j < DILITHIUM_N / 4; j++) {
 5769        sword32 r0 = r[j +   0];
 5770        sword32 r2 = r[j +  64];
 5771        sword32 r4 = r[j + 128];
 5772        sword32 r6 = r[j + 192];
 5773
 5774        t0 = dilithium_mont_red((sword64)zeta128 * r4);
 5775        t2 = dilithium_mont_red((sword64)zeta128 * r6);
 5776        r4 = r0 - t0;
 5777        r6 = r2 - t2;
 5778        r0 += t0;
 5779        r2 += t2;
 5780
 5781        t0 = dilithium_mont_red((sword64)zeta640 * r2);
 5782        t2 = dilithium_mont_red((sword64)zeta641 * r6);
 5783        r2 = r0 - t0;
 5784        r6 = r4 - t2;
 5785        r0 += t0;
 5786        r4 += t2;
 5787
 5788        r[j +   0] = r0;
 5789        r[j +  64] = r2;
 5790        r[j + 128] = r4;
 5791        r[j + 192] = r6;
 5792    }
 5793
 5794    for (j = 0; j < DILITHIUM_N; j += 64) {
 5795        unsigned int i;
 5796        sword32 zeta32  = zetas[ 4 + j / 64 + 0];
 5797        sword32 zeta160 = zetas[ 8 + j / 32 + 0];
 5798        sword32 zeta161 = zetas[ 8 + j / 32 + 1];
 5799        for (i = 0; i < 16; i++) {
 5800            sword32 r0 = r[j + i +  0];
 5801            sword32 r2 = r[j + i + 16];
 5802            sword32 r4 = r[j + i + 32];
 5803            sword32 r6 = r[j + i + 48];
 5804
 5805            t0 = dilithium_mont_red((sword64)zeta32 * r4);
 5806            t2 = dilithium_mont_red((sword64)zeta32 * r6);
 5807            r4 = r0 - t0;
 5808            r6 = r2 - t2;
 5809            r0 += t0;
 5810            r2 += t2;
 5811
 5812            t0 = dilithium_mont_red((sword64)zeta160 * r2);
 5813            t2 = dilithium_mont_red((sword64)zeta161 * r6);
 5814            r2 = r0 - t0;
 5815            r6 = r4 - t2;
 5816            r0 += t0;
 5817            r4 += t2;
 5818
 5819            r[j + i +  0] = r0;
 5820            r[j + i + 16] = r2;
 5821            r[j + i + 32] = r4;
 5822            r[j + i + 48] = r6;
 5823        }
 5824    }
 5825
 5826    for (j = 0; j < DILITHIUM_N; j += 16) {
 5827        unsigned int i;
 5828        sword32 zeta8   = zetas[16 + j / 16];
 5829        sword32 zeta40  = zetas[32 + j / 8 + 0];
 5830        sword32 zeta41  = zetas[32 + j / 8 + 1];
 5831        for (i = 0; i < 4; i++) {
 5832            sword32 r0 = r[j + i +  0];
 5833            sword32 r2 = r[j + i +  4];
 5834            sword32 r4 = r[j + i +  8];
 5835            sword32 r6 = r[j + i + 12];
 5836
 5837            t0 = dilithium_mont_red((sword64)zeta8 * r4);
 5838            t2 = dilithium_mont_red((sword64)zeta8 * r6);
 5839            r4 = r0 - t0;
 5840            r6 = r2 - t2;
 5841            r0 += t0;
 5842            r2 += t2;
 5843
 5844            t0 = dilithium_mont_red((sword64)zeta40 * r2);
 5845            t2 = dilithium_mont_red((sword64)zeta41 * r6);
 5846            r2 = r0 - t0;
 5847            r6 = r4 - t2;
 5848            r0 += t0;
 5849            r4 += t2;
 5850
 5851            r[j + i +  0] = r0;
 5852            r[j + i +  4] = r2;
 5853            r[j + i +  8] = r4;
 5854            r[j + i + 12] = r6;
 5855        }
 5856    }
 5857
 5858    k = 128;
 5859    for (j = 0; j < DILITHIUM_N; j += 4) {
 5860        sword32 zeta2 = zetas[64 + j / 4];
 5861        sword32 r0 = r[j + 0];
 5862        sword32 r2 = r[j + 1];
 5863        sword32 r4 = r[j + 2];
 5864        sword32 r6 = r[j + 3];
 5865
 5866        t0 = dilithium_mont_red((sword64)zeta2 * r4);
 5867        t2 = dilithium_mont_red((sword64)zeta2 * r6);
 5868        r4 = r0 - t0;
 5869        r6 = r2 - t2;
 5870        r0 += t0;
 5871        r2 += t2;
 5872
 5873        t0 = dilithium_mont_red((sword64)zetas[k++] * r2);
 5874        t2 = dilithium_mont_red((sword64)zetas[k++] * r6);
 5875        r2 = r0 - t0;
 5876        r6 = r4 - t2;
 5877        r0 += t0;
 5878        r4 += t2;
 5879
 5880        r[j + 0] = r0;
 5881        r[j + 1] = r2;
 5882        r[j + 2] = r4;
 5883        r[j + 3] = r6;
 5884    }
 5885#else
 5886    unsigned int j;
 5887    unsigned int k;
 5888    sword32 t0;
 5889    sword32 t1;
 5890    sword32 t2;
 5891    sword32 t3;
 5892
 5893    sword32 zeta128 = zetas[1];
 5894    sword32 zeta640 = zetas[2];
 5895    sword32 zeta641 = zetas[3];
 5896    for (j = 0; j < DILITHIUM_N / 8; j++) {
 5897        sword32 r0 = r[j +   0];
 5898        sword32 r1 = r[j +  32];
 5899        sword32 r2 = r[j +  64];
 5900        sword32 r3 = r[j +  96];
 5901        sword32 r4 = r[j + 128];
 5902        sword32 r5 = r[j + 160];
 5903        sword32 r6 = r[j + 192];
 5904        sword32 r7 = r[j + 224];
 5905
 5906        t0 = dilithium_mont_red((sword64)zeta128 * r4);
 5907        t1 = dilithium_mont_red((sword64)zeta128 * r5);
 5908        t2 = dilithium_mont_red((sword64)zeta128 * r6);
 5909        t3 = dilithium_mont_red((sword64)zeta128 * r7);
 5910        r4 = r0 - t0;
 5911        r5 = r1 - t1;
 5912        r6 = r2 - t2;
 5913        r7 = r3 - t3;
 5914        r0 += t0;
 5915        r1 += t1;
 5916        r2 += t2;
 5917        r3 += t3;
 5918
 5919        t0 = dilithium_mont_red((sword64)zeta640 * r2);
 5920        t1 = dilithium_mont_red((sword64)zeta640 * r3);
 5921        t2 = dilithium_mont_red((sword64)zeta641 * r6);
 5922        t3 = dilithium_mont_red((sword64)zeta641 * r7);
 5923        r2 = r0 - t0;
 5924        r3 = r1 - t1;
 5925        r6 = r4 - t2;
 5926        r7 = r5 - t3;
 5927        r0 += t0;
 5928        r1 += t1;
 5929        r4 += t2;
 5930        r5 += t3;
 5931
 5932        r[j +   0] = r0;
 5933        r[j +  32] = r1;
 5934        r[j +  64] = r2;
 5935        r[j +  96] = r3;
 5936        r[j + 128] = r4;
 5937        r[j + 160] = r5;
 5938        r[j + 192] = r6;
 5939        r[j + 224] = r7;
 5940    }
 5941
 5942    for (j = 0; j < DILITHIUM_N; j += 64) {
 5943        unsigned int i;
 5944        sword32 zeta32  = zetas[ 4 + j / 64 + 0];
 5945        sword32 zeta160 = zetas[ 8 + j / 32 + 0];
 5946        sword32 zeta161 = zetas[ 8 + j / 32 + 1];
 5947        sword32 zeta80  = zetas[16 + j / 16 + 0];
 5948        sword32 zeta81  = zetas[16 + j / 16 + 1];
 5949        sword32 zeta82  = zetas[16 + j / 16 + 2];
 5950        sword32 zeta83  = zetas[16 + j / 16 + 3];
 5951        for (i = 0; i < 8; i++) {
 5952            sword32 r0 = r[j + i +  0];
 5953            sword32 r1 = r[j + i +  8];
 5954            sword32 r2 = r[j + i + 16];
 5955            sword32 r3 = r[j + i + 24];
 5956            sword32 r4 = r[j + i + 32];
 5957            sword32 r5 = r[j + i + 40];
 5958            sword32 r6 = r[j + i + 48];
 5959            sword32 r7 = r[j + i + 56];
 5960
 5961            t0 = dilithium_mont_red((sword64)zeta32 * r4);
 5962            t1 = dilithium_mont_red((sword64)zeta32 * r5);
 5963            t2 = dilithium_mont_red((sword64)zeta32 * r6);
 5964            t3 = dilithium_mont_red((sword64)zeta32 * r7);
 5965            r4 = r0 - t0;
 5966            r5 = r1 - t1;
 5967            r6 = r2 - t2;
 5968            r7 = r3 - t3;
 5969            r0 += t0;
 5970            r1 += t1;
 5971            r2 += t2;
 5972            r3 += t3;
 5973
 5974            t0 = dilithium_mont_red((sword64)zeta160 * r2);
 5975            t1 = dilithium_mont_red((sword64)zeta160 * r3);
 5976            t2 = dilithium_mont_red((sword64)zeta161 * r6);
 5977            t3 = dilithium_mont_red((sword64)zeta161 * r7);
 5978            r2 = r0 - t0;
 5979            r3 = r1 - t1;
 5980            r6 = r4 - t2;
 5981            r7 = r5 - t3;
 5982            r0 += t0;
 5983            r1 += t1;
 5984            r4 += t2;
 5985            r5 += t3;
 5986
 5987            t0 = dilithium_mont_red((sword64)zeta80 * r1);
 5988            t1 = dilithium_mont_red((sword64)zeta81 * r3);
 5989            t2 = dilithium_mont_red((sword64)zeta82 * r5);
 5990            t3 = dilithium_mont_red((sword64)zeta83 * r7);
 5991            r1 = r0 - t0;
 5992            r3 = r2 - t1;
 5993            r5 = r4 - t2;
 5994            r7 = r6 - t3;
 5995            r0 += t0;
 5996            r2 += t1;
 5997            r4 += t2;
 5998            r6 += t3;
 5999
 6000            r[j + i +  0] = r0;
 6001            r[j + i +  8] = r1;
 6002            r[j + i + 16] = r2;
 6003            r[j + i + 24] = r3;
 6004            r[j + i + 32] = r4;
 6005            r[j + i + 40] = r5;
 6006            r[j + i + 48] = r6;
 6007            r[j + i + 56] = r7;
 6008        }
 6009    }
 6010
 6011    k = 128;
 6012    for (j = 0; j < DILITHIUM_N; j += 8) {
 6013        sword32 zeta4  = zetas[32 + j / 8 + 0];
 6014        sword32 zeta20 = zetas[64 + j / 4 + 0];
 6015        sword32 zeta21 = zetas[64 + j / 4 + 1];
 6016        sword32 r0 = r[j + 0];
 6017        sword32 r1 = r[j + 1];
 6018        sword32 r2 = r[j + 2];
 6019        sword32 r3 = r[j + 3];
 6020        sword32 r4 = r[j + 4];
 6021        sword32 r5 = r[j + 5];
 6022        sword32 r6 = r[j + 6];
 6023        sword32 r7 = r[j + 7];
 6024
 6025        t0 = dilithium_mont_red((sword64)zeta4 * r4);
 6026        t1 = dilithium_mont_red((sword64)zeta4 * r5);
 6027        t2 = dilithium_mont_red((sword64)zeta4 * r6);
 6028        t3 = dilithium_mont_red((sword64)zeta4 * r7);
 6029        r4 = r0 - t0;
 6030        r5 = r1 - t1;
 6031        r6 = r2 - t2;
 6032        r7 = r3 - t3;
 6033        r0 += t0;
 6034        r1 += t1;
 6035        r2 += t2;
 6036        r3 += t3;
 6037
 6038        t0 = dilithium_mont_red((sword64)zeta20 * r2);
 6039        t1 = dilithium_mont_red((sword64)zeta20 * r3);
 6040        t2 = dilithium_mont_red((sword64)zeta21 * r6);
 6041        t3 = dilithium_mont_red((sword64)zeta21 * r7);
 6042        r2 = r0 - t0;
 6043        r3 = r1 - t1;
 6044        r6 = r4 - t2;
 6045        r7 = r5 - t3;
 6046        r0 += t0;
 6047        r1 += t1;
 6048        r4 += t2;
 6049        r5 += t3;
 6050
 6051        t0 = dilithium_mont_red((sword64)zetas[k++] * r1);
 6052        t1 = dilithium_mont_red((sword64)zetas[k++] * r3);
 6053        t2 = dilithium_mont_red((sword64)zetas[k++] * r5);
 6054        t3 = dilithium_mont_red((sword64)zetas[k++] * r7);
 6055        r1 = r0 - t0;
 6056        r3 = r2 - t1;
 6057        r5 = r4 - t2;
 6058        r7 = r6 - t3;
 6059        r0 += t0;
 6060        r2 += t1;
 6061        r4 += t2;
 6062        r6 += t3;
 6063
 6064        r[j + 0] = r0;
 6065        r[j + 1] = r1;
 6066        r[j + 2] = r2;
 6067        r[j + 3] = r3;
 6068        r[j + 4] = r4;
 6069        r[j + 5] = r5;
 6070        r[j + 6] = r6;
 6071        r[j + 7] = r7;
 6072    }
 6073#endif
 6074}
 6075
 6076#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 6077     defined(WC_DILITHIUM_CACHE_PRIV_VECTORS)
 6078/* Number-Theoretic Transform.
 6079 *
 6080 * @param [in, out] r  Polynomial to transform.
 6081 */
 6082static void dilithium_ntt(sword32* r)
 6083{
 6084#ifdef USE_INTEL_SPEEDUP
 6085    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 6086        wc_mldsa_ntt_avx2(r);
 6087        RESTORE_VECTOR_REGISTERS();
 6088    }
 6089    else
 6090#endif
 6091    {
 6092        dilithium_ntt_c(r);
 6093    }
 6094}
 6095#endif
 6096
 6097#if !defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 6098    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 6099     (!defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 6100      defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC))) || \
 6101    (defined(WOLFSSL_DILITHIUM_SMALL) && \
 6102     (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 6103      defined(WOLFSSL_DILITHIUM_CHECK_KEY)))
 6104/* Number-Theoretic Transform.
 6105 *
 6106 * @param [in, out] r  Polynomial to transform.
 6107 */
 6108static void dilithium_ntt_full(sword32* r)
 6109{
 6110#ifdef USE_INTEL_SPEEDUP
 6111    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 6112        wc_mldsa_ntt_full_avx2(r);
 6113        RESTORE_VECTOR_REGISTERS();
 6114    }
 6115    else
 6116#endif
 6117    {
 6118        dilithium_ntt_c(r);
 6119    }
 6120}
 6121#endif
 6122
 6123#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 6124    (!defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 6125     defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC) || \
 6126     defined(WC_DILITHIUM_CACHE_PRIV_VECTORS))
 6127/* Number-Theoretic Transform.
 6128 *
 6129 * @param [in, out]  r  Vector of polynomials to transform.
 6130 * @param [in]       l  Dimension of polynomial.
 6131 */
 6132static void dilithium_vec_ntt(sword32* r, byte l)
 6133{
 6134    unsigned int i;
 6135
 6136    for (i = 0; i < l; i++) {
 6137        dilithium_ntt(r);
 6138        r += DILITHIUM_N;
 6139    }
 6140}
 6141#endif
 6142#endif
 6143
 6144#if (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 6145     (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 6146      (!defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 6147       defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC)))) || \
 6148    (defined(WOLFSSL_DILITHIUM_SMALL) && \
 6149     (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 6150      defined(WOLFSSL_DILITHIUM_CHECK_KEY)))
 6151/* Number-Theoretic Transform.
 6152 *
 6153 * @param [in, out]  r  Vector of polynomials to transform.
 6154 * @param [in]       l  Dimension of polynomial.
 6155 */
 6156static void dilithium_vec_ntt_full(sword32* r, byte l)
 6157{
 6158    unsigned int i;
 6159
 6160    for (i = 0; i < l; i++) {
 6161        dilithium_ntt_full(r);
 6162        r += DILITHIUM_N;
 6163    }
 6164}
 6165#endif
 6166
 6167#ifndef WOLFSSL_DILITHIUM_SMALL
 6168
 6169/* Zeta index value 1 not in montgomery form. */
 6170#define DILITHIUM_NTT_ZETA_1    ((sword32)-3572223)
 6171
 6172/* Number-Theoretic Transform with small initial values.
 6173 *
 6174 * @param [in, out] r  Polynomial to transform.
 6175 */
 6176static void dilithium_ntt_small_c(sword32* r)
 6177{
 6178    unsigned int k;
 6179    unsigned int j;
 6180#ifdef WOLFSSL_DILITHIUM_NO_LARGE_CODE
 6181    unsigned int start;
 6182    sword32 zeta;
 6183
 6184    for (j = 0; j < DILITHIUM_N / 2; ++j) {
 6185        sword32 t = dilithium_red(DILITHIUM_NTT_ZETA_1 *
 6186                                  r[j + DILITHIUM_N / 2]);
 6187        sword32 rj = r[j];
 6188        r[j + DILITHIUM_N / 2] = rj - t;
 6189        r[j] = rj + t;
 6190    }
 6191
 6192    k = 1;
 6193    NTT(64);
 6194    NTT(32);
 6195    NTT(16);
 6196    NTT(8);
 6197    NTT(4);
 6198    NTT(2);
 6199
 6200    for (j = 0; j < DILITHIUM_N; j += 2) {
 6201        sword32 t = dilithium_mont_red((sword64)zetas[++k] * r[j + 1]);
 6202        sword32 rj = r[j];
 6203        r[j + 1] = rj - t;
 6204        r[j] = rj + t;
 6205    }
 6206#elif defined(WC_32BIT_CPU)
 6207    sword32 t0;
 6208    sword32 t2;
 6209
 6210    sword32 zeta640 = zetas[2];
 6211    sword32 zeta641 = zetas[3];
 6212    for (j = 0; j < DILITHIUM_N / 4; j++) {
 6213        sword32 r0 = r[j +   0];
 6214        sword32 r2 = r[j +  64];
 6215        sword32 r4 = r[j + 128];
 6216        sword32 r6 = r[j + 192];
 6217
 6218        t0 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r4);
 6219        t2 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r6);
 6220        r4 = r0 - t0;
 6221        r6 = r2 - t2;
 6222        r0 += t0;
 6223        r2 += t2;
 6224
 6225        t0 = dilithium_mont_red((sword64)zeta640 * r2);
 6226        t2 = dilithium_mont_red((sword64)zeta641 * r6);
 6227        r2 = r0 - t0;
 6228        r6 = r4 - t2;
 6229        r0 += t0;
 6230        r4 += t2;
 6231
 6232        r[j +   0] = r0;
 6233        r[j +  64] = r2;
 6234        r[j + 128] = r4;
 6235        r[j + 192] = r6;
 6236    }
 6237
 6238    for (j = 0; j < DILITHIUM_N; j += 64) {
 6239        unsigned int i;
 6240        sword32 zeta32  = zetas[ 4 + j / 64 + 0];
 6241        sword32 zeta160 = zetas[ 8 + j / 32 + 0];
 6242        sword32 zeta161 = zetas[ 8 + j / 32 + 1];
 6243        for (i = 0; i < 16; i++) {
 6244            sword32 r0 = r[j + i +  0];
 6245            sword32 r2 = r[j + i + 16];
 6246            sword32 r4 = r[j + i + 32];
 6247            sword32 r6 = r[j + i + 48];
 6248
 6249            t0 = dilithium_mont_red((sword64)zeta32 * r4);
 6250            t2 = dilithium_mont_red((sword64)zeta32 * r6);
 6251            r4 = r0 - t0;
 6252            r6 = r2 - t2;
 6253            r0 += t0;
 6254            r2 += t2;
 6255
 6256            t0 = dilithium_mont_red((sword64)zeta160 * r2);
 6257            t2 = dilithium_mont_red((sword64)zeta161 * r6);
 6258            r2 = r0 - t0;
 6259            r6 = r4 - t2;
 6260            r0 += t0;
 6261            r4 += t2;
 6262
 6263            r[j + i +  0] = r0;
 6264            r[j + i + 16] = r2;
 6265            r[j + i + 32] = r4;
 6266            r[j + i + 48] = r6;
 6267        }
 6268    }
 6269
 6270    for (j = 0; j < DILITHIUM_N; j += 16) {
 6271        unsigned int i;
 6272        sword32 zeta8   = zetas[16 + j / 16];
 6273        sword32 zeta40  = zetas[32 + j / 8 + 0];
 6274        sword32 zeta41  = zetas[32 + j / 8 + 1];
 6275        for (i = 0; i < 4; i++) {
 6276            sword32 r0 = r[j + i +  0];
 6277            sword32 r2 = r[j + i +  4];
 6278            sword32 r4 = r[j + i +  8];
 6279            sword32 r6 = r[j + i + 12];
 6280
 6281            t0 = dilithium_mont_red((sword64)zeta8 * r4);
 6282            t2 = dilithium_mont_red((sword64)zeta8 * r6);
 6283            r4 = r0 - t0;
 6284            r6 = r2 - t2;
 6285            r0 += t0;
 6286            r2 += t2;
 6287
 6288            t0 = dilithium_mont_red((sword64)zeta40 * r2);
 6289            t2 = dilithium_mont_red((sword64)zeta41 * r6);
 6290            r2 = r0 - t0;
 6291            r6 = r4 - t2;
 6292            r0 += t0;
 6293            r4 += t2;
 6294
 6295            r[j + i +  0] = r0;
 6296            r[j + i +  4] = r2;
 6297            r[j + i +  8] = r4;
 6298            r[j + i + 12] = r6;
 6299        }
 6300    }
 6301
 6302    k = 128;
 6303    for (j = 0; j < DILITHIUM_N; j += 4) {
 6304        sword32 zeta2 = zetas[64 + j / 4];
 6305        sword32 r0 = r[j + 0];
 6306        sword32 r2 = r[j + 1];
 6307        sword32 r4 = r[j + 2];
 6308        sword32 r6 = r[j + 3];
 6309
 6310        t0 = dilithium_mont_red((sword64)zeta2 * r4);
 6311        t2 = dilithium_mont_red((sword64)zeta2 * r6);
 6312        r4 = r0 - t0;
 6313        r6 = r2 - t2;
 6314        r0 += t0;
 6315        r2 += t2;
 6316
 6317        t0 = dilithium_mont_red((sword64)zetas[k++] * r2);
 6318        t2 = dilithium_mont_red((sword64)zetas[k++] * r6);
 6319        r2 = r0 - t0;
 6320        r6 = r4 - t2;
 6321        r0 += t0;
 6322        r4 += t2;
 6323
 6324        r[j + 0] = r0;
 6325        r[j + 1] = r2;
 6326        r[j + 2] = r4;
 6327        r[j + 3] = r6;
 6328    }
 6329#else
 6330    sword32 t0;
 6331    sword32 t1;
 6332    sword32 t2;
 6333    sword32 t3;
 6334    sword32 zeta640 = zetas[2];
 6335    sword32 zeta641 = zetas[3];
 6336    for (j = 0; j < DILITHIUM_N / 8; j++) {
 6337        sword32 r0 = r[j +   0];
 6338        sword32 r1 = r[j +  32];
 6339        sword32 r2 = r[j +  64];
 6340        sword32 r3 = r[j +  96];
 6341        sword32 r4 = r[j + 128];
 6342        sword32 r5 = r[j + 160];
 6343        sword32 r6 = r[j + 192];
 6344        sword32 r7 = r[j + 224];
 6345
 6346        t0 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r4);
 6347        t1 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r5);
 6348        t2 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r6);
 6349        t3 = dilithium_red(DILITHIUM_NTT_ZETA_1 * r7);
 6350        r4 = r0 - t0;
 6351        r5 = r1 - t1;
 6352        r6 = r2 - t2;
 6353        r7 = r3 - t3;
 6354        r0 += t0;
 6355        r1 += t1;
 6356        r2 += t2;
 6357        r3 += t3;
 6358
 6359        t0 = dilithium_mont_red((sword64)zeta640 * r2);
 6360        t1 = dilithium_mont_red((sword64)zeta640 * r3);
 6361        t2 = dilithium_mont_red((sword64)zeta641 * r6);
 6362        t3 = dilithium_mont_red((sword64)zeta641 * r7);
 6363        r2 = r0 - t0;
 6364        r3 = r1 - t1;
 6365        r6 = r4 - t2;
 6366        r7 = r5 - t3;
 6367        r0 += t0;
 6368        r1 += t1;
 6369        r4 += t2;
 6370        r5 += t3;
 6371
 6372        r[j +   0] = r0;
 6373        r[j +  32] = r1;
 6374        r[j +  64] = r2;
 6375        r[j +  96] = r3;
 6376        r[j + 128] = r4;
 6377        r[j + 160] = r5;
 6378        r[j + 192] = r6;
 6379        r[j + 224] = r7;
 6380    }
 6381
 6382    for (j = 0; j < DILITHIUM_N; j += 64) {
 6383        unsigned int i;
 6384        sword32 zeta32  = zetas[ 4 + j / 64 + 0];
 6385        sword32 zeta160 = zetas[ 8 + j / 32 + 0];
 6386        sword32 zeta161 = zetas[ 8 + j / 32 + 1];
 6387        sword32 zeta80  = zetas[16 + j / 16 + 0];
 6388        sword32 zeta81  = zetas[16 + j / 16 + 1];
 6389        sword32 zeta82  = zetas[16 + j / 16 + 2];
 6390        sword32 zeta83  = zetas[16 + j / 16 + 3];
 6391        for (i = 0; i < 8; i++) {
 6392            sword32 r0 = r[j + i +  0];
 6393            sword32 r1 = r[j + i +  8];
 6394            sword32 r2 = r[j + i + 16];
 6395            sword32 r3 = r[j + i + 24];
 6396            sword32 r4 = r[j + i + 32];
 6397            sword32 r5 = r[j + i + 40];
 6398            sword32 r6 = r[j + i + 48];
 6399            sword32 r7 = r[j + i + 56];
 6400
 6401            t0 = dilithium_mont_red((sword64)zeta32 * r4);
 6402            t1 = dilithium_mont_red((sword64)zeta32 * r5);
 6403            t2 = dilithium_mont_red((sword64)zeta32 * r6);
 6404            t3 = dilithium_mont_red((sword64)zeta32 * r7);
 6405            r4 = r0 - t0;
 6406            r5 = r1 - t1;
 6407            r6 = r2 - t2;
 6408            r7 = r3 - t3;
 6409            r0 += t0;
 6410            r1 += t1;
 6411            r2 += t2;
 6412            r3 += t3;
 6413
 6414            t0 = dilithium_mont_red((sword64)zeta160 * r2);
 6415            t1 = dilithium_mont_red((sword64)zeta160 * r3);
 6416            t2 = dilithium_mont_red((sword64)zeta161 * r6);
 6417            t3 = dilithium_mont_red((sword64)zeta161 * r7);
 6418            r2 = r0 - t0;
 6419            r3 = r1 - t1;
 6420            r6 = r4 - t2;
 6421            r7 = r5 - t3;
 6422            r0 += t0;
 6423            r1 += t1;
 6424            r4 += t2;
 6425            r5 += t3;
 6426
 6427            t0 = dilithium_mont_red((sword64)zeta80 * r1);
 6428            t1 = dilithium_mont_red((sword64)zeta81 * r3);
 6429            t2 = dilithium_mont_red((sword64)zeta82 * r5);
 6430            t3 = dilithium_mont_red((sword64)zeta83 * r7);
 6431            r1 = r0 - t0;
 6432            r3 = r2 - t1;
 6433            r5 = r4 - t2;
 6434            r7 = r6 - t3;
 6435            r0 += t0;
 6436            r2 += t1;
 6437            r4 += t2;
 6438            r6 += t3;
 6439
 6440            r[j + i +  0] = r0;
 6441            r[j + i +  8] = r1;
 6442            r[j + i + 16] = r2;
 6443            r[j + i + 24] = r3;
 6444            r[j + i + 32] = r4;
 6445            r[j + i + 40] = r5;
 6446            r[j + i + 48] = r6;
 6447            r[j + i + 56] = r7;
 6448        }
 6449    }
 6450
 6451    k = 128;
 6452    for (j = 0; j < DILITHIUM_N; j += 8) {
 6453        sword32 zeta4  = zetas[32 + j / 8 + 0];
 6454        sword32 zeta20 = zetas[64 + j / 4 + 0];
 6455        sword32 zeta21 = zetas[64 + j / 4 + 1];
 6456        sword32 r0 = r[j + 0];
 6457        sword32 r1 = r[j + 1];
 6458        sword32 r2 = r[j + 2];
 6459        sword32 r3 = r[j + 3];
 6460        sword32 r4 = r[j + 4];
 6461        sword32 r5 = r[j + 5];
 6462        sword32 r6 = r[j + 6];
 6463        sword32 r7 = r[j + 7];
 6464
 6465        t0 = dilithium_mont_red((sword64)zeta4 * r4);
 6466        t1 = dilithium_mont_red((sword64)zeta4 * r5);
 6467        t2 = dilithium_mont_red((sword64)zeta4 * r6);
 6468        t3 = dilithium_mont_red((sword64)zeta4 * r7);
 6469        r4 = r0 - t0;
 6470        r5 = r1 - t1;
 6471        r6 = r2 - t2;
 6472        r7 = r3 - t3;
 6473        r0 += t0;
 6474        r1 += t1;
 6475        r2 += t2;
 6476        r3 += t3;
 6477
 6478        t0 = dilithium_mont_red((sword64)zeta20 * r2);
 6479        t1 = dilithium_mont_red((sword64)zeta20 * r3);
 6480        t2 = dilithium_mont_red((sword64)zeta21 * r6);
 6481        t3 = dilithium_mont_red((sword64)zeta21 * r7);
 6482        r2 = r0 - t0;
 6483        r3 = r1 - t1;
 6484        r6 = r4 - t2;
 6485        r7 = r5 - t3;
 6486        r0 += t0;
 6487        r1 += t1;
 6488        r4 += t2;
 6489        r5 += t3;
 6490
 6491        t0 = dilithium_mont_red((sword64)zetas[k++] * r1);
 6492        t1 = dilithium_mont_red((sword64)zetas[k++] * r3);
 6493        t2 = dilithium_mont_red((sword64)zetas[k++] * r5);
 6494        t3 = dilithium_mont_red((sword64)zetas[k++] * r7);
 6495        r1 = r0 - t0;
 6496        r3 = r2 - t1;
 6497        r5 = r4 - t2;
 6498        r7 = r6 - t3;
 6499        r0 += t0;
 6500        r2 += t1;
 6501        r4 += t2;
 6502        r6 += t3;
 6503
 6504        r[j + 0] = r0;
 6505        r[j + 1] = r1;
 6506        r[j + 2] = r2;
 6507        r[j + 3] = r3;
 6508        r[j + 4] = r4;
 6509        r[j + 5] = r5;
 6510        r[j + 6] = r6;
 6511        r[j + 7] = r7;
 6512    }
 6513#endif
 6514}
 6515
 6516#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 6517     defined(WC_DILITHIUM_CACHE_PRIV_VECTORS)
 6518/* Number-Theoretic Transform.
 6519 *
 6520 * @param [in, out] r  Polynomial to transform.
 6521 */
 6522static void dilithium_ntt_small(sword32* r)
 6523{
 6524#ifdef USE_INTEL_SPEEDUP
 6525    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 6526        wc_mldsa_ntt_avx2(r);
 6527        RESTORE_VECTOR_REGISTERS();
 6528    }
 6529    else
 6530#endif
 6531    {
 6532        dilithium_ntt_small_c(r);
 6533    }
 6534}
 6535#endif
 6536
 6537#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 6538    !defined(WOLFSSL_DILITHIUM_NO_VERIFY) || \
 6539     defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 6540/* Number-Theoretic Transform.
 6541 *
 6542 * @param [in, out] r  Polynomial to transform.
 6543 */
 6544static void dilithium_ntt_small_full(sword32* r)
 6545{
 6546#ifdef USE_INTEL_SPEEDUP
 6547    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 6548        wc_mldsa_ntt_full_avx2(r);
 6549        RESTORE_VECTOR_REGISTERS();
 6550    }
 6551    else
 6552#endif
 6553    {
 6554        dilithium_ntt_small_c(r);
 6555    }
 6556}
 6557#endif
 6558
 6559#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 6560    (!defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 6561     defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC) || \
 6562     defined(WC_DILITHIUM_CACHE_PRIV_VECTORS))
 6563/* Number-Theoretic Transform with small initial values.
 6564 *
 6565 * @param [in, out]  r  Vector of polynomials to transform.
 6566 * @param [in]       l  Dimension of polynomial.
 6567 */
 6568static void dilithium_vec_ntt_small(sword32* r, byte l)
 6569{
 6570    unsigned int i;
 6571
 6572    for (i = 0; i < l; i++) {
 6573        dilithium_ntt_small(r);
 6574        r += DILITHIUM_N;
 6575    }
 6576}
 6577#endif
 6578
 6579#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 6580     defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 6581/* Number-Theoretic Transform with small initial values.
 6582 *
 6583 * @param [in, out]  r  Vector of polynomials to transform.
 6584 * @param [in]       l  Dimension of polynomial.
 6585 */
 6586static void dilithium_vec_ntt_small_full(sword32* r, byte l)
 6587{
 6588    unsigned int i;
 6589
 6590    for (i = 0; i < l; i++) {
 6591        dilithium_ntt_small_full(r);
 6592        r += DILITHIUM_N;
 6593    }
 6594}
 6595#endif
 6596
 6597#else
 6598
 6599/* Number-Theoretic Transform with small initial values.
 6600 *
 6601 * @param [in, out] r  Polynomial to transform.
 6602 */
 6603#define dilithium_ntt_small          dilithium_ntt
 6604/* Number-Theoretic Transform with small initial values.
 6605 *
 6606 * @param [in, out] r  Polynomial to transform.
 6607 */
 6608#define dilithium_ntt_small_full     dilithium_ntt_full
 6609/* Number-Theoretic Transform with small initial values.
 6610 *
 6611 * @param [in, out]  r  Vector of polynomials to transform.
 6612 * @param [in]       l  Dimension of polynomial.
 6613 */
 6614#define dilithium_vec_ntt_small      dilithium_vec_ntt
 6615/* Number-Theoretic Transform with small initial values.
 6616 *
 6617 * @param [in, out]  r  Vector of polynomials to transform.
 6618 * @param [in]       l  Dimension of polynomial.
 6619 */
 6620#define dilithium_vec_ntt_small_full dilithium_vec_ntt_full
 6621
 6622
 6623#endif /* WOLFSSL_DILITHIUM_SMALL */
 6624
 6625
 6626/* One iteration of Inverse Number-Theoretic Transform.
 6627 *
 6628 * @param [in] len  Length of sequence.
 6629 */
 6630#define INVNTT(len)                                                         \
 6631do {                                                                        \
 6632    for (start = 0; start < DILITHIUM_N; start += 2 * (len)) {              \
 6633        zeta = zetas_inv[k++];                                              \
 6634        for (j = 0; j < (len); ++j) {                                       \
 6635            sword32 rj = r[start + j];                                      \
 6636            sword32 rjl = r[start + j + (len)];                             \
 6637            sword32 t = rj + rjl;                                           \
 6638            r[start + j] = t;                                               \
 6639            rjl = rj - rjl;                                                 \
 6640            r[start + j + (len)] = dilithium_mont_red((sword64)zeta * rjl); \
 6641        }                                                                   \
 6642    }                                                                       \
 6643}                                                                           \
 6644while (0)
 6645
 6646/* Inverse Number-Theoretic Transform.
 6647 *
 6648 * @param [in, out] r  Polynomial to transform.
 6649 */
 6650static void dilithium_invntt_c(sword32* r)
 6651{
 6652#ifdef WOLFSSL_DILITHIUM_SMALL
 6653    unsigned int len;
 6654    unsigned int k;
 6655    unsigned int j;
 6656    sword32 zeta;
 6657
 6658    k = 256;
 6659    for (len = 1; len <= DILITHIUM_N / 2; len <<= 1) {
 6660        unsigned int start;
 6661        for (start = 0; start < DILITHIUM_N; start = j + len) {
 6662            zeta = -zetas[--k];
 6663            for (j = start; j < start + len; ++j) {
 6664                sword32 rj = r[j];
 6665                sword32 rjl = r[j + len];
 6666                sword32 t = rj + rjl;
 6667                r[j] = t;
 6668                rjl = rj - rjl;
 6669                r[j + len] = dilithium_mont_red((sword64)zeta * rjl);
 6670            }
 6671        }
 6672    }
 6673
 6674    zeta = -zetas[0];
 6675    for (j = 0; j < DILITHIUM_N; ++j) {
 6676        r[j] = dilithium_mont_red((sword64)zeta * r[j]);
 6677    }
 6678#elif defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE)
 6679    unsigned int j;
 6680    unsigned int k = 0;
 6681    unsigned int start;
 6682    sword32 zeta;
 6683
 6684    for (j = 0; j < DILITHIUM_N; j += 2) {
 6685        sword32 rj = r[j];
 6686        sword32 rjl = r[j + 1];
 6687        sword32 t = rj + rjl;
 6688        r[j] = t;
 6689        rjl = rj - rjl;
 6690        r[j + 1] = dilithium_mont_red((sword64)zetas_inv[k++] * rjl);
 6691    }
 6692
 6693    INVNTT(2);
 6694    INVNTT(4);
 6695    INVNTT(8);
 6696    INVNTT(16);
 6697    INVNTT(32);
 6698    INVNTT(64);
 6699    INVNTT(128);
 6700
 6701    zeta = zetas_inv[255];
 6702    for (j = 0; j < DILITHIUM_N; ++j) {
 6703        r[j] = dilithium_mont_red((sword64)zeta * r[j]);
 6704    }
 6705#elif defined(WC_32BIT_CPU)
 6706    unsigned int j;
 6707    unsigned int k = 0;
 6708    sword32 t0;
 6709    sword32 t2;
 6710
 6711    sword32 zeta640;
 6712    sword32 zeta641;
 6713    sword32 zeta128;
 6714    sword32 zeta256;
 6715    for (j = 0; j < DILITHIUM_N; j += 4) {
 6716        sword32 zeta2 = zetas_inv[128 + j / 4];
 6717        sword32 r0 = r[j + 0];
 6718        sword32 r2 = r[j + 1];
 6719        sword32 r4 = r[j + 2];
 6720        sword32 r6 = r[j + 3];
 6721
 6722        t0 = dilithium_mont_red((sword64)zetas_inv[k++] * (r0 - r2));
 6723        t2 = dilithium_mont_red((sword64)zetas_inv[k++] * (r4 - r6));
 6724        r0 += r2;
 6725        r4 += r6;
 6726        r2 = t0;
 6727        r6 = t2;
 6728
 6729        t0 = dilithium_mont_red((sword64)zeta2 * (r0 - r4));
 6730        t2 = dilithium_mont_red((sword64)zeta2 * (r2 - r6));
 6731        r0 += r4;
 6732        r2 += r6;
 6733        r4 = t0;
 6734        r6 = t2;
 6735
 6736        r[j + 0] = r0;
 6737        r[j + 1] = r2;
 6738        r[j + 2] = r4;
 6739        r[j + 3] = r6;
 6740    }
 6741
 6742    for (j = 0; j < DILITHIUM_N; j += 16) {
 6743        unsigned int i;
 6744        sword32 zeta40 = zetas_inv[192 + j / 8 + 0];
 6745        sword32 zeta41 = zetas_inv[192 + j / 8 + 1];
 6746        sword32 zeta8  = zetas_inv[224 + j / 16 + 0];
 6747        for (i = 0; i < 4; i++) {
 6748            sword32 r0 = r[j + i +  0];
 6749            sword32 r2 = r[j + i +  4];
 6750            sword32 r4 = r[j + i +  8];
 6751            sword32 r6 = r[j + i + 12];
 6752
 6753            t0 = dilithium_mont_red((sword64)zeta40 * (r0 - r2));
 6754            t2 = dilithium_mont_red((sword64)zeta41 * (r4 - r6));
 6755            r0 += r2;
 6756            r4 += r6;
 6757            r2 = t0;
 6758            r6 = t2;
 6759
 6760            t0 = dilithium_mont_red((sword64)zeta8 * (r0 - r4));
 6761            t2 = dilithium_mont_red((sword64)zeta8 * (r2 - r6));
 6762            r0 += r4;
 6763            r2 += r6;
 6764            r4 = t0;
 6765            r6 = t2;
 6766
 6767            r[j + i +  0] = r0;
 6768            r[j + i +  4] = r2;
 6769            r[j + i +  8] = r4;
 6770            r[j + i + 12] = r6;
 6771        }
 6772    }
 6773
 6774    for (j = 0; j < DILITHIUM_N; j += 64) {
 6775        unsigned int i;
 6776        sword32 zeta160 = zetas_inv[240 + j / 32 + 0];
 6777        sword32 zeta161 = zetas_inv[240 + j / 32 + 1];
 6778        sword32 zeta32  = zetas_inv[248 + j / 64 + 0];
 6779        for (i = 0; i < 16; i++) {
 6780            sword32 r0 = r[j + i +  0];
 6781            sword32 r2 = r[j + i + 16];
 6782            sword32 r4 = r[j + i + 32];
 6783            sword32 r6 = r[j + i + 48];
 6784
 6785            t0 = dilithium_mont_red((sword64)zeta160 * (r0 - r2));
 6786            t2 = dilithium_mont_red((sword64)zeta161 * (r4 - r6));
 6787            r0 += r2;
 6788            r4 += r6;
 6789            r2 = t0;
 6790            r6 = t2;
 6791
 6792            t0 = dilithium_mont_red((sword64)zeta32 * (r0 - r4));
 6793            t2 = dilithium_mont_red((sword64)zeta32 * (r2 - r6));
 6794            r0 += r4;
 6795            r2 += r6;
 6796            r4 = t0;
 6797            r6 = t2;
 6798
 6799            r[j + i +  0] = r0;
 6800            r[j + i + 16] = r2;
 6801            r[j + i + 32] = r4;
 6802            r[j + i + 48] = r6;
 6803        }
 6804    }
 6805
 6806    zeta640 = zetas_inv[252];
 6807    zeta641 = zetas_inv[253];
 6808    zeta128 = zetas_inv[254];
 6809    zeta256 = zetas_inv[255];
 6810    for (j = 0; j < DILITHIUM_N / 4; j++) {
 6811        sword32 r0 = r[j +   0];
 6812        sword32 r2 = r[j +  64];
 6813        sword32 r4 = r[j + 128];
 6814        sword32 r6 = r[j + 192];
 6815
 6816        t0 = dilithium_mont_red((sword64)zeta640 * (r0 - r2));
 6817        t2 = dilithium_mont_red((sword64)zeta641 * (r4 - r6));
 6818        r0 += r2;
 6819        r4 += r6;
 6820        r2 = t0;
 6821        r6 = t2;
 6822
 6823        t0 = dilithium_mont_red((sword64)zeta128 * (r0 - r4));
 6824        t2 = dilithium_mont_red((sword64)zeta128 * (r2 - r6));
 6825        r0 += r4;
 6826        r2 += r6;
 6827        r4 = t0;
 6828        r6 = t2;
 6829
 6830        r0 = dilithium_mont_red((sword64)zeta256 * r0);
 6831        r2 = dilithium_mont_red((sword64)zeta256 * r2);
 6832        r4 = dilithium_mont_red((sword64)zeta256 * r4);
 6833        r6 = dilithium_mont_red((sword64)zeta256 * r6);
 6834
 6835        r[j +   0] = r0;
 6836        r[j +  64] = r2;
 6837        r[j + 128] = r4;
 6838        r[j + 192] = r6;
 6839    }
 6840#else
 6841    unsigned int j;
 6842    unsigned int k = 0;
 6843    sword32 t0;
 6844    sword32 t1;
 6845    sword32 t2;
 6846    sword32 t3;
 6847
 6848    sword32 zeta640;
 6849    sword32 zeta641;
 6850    sword32 zeta128;
 6851    sword32 zeta256;
 6852    for (j = 0; j < DILITHIUM_N; j += 8) {
 6853        sword32 zeta20 = zetas_inv[128 + j / 4 + 0];
 6854        sword32 zeta21 = zetas_inv[128 + j / 4 + 1];
 6855        sword32 zeta4  = zetas_inv[192 + j / 8 + 0];
 6856        sword32 r0 = r[j + 0];
 6857        sword32 r1 = r[j + 1];
 6858        sword32 r2 = r[j + 2];
 6859        sword32 r3 = r[j + 3];
 6860        sword32 r4 = r[j + 4];
 6861        sword32 r5 = r[j + 5];
 6862        sword32 r6 = r[j + 6];
 6863        sword32 r7 = r[j + 7];
 6864
 6865        t0 = dilithium_mont_red((sword64)zetas_inv[k++] * (r0 - r1));
 6866        t1 = dilithium_mont_red((sword64)zetas_inv[k++] * (r2 - r3));
 6867        t2 = dilithium_mont_red((sword64)zetas_inv[k++] * (r4 - r5));
 6868        t3 = dilithium_mont_red((sword64)zetas_inv[k++] * (r6 - r7));
 6869        r0 += r1;
 6870        r2 += r3;
 6871        r4 += r5;
 6872        r6 += r7;
 6873        r1 = t0;
 6874        r3 = t1;
 6875        r5 = t2;
 6876        r7 = t3;
 6877
 6878        t0 = dilithium_mont_red((sword64)zeta20 * (r0 - r2));
 6879        t1 = dilithium_mont_red((sword64)zeta20 * (r1 - r3));
 6880        t2 = dilithium_mont_red((sword64)zeta21 * (r4 - r6));
 6881        t3 = dilithium_mont_red((sword64)zeta21 * (r5 - r7));
 6882        r0 += r2;
 6883        r1 += r3;
 6884        r4 += r6;
 6885        r5 += r7;
 6886        r2 = t0;
 6887        r3 = t1;
 6888        r6 = t2;
 6889        r7 = t3;
 6890
 6891        t0 = dilithium_mont_red((sword64)zeta4 * (r0 - r4));
 6892        t1 = dilithium_mont_red((sword64)zeta4 * (r1 - r5));
 6893        t2 = dilithium_mont_red((sword64)zeta4 * (r2 - r6));
 6894        t3 = dilithium_mont_red((sword64)zeta4 * (r3 - r7));
 6895        r0 += r4;
 6896        r1 += r5;
 6897        r2 += r6;
 6898        r3 += r7;
 6899        r4 = t0;
 6900        r5 = t1;
 6901        r6 = t2;
 6902        r7 = t3;
 6903
 6904        r[j + 0] = r0;
 6905        r[j + 1] = r1;
 6906        r[j + 2] = r2;
 6907        r[j + 3] = r3;
 6908        r[j + 4] = r4;
 6909        r[j + 5] = r5;
 6910        r[j + 6] = r6;
 6911        r[j + 7] = r7;
 6912    }
 6913
 6914    for (j = 0; j < DILITHIUM_N; j += 64) {
 6915        unsigned int i;
 6916        sword32 zeta80  = zetas_inv[224 + j / 16 + 0];
 6917        sword32 zeta81  = zetas_inv[224 + j / 16 + 1];
 6918        sword32 zeta82  = zetas_inv[224 + j / 16 + 2];
 6919        sword32 zeta83  = zetas_inv[224 + j / 16 + 3];
 6920        sword32 zeta160 = zetas_inv[240 + j / 32 + 0];
 6921        sword32 zeta161 = zetas_inv[240 + j / 32 + 1];
 6922        sword32 zeta32  = zetas_inv[248 + j / 64 + 0];
 6923        for (i = 0; i < 8; i++) {
 6924            sword32 r0 = r[j + i +  0];
 6925            sword32 r1 = r[j + i +  8];
 6926            sword32 r2 = r[j + i + 16];
 6927            sword32 r3 = r[j + i + 24];
 6928            sword32 r4 = r[j + i + 32];
 6929            sword32 r5 = r[j + i + 40];
 6930            sword32 r6 = r[j + i + 48];
 6931            sword32 r7 = r[j + i + 56];
 6932
 6933            t0 = dilithium_mont_red((sword64)zeta80 * (r0 - r1));
 6934            t1 = dilithium_mont_red((sword64)zeta81 * (r2 - r3));
 6935            t2 = dilithium_mont_red((sword64)zeta82 * (r4 - r5));
 6936            t3 = dilithium_mont_red((sword64)zeta83 * (r6 - r7));
 6937            r0 += r1;
 6938            r2 += r3;
 6939            r4 += r5;
 6940            r6 += r7;
 6941            r1 = t0;
 6942            r3 = t1;
 6943            r5 = t2;
 6944            r7 = t3;
 6945
 6946            t0 = dilithium_mont_red((sword64)zeta160 * (r0 - r2));
 6947            t1 = dilithium_mont_red((sword64)zeta160 * (r1 - r3));
 6948            t2 = dilithium_mont_red((sword64)zeta161 * (r4 - r6));
 6949            t3 = dilithium_mont_red((sword64)zeta161 * (r5 - r7));
 6950            r0 += r2;
 6951            r1 += r3;
 6952            r4 += r6;
 6953            r5 += r7;
 6954            r2 = t0;
 6955            r3 = t1;
 6956            r6 = t2;
 6957            r7 = t3;
 6958
 6959            t0 = dilithium_mont_red((sword64)zeta32 * (r0 - r4));
 6960            t1 = dilithium_mont_red((sword64)zeta32 * (r1 - r5));
 6961            t2 = dilithium_mont_red((sword64)zeta32 * (r2 - r6));
 6962            t3 = dilithium_mont_red((sword64)zeta32 * (r3 - r7));
 6963            r0 += r4;
 6964            r1 += r5;
 6965            r2 += r6;
 6966            r3 += r7;
 6967            r4 = t0;
 6968            r5 = t1;
 6969            r6 = t2;
 6970            r7 = t3;
 6971
 6972            r[j + i +  0] = r0;
 6973            r[j + i +  8] = r1;
 6974            r[j + i + 16] = r2;
 6975            r[j + i + 24] = r3;
 6976            r[j + i + 32] = r4;
 6977            r[j + i + 40] = r5;
 6978            r[j + i + 48] = r6;
 6979            r[j + i + 56] = r7;
 6980        }
 6981    }
 6982
 6983    zeta640 = zetas_inv[252];
 6984    zeta641 = zetas_inv[253];
 6985    zeta128 = zetas_inv[254];
 6986    zeta256 = zetas_inv[255];
 6987    for (j = 0; j < DILITHIUM_N / 8; j++) {
 6988        sword32 r0 = r[j +   0];
 6989        sword32 r1 = r[j +  32];
 6990        sword32 r2 = r[j +  64];
 6991        sword32 r3 = r[j +  96];
 6992        sword32 r4 = r[j + 128];
 6993        sword32 r5 = r[j + 160];
 6994        sword32 r6 = r[j + 192];
 6995        sword32 r7 = r[j + 224];
 6996
 6997        t0 = dilithium_mont_red((sword64)zeta640 * (r0 - r2));
 6998        t1 = dilithium_mont_red((sword64)zeta640 * (r1 - r3));
 6999        t2 = dilithium_mont_red((sword64)zeta641 * (r4 - r6));
 7000        t3 = dilithium_mont_red((sword64)zeta641 * (r5 - r7));
 7001        r0 += r2;
 7002        r1 += r3;
 7003        r4 += r6;
 7004        r5 += r7;
 7005        r2 = t0;
 7006        r3 = t1;
 7007        r6 = t2;
 7008        r7 = t3;
 7009
 7010        t0 = dilithium_mont_red((sword64)zeta128 * (r0 - r4));
 7011        t1 = dilithium_mont_red((sword64)zeta128 * (r1 - r5));
 7012        t2 = dilithium_mont_red((sword64)zeta128 * (r2 - r6));
 7013        t3 = dilithium_mont_red((sword64)zeta128 * (r3 - r7));
 7014        r0 += r4;
 7015        r1 += r5;
 7016        r2 += r6;
 7017        r3 += r7;
 7018        r4 = t0;
 7019        r5 = t1;
 7020        r6 = t2;
 7021        r7 = t3;
 7022
 7023        r0 = dilithium_mont_red((sword64)zeta256 * r0);
 7024        r1 = dilithium_mont_red((sword64)zeta256 * r1);
 7025        r2 = dilithium_mont_red((sword64)zeta256 * r2);
 7026        r3 = dilithium_mont_red((sword64)zeta256 * r3);
 7027        r4 = dilithium_mont_red((sword64)zeta256 * r4);
 7028        r5 = dilithium_mont_red((sword64)zeta256 * r5);
 7029        r6 = dilithium_mont_red((sword64)zeta256 * r6);
 7030        r7 = dilithium_mont_red((sword64)zeta256 * r7);
 7031
 7032        r[j +   0] = r0;
 7033        r[j +  32] = r1;
 7034        r[j +  64] = r2;
 7035        r[j +  96] = r3;
 7036        r[j + 128] = r4;
 7037        r[j + 160] = r5;
 7038        r[j + 192] = r6;
 7039        r[j + 224] = r7;
 7040    }
 7041#endif
 7042}
 7043
 7044#if !defined(WOLFSSL_DILITHIUM_NO_SIGN)
 7045/* Inverse Number-Theoretic Transform.
 7046 *
 7047 * @param [in, out] r  Polynomial to transform.
 7048 */
 7049static void dilithium_invntt(sword32* r)
 7050{
 7051#ifdef USE_INTEL_SPEEDUP
 7052    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7053        wc_mldsa_invntt_avx2(r);
 7054        RESTORE_VECTOR_REGISTERS();
 7055    }
 7056    else
 7057#endif
 7058    {
 7059        dilithium_invntt_c(r);
 7060    }
 7061}
 7062#endif
 7063
 7064/* Inverse Number-Theoretic Transform.
 7065 *
 7066 * @param [in, out] r  Polynomial to transform.
 7067 */
 7068static void dilithium_invntt_full(sword32* r)
 7069{
 7070#ifdef USE_INTEL_SPEEDUP
 7071    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7072        wc_mldsa_invntt_full_avx2(r);
 7073        RESTORE_VECTOR_REGISTERS();
 7074    }
 7075    else
 7076#endif
 7077    {
 7078        dilithium_invntt_c(r);
 7079    }
 7080}
 7081
 7082#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7083     defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 7084    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7085     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 7086    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 7087     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 7088/* Inverse Number-Theoretic Transform.
 7089 *
 7090 * @param [in, out]  r  Vector of polynomials to transform.
 7091 * @param [in]       l  Dimension of polynomial.
 7092 */
 7093static void dilithium_vec_invntt_full(sword32* r, byte l)
 7094{
 7095    unsigned int i;
 7096
 7097    for (i = 0; i < l; i++) {
 7098        dilithium_invntt_full(r);
 7099        r += DILITHIUM_N;
 7100    }
 7101}
 7102#endif
 7103
 7104#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7105     defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 7106    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7107     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 7108    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 7109     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 7110/* Matrix multiplication.
 7111 *
 7112 * @param [out] r  Vector of polynomials that is result.
 7113 * @param [in]  m  Matrix of polynomials.
 7114 * @param [in]  v  Vector of polynomials.
 7115 * @param [in]  k  First dimension of matrix and dimension of result.
 7116 * @param [in]  l  Second dimension of matrix and dimension of v.
 7117 */
 7118static void dilithium_matrix_mul_c(sword32* r, const sword32* m,
 7119    const sword32* v, byte k, byte l)
 7120{
 7121    byte i;
 7122
 7123    for (i = 0; i < k; i++) {
 7124        byte j;
 7125        unsigned int e;
 7126        const sword32* vt = v;
 7127
 7128#ifdef WOLFSSL_DILITHIUM_SMALL
 7129        for (e = 0; e < DILITHIUM_N; e++) {
 7130            r[e] = dilithium_mont_red((sword64)m[e] * vt[e]);
 7131        }
 7132        m += DILITHIUM_N;
 7133        vt += DILITHIUM_N;
 7134        for (j = 1; j < l; j++) {
 7135            for (e = 0; e < DILITHIUM_N; e++) {
 7136                r[e] += dilithium_mont_red((sword64)m[e] * vt[e]);
 7137            }
 7138            m += DILITHIUM_N;
 7139            vt += DILITHIUM_N;
 7140        }
 7141#elif defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE)
 7142        (void)j;
 7143        if (l == 4) {
 7144            for (e = 0; e < DILITHIUM_N; e++) {
 7145                sword64 t = ((sword64)m[e + 0 * 256] * vt[e + 0 * 256]) +
 7146                            ((sword64)m[e + 1 * 256] * vt[e + 1 * 256]) +
 7147                            ((sword64)m[e + 2 * 256] * vt[e + 2 * 256]) +
 7148                            ((sword64)m[e + 3 * 256] * vt[e + 3 * 256]);
 7149                r[e] = dilithium_mont_red(t);
 7150            }
 7151            m += DILITHIUM_N * 4;
 7152        }
 7153        if (l == 5) {
 7154            for (e = 0; e < DILITHIUM_N; e++) {
 7155                sword64 t = ((sword64)m[e + 0 * 256] * vt[e + 0 * 256]) +
 7156                            ((sword64)m[e + 1 * 256] * vt[e + 1 * 256]) +
 7157                            ((sword64)m[e + 2 * 256] * vt[e + 2 * 256]) +
 7158                            ((sword64)m[e + 3 * 256] * vt[e + 3 * 256]) +
 7159                            ((sword64)m[e + 4 * 256] * vt[e + 4 * 256]);
 7160                r[e] = dilithium_mont_red(t);
 7161            }
 7162            m += DILITHIUM_N * 5;
 7163        }
 7164        if (l == 7) {
 7165            for (e = 0; e < DILITHIUM_N; e++) {
 7166                sword64 t = ((sword64)m[e + 0 * 256] * vt[e + 0 * 256]) +
 7167                            ((sword64)m[e + 1 * 256] * vt[e + 1 * 256]) +
 7168                            ((sword64)m[e + 2 * 256] * vt[e + 2 * 256]) +
 7169                            ((sword64)m[e + 3 * 256] * vt[e + 3 * 256]) +
 7170                            ((sword64)m[e + 4 * 256] * vt[e + 4 * 256]) +
 7171                            ((sword64)m[e + 5 * 256] * vt[e + 5 * 256]) +
 7172                            ((sword64)m[e + 6 * 256] * vt[e + 6 * 256]);
 7173                r[e] = dilithium_mont_red(t);
 7174            }
 7175            m += DILITHIUM_N * 7;
 7176        }
 7177#else
 7178        sword64 t0;
 7179        sword64 t1;
 7180#if !defined(WOLFSSL_NO_ML_DSA_44) || !defined(WOLFSSL_NO_ML_DSA_65)
 7181        sword64 t2;
 7182        sword64 t3;
 7183#endif
 7184
 7185        (void)j;
 7186#ifndef WOLFSSL_NO_ML_DSA_44
 7187        if (l == 4) {
 7188            for (e = 0; e < DILITHIUM_N; e += 4) {
 7189                t0 = ((sword64)m[e + 0 + 0 * 256] * vt[e + 0 + 0 * 256]) +
 7190                     ((sword64)m[e + 0 + 1 * 256] * vt[e + 0 + 1 * 256]) +
 7191                     ((sword64)m[e + 0 + 2 * 256] * vt[e + 0 + 2 * 256]) +
 7192                     ((sword64)m[e + 0 + 3 * 256] * vt[e + 0 + 3 * 256]);
 7193                t1 = ((sword64)m[e + 1 + 0 * 256] * vt[e + 1 + 0 * 256]) +
 7194                     ((sword64)m[e + 1 + 1 * 256] * vt[e + 1 + 1 * 256]) +
 7195                     ((sword64)m[e + 1 + 2 * 256] * vt[e + 1 + 2 * 256]) +
 7196                     ((sword64)m[e + 1 + 3 * 256] * vt[e + 1 + 3 * 256]);
 7197                t2 = ((sword64)m[e + 2 + 0 * 256] * vt[e + 2 + 0 * 256]) +
 7198                     ((sword64)m[e + 2 + 1 * 256] * vt[e + 2 + 1 * 256]) +
 7199                     ((sword64)m[e + 2 + 2 * 256] * vt[e + 2 + 2 * 256]) +
 7200                     ((sword64)m[e + 2 + 3 * 256] * vt[e + 2 + 3 * 256]);
 7201                t3 = ((sword64)m[e + 3 + 0 * 256] * vt[e + 3 + 0 * 256]) +
 7202                     ((sword64)m[e + 3 + 1 * 256] * vt[e + 3 + 1 * 256]) +
 7203                     ((sword64)m[e + 3 + 2 * 256] * vt[e + 3 + 2 * 256]) +
 7204                     ((sword64)m[e + 3 + 3 * 256] * vt[e + 3 + 3 * 256]);
 7205                r[e + 0] = dilithium_mont_red(t0);
 7206                r[e + 1] = dilithium_mont_red(t1);
 7207                r[e + 2] = dilithium_mont_red(t2);
 7208                r[e + 3] = dilithium_mont_red(t3);
 7209            }
 7210            m += DILITHIUM_N * 4;
 7211        }
 7212#endif
 7213#ifndef WOLFSSL_NO_ML_DSA_65
 7214        if (l == 5) {
 7215            for (e = 0; e < DILITHIUM_N; e += 4) {
 7216                t0 = ((sword64)m[e + 0 + 0 * 256] * vt[e + 0 + 0 * 256]) +
 7217                     ((sword64)m[e + 0 + 1 * 256] * vt[e + 0 + 1 * 256]) +
 7218                     ((sword64)m[e + 0 + 2 * 256] * vt[e + 0 + 2 * 256]) +
 7219                     ((sword64)m[e + 0 + 3 * 256] * vt[e + 0 + 3 * 256]) +
 7220                     ((sword64)m[e + 0 + 4 * 256] * vt[e + 0 + 4 * 256]);
 7221                t1 = ((sword64)m[e + 1 + 0 * 256] * vt[e + 1 + 0 * 256]) +
 7222                     ((sword64)m[e + 1 + 1 * 256] * vt[e + 1 + 1 * 256]) +
 7223                     ((sword64)m[e + 1 + 2 * 256] * vt[e + 1 + 2 * 256]) +
 7224                     ((sword64)m[e + 1 + 3 * 256] * vt[e + 1 + 3 * 256]) +
 7225                     ((sword64)m[e + 1 + 4 * 256] * vt[e + 1 + 4 * 256]);
 7226                t2 = ((sword64)m[e + 2 + 0 * 256] * vt[e + 2 + 0 * 256]) +
 7227                     ((sword64)m[e + 2 + 1 * 256] * vt[e + 2 + 1 * 256]) +
 7228                     ((sword64)m[e + 2 + 2 * 256] * vt[e + 2 + 2 * 256]) +
 7229                     ((sword64)m[e + 2 + 3 * 256] * vt[e + 2 + 3 * 256]) +
 7230                     ((sword64)m[e + 2 + 4 * 256] * vt[e + 2 + 4 * 256]);
 7231                t3 = ((sword64)m[e + 3 + 0 * 256] * vt[e + 3 + 0 * 256]) +
 7232                     ((sword64)m[e + 3 + 1 * 256] * vt[e + 3 + 1 * 256]) +
 7233                     ((sword64)m[e + 3 + 2 * 256] * vt[e + 3 + 2 * 256]) +
 7234                     ((sword64)m[e + 3 + 3 * 256] * vt[e + 3 + 3 * 256]) +
 7235                     ((sword64)m[e + 3 + 4 * 256] * vt[e + 3 + 4 * 256]);
 7236                r[e + 0] = dilithium_mont_red(t0);
 7237                r[e + 1] = dilithium_mont_red(t1);
 7238                r[e + 2] = dilithium_mont_red(t2);
 7239                r[e + 3] = dilithium_mont_red(t3);
 7240            }
 7241            m += DILITHIUM_N * 5;
 7242        }
 7243#endif
 7244#ifndef WOLFSSL_NO_ML_DSA_87
 7245        if (l == 7) {
 7246            for (e = 0; e < DILITHIUM_N; e += 2) {
 7247                t0 = ((sword64)m[e + 0 + 0 * 256] * vt[e + 0 + 0 * 256]) +
 7248                     ((sword64)m[e + 0 + 1 * 256] * vt[e + 0 + 1 * 256]) +
 7249                     ((sword64)m[e + 0 + 2 * 256] * vt[e + 0 + 2 * 256]) +
 7250                     ((sword64)m[e + 0 + 3 * 256] * vt[e + 0 + 3 * 256]) +
 7251                     ((sword64)m[e + 0 + 4 * 256] * vt[e + 0 + 4 * 256]) +
 7252                     ((sword64)m[e + 0 + 5 * 256] * vt[e + 0 + 5 * 256]) +
 7253                     ((sword64)m[e + 0 + 6 * 256] * vt[e + 0 + 6 * 256]);
 7254                t1 = ((sword64)m[e + 1 + 0 * 256] * vt[e + 1 + 0 * 256]) +
 7255                     ((sword64)m[e + 1 + 1 * 256] * vt[e + 1 + 1 * 256]) +
 7256                     ((sword64)m[e + 1 + 2 * 256] * vt[e + 1 + 2 * 256]) +
 7257                     ((sword64)m[e + 1 + 3 * 256] * vt[e + 1 + 3 * 256]) +
 7258                     ((sword64)m[e + 1 + 4 * 256] * vt[e + 1 + 4 * 256]) +
 7259                     ((sword64)m[e + 1 + 5 * 256] * vt[e + 1 + 5 * 256]) +
 7260                     ((sword64)m[e + 1 + 6 * 256] * vt[e + 1 + 6 * 256]);
 7261                r[e + 0] = dilithium_mont_red(t0);
 7262                r[e + 1] = dilithium_mont_red(t1);
 7263            }
 7264            m += DILITHIUM_N * 7;
 7265        }
 7266#endif
 7267#endif
 7268        r += DILITHIUM_N;
 7269    }
 7270}
 7271
 7272/* Matrix multiplication.
 7273 *
 7274 * @param [out] r  Vector of polynomials that is result.
 7275 * @param [in]  m  Matrix of polynomials.
 7276 * @param [in]  v  Vector of polynomials.
 7277 * @param [in]  k  First dimension of matrix and dimension of result.
 7278 * @param [in]  l  Second dimension of matrix and dimension of v.
 7279 */
 7280static void dilithium_matrix_mul(sword32* r, const sword32* m, const sword32* v,
 7281     byte k, byte l)
 7282{
 7283#ifdef USE_INTEL_SPEEDUP
 7284    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7285        int i;
 7286        if (l == 4) {
 7287            for (i = 0; i < k; i++) {
 7288                wc_mldsa_mul_vec_4_avx2(r, m, v);
 7289                m += l * DILITHIUM_N;
 7290                r += DILITHIUM_N;
 7291            }
 7292        }
 7293        else if (l == 5) {
 7294            for (i = 0; i < k; i++) {
 7295                wc_mldsa_mul_vec_5_avx2(r, m, v);
 7296                m += l * DILITHIUM_N;
 7297                r += DILITHIUM_N;
 7298            }
 7299        }
 7300        else {
 7301            for (i = 0; i < k; i++) {
 7302                wc_mldsa_mul_vec_7_avx2(r, m, v);
 7303                m += l * DILITHIUM_N;
 7304                r += DILITHIUM_N;
 7305            }
 7306        }
 7307        RESTORE_VECTOR_REGISTERS();
 7308    }
 7309    else
 7310#endif
 7311    {
 7312        dilithium_matrix_mul_c(r, m, v, k, l);
 7313    }
 7314}
 7315#endif
 7316
 7317#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 7318    (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7319     !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))
 7320/* Polynomial multiplication.
 7321 *
 7322 * @param [out] r  Polynomial result.
 7323 * @param [in]  a  Polynomial
 7324 * @param [in]  b  Polynomial.
 7325 */
 7326static void dilithium_mul_c(sword32* r, sword32* a, sword32* b)
 7327{
 7328    unsigned int e;
 7329#ifdef WOLFSSL_DILITHIUM_SMALL
 7330    for (e = 0; e < DILITHIUM_N; e++) {
 7331        r[e] = dilithium_mont_red((sword64)a[e] * b[e]);
 7332    }
 7333#elif defined(WOLFSSL_DILITHIUM_NO_LARGE_CODE)
 7334    for (e = 0; e < DILITHIUM_N; e += 8) {
 7335        r[e+0] = dilithium_mont_red((sword64)a[e+0] * b[e+0]);
 7336        r[e+1] = dilithium_mont_red((sword64)a[e+1] * b[e+1]);
 7337        r[e+2] = dilithium_mont_red((sword64)a[e+2] * b[e+2]);
 7338        r[e+3] = dilithium_mont_red((sword64)a[e+3] * b[e+3]);
 7339        r[e+4] = dilithium_mont_red((sword64)a[e+4] * b[e+4]);
 7340        r[e+5] = dilithium_mont_red((sword64)a[e+5] * b[e+5]);
 7341        r[e+6] = dilithium_mont_red((sword64)a[e+6] * b[e+6]);
 7342        r[e+7] = dilithium_mont_red((sword64)a[e+7] * b[e+7]);
 7343    }
 7344#else
 7345    for (e = 0; e < DILITHIUM_N; e += 16) {
 7346        r[e+ 0] = dilithium_mont_red((sword64)a[e+ 0] * b[e+ 0]);
 7347        r[e+ 1] = dilithium_mont_red((sword64)a[e+ 1] * b[e+ 1]);
 7348        r[e+ 2] = dilithium_mont_red((sword64)a[e+ 2] * b[e+ 2]);
 7349        r[e+ 3] = dilithium_mont_red((sword64)a[e+ 3] * b[e+ 3]);
 7350        r[e+ 4] = dilithium_mont_red((sword64)a[e+ 4] * b[e+ 4]);
 7351        r[e+ 5] = dilithium_mont_red((sword64)a[e+ 5] * b[e+ 5]);
 7352        r[e+ 6] = dilithium_mont_red((sword64)a[e+ 6] * b[e+ 6]);
 7353        r[e+ 7] = dilithium_mont_red((sword64)a[e+ 7] * b[e+ 7]);
 7354        r[e+ 8] = dilithium_mont_red((sword64)a[e+ 8] * b[e+ 8]);
 7355        r[e+ 9] = dilithium_mont_red((sword64)a[e+ 9] * b[e+ 9]);
 7356        r[e+10] = dilithium_mont_red((sword64)a[e+10] * b[e+10]);
 7357        r[e+11] = dilithium_mont_red((sword64)a[e+11] * b[e+11]);
 7358        r[e+12] = dilithium_mont_red((sword64)a[e+12] * b[e+12]);
 7359        r[e+13] = dilithium_mont_red((sword64)a[e+13] * b[e+13]);
 7360        r[e+14] = dilithium_mont_red((sword64)a[e+14] * b[e+14]);
 7361        r[e+15] = dilithium_mont_red((sword64)a[e+15] * b[e+15]);
 7362    }
 7363#endif
 7364}
 7365
 7366#if !defined(WOLFSSL_DILITHIUM_NO_SIGN)
 7367/* Polynomial multiplication.
 7368 *
 7369 * @param [out] r  Polynomial result.
 7370 * @param [in]  a  Polynomial
 7371 * @param [in]  b  Polynomial.
 7372 */
 7373static void dilithium_mul(sword32* r, sword32* a, sword32* b)
 7374{
 7375#ifdef USE_INTEL_SPEEDUP
 7376    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7377        wc_mldsa_mul_avx2(r, a, b);
 7378        RESTORE_VECTOR_REGISTERS();
 7379    }
 7380    else
 7381#endif
 7382    {
 7383        dilithium_mul_c(r, a, b);
 7384    }
 7385}
 7386#endif
 7387
 7388#if !defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7389    !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)
 7390/* Vector multiplication.
 7391 *
 7392 * @param [out] r  Vector of polynomials that is result.
 7393 * @param [in]  a  Polynomials
 7394 * @param [in]  b  Vector of polynomials.
 7395 * @param [in]  l  Dimension of vectors.
 7396 */
 7397static void dilithium_vec_mul(sword32* r, sword32* a, sword32* b, byte l)
 7398{
 7399    byte i;
 7400
 7401#ifdef USE_INTEL_SPEEDUP
 7402    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7403        for (i = 0; i < l; i++) {
 7404            wc_mldsa_mul_avx2(r, a, b);
 7405            r += DILITHIUM_N;
 7406            b += DILITHIUM_N;
 7407        }
 7408        RESTORE_VECTOR_REGISTERS();
 7409    }
 7410    else
 7411#endif
 7412    {
 7413        for (i = 0; i < l; i++) {
 7414            dilithium_mul_c(r, a, b);
 7415            r += DILITHIUM_N;
 7416            b += DILITHIUM_N;
 7417        }
 7418    }
 7419}
 7420#endif
 7421#endif
 7422
 7423#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 7424    (defined(WOLFSSL_DILITHIUM_SMALL) && \
 7425     (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7426      (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7427       !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 7428      defined(WOLFSSL_DILITHIUM_CHECK_KEY)))
 7429/* Modulo reduce values in polynomial. Range (-2^31)..(2^31-1).
 7430 *
 7431 * @param [in, out] a  Polynomial.
 7432 */
 7433static void dilithium_poly_red_c(sword32* a)
 7434{
 7435    unsigned int j;
 7436#ifdef WOLFSSL_DILITHIUM_SMALL
 7437    for (j = 0; j < DILITHIUM_N; j++) {
 7438        a[j] = dilithium_red(a[j]);
 7439    }
 7440#else
 7441    for (j = 0; j < DILITHIUM_N; j += 8) {
 7442        a[j+0] = dilithium_red(a[j+0]);
 7443        a[j+1] = dilithium_red(a[j+1]);
 7444        a[j+2] = dilithium_red(a[j+2]);
 7445        a[j+3] = dilithium_red(a[j+3]);
 7446        a[j+4] = dilithium_red(a[j+4]);
 7447        a[j+5] = dilithium_red(a[j+5]);
 7448        a[j+6] = dilithium_red(a[j+6]);
 7449        a[j+7] = dilithium_red(a[j+7]);
 7450    }
 7451#endif
 7452}
 7453
 7454/* Modulo reduce values in polynomial. Range (-2^31)..(2^31-1).
 7455 *
 7456 * @param [in, out] a  Polynomial.
 7457 */
 7458static void dilithium_poly_red(sword32* a)
 7459{
 7460#ifdef USE_INTEL_SPEEDUP
 7461    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7462        wc_mldsa_poly_red_avx2(a);
 7463        RESTORE_VECTOR_REGISTERS();
 7464    }
 7465    else
 7466#endif
 7467    {
 7468        dilithium_poly_red_c(a);
 7469    }
 7470}
 7471
 7472#if (defined(WOLFSSL_DILITHIUM_SMALL) && \
 7473     (!defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7474      (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7475       !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM)) || \
 7476      defined(WOLFSSL_DILITHIUM_CHECK_KEY))) || \
 7477    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 7478     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 7479/* Modulo reduce values in polynomials of vector. Range (-2^31)..(2^31-1).
 7480 *
 7481 * @param [in, out] a  Vector of polynomials.
 7482 * @param [in]      l  Dimension of vector.
 7483 */
 7484static void dilithium_vec_red(sword32* a, byte l)
 7485{
 7486    byte i;
 7487
 7488    for (i = 0; i < l; i++) {
 7489        dilithium_poly_red(a);
 7490        a += DILITHIUM_N;
 7491    }
 7492}
 7493#endif
 7494#endif
 7495
 7496#if (!defined(WOLFSSL_DILITHIUM_NO_SIGN) || \
 7497     (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7498      !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))) || \
 7499    defined(WOLFSSL_DILITHIUM_CHECK_KEY)
 7500/* Subtract polynomials a from r. r -= a.
 7501 *
 7502 * @param [out] r  Polynomial to subtract from.
 7503 * @param [in]  a  Polynomial to subtract.
 7504 */
 7505static void dilithium_sub_c(sword32* r, const sword32* a)
 7506{
 7507    unsigned int j;
 7508#ifdef WOLFSSL_DILITHIUM_SMALL
 7509    for (j = 0; j < DILITHIUM_N; j++) {
 7510        r[j] -= a[j];
 7511    }
 7512#else
 7513    for (j = 0; j < DILITHIUM_N; j += 8) {
 7514        r[j+0] -= a[j+0];
 7515        r[j+1] -= a[j+1];
 7516        r[j+2] -= a[j+2];
 7517        r[j+3] -= a[j+3];
 7518        r[j+4] -= a[j+4];
 7519        r[j+5] -= a[j+5];
 7520        r[j+6] -= a[j+6];
 7521        r[j+7] -= a[j+7];
 7522    }
 7523#endif
 7524}
 7525
 7526/* Subtract polynomials a from r. r -= a.
 7527 *
 7528 * @param [out] r  Polynomial to subtract from.
 7529 * @param [in]  a  Polynomial to subtract.
 7530 */
 7531static void dilithium_sub(sword32* r, const sword32* a)
 7532{
 7533#ifdef USE_INTEL_SPEEDUP
 7534    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7535        wc_mldsa_poly_sub_avx2(r, a);
 7536        RESTORE_VECTOR_REGISTERS();
 7537    }
 7538    else
 7539#endif
 7540    {
 7541        dilithium_sub_c(r, a);
 7542    }
 7543}
 7544
 7545#if defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 7546   (!defined(WOLFSSL_DILITHIUM_NO_VERIFY) && \
 7547    !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM))
 7548/* Subtract vector a from r. r -= a.
 7549 *
 7550 * @param [out] r  Vector of polynomials that is result.
 7551 * @param [in]  a  Vector of polynomials to subtract.
 7552 * @param [in]  l  Dimension of vectors.
 7553 */
 7554static void dilithium_vec_sub(sword32* r, const sword32* a, byte l)
 7555{
 7556    byte i;
 7557
 7558    for (i = 0; i < l; i++) {
 7559        dilithium_sub(r, a);
 7560        r += DILITHIUM_N;
 7561        a += DILITHIUM_N;
 7562    }
 7563}
 7564#endif
 7565#endif
 7566
 7567#ifndef WOLFSSL_DILITHIUM_VERIFY_ONLY
 7568/* Add polynomials a to r. r += a.
 7569 *
 7570 * @param [out] r  Polynomial to add to.
 7571 * @param [in]  a  Polynomial to add.
 7572 */
 7573static void dilithium_add_c(sword32* r, const sword32* a)
 7574{
 7575    unsigned int j;
 7576#ifdef WOLFSSL_DILITHIUM_SMALL
 7577    for (j = 0; j < DILITHIUM_N; j++) {
 7578        r[j] += a[j];
 7579    }
 7580#else
 7581    for (j = 0; j < DILITHIUM_N; j += 8) {
 7582        r[j+0] += a[j+0];
 7583        r[j+1] += a[j+1];
 7584        r[j+2] += a[j+2];
 7585        r[j+3] += a[j+3];
 7586        r[j+4] += a[j+4];
 7587        r[j+5] += a[j+5];
 7588        r[j+6] += a[j+6];
 7589        r[j+7] += a[j+7];
 7590    }
 7591#endif
 7592}
 7593
 7594/* Add polynomials a to r. r += a.
 7595 *
 7596 * @param [out] r  Polynomial to add to.
 7597 * @param [in]  a  Polynomial to add.
 7598 */
 7599static void dilithium_add(sword32* r, const sword32* a)
 7600{
 7601#ifdef USE_INTEL_SPEEDUP
 7602    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7603        wc_mldsa_poly_add_avx2(r, a);
 7604        RESTORE_VECTOR_REGISTERS();
 7605    }
 7606    else
 7607#endif
 7608    {
 7609        dilithium_add_c(r, a);
 7610    }
 7611}
 7612
 7613#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7614    defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 7615    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 7616     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 7617/* Add vector a to r. r += a.
 7618 *
 7619 * @param [out] r  Vector of polynomials that is result.
 7620 * @param [in]  a  Vector of polynomials to add.
 7621 * @param [in]  l  Dimension of vectors.
 7622 */
 7623static void dilithium_vec_add(sword32* r, const sword32* a, byte l)
 7624{
 7625    byte i;
 7626
 7627    for (i = 0; i < l; i++) {
 7628        dilithium_add(r, a);
 7629        r += DILITHIUM_N;
 7630        a += DILITHIUM_N;
 7631    }
 7632}
 7633#endif
 7634
 7635/* If v is negative, evaluate to Q. Otherwise evaluate to 0. Constant-time.
 7636 * (word32)v >> 31 is 1 when v's sign bit is set. 0U - 1 = all-ones,
 7637 * 0U - 0 = 0. & Q yields Q or 0. */
 7638#define DILITHIUM_POS_OFFSET(v) \
 7639    ((sword32)((0U - (((word32)(v)) >> 31)) & DILITHIUM_Q))
 7640
 7641/* Make values in polynomial be in positive range.
 7642 *
 7643 * @param [in, out] a  Polynomial.
 7644 */
 7645static void dilithium_make_pos_c(sword32* a)
 7646{
 7647    unsigned int j;
 7648#ifdef WOLFSSL_DILITHIUM_SMALL
 7649    for (j = 0; j < DILITHIUM_N; j++) {
 7650        a[j] += DILITHIUM_POS_OFFSET(a[j]);
 7651    }
 7652#else
 7653    for (j = 0; j < DILITHIUM_N; j += 8) {
 7654        a[j+0] += DILITHIUM_POS_OFFSET(a[j+0]);
 7655        a[j+1] += DILITHIUM_POS_OFFSET(a[j+1]);
 7656        a[j+2] += DILITHIUM_POS_OFFSET(a[j+2]);
 7657        a[j+3] += DILITHIUM_POS_OFFSET(a[j+3]);
 7658        a[j+4] += DILITHIUM_POS_OFFSET(a[j+4]);
 7659        a[j+5] += DILITHIUM_POS_OFFSET(a[j+5]);
 7660        a[j+6] += DILITHIUM_POS_OFFSET(a[j+6]);
 7661        a[j+7] += DILITHIUM_POS_OFFSET(a[j+7]);
 7662    }
 7663#endif
 7664}
 7665
 7666/* Make values in polynomial be in positive range.
 7667 *
 7668 * @param [in, out] a  Polynomial.
 7669 */
 7670static void dilithium_make_pos(sword32* a)
 7671{
 7672#ifdef USE_INTEL_SPEEDUP
 7673    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
 7674        wc_mldsa_poly_make_pos_avx2(a);
 7675        RESTORE_VECTOR_REGISTERS();
 7676    }
 7677    else
 7678#endif
 7679    {
 7680        dilithium_make_pos_c(a);
 7681    }
 7682}
 7683
 7684#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY) || \
 7685    defined(WOLFSSL_DILITHIUM_CHECK_KEY) || \
 7686    (!defined(WOLFSSL_DILITHIUM_NO_SIGN) && \
 7687     !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM))
 7688/* Make values in polynomials of vector be in positive range.
 7689 *
 7690 * @param [in, out] a  Vector of polynomials.
 7691 * @param [in]      l  Dimension of vector.
 7692 */
 7693static void dilithium_vec_make_pos(sword32* a, byte l)
 7694{
 7695    byte i;
 7696
 7697    for (i = 0; i < l; i++) {
 7698        dilithium_make_pos(a);
 7699        a += DILITHIUM_N;
 7700    }
 7701}
 7702#endif
 7703
 7704#endif /* !WOLFSSL_DILITHIUM_VERIFY_ONLY */
 7705
 7706/******************************************************************************/
 7707
 7708#ifndef WOLFSSL_DILITHIUM_NO_MAKE_KEY
 7709
 7710/* Make a key from a random seed.
 7711 *
 7712 * xi is seed passed in.
 7713 * FIPS 204. 6.1: Algorithm 6 ML-DSA.KeyGen_internal(xi)
 7714 *   1: (rho, rho', K) E B32 x B64 x B32 <- H(xi||k||l, 1024)
 7715 *   2:
 7716 *   3: A_circum <- ExpandA(rho)
 7717 *   4: (s1,s2) <- ExpandS(rho')
 7718 *   5: t <- NTT-1(A_circum o NTT(s1)) + s2
 7719 *   6: (t1, t0) <- Power2Round(t, d)
 7720 *   7: pk <- pkEncode(rho, t1)
 7721 *   8: tr <- H(pk, 64)
 7722 *   9: sk <- skEncode(rho, K, tr, s1, s2, t0)
 7723 *  10: return (pk, sk)
 7724 *
 7725 * FIPS 204. 7.2: Algorithm 22 pkEncode(rho, t1)
 7726 *   1: pk <- rho
 7727 *   2: for i from 0 to k - 1 do
 7728 *   3:     pk <- pk || SimpleBitPack(t1[i], 2^(bitlen(q-1)-d) - 1)
 7729 *   4: end for
 7730 *   5: return pk
 7731 *
 7732 * FIPS 204. 7.2: Algorithm 24 skEncode(rho, K, tr, s, s2, t0)
 7733 *   1: sk <- rho || K || tr
 7734 *   2: for i from 0 to l - 1 do
 7735 *   3:     sk <- sk || BitPack(s1[i], eta, eta)
 7736 *   4: end for
 7737 *   5: for i from 0 to k - 1 do
 7738 *   6:     sk <- sk || BitPack(s2[i], eta, eta)
 7739 *   7: end for
 7740 *   8: for i from 0 to k - 1 do
 7741 *   9:     sk <- sk || BitPack(t0[i], 2^(d-1)-1, 2^(d-1))
 7742 *  10: end for
 7743 *  11: return sk
 7744 *
 7745 * Public and private key store in key.
 7746 *
 7747 * @param [in, out] key   Dilithium key.
 7748 * @param [in]      seed  Seed to hash to generate values.
 7749 * @return  0 on success.
 7750 * @return  MEMORY_E when memory allocation fails.
 7751 * @return  Other negative when an error occurs.
 7752 */
 7753static int dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
 7754{
 7755#ifndef WOLFSSL_DILITHIUM_MAKE_KEY_SMALL_MEM
 7756    int ret = 0;
 7757    const wc_dilithium_params* params = key->params;
 7758    sword32* a = NULL;
 7759    sword32* s1 = NULL;
 7760    sword32* s2 = NULL;
 7761    sword32* t = NULL;
 7762    byte* pub_seed = NULL;
 7763    byte kl[2];
 7764
 7765#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
 7766    ret = dilithium_alloc_priv_buf(key);
 7767    if (ret == 0) {
 7768        ret = dilithium_alloc_pub_buf(key);
 7769    }
 7770#endif
 7771
 7772    if (ret == 0) {
 7773        pub_seed = key->k;
 7774    }
 7775
 7776    /* Allocate memory for large intermediates. */
 7777#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 7778#ifndef WC_DILITHIUM_FIXED_ARRAY
 7779    if (key->a == NULL) {
 7780        key->a = (sword32*)XMALLOC(params->aSz, key->heap,
 7781            DYNAMIC_TYPE_DILITHIUM);
 7782        if (key->a == NULL) {
 7783            ret = MEMORY_E;
 7784        }
 7785    }
 7786#endif
 7787    if (ret == 0) {
 7788        a = key->a;
 7789    }
 7790#endif
 7791#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
 7792#ifndef WC_DILITHIUM_FIXED_ARRAY
 7793    if ((ret == 0) && (key->s1 == NULL)) {
 7794        key->s1 = (sword32*)XMALLOC(params->aSz, key->heap,
 7795            DYNAMIC_TYPE_DILITHIUM);
 7796        if (key->s1 == NULL) {
 7797            ret = MEMORY_E;
 7798        }
 7799        else {
 7800            key->s2 = key->s1  + params->s1Sz / sizeof(*s1);
 7801            key->t0 = key->s2  + params->s2Sz / sizeof(*s2);
 7802        }
 7803    }
 7804#endif
 7805    if (ret == 0) {
 7806        s1 = key->s1;
 7807        s2 = key->s2;
 7808        t  = key->t0;
 7809    }
 7810#else
 7811    if (ret == 0) {
 7812        unsigned int allocSz;
 7813
 7814        allocSz = (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz;
 7815#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 7816        allocSz += params->aSz;
 7817#endif
 7818
 7819        /* s1, s2, t, a */
 7820        s1 = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 7821        if (s1 == NULL) {
 7822            ret = MEMORY_E;
 7823        }
 7824        else {
 7825            s2 = s1 + params->s1Sz / sizeof(*s1);
 7826            t  = s2 + params->s2Sz / sizeof(*s2);
 7827#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 7828            a  = t  + params->s2Sz / sizeof(*t);
 7829#endif
 7830        }
 7831    }
 7832#endif
 7833
 7834    if (ret == 0) {
 7835#ifdef WOLFSSL_DILITHIUM_FIPS204_DRAFT
 7836        if (key->params->level >= WC_ML_DSA_DRAFT) {
 7837            /* Step 2: Create public seed, private seed and K from seed.
 7838             * Step 9; Alg 18, Step 1: Public seed is placed into private key.
 7839             */
 7840            ret = dilithium_shake256(&key->shake, seed, DILITHIUM_SEED_SZ,
 7841                pub_seed, DILITHIUM_SEEDS_SZ);
 7842        }
 7843        else
 7844#endif
 7845        {
 7846            kl[0] = params->k;
 7847            kl[1] = params->l;
 7848            /* Step 1: Create public seed, private seed and K from seed.
 7849             * Step 9; Alg 24, Step 1: Public seed is placed into private key.
 7850             */
 7851            ret = dilithium_hash256(&key->shake, seed, DILITHIUM_SEED_SZ, kl, 2,
 7852                pub_seed, DILITHIUM_SEEDS_SZ);
 7853        }
 7854    }
 7855    if (ret == 0) {
 7856        /* Step 7; Alg 22 Step 1: Copy public seed into public key. */
 7857        XMEMCPY(key->p, pub_seed, DILITHIUM_PUB_SEED_SZ);
 7858
 7859        /* Step 3: Expand public seed into a matrix of polynomials. */
 7860        ret = dilithium_expand_a(&key->shake, pub_seed, params->k, params->l,
 7861            a, key->heap);
 7862    }
 7863    if (ret == 0) {
 7864        byte* priv_seed = key->k + DILITHIUM_PUB_SEED_SZ;
 7865
 7866        /* Step 4: Expand private seed into to vectors of polynomials. */
 7867        ret = dilithium_expand_s(&key->shake, priv_seed, params->eta, s1,
 7868            params->l, s2, params->k);
 7869    }
 7870    if (ret == 0) {
 7871        byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 7872        byte* tr = k + DILITHIUM_K_SZ;
 7873        byte* s1p = tr + DILITHIUM_TR_SZ;
 7874        byte* s2p = s1p + params->s1EncSz;
 7875        byte* t0 = s2p + params->s2EncSz;
 7876        byte* t1 = key->p + DILITHIUM_PUB_SEED_SZ;
 7877
 7878        /* Step 9: Move k down to after public seed. */
 7879        XMEMCPY(k, k + DILITHIUM_PRIV_SEED_SZ, DILITHIUM_K_SZ);
 7880        /* Step 9. Alg 24 Steps 2-4: Encode s1 into private key. */
 7881        dilthium_vec_encode_eta_bits(s1, params->l, params->eta, s1p);
 7882        /* Step 9. Alg 24 Steps 5-7: Encode s2 into private key. */
 7883        dilthium_vec_encode_eta_bits(s2, params->k, params->eta, s2p);
 7884
 7885        /* Step 5: t <- NTT-1(A_circum o NTT(s1)) + s2 */
 7886        dilithium_vec_ntt_small_full(s1, params->l);
 7887        dilithium_matrix_mul(t, a, s1, params->k, params->l);
 7888    #ifdef WOLFSSL_DILITHIUM_SMALL
 7889        dilithium_vec_red(t, params->k);
 7890    #endif
 7891        dilithium_vec_invntt_full(t, params->k);
 7892        dilithium_vec_add(t, s2, params->k);
 7893
 7894        /* Make positive for decomposing. */
 7895        dilithium_vec_make_pos(t, params->k);
 7896        /* Step 6, Step 7, Step 9. Alg 22 Steps 2-4, Alg 24 Steps 8-10.
 7897         * Decompose t in t0 and t1 and encode into public and private key.
 7898         */
 7899        dilithium_vec_encode_t0_t1(t, params->k, t0, t1);
 7900        /* Step 8. Alg 24, Step 1: Hash public key into private key. */
 7901        ret = dilithium_shake256(&key->shake, key->p, params->pkSz, tr,
 7902            DILITHIUM_TR_SZ);
 7903    }
 7904    if (ret == 0) {
 7905        /* Public key and private key are available. */
 7906        key->prvKeySet = 1;
 7907        key->pubKeySet = 1;
 7908#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 7909        /* Matrix A is available. */
 7910        key->aSet = 1;
 7911#endif
 7912#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
 7913        /* Private vectors are not available as they were overwritten. */
 7914        key->privVecsSet = 0;
 7915#endif
 7916#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
 7917        /* Public vector, t1, is not available as it was not created. */
 7918        key->pubVecSet = 0;
 7919#endif
 7920    }
 7921
 7922#ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
 7923    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
 7924#endif
 7925    return ret;
 7926#else
 7927    int ret = 0;
 7928    const wc_dilithium_params* params = key->params;
 7929    sword32* a = NULL;
 7930    sword32* s1 = NULL;
 7931    sword32* s2 = NULL;
 7932    sword32* t = NULL;
 7933#ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 7934    sword64* t64 = NULL;
 7935#endif
 7936    byte* h = NULL;
 7937    byte* pub_seed = NULL;
 7938    unsigned int r;
 7939    unsigned int s;
 7940    byte kl[2];
 7941
 7942#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
 7943    ret = dilithium_alloc_priv_buf(key);
 7944    if (ret == 0) {
 7945        ret = dilithium_alloc_pub_buf(key);
 7946    }
 7947#endif
 7948
 7949    if (ret == 0) {
 7950        pub_seed = key->k;
 7951    }
 7952
 7953    /* Allocate memory for large intermediates. */
 7954    if (ret == 0) {
 7955        unsigned int allocSz;
 7956
 7957        /* s1-l, s2-k, t-k, a-1 */
 7958        allocSz  = (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz +
 7959                   (unsigned int)DILITHIUM_REJ_NTT_POLY_H_SIZE +
 7960                   (unsigned int)DILITHIUM_POLY_SIZE;
 7961    #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 7962        /* t64 */
 7963        allocSz += (unsigned int)DILITHIUM_POLY_SIZE * 2U;
 7964    #endif
 7965        s1 = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 7966        if (s1 == NULL) {
 7967            ret = MEMORY_E;
 7968        }
 7969        else {
 7970            s2 = s1 + params->s1Sz / sizeof(*s1);
 7971            t  = s2 + params->s2Sz / sizeof(*s2);
 7972            h  = (byte*)(t  + params->s2Sz / sizeof(*t));
 7973            a  = (sword32*)(h + DILITHIUM_REJ_NTT_POLY_H_SIZE);
 7974        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 7975            t64 = (sword64*)(a + DILITHIUM_N);
 7976        #endif
 7977        }
 7978    }
 7979
 7980    if (ret == 0) {
 7981#ifdef WOLFSSL_DILITHIUM_FIPS204_DRAFT
 7982        if (key->params->level >= WC_ML_DSA_DRAFT) {
 7983            /* Step 2: Create public seed, private seed and K from seed.
 7984             * Step 9; Alg 18, Step 1: Public seed is placed into private key.
 7985             */
 7986            ret = dilithium_shake256(&key->shake, seed, DILITHIUM_SEED_SZ,
 7987                pub_seed, DILITHIUM_SEEDS_SZ);
 7988        }
 7989        else
 7990#endif
 7991        {
 7992            kl[0] = params->k;
 7993            kl[1] = params->l;
 7994            /* Step 1: Create public seed, private seed and K from seed.
 7995             * Step 9; Alg 24, Step 1: Public seed is placed into private key.
 7996             */
 7997            ret = dilithium_hash256(&key->shake, seed, DILITHIUM_SEED_SZ, kl, 2,
 7998                pub_seed, DILITHIUM_SEEDS_SZ);
 7999        }
 8000    }
 8001    if (ret == 0) {
 8002        byte* priv_seed = key->k + DILITHIUM_PUB_SEED_SZ;
 8003
 8004        /* Step 7; Alg 22 Step 1: Copy public seed into public key. */
 8005        XMEMCPY(key->p, pub_seed, DILITHIUM_PUB_SEED_SZ);
 8006
 8007        /* Step 4: Expand private seed into to vectors of polynomials. */
 8008        ret = dilithium_expand_s(&key->shake, priv_seed, params->eta, s1,
 8009            params->l, s2, params->k);
 8010    }
 8011    if (ret == 0) {
 8012        byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 8013        byte* tr = k + DILITHIUM_K_SZ;
 8014        byte* s1p = tr + DILITHIUM_TR_SZ;
 8015        byte* s2p = s1p + params->s1EncSz;
 8016        byte* t0 = s2p + params->s2EncSz;
 8017        byte* t1 = key->p + DILITHIUM_PUB_SEED_SZ;
 8018        byte aseed[DILITHIUM_GEN_A_SEED_SZ];
 8019        sword32* s2t = s2;
 8020        sword32* tt = t;
 8021
 8022        /* Step 9: Move k down to after public seed. */
 8023        XMEMCPY(k, k + DILITHIUM_PRIV_SEED_SZ, DILITHIUM_K_SZ);
 8024        /* Step 9. Alg 24 Steps 2-4: Encode s1 into private key. */
 8025        dilthium_vec_encode_eta_bits(s1, params->l, params->eta, s1p);
 8026        /* Step 9. Alg 24 Steps 5-7: Encode s2 into private key. */
 8027        dilthium_vec_encode_eta_bits(s2, params->k, params->eta, s2p);
 8028
 8029        /* Step 5: NTT(s1) */
 8030        dilithium_vec_ntt_small_full(s1, params->l);
 8031        /* Step 5: t <- NTT-1(A_circum o NTT(s1)) + s2 */
 8032        XMEMCPY(aseed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 8033        for (r = 0; (ret == 0) && (r < params->k); r++) {
 8034            sword32* s1t = s1;
 8035            unsigned int e;
 8036
 8037            /* Put r/i into buffer to be hashed. */
 8038            aseed[DILITHIUM_PUB_SEED_SZ + 1] = (byte)r;
 8039            for (s = 0; (ret == 0) && (s < params->l); s++) {
 8040                /* Put s into buffer to be hashed. */
 8041                aseed[DILITHIUM_PUB_SEED_SZ + 0] = (byte)s;
 8042                /* Step 3: Expand public seed into a matrix of polynomials. */
 8043                ret = dilithium_rej_ntt_poly_ex(&key->shake, aseed, a, h);
 8044                if (ret != 0) {
 8045                    break;
 8046                }
 8047                /* Matrix multiply. */
 8048            #ifndef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8049                if (s == 0) {
 8050                #ifdef WOLFSSL_DILITHIUM_SMALL
 8051                    for (e = 0; e < DILITHIUM_N; e++) {
 8052                        tt[e] = dilithium_mont_red((sword64)a[e] * s1t[e]);
 8053                    }
 8054                #else
 8055                    for (e = 0; e < DILITHIUM_N; e += 8) {
 8056                        tt[e+0] = dilithium_mont_red((sword64)a[e+0]*s1t[e+0]);
 8057                        tt[e+1] = dilithium_mont_red((sword64)a[e+1]*s1t[e+1]);
 8058                        tt[e+2] = dilithium_mont_red((sword64)a[e+2]*s1t[e+2]);
 8059                        tt[e+3] = dilithium_mont_red((sword64)a[e+3]*s1t[e+3]);
 8060                        tt[e+4] = dilithium_mont_red((sword64)a[e+4]*s1t[e+4]);
 8061                        tt[e+5] = dilithium_mont_red((sword64)a[e+5]*s1t[e+5]);
 8062                        tt[e+6] = dilithium_mont_red((sword64)a[e+6]*s1t[e+6]);
 8063                        tt[e+7] = dilithium_mont_red((sword64)a[e+7]*s1t[e+7]);
 8064                    }
 8065                #endif
 8066                }
 8067                else {
 8068                #ifdef WOLFSSL_DILITHIUM_SMALL
 8069                    for (e = 0; e < DILITHIUM_N; e++) {
 8070                        tt[e] += dilithium_mont_red((sword64)a[e] * s1t[e]);
 8071                    }
 8072                #else
 8073                    for (e = 0; e < DILITHIUM_N; e += 8) {
 8074                        tt[e+0] += dilithium_mont_red((sword64)a[e+0]*s1t[e+0]);
 8075                        tt[e+1] += dilithium_mont_red((sword64)a[e+1]*s1t[e+1]);
 8076                        tt[e+2] += dilithium_mont_red((sword64)a[e+2]*s1t[e+2]);
 8077                        tt[e+3] += dilithium_mont_red((sword64)a[e+3]*s1t[e+3]);
 8078                        tt[e+4] += dilithium_mont_red((sword64)a[e+4]*s1t[e+4]);
 8079                        tt[e+5] += dilithium_mont_red((sword64)a[e+5]*s1t[e+5]);
 8080                        tt[e+6] += dilithium_mont_red((sword64)a[e+6]*s1t[e+6]);
 8081                        tt[e+7] += dilithium_mont_red((sword64)a[e+7]*s1t[e+7]);
 8082                    }
 8083                #endif
 8084                }
 8085            #else
 8086                if (s == 0) {
 8087                #ifdef WOLFSSL_DILITHIUM_SMALL
 8088                    for (e = 0; e < DILITHIUM_N; e++) {
 8089                        t64[e] = (sword64)a[e] * s1t[e];
 8090                    }
 8091                #else
 8092                    for (e = 0; e < DILITHIUM_N; e += 8) {
 8093                        t64[e+0] = (sword64)a[e+0] * s1t[e+0];
 8094                        t64[e+1] = (sword64)a[e+1] * s1t[e+1];
 8095                        t64[e+2] = (sword64)a[e+2] * s1t[e+2];
 8096                        t64[e+3] = (sword64)a[e+3] * s1t[e+3];
 8097                        t64[e+4] = (sword64)a[e+4] * s1t[e+4];
 8098                        t64[e+5] = (sword64)a[e+5] * s1t[e+5];
 8099                        t64[e+6] = (sword64)a[e+6] * s1t[e+6];
 8100                        t64[e+7] = (sword64)a[e+7] * s1t[e+7];
 8101                    }
 8102                #endif
 8103                }
 8104                else {
 8105                #ifdef WOLFSSL_DILITHIUM_SMALL
 8106                    for (e = 0; e < DILITHIUM_N; e++) {
 8107                        t64[e] += (sword64)a[e] * s1t[e];
 8108                    }
 8109                #else
 8110                    for (e = 0; e < DILITHIUM_N; e += 8) {
 8111                        t64[e+0] += (sword64)a[e+0] * s1t[e+0];
 8112                        t64[e+1] += (sword64)a[e+1] * s1t[e+1];
 8113                        t64[e+2] += (sword64)a[e+2] * s1t[e+2];
 8114                        t64[e+3] += (sword64)a[e+3] * s1t[e+3];
 8115                        t64[e+4] += (sword64)a[e+4] * s1t[e+4];
 8116                        t64[e+5] += (sword64)a[e+5] * s1t[e+5];
 8117                        t64[e+6] += (sword64)a[e+6] * s1t[e+6];
 8118                        t64[e+7] += (sword64)a[e+7] * s1t[e+7];
 8119                    }
 8120                #endif
 8121                }
 8122            #endif
 8123                /* Next polynomial. */
 8124                s1t += DILITHIUM_N;
 8125            }
 8126        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8127            for (e = 0; e < DILITHIUM_N; e++) {
 8128                tt[e] = dilithium_mont_red(t64[e]);
 8129            }
 8130        #endif
 8131            dilithium_invntt_full(tt);
 8132            dilithium_add(tt, s2t);
 8133            /* Make positive for decomposing. */
 8134            dilithium_make_pos(tt);
 8135
 8136            tt += DILITHIUM_N;
 8137            s2t += DILITHIUM_N;
 8138        }
 8139
 8140        /* Step 6, Step 7, Step 9. Alg 22 Steps 2-4, Alg 24 Steps 8-10.
 8141         * Decompose t in t0 and t1 and encode into public and private key.
 8142         */
 8143        dilithium_vec_encode_t0_t1(t, params->k, t0, t1);
 8144        /* Step 8. Alg 24, Step 1: Hash public key into private key. */
 8145        ret = dilithium_shake256(&key->shake, key->p, params->pkSz, tr,
 8146            DILITHIUM_TR_SZ);
 8147    }
 8148    if (ret == 0) {
 8149        /* Public key and private key are available. */
 8150        key->prvKeySet = 1;
 8151        key->pubKeySet = 1;
 8152    }
 8153
 8154    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8155    return ret;
 8156#endif
 8157}
 8158
 8159/* Make a key from a random seed.
 8160 *
 8161 * FIPS 204. 5.1: Algorithm 1 ML-DSA.KeyGen()
 8162 *   1: xi <- B32  [Choose random seed]
 8163 *   2: if xi = NULL then
 8164 *   3:   return falsam
 8165 *   4: end if
 8166 *   5: return ML-DSA.KeyGen_internal(xi)
 8167 *
 8168 * @param [in, out] key  Dilithium key.
 8169 * @param [in]      rng  Random number generator.
 8170 * @return  0 on success.
 8171 * @return  MEMORY_E when memory allocation fails.
 8172 * @return  Other negative when an error occurs.
 8173 */
 8174static int dilithium_make_key(dilithium_key* key, WC_RNG* rng)
 8175{
 8176    int ret;
 8177    byte seed[DILITHIUM_SEED_SZ];
 8178
 8179    /* Step 1: Generate a 32 byte random seed. */
 8180    ret = wc_RNG_GenerateBlock(rng, seed, DILITHIUM_SEED_SZ);
 8181    /* Step 2: Check for error. */
 8182    if (ret == 0) {
 8183        /* Step 5: Make key with random seed. */
 8184        ret = wc_dilithium_make_key_from_seed(key, seed);
 8185    }
 8186
 8187    ForceZero(seed, sizeof(seed));
 8188    return ret;
 8189}
 8190#endif /* !WOLFSSL_DILITHIUM_NO_MAKE_KEY */
 8191
 8192#ifndef WOLFSSL_DILITHIUM_NO_SIGN
 8193
 8194#if !defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM) || \
 8195    defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC) || \
 8196    defined(WC_DILITHIUM_CACHE_PRIV_VECTORS)
 8197/* Decode, from private key, and NTT private key vectors s1, s2, and t0.
 8198 *
 8199 * FIPS 204. 6: Algorithm 2 ML-DSA.Sign(sk, M)
 8200 *   1: (rho, K, tr, s1, s2, t0) <- skDecode(sk)
 8201 *   2: s1_circum <- NTT(s1)
 8202 *   3: s2_circum <- NTT(s2)
 8203 *   4: t0_circum <- NTT(t0)
 8204 *
 8205 * @param [in, out] key  Dilithium key.
 8206 * @param [out]     s1   Vector of polynomials s1.
 8207 * @param [out]     s2   Vector of polynomials s2.
 8208 * @param [out]     t0   Vector of polynomials t0.
 8209 */
 8210static void dilithium_make_priv_vecs(dilithium_key* key, sword32* s1,
 8211    sword32* s2, sword32* t0)
 8212{
 8213    const wc_dilithium_params* params = key->params;
 8214    const byte* pubSeed = key->k;
 8215    const byte* k = pubSeed + DILITHIUM_PUB_SEED_SZ;
 8216    const byte* tr = k + DILITHIUM_K_SZ;
 8217    const byte* s1p = tr + DILITHIUM_TR_SZ;
 8218    const byte* s2p = s1p + params->s1EncSz;
 8219    const byte* t0p = s2p + params->s2EncSz;
 8220
 8221    /* Step 1: Decode s1, s2, t0. */
 8222    dilithium_vec_decode_eta_bits(s1p, params->eta, s1, params->l);
 8223    dilithium_vec_decode_eta_bits(s2p, params->eta, s2, params->k);
 8224    dilithium_vec_decode_t0(t0p, params->k, t0);
 8225
 8226    /* Step 2: NTT s1. */
 8227    dilithium_vec_ntt_small(s1, params->l);
 8228    /* Step 3: NTT s2. */
 8229    dilithium_vec_ntt_small(s2, params->k);
 8230    /* Step 4: NTT t0. */
 8231    dilithium_vec_ntt(t0, params->k);
 8232
 8233#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
 8234    /* Private key vectors have been created. */
 8235    key->privVecsSet = 1;
 8236#endif
 8237}
 8238#endif
 8239
 8240/* Sign a message with the key and a seed.
 8241 *
 8242 * FIPS 204. 5.2: Algorithm 2 ML-DSA.sign(sk, M, ctx)
 8243 *   ...
 8244 *  10: M' <- ByyesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
 8245 *                        ctx) || M
 8246 *   ...
 8247 *
 8248 * FIPS 204. 6: Algorithm 2 ML-DSA.Sign(sk, M)
 8249 *   1: (rho, K, tr, s1, s2, t0) <- skDecode(sk)
 8250 *   2: s1_circum <- NTT(s1)
 8251 *   3: s2_circum <- NTT(s2)
 8252 *   4: t0_circum <- NTT(t0)
 8253 *   5: A_circum <- ExpandA(rho)
 8254 *   6: mu <- H(tr||M, 512)
 8255 *   7: rnd <- {0,1}256
 8256 *   8: rho' <- H(K||rnd||mu, 512)
 8257 *   9: kappa <- 0
 8258 *  10: (z, h) <- falsam
 8259 *  11: while (z, h) = falsam do
 8260 *  12:    y <- ExpandMask(rho', kappa)
 8261 *  13:    w <- NTT-1(A_circum o NTT(y))
 8262 *  14:    w1 <- HighBits(w)
 8263 *  15:    c_tilde E {0,1}2*lambda <- H(mu|w1Encode(w1), 2 * lambda)
 8264 *  16:    (c1_tilde, c2_tilde) E {0,1}256 x {0,1}2*lambda-256 <- c_tilde
 8265 *  17:     c < SampleInBall(c1_tilde)
 8266 *  18:     c_circum <- NTT(c)
 8267 *  19:     <<cs1>> <- NTT-1(c_circum o s1_circum)
 8268 *  20:     <<cs2>> <- NTT-1(c_circum o s2_circum)
 8269 *  21:     z <- y + <<cs1>>
 8270 *  22:     r0 <- LowBits(w - <<cs2>>
 8271 *  23:     if ||z||inf >= GAMMA1 - BETA or ||r0||inf GAMMA2 - BETA then
 8272 *                                                             (z, h) <- falsam
 8273 *  24:     else
 8274 *  25:         <<ct0>> <- NTT-1(c_circum o t0_circum)
 8275 *  26:         h < MakeHint(-<<ct0>>, w - <<sc2>> + <<ct0>>)
 8276 *  27:         if (||<<ct>>||inf >= GAMMMA1 or
 8277 *                 the number of 1's in h is greater than OMEGA, then
 8278 *                                                             (z, h) <- falsam
 8279 *  28:         end if
 8280 *  29:     end if
 8281 *  30:     kappa <- kappa + l
 8282 *  31: end while
 8283 *  32: sigma <- sigEncode(c_tilde, z mod +/- q, h)
 8284 *  33: return sigma
 8285 *
 8286 * @param [in, out] key     Dilithium key.
 8287 * @param [in]      seedMu  Random seed || mu.
 8288 * @param [out]     sig     Buffer to hold signature.
 8289 * @param [in, out] sigLen  On in, length of buffer in bytes.
 8290 *                          On out, the length of the signature in bytes.
 8291 * @return  0 on success.
 8292 * @return  BAD_FUNC_ARG when context length is greater than 255.
 8293 * @return  BUFFER_E when the signature buffer is too small.
 8294 * @return  MEMORY_E when memory allocation fails.
 8295 * @return  Other negative when an error occurs.
 8296 */
 8297static int dilithium_sign_with_seed_mu(dilithium_key* key,
 8298    const byte* seedMu, byte* sig, word32 *sigLen)
 8299{
 8300#ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM
 8301    int ret = 0;
 8302    const wc_dilithium_params* params = key->params;
 8303    const byte* pub_seed = key->k;
 8304    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 8305    const byte* mu = seedMu + DILITHIUM_RND_SZ;
 8306    sword32* a = NULL;
 8307    sword32* s1 = NULL;
 8308    sword32* s2 = NULL;
 8309    sword32* t0 = NULL;
 8310    sword32* y = NULL;
 8311    sword32* w0 = NULL;
 8312    sword32* w1 = NULL;
 8313    sword32* c = NULL;
 8314    sword32* z = NULL;
 8315    sword32* ct0 = NULL;
 8316    byte priv_rand_seed[DILITHIUM_Y_SEED_SZ];
 8317    byte* h = sig + params->lambda / 4 + params->zEncSz;
 8318    unsigned int allocSz = 0;
 8319#ifdef WC_MLDSA_FAULT_HARDEN
 8320   sword32* y_check;
 8321#endif
 8322
 8323    /* Check the signature buffer isn't too small. */
 8324    if (*sigLen < params->sigSz) {
 8325        ret = BUFFER_E;
 8326    }
 8327    if (ret == 0) {
 8328        /* Return the size of the signature. */
 8329        *sigLen = params->sigSz;
 8330    }
 8331
 8332    /* Allocate memory for large intermediates. */
 8333#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 8334#ifndef WC_DILITHIUM_FIXED_ARRAY
 8335    if ((ret == 0) && (key->a == NULL)) {
 8336        key->a = (sword32*)XMALLOC((size_t)params->aSz, key->heap,
 8337            DYNAMIC_TYPE_DILITHIUM);
 8338        if (key->a == NULL) {
 8339            ret = MEMORY_E;
 8340        }
 8341        else {
 8342            XMEMSET(key->a, 0, (size_t)params->aSz);
 8343        }
 8344    }
 8345#endif
 8346    if (ret == 0) {
 8347        a = key->a;
 8348    }
 8349#endif
 8350#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
 8351#ifndef WC_DILITHIUM_FIXED_ARRAY
 8352    if ((ret == 0) && (key->s1 == NULL)) {
 8353        key->s1 = (sword32*)XMALLOC(params->aSz, key->heap,
 8354            DYNAMIC_TYPE_DILITHIUM);
 8355        if (key->s1 == NULL) {
 8356            ret = MEMORY_E;
 8357        }
 8358        else {
 8359            XMEMSET(key->s1, 0, params->aSz);
 8360            key->s2 = key->s1  + params->s1Sz / sizeof(*s1);
 8361            key->t0 = key->s2  + params->s2Sz / sizeof(*s2);
 8362        }
 8363    }
 8364#endif
 8365    if (ret == 0) {
 8366        s1 = key->s1;
 8367        s2 = key->s2;
 8368        t0 = key->t0;
 8369    }
 8370#endif
 8371    if (ret == 0) {
 8372        /* y-l, w0-k, w1-k, c-1, z-l, ct0-k */
 8373        allocSz = (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz +
 8374            (unsigned int)DILITHIUM_POLY_SIZE + params->s1Sz + params->s2Sz;
 8375#ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
 8376        /* s1-l, s2-k, t0-k */
 8377        allocSz += (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz;
 8378#endif
 8379#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 8380        /* A */
 8381        allocSz += params->aSz;
 8382#endif
 8383        y = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8384        if (y == NULL) {
 8385            ret = MEMORY_E;
 8386        }
 8387        else {
 8388        #ifdef WC_MLDSA_FAULT_HARDEN
 8389            y_check = y;
 8390        #endif
 8391            w0  = y   + params->s1Sz / sizeof(*y);
 8392            w1  = w0  + params->s2Sz / sizeof(*w0);
 8393            c   = w1  + params->s2Sz / sizeof(*w1);
 8394            z   = c   + DILITHIUM_N;
 8395            ct0 = z   + params->s1Sz / sizeof(*z);
 8396#ifndef WC_DILITHIUM_CACHE_PRIV_VECTORS
 8397            s1  = ct0 + params->s2Sz / sizeof(*ct0);
 8398            s2  = s1  + params->s1Sz / sizeof(*s1);
 8399            t0  = s2  + params->s2Sz / sizeof(*s2);
 8400#endif
 8401#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 8402            a   = t0  + params->s2Sz / sizeof(*s2);
 8403#endif
 8404        }
 8405    }
 8406
 8407    if (ret == 0) {
 8408#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
 8409        /* Check that we haven't already cached the private vectors. */
 8410        if (!key->privVecsSet)
 8411#endif
 8412        {
 8413            /* Steps 1-4: Decode and NTT vectors s1, s2, and t0. */
 8414            dilithium_make_priv_vecs(key, s1, s2, t0);
 8415        }
 8416
 8417#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 8418        /* Check that we haven't already cached the matrix A. */
 8419        if (!key->aSet)
 8420#endif
 8421        {
 8422            /* Step 5: Create the matrix A from the public seed. */
 8423            ret = dilithium_expand_a(&key->shake, pub_seed, params->k,
 8424                params->l, a, key->heap);
 8425#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 8426            key->aSet = (ret == 0);
 8427#endif
 8428        }
 8429    }
 8430    if (ret == 0) {
 8431        /* Step 9: Compute private random using hash. */
 8432        ret = dilithium_hash256(&key->shake, k, DILITHIUM_K_SZ, seedMu,
 8433            DILITHIUM_RND_SZ + DILITHIUM_MU_SZ, priv_rand_seed,
 8434            DILITHIUM_PRIV_RAND_SEED_SZ);
 8435    }
 8436    if (ret == 0) {
 8437        word16 kappa = 0;
 8438        int valid = 0;
 8439
 8440        /* Step 11: Start rejection sampling loop */
 8441        do {
 8442            WC_DECLARE_VAR(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ, 0);
 8443            sword32* w = w1;
 8444            sword32* y_ntt = z;
 8445            sword32* cs2 = ct0;
 8446            byte* commit = sig;
 8447
 8448            /* Step 12: Compute vector y from private random seed and kappa. */
 8449            dilithium_vec_expand_mask(&key->shake, priv_rand_seed, kappa,
 8450                params->gamma1_bits, y, params->l);
 8451        #ifdef WOLFSSL_DILITHIUM_SIGN_CHECK_Y
 8452            valid = dilithium_vec_check_low(y, params->l,
 8453                ((sword32)1 << params->gamma1_bits) - params->beta);
 8454            if (valid)
 8455        #endif
 8456            {
 8457                /* Step 13: NTT-1(A o NTT(y)) */
 8458                XMEMCPY(y_ntt, y, params->s1Sz);
 8459            #ifdef WC_MLDSA_FAULT_HARDEN
 8460                if (y_check != y) {
 8461                    valid = 0;
 8462                    ret = BAD_COND_E;
 8463                }
 8464            }
 8465            if (ret == 0) {
 8466            #endif
 8467                dilithium_vec_ntt_full(y_ntt, params->l);
 8468                dilithium_matrix_mul(w, a, y_ntt, params->k, params->l);
 8469            #ifdef WOLFSSL_DILITHIUM_SMALL
 8470                dilithium_vec_red(w, params->k);
 8471            #endif
 8472                dilithium_vec_invntt_full(w, params->k);
 8473                /* Step 14, Step 22: Make values positive and decompose. */
 8474                dilithium_vec_make_pos(w, params->k);
 8475                dilithium_vec_decompose(w, params->k, params->gamma2, w0, w1);
 8476        #ifdef WOLFSSL_DILITHIUM_SIGN_CHECK_W0
 8477                valid = dilithium_vec_check_low(w0, params->k,
 8478                    params->gamma2 - params->beta);
 8479            }
 8480            if (valid) {
 8481        #endif
 8482                /* Step 15: Encode w1. */
 8483                WC_ALLOC_VAR_EX(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ, key->heap,
 8484                    DYNAMIC_TYPE_DILITHIUM, ret=MEMORY_E);
 8485                if (WC_VAR_OK(w1e))
 8486                {
 8487                    dilithium_vec_encode_w1(w1, params->k, params->gamma2, w1e);
 8488                    /* Step 15: Hash mu and encoded w1.
 8489                     * Step 32: Hash is stored in signature. */
 8490                    ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ,
 8491                        w1e, params->w1EncSz, commit, params->lambda / 4);
 8492                }
 8493                if (ret == 0) {
 8494                    /* Step 17: Compute c from first 256 bits of commit. */
 8495                    ret = dilithium_sample_in_ball(params->level, &key->shake,
 8496                        commit, params->lambda / 4, params->tau, c, key->heap);
 8497                }
 8498                if (ret == 0) {
 8499                    sword32 hi;
 8500                    byte i;
 8501
 8502                    valid = 1;
 8503                    /* Step 18: NTT(c). */
 8504                    dilithium_ntt_small(c);
 8505                    hi = params->gamma2 - params->beta;
 8506                    for (i = 0; valid && i < params->k; i++) {
 8507                        /* Step 20: cs2 = NTT-1(c o s2) */
 8508                        dilithium_mul(cs2 + i * DILITHIUM_N, c,
 8509                            s2 + i * DILITHIUM_N);
 8510                        dilithium_invntt(cs2 + i * DILITHIUM_N);
 8511                        /* Step 22: w0 - cs2 */
 8512                        dilithium_sub(w0 + i * DILITHIUM_N,
 8513                            cs2 + i * DILITHIUM_N);
 8514                        /* Step 23: Check w0 - cs2 has low enough values. */
 8515                        valid = dilithium_vec_check_low(w0 + i * DILITHIUM_N, 1,
 8516                            hi);
 8517                    }
 8518                    hi = ((sword32)1 << params->gamma1_bits) - params->beta;
 8519                    for (i = 0; valid && i < params->l; i++) {
 8520                        /* Step 19: cs1 = NTT-1(c o s1) */
 8521                        dilithium_mul(z + i * DILITHIUM_N, c,
 8522                            s1 + i * DILITHIUM_N);
 8523                        dilithium_invntt(z + i * DILITHIUM_N);
 8524                        /* Step 21: z = y + cs1 */
 8525                        dilithium_add(z + i * DILITHIUM_N, y + i * DILITHIUM_N);
 8526                        dilithium_poly_red(z + i * DILITHIUM_N);
 8527                        /* Step 23: Check z has low enough values. */
 8528                        valid = dilithium_vec_check_low(z + i * DILITHIUM_N, 1,
 8529                            hi);
 8530                    }
 8531                    for (i = 0; valid && i < params->k; i++) {
 8532                        /* Step 25: ct0 = NTT-1(c o t0) */
 8533                        dilithium_mul(ct0 + i * DILITHIUM_N, c,
 8534                            t0 + i * DILITHIUM_N);
 8535                        dilithium_invntt(ct0 + i * DILITHIUM_N);
 8536                        /* Step 27: Check ct0 has low enough values. */
 8537                        hi = params->gamma2;
 8538                        valid = dilithium_vec_check_low(ct0 + i * DILITHIUM_N,
 8539                            1, hi);
 8540                    }
 8541                    if (valid) {
 8542                        /* Step 26: ct0 = ct0 + w0 */
 8543                        dilithium_vec_add(ct0, w0, params->k);
 8544                        dilithium_vec_red(ct0, params->k);
 8545                        /* Step 26, 27: Make hint from ct0 and w1 and check
 8546                         * number of hints is valid.
 8547                         * Step 32: h is encoded into signature.
 8548                         */
 8549                        valid = (dilithium_make_hint(ct0, w1, params->k,
 8550                            params->gamma2, params->omega, h) >= 0);
 8551                    }
 8552                }
 8553
 8554                WC_FREE_VAR_EX(w1e, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8555            }
 8556
 8557            if (!valid) {
 8558                /* Too many attempts - something wrong with implementation. */
 8559                if ((kappa > (word16)(kappa + params->l))) {
 8560                    ret = BAD_COND_E;
 8561                }
 8562
 8563                /* Step 30: increment value to append to seed to unique value.
 8564                 */
 8565                kappa = (word16)(kappa + params->l);
 8566            }
 8567        }
 8568        /* Step 11: Check we have a valid signature. */
 8569        while ((ret == 0) && (!valid));
 8570    }
 8571    if (ret == 0) {
 8572        byte* ze = sig + params->lambda / 4;
 8573        /* Step 32: Encode z into signature.
 8574         * Commit (c) and h already encoded into signature. */
 8575        dilithium_vec_encode_gamma1(z, params->l, params->gamma1_bits, ze);
 8576    }
 8577
 8578    ForceZero(priv_rand_seed, sizeof(priv_rand_seed));
 8579    if (y != NULL) {
 8580        ForceZero(y, allocSz);
 8581    }
 8582    XFREE(y, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8583    return ret;
 8584#else
 8585    int ret = 0;
 8586    const wc_dilithium_params* params = key->params;
 8587    const byte* pub_seed = key->k;
 8588    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 8589    const byte* tr = k + DILITHIUM_K_SZ;
 8590    const byte* s1p = tr + DILITHIUM_TR_SZ;
 8591    const byte* s2p = s1p + params->s1EncSz;
 8592    const byte* t0p = s2p + params->s2EncSz;
 8593    const byte* mu = seedMu + DILITHIUM_RND_SZ;
 8594    sword32* a = NULL;
 8595    sword32* s1 = NULL;
 8596    sword32* s2 = NULL;
 8597    sword32* t0 = NULL;
 8598    sword32* y = NULL;
 8599    sword32* y_ntt = NULL;
 8600    sword32* w0 = NULL;
 8601    sword32* w1 = NULL;
 8602    sword32* c = NULL;
 8603    sword32* z = NULL;
 8604    sword32* ct0 = NULL;
 8605#ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8606    sword64* t64 = NULL;
 8607#endif
 8608    byte* blocks = NULL;
 8609    byte priv_rand_seed[DILITHIUM_Y_SEED_SZ];
 8610    byte* h = sig + params->lambda / 4 + params->zEncSz;
 8611    unsigned int allocSz = 0;
 8612#ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8613    byte maxK = (byte)min(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A,
 8614        params->k);
 8615#endif
 8616#ifdef WC_MLDSA_FAULT_HARDEN
 8617    sword32* y_check;
 8618#endif
 8619
 8620    /* Check the signature buffer isn't too small. */
 8621    if ((ret == 0) && (*sigLen < params->sigSz)) {
 8622        ret = BUFFER_E;
 8623    }
 8624    if (ret == 0) {
 8625        /* Return the size of the signature. */
 8626        *sigLen = params->sigSz;
 8627    }
 8628
 8629    /* Allocate memory for large intermediates. */
 8630    if (ret == 0) {
 8631        /* y-l, w0-k, w1-k, blocks, c-1, z-1, A-1 */
 8632        allocSz  = (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz +
 8633                   (unsigned int)DILITHIUM_REJ_NTT_POLY_H_SIZE +
 8634                   (unsigned int)DILITHIUM_POLY_SIZE +
 8635                   (unsigned int)DILITHIUM_POLY_SIZE +
 8636                   (unsigned int)DILITHIUM_POLY_SIZE;
 8637    #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 8638        allocSz += (unsigned int)params->s1Sz + params->s2Sz + params->s2Sz;
 8639    #elif defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A)
 8640        allocSz += (unsigned int)maxK * params->l *
 8641                   (unsigned int)DILITHIUM_POLY_SIZE;
 8642    #endif
 8643    #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8644        allocSz += (unsigned int)DILITHIUM_POLY_SIZE * 2U;
 8645    #endif
 8646        y = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8647        if (y == NULL) {
 8648            ret = MEMORY_E;
 8649        }
 8650        else {
 8651        #ifdef WC_MLDSA_FAULT_HARDEN
 8652            y_check = y;
 8653        #endif
 8654            w0     = y  + params->s1Sz / sizeof(*y_ntt);
 8655            w1     = w0 + params->s2Sz / sizeof(*w0);
 8656            blocks = (byte*)(w1 + params->s2Sz / sizeof(*w1));
 8657            c      = (sword32*)(blocks + DILITHIUM_REJ_NTT_POLY_H_SIZE);
 8658            z      = c  + DILITHIUM_N;
 8659            a      = z  + DILITHIUM_N;
 8660            ct0    = z;
 8661    #if defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A)
 8662            y_ntt  = w0;
 8663            s1     = z;
 8664            s2     = z;
 8665            t0     = z;
 8666        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8667            t64    = (sword64*)(a + (1 + maxK * params->l) * DILITHIUM_N);
 8668        #endif
 8669    #elif defined(WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC)
 8670            y_ntt  = z;
 8671            s1     = a  + DILITHIUM_N;
 8672            s2     = s1 + params->s1Sz / sizeof(*s1);
 8673            t0     = s2 + params->s2Sz / sizeof(*s2);
 8674        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8675            t64    = (sword64*)(t0 + params->s2Sz / sizeof(*t0));
 8676        #endif
 8677    #else
 8678            y_ntt  = z;
 8679            s1     = z;
 8680            s2     = z;
 8681            t0     = z;
 8682        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8683            t64    = (sword64*)(a + DILITHIUM_N);
 8684        #endif
 8685    #endif
 8686        }
 8687    }
 8688
 8689    if (ret == 0) {
 8690        /* Step 9: Compute private random using hash. */
 8691        ret = dilithium_hash256(&key->shake, k, DILITHIUM_K_SZ, seedMu,
 8692            DILITHIUM_RND_SZ + DILITHIUM_MU_SZ, priv_rand_seed,
 8693            DILITHIUM_PRIV_RAND_SEED_SZ);
 8694    }
 8695#ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 8696    if (ret == 0) {
 8697        dilithium_make_priv_vecs(key, s1, s2, t0);
 8698    }
 8699#endif
 8700#ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8701    if (ret == 0) {
 8702        /* Step 5: Create the matrix A from the public seed. */
 8703        ret = dilithium_expand_a(&key->shake, pub_seed, maxK, params->l, a,
 8704            key->heap);
 8705    }
 8706#endif
 8707    if (ret == 0) {
 8708        word16 kappa = 0;
 8709        int valid;
 8710
 8711        /* Step 11: Start rejection sampling loop */
 8712        do {
 8713            byte aseed[DILITHIUM_GEN_A_SEED_SZ];
 8714            WC_DECLARE_VAR(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ, 0);
 8715            sword32* w = w1;
 8716            byte* commit = sig;
 8717            byte r;
 8718            byte s;
 8719            sword32 hi;
 8720            sword32* wt = w;
 8721            sword32* w0t = w0;
 8722            sword32* w1t = w1;
 8723            sword32* at = a;
 8724
 8725        #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8726            w0t += WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A * DILITHIUM_N;
 8727            w1t += WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A * DILITHIUM_N;
 8728            wt += WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A * DILITHIUM_N;
 8729            at += WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A * params->l *
 8730                DILITHIUM_N;
 8731        #endif
 8732
 8733            valid = 1;
 8734            /* Step 12: Compute vector y from private random seed and kappa. */
 8735            dilithium_vec_expand_mask(&key->shake, priv_rand_seed, kappa,
 8736                params->gamma1_bits, y, params->l);
 8737        #ifdef WOLFSSL_DILITHIUM_SIGN_CHECK_Y
 8738            valid = dilithium_vec_check_low(y, params->l,
 8739                ((sword32)1 << params->gamma1_bits) - params->beta);
 8740        #endif
 8741
 8742        #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8743            /* Step 13: NTT-1(A o NTT(y)) */
 8744            XMEMCPY(y_ntt, y, params->s1Sz);
 8745            dilithium_vec_ntt_full(y_ntt, params->l);
 8746            dilithium_matrix_mul(w, a, y_ntt, maxK, params->l);
 8747        #ifdef WOLFSSL_DILITHIUM_SMALL
 8748            dilithium_vec_red(w, params->k);
 8749        #endif
 8750            dilithium_vec_invntt_full(w, maxK);
 8751            /* Step 14, Step 22: Make values positive and decompose. */
 8752            dilithium_vec_make_pos(w, maxK);
 8753            dilithium_vec_decompose(w, maxK, params->gamma2, w0, w1);
 8754        #endif
 8755            /* Step 5: Create the matrix A from the public seed. */
 8756            /* Copy the seed into a buffer that has space for s and r. */
 8757            XMEMCPY(aseed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 8758        #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8759            r = WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A;
 8760        #else
 8761            r = 0;
 8762        #endif
 8763            /* Alg 26. Step 1: Loop over first dimension of matrix. */
 8764            for (; (ret == 0) && valid && (r < params->k); r++) {
 8765                unsigned int e;
 8766                sword32* yt = y;
 8767            #ifdef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC_A
 8768                sword32* y_ntt_t = z;
 8769            #else
 8770                sword32* y_ntt_t = y_ntt;
 8771            #endif
 8772            #ifdef WC_MLDSA_FAULT_HARDEN
 8773                sword32* yt_check = yt;
 8774            #endif
 8775            #ifdef WC_MLDSA_FAULT_HARDEN
 8776                if (y_check != y) {
 8777                    valid = 0;
 8778                    ret = BAD_COND_E;
 8779                    break;
 8780                }
 8781            #endif
 8782
 8783                /* Put r/i into buffer to be hashed. */
 8784                aseed[DILITHIUM_PUB_SEED_SZ + 1] = r;
 8785                /* Alg 26. Step 2: Loop over second dimension of matrix. */
 8786                for (s = 0; (ret == 0) && (s < params->l); s++) {
 8787                    /* Put s into buffer to be hashed. */
 8788                    aseed[DILITHIUM_PUB_SEED_SZ + 0] = s;
 8789                    /* Alg 26. Step 3: Create polynomial from hashing seed. */
 8790                    ret = dilithium_rej_ntt_poly_ex(&key->shake, aseed, at,
 8791                        blocks);
 8792                    if (ret != 0) {
 8793                        break;
 8794                    }
 8795                    XMEMCPY(y_ntt_t, yt, DILITHIUM_POLY_SIZE);
 8796                #ifdef WC_MLDSA_FAULT_HARDEN
 8797                    if (yt_check + s * DILITHIUM_N != yt) {
 8798                        ret = BAD_COND_E;
 8799                        break;
 8800                    }
 8801                #endif
 8802                    dilithium_ntt_full(y_ntt_t);
 8803                    /* Matrix multiply. */
 8804                #ifndef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8805                    if (s == 0) {
 8806                    #ifdef WOLFSSL_DILITHIUM_SMALL
 8807                        for (e = 0; e < DILITHIUM_N; e++) {
 8808                            wt[e] = dilithium_mont_red((sword64)at[e] *
 8809                                y_ntt_t[e]);
 8810                        }
 8811                    #else
 8812                        for (e = 0; e < DILITHIUM_N; e += 8) {
 8813                            wt[e + 0] = dilithium_mont_red((sword64)at[e + 0] *
 8814                                y_ntt_t[e + 0]);
 8815                            wt[e + 1] = dilithium_mont_red((sword64)at[e + 1] *
 8816                                y_ntt_t[e + 1]);
 8817                            wt[e + 2] = dilithium_mont_red((sword64)at[e + 2] *
 8818                                y_ntt_t[e + 2]);
 8819                            wt[e + 3] = dilithium_mont_red((sword64)at[e + 3] *
 8820                                y_ntt_t[e + 3]);
 8821                            wt[e + 4] = dilithium_mont_red((sword64)at[e + 4] *
 8822                                y_ntt_t[e + 4]);
 8823                            wt[e + 5] = dilithium_mont_red((sword64)at[e + 5] *
 8824                                y_ntt_t[e + 5]);
 8825                            wt[e + 6] = dilithium_mont_red((sword64)at[e + 6] *
 8826                                y_ntt_t[e + 6]);
 8827                            wt[e + 7] = dilithium_mont_red((sword64)at[e + 7] *
 8828                                y_ntt_t[e + 7]);
 8829                        }
 8830                    #endif
 8831                    }
 8832                    else {
 8833                    #ifdef WOLFSSL_DILITHIUM_SMALL
 8834                        for (e = 0; e < DILITHIUM_N; e++) {
 8835                            wt[e] += dilithium_mont_red((sword64)at[e] *
 8836                                y_ntt_t[e]);
 8837                        }
 8838                    #else
 8839                        for (e = 0; e < DILITHIUM_N; e += 8) {
 8840                            wt[e + 0] += dilithium_mont_red((sword64)at[e + 0] *
 8841                                y_ntt_t[e + 0]);
 8842                            wt[e + 1] += dilithium_mont_red((sword64)at[e + 1] *
 8843                                y_ntt_t[e + 1]);
 8844                            wt[e + 2] += dilithium_mont_red((sword64)at[e + 2] *
 8845                                y_ntt_t[e + 2]);
 8846                            wt[e + 3] += dilithium_mont_red((sword64)at[e + 3] *
 8847                                y_ntt_t[e + 3]);
 8848                            wt[e + 4] += dilithium_mont_red((sword64)at[e + 4] *
 8849                                y_ntt_t[e + 4]);
 8850                            wt[e + 5] += dilithium_mont_red((sword64)at[e + 5] *
 8851                                y_ntt_t[e + 5]);
 8852                            wt[e + 6] += dilithium_mont_red((sword64)at[e + 6] *
 8853                                y_ntt_t[e + 6]);
 8854                            wt[e + 7] += dilithium_mont_red((sword64)at[e + 7] *
 8855                                y_ntt_t[e + 7]);
 8856                        }
 8857                    #endif
 8858                    }
 8859                #else
 8860                    if (s == 0) {
 8861                    #ifdef WOLFSSL_DILITHIUM_SMALL
 8862                        for (e = 0; e < DILITHIUM_N; e++) {
 8863                            t64[e] = (sword64)at[e] * y_ntt_t[e];
 8864                        }
 8865                    #else
 8866                        for (e = 0; e < DILITHIUM_N; e += 8) {
 8867                            t64[e+0] = (sword64)at[e+0] * y_ntt_t[e+0];
 8868                            t64[e+1] = (sword64)at[e+1] * y_ntt_t[e+1];
 8869                            t64[e+2] = (sword64)at[e+2] * y_ntt_t[e+2];
 8870                            t64[e+3] = (sword64)at[e+3] * y_ntt_t[e+3];
 8871                            t64[e+4] = (sword64)at[e+4] * y_ntt_t[e+4];
 8872                            t64[e+5] = (sword64)at[e+5] * y_ntt_t[e+5];
 8873                            t64[e+6] = (sword64)at[e+6] * y_ntt_t[e+6];
 8874                            t64[e+7] = (sword64)at[e+7] * y_ntt_t[e+7];
 8875                        }
 8876                    #endif
 8877                    }
 8878                    else {
 8879                    #ifdef WOLFSSL_DILITHIUM_SMALL
 8880                        for (e = 0; e < DILITHIUM_N; e++) {
 8881                            t64[e] += (sword64)at[e] * y_ntt_t[e];
 8882                        }
 8883                    #else
 8884                        for (e = 0; e < DILITHIUM_N; e += 8) {
 8885                            t64[e+0] += (sword64)at[e+0] * y_ntt_t[e+0];
 8886                            t64[e+1] += (sword64)at[e+1] * y_ntt_t[e+1];
 8887                            t64[e+2] += (sword64)at[e+2] * y_ntt_t[e+2];
 8888                            t64[e+3] += (sword64)at[e+3] * y_ntt_t[e+3];
 8889                            t64[e+4] += (sword64)at[e+4] * y_ntt_t[e+4];
 8890                            t64[e+5] += (sword64)at[e+5] * y_ntt_t[e+5];
 8891                            t64[e+6] += (sword64)at[e+6] * y_ntt_t[e+6];
 8892                            t64[e+7] += (sword64)at[e+7] * y_ntt_t[e+7];
 8893                        }
 8894                    #endif
 8895                    }
 8896                #endif
 8897                    /* Next polynomial. */
 8898                    yt += DILITHIUM_N;
 8899                }
 8900                if (ret != 0) {
 8901                    break;
 8902                }
 8903            #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 8904                for (e = 0; e < DILITHIUM_N; e++) {
 8905                    wt[e] = dilithium_mont_red(t64[e]);
 8906                }
 8907            #endif
 8908                dilithium_invntt_full(wt);
 8909                /* Step 14, Step 22: Make values positive and decompose. */
 8910                dilithium_make_pos(wt);
 8911            #ifndef WOLFSSL_NO_ML_DSA_44
 8912                if (params->gamma2 == DILITHIUM_Q_LOW_88) {
 8913                    /* For each value of polynomial. */
 8914                    for (e = 0; e < DILITHIUM_N; e++) {
 8915                        /* Decompose value into two vectors. */
 8916                        dilithium_decompose_q88(wt[e], &w0t[e], &w1t[e]);
 8917                    }
 8918                }
 8919            #endif
 8920            #if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 8921                if (params->gamma2 == DILITHIUM_Q_LOW_32) {
 8922                    /* For each value of polynomial. */
 8923                    for (e = 0; e < DILITHIUM_N; e++) {
 8924                        /* Decompose value into two vectors. */
 8925                        dilithium_decompose_q32(wt[e], &w0t[e], &w1t[e]);
 8926                    }
 8927                }
 8928            #endif
 8929            #ifdef WOLFSSL_DILITHIUM_SIGN_CHECK_W0
 8930                valid = dilithium_vec_check_low(w0t,
 8931                    params->gamma2 - params->beta);
 8932            #endif
 8933                wt  += DILITHIUM_N;
 8934                w0t += DILITHIUM_N;
 8935                w1t += DILITHIUM_N;
 8936            }
 8937            if ((ret == 0) && valid) {
 8938                sword32* yt = y;
 8939            #ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 8940                const byte* s1pt = s1p;
 8941            #endif
 8942                byte* ze = sig + params->lambda / 4;
 8943
 8944                /* Step 15: Encode w1. */
 8945                WC_ALLOC_VAR_EX(w1e, byte, DILITHIUM_MAX_W1_ENC_SZ,
 8946                    key->heap, DYNAMIC_TYPE_DILITHIUM, ret=MEMORY_E);
 8947                if (WC_VAR_OK(w1e)) {
 8948                    dilithium_vec_encode_w1(w1, params->k, params->gamma2,
 8949                        w1e);
 8950                    /* Step 15: Hash mu and encoded w1.
 8951                     * Step 32: Hash is stored in signature. */
 8952                    ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ,
 8953                        w1e, params->w1EncSz, commit, params->lambda / 4);
 8954                }
 8955                WC_FREE_VAR_EX(w1e, key->heap, DYNAMIC_TYPE_DILITHIUM);
 8956                if (ret == 0) {
 8957                    /* Step 17: Compute c from first 256 bits of commit. */
 8958                    ret = dilithium_sample_in_ball_ex(params->level,
 8959                        &key->shake, commit, params->lambda / 4, params->tau, c,
 8960                        blocks);
 8961                }
 8962                if (ret == 0) {
 8963                    /* Step 18: NTT(c). */
 8964                    dilithium_ntt_small(c);
 8965                }
 8966
 8967                for (s = 0; (ret == 0) && valid && (s < params->l); s++) {
 8968            #ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 8969                #if !defined(WOLFSSL_NO_ML_DSA_44) || \
 8970                    !defined(WOLFSSL_NO_ML_DSA_87)
 8971                    /* -2..2 */
 8972                    if (params->eta == DILITHIUM_ETA_2) {
 8973                        dilithium_decode_eta_2_bits(s1pt, s1);
 8974                        s1pt += DILITHIUM_ETA_2_BITS * DILITHIUM_N / 8;
 8975                    }
 8976                #endif
 8977                #ifndef WOLFSSL_NO_ML_DSA_65
 8978                    /* -4..4 */
 8979                    if (params->eta == DILITHIUM_ETA_4) {
 8980                        dilithium_decode_eta_4_bits(s1pt, s1);
 8981                        s1pt += DILITHIUM_N / 2;
 8982                    }
 8983                #endif
 8984                    dilithium_ntt_small(s1);
 8985                    dilithium_mul(z, c, s1);
 8986            #else
 8987                    dilithium_mul(z, c, s1 + s * DILITHIUM_N);
 8988            #endif
 8989                    /* Step 19: cs1 = NTT-1(c o s1) */
 8990                    dilithium_invntt(z);
 8991                    /* Step 21: z = y + cs1 */
 8992                    dilithium_add(z, yt);
 8993                    dilithium_poly_red(z);
 8994                    /* Step 23: Check z has low enough values. */
 8995                    hi = ((sword32)1 << params->gamma1_bits) - params->beta;
 8996                    valid = dilithium_check_low(z, hi);
 8997                    if (valid) {
 8998                        /* Step 32: Encode z into signature.
 8999                         * Commit (c) and h already encoded into signature. */
 9000                    #if !defined(WOLFSSL_NO_ML_DSA_44)
 9001                        if (params->gamma1_bits == DILITHIUM_GAMMA1_BITS_17) {
 9002                            dilithium_encode_gamma1_17_bits(z, ze);
 9003                            /* Move to next place to encode to. */
 9004                            ze += DILITHIUM_GAMMA1_17_ENC_BITS / 2 *
 9005                                  DILITHIUM_N / 4;
 9006                        }
 9007                    #endif
 9008                    #if !defined(WOLFSSL_NO_ML_DSA_65) || \
 9009                        !defined(WOLFSSL_NO_ML_DSA_87)
 9010                        if (params->gamma1_bits == DILITHIUM_GAMMA1_BITS_19) {
 9011                            dilithium_encode_gamma1_19_bits(z, ze);
 9012                            /* Move to next place to encode to. */
 9013                            ze += DILITHIUM_GAMMA1_19_ENC_BITS / 2 *
 9014                                  DILITHIUM_N / 4;
 9015                        }
 9016                    #endif
 9017                    }
 9018
 9019                    yt += DILITHIUM_N;
 9020                }
 9021            }
 9022            if ((ret == 0) && valid) {
 9023                const byte* t0pt = t0p;
 9024            #ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 9025                const byte* s2pt = s2p;
 9026            #endif
 9027                sword32* cs2 = ct0;
 9028                byte idx = 0;
 9029                w0t = w0;
 9030                w1t = w1;
 9031
 9032                for (r = 0; valid && (r < params->k); r++) {
 9033            #ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 9034                #if !defined(WOLFSSL_NO_ML_DSA_44) || \
 9035                    !defined(WOLFSSL_NO_ML_DSA_87)
 9036                    /* -2..2 */
 9037                    if (params->eta == DILITHIUM_ETA_2) {
 9038                        dilithium_decode_eta_2_bits(s2pt, s2);
 9039                        s2pt += DILITHIUM_ETA_2_BITS * DILITHIUM_N / 8;
 9040                    }
 9041                #endif
 9042                #ifndef WOLFSSL_NO_ML_DSA_65
 9043                    /* -4..4 */
 9044                    if (params->eta == DILITHIUM_ETA_4) {
 9045                        dilithium_decode_eta_4_bits(s2pt, s2);
 9046                        s2pt += DILITHIUM_N / 2;
 9047                    }
 9048                #endif
 9049                    dilithium_ntt_small(s2);
 9050                    /* Step 20: cs2 = NTT-1(c o s2) */
 9051                    dilithium_mul(cs2, c, s2);
 9052            #else
 9053                    /* Step 20: cs2 = NTT-1(c o s2) */
 9054                    dilithium_mul(cs2, c, s2 + r * DILITHIUM_N);
 9055            #endif
 9056                    dilithium_invntt(cs2);
 9057                    /* Step 22: w0 - cs2 */
 9058                    dilithium_sub(w0t, cs2);
 9059                    dilithium_poly_red(w0t);
 9060                    /* Step 23: Check w0 - cs2 has low enough values. */
 9061                    hi = params->gamma2 - params->beta;
 9062                    valid = dilithium_check_low(w0t, hi);
 9063                    if (valid) {
 9064                    #ifndef WOLFSSL_DILITHIUM_SIGN_SMALL_MEM_PRECALC
 9065                        dilithium_decode_t0(t0pt, t0);
 9066                        dilithium_ntt(t0);
 9067
 9068                        /* Step 25: ct0 = NTT-1(c o t0) */
 9069                        dilithium_mul(ct0, c, t0);
 9070                    #else
 9071                        /* Step 25: ct0 = NTT-1(c o t0) */
 9072                        dilithium_mul(ct0, c, t0 + r * DILITHIUM_N);
 9073                    #endif
 9074                        dilithium_invntt(ct0);
 9075                        /* Step 27: Check ct0 has low enough values. */
 9076                        valid = dilithium_check_low(ct0, params->gamma2);
 9077                    }
 9078                    if (valid) {
 9079                        /* Step 26: ct0 = ct0 + w0 */
 9080                        dilithium_add(ct0, w0t);
 9081                        dilithium_poly_red(ct0);
 9082
 9083                        /* Step 26, 27: Make hint from ct0 and w1 and check
 9084                         * number of hints is valid.
 9085                         * Step 32: h is encoded into signature.
 9086                         */
 9087                    #ifndef WOLFSSL_NO_ML_DSA_44
 9088                        if (params->gamma2 == DILITHIUM_Q_LOW_88) {
 9089                            valid = (dilithium_make_hint_88(ct0, w1t, h,
 9090                                &idx) == 0);
 9091                            /* Alg 14, Step 10: Store count of hints for
 9092                             *                  polynomial at end of list. */
 9093                            h[PARAMS_ML_DSA_44_OMEGA + r] = idx;
 9094                        }
 9095                    #endif
 9096                    #if !defined(WOLFSSL_NO_ML_DSA_65) || \
 9097                        !defined(WOLFSSL_NO_ML_DSA_87)
 9098                        if (params->gamma2 == DILITHIUM_Q_LOW_32) {
 9099                            valid = (dilithium_make_hint_32(ct0, w1t,
 9100                                params->omega, h, &idx) == 0);
 9101                            /* Alg 14, Step 10: Store count of hints for
 9102                             *                  polynomial at end of list. */
 9103                            h[params->omega + r] = idx;
 9104                        }
 9105                    #endif
 9106                    }
 9107
 9108                    t0pt += DILITHIUM_D * DILITHIUM_N / 8;
 9109                    w0t += DILITHIUM_N;
 9110                    w1t += DILITHIUM_N;
 9111                }
 9112                /* Set remaining hints to zero. */
 9113                XMEMSET(h + idx, 0, (size_t)(params->omega - idx));
 9114            }
 9115
 9116            if (!valid) {
 9117                /* Too many attempts - something wrong with implementation. */
 9118                if ((kappa > (word16)(kappa + params->l))) {
 9119                    ret = BAD_COND_E;
 9120                }
 9121
 9122                /* Step 30: increment value to append to seed to unique value.
 9123                 */
 9124                kappa = (word16)(kappa + params->l);
 9125            }
 9126        }
 9127        /* Step 11: Check we have a valid signature. */
 9128        while ((ret == 0) && (!valid));
 9129    }
 9130
 9131    ForceZero(priv_rand_seed, sizeof(priv_rand_seed));
 9132    if (y != NULL) {
 9133        ForceZero(y, allocSz);
 9134    }
 9135    XFREE(y, key->heap, DYNAMIC_TYPE_DILITHIUM);
 9136    return ret;
 9137#endif
 9138}
 9139
 9140/* Sign a message with the key and a seed.
 9141 *
 9142 * FIPS 204. 5.2: Algorithm 2 ML-DSA.Sign(sk, M, ctx)
 9143 *   ...
 9144 *  10: M' <- BytesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
 9145 *                        ctx || M)
 9146 *  11: sigma <- ML-DSA.Sign_internal(sk, M', rnd)
 9147 *  12: return sigma
 9148 *
 9149 * FIPS 204. 6.2: Algorithm 7 ML-DSA.SignInternal(sk, M', rnd)
 9150 *   ...
 9151 *   6: mu <- H(BytesToBits(tr)||M', 64)
 9152 *   ...
 9153 *
 9154 * @param [in, out] key     Dilithium key.
 9155 * @param [in]      seed    Random seed.
 9156 * @param [in]      ctx     Context of signature.
 9157 * @param [in]      ctxLen  Length of context in bytes.
 9158 * @param [in]      msg     Message data to sign.
 9159 * @param [in]      msgLen  Length of message data in bytes.
 9160 * @param [out]     sig     Buffer to hold signature.
 9161 * @param [in, out] sigLen  On in, length of buffer in bytes.
 9162 *                          On out, the length of the signature in bytes.
 9163 * @return  0 on success.
 9164 * @return  BAD_FUNC_ARG when context length is greater than 255.
 9165 * @return  BUFFER_E when the signature buffer is too small.
 9166 * @return  MEMORY_E when memory allocation fails.
 9167 * @return  Other negative when an error occurs.
 9168 */
 9169static int dilithium_sign_ctx_msg_with_seed(dilithium_key* key,
 9170    const byte* seed, const byte* ctx, byte ctxLen, const byte* msg,
 9171    word32 msgLen, byte* sig, word32 *sigLen)
 9172{
 9173    int ret;
 9174    const byte* pub_seed = key->k;
 9175    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9176    const byte* tr = k + DILITHIUM_K_SZ;
 9177    byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
 9178    byte* mu = seedMu + DILITHIUM_RND_SZ;
 9179
 9180    XMEMCPY(seedMu, seed, DILITHIUM_RND_SZ);
 9181    /* Step 6. Calculate mu. */
 9182    ret = dilithium_hash256_ctx_msg(&key->shake, tr, DILITHIUM_TR_SZ, 0,
 9183        ctx, ctxLen, msg, msgLen, mu, DILITHIUM_MU_SZ);
 9184    if (ret == 0) {
 9185        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
 9186    }
 9187
 9188    ForceZero(seedMu, sizeof(seedMu));
 9189    return ret;
 9190}
 9191
 9192/* Sign a message with the key and a seed.
 9193 *
 9194 * FIPS 204. 5.2: Algorithm 2 ML-DSA.Sign(sk, M, ctx)
 9195 *   ...
 9196 *  10: M' <- BytesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
 9197 *                        ctx || M)
 9198 *  11: sigma <- ML-DSA.Sign_internal(sk, M', rnd)
 9199 *  12: return sigma
 9200 *
 9201 * FIPS 204. 6.2: Algorithm 7 ML-DSA.SignInternal(sk, M', rnd)
 9202 *   ...
 9203 *   6: mu <- H(BytesToBits(tr)||M', 64)
 9204 *   ...
 9205 *
 9206 * @param [in, out] key     Dilithium key.
 9207 * @param [in]      seed    Random seed.
 9208 * @param [in]      msg     Message data to sign.
 9209 * @param [in]      msgLen  Length of message data in bytes.
 9210 * @param [out]     sig     Buffer to hold signature.
 9211 * @param [in, out] sigLen  On in, length of buffer in bytes.
 9212 *                          On out, the length of the signature in bytes.
 9213 * @return  0 on success.
 9214 * @return  BAD_FUNC_ARG when context length is greater than 255.
 9215 * @return  BUFFER_E when the signature buffer is too small.
 9216 * @return  MEMORY_E when memory allocation fails.
 9217 * @return  Other negative when an error occurs.
 9218 */
 9219#ifdef WOLFSSL_DILITHIUM_NO_CTX
 9220static int dilithium_sign_msg_with_seed(dilithium_key* key, const byte* seed,
 9221    const byte* msg, word32 msgLen, byte* sig, word32 *sigLen)
 9222{
 9223    int ret;
 9224    const byte* pub_seed = key->k;
 9225    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9226    const byte* tr = k + DILITHIUM_K_SZ;
 9227    byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
 9228    byte* mu = seedMu + DILITHIUM_RND_SZ;
 9229
 9230    XMEMCPY(seedMu, seed, DILITHIUM_RND_SZ);
 9231    /* Step 6. Calculate mu. */
 9232    ret = dilithium_hash256(&key->shake, tr, DILITHIUM_TR_SZ, msg, msgLen, mu,
 9233        DILITHIUM_MU_SZ);
 9234    if (ret == 0) {
 9235        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
 9236    }
 9237
 9238    ForceZero(seedMu, sizeof(seedMu));
 9239    return ret;
 9240}
 9241#endif /* WOLFSSL_DILITHIUM_NO_CTX */
 9242
 9243/* Sign a message with the key and a random number generator.
 9244 *
 9245 * FIPS 204. 5.2: Algorithm 2 ML-DSA.Sign(sk, M, ctx)
 9246 *   ...
 9247 *   5: rnd <- B32  [Randomly generated.]
 9248 *   6: if rnd = NULL then
 9249 *   7:     return falsam
 9250 *   8: end if
 9251 *   9:
 9252 *  10: M' <- BytesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
 9253 *                        ctx || M)
 9254 *   ...
 9255 *
 9256 * FIPS 204. 6.2: Algorithm 7 ML-DSA.SignInternal(sk, M', rnd)
 9257 *   ...
 9258 *   6: mu <- H(BytesToBits(tr)||M', 64)
 9259 *   ...
 9260 *
 9261 * @param [in, out] key     Dilithium key.
 9262 * @param [in, out] rng     Random number generator.
 9263 * @param [in]      ctx     Context of signature.
 9264 * @param [in]      ctxLen  Length of context.
 9265 * @param [in]      msg     Message data to sign.
 9266 * @param [in]      msgLen  Length of message data in bytes.
 9267 * @param [out]     sig     Buffer to hold signature.
 9268 * @param [in, out] sigLen  On in, length of buffer in bytes.
 9269 *                          On out, the length of the signature in bytes.
 9270 * @return  0 on success.
 9271 * @return  BUFFER_E when the signature buffer is too small.
 9272 * @return  MEMORY_E when memory allocation fails.
 9273 * @return  Other negative when an error occurs.
 9274 */
 9275static int dilithium_sign_ctx_msg(dilithium_key* key, WC_RNG* rng,
 9276    const byte* ctx, byte ctxLen, const byte* msg, word32 msgLen, byte* sig,
 9277    word32 *sigLen)
 9278{
 9279    int ret = 0;
 9280    const byte* pub_seed = key->k;
 9281    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9282    const byte* tr = k + DILITHIUM_K_SZ;
 9283    byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
 9284    byte* mu = seedMu + DILITHIUM_RND_SZ;
 9285
 9286    /* Must have a random number generator. */
 9287    if (rng == NULL) {
 9288        ret = BAD_FUNC_ARG;
 9289    }
 9290
 9291    if (ret == 0) {
 9292        /* Step 7: Generate random seed. */
 9293        ret = wc_RNG_GenerateBlock(rng, seedMu, DILITHIUM_RND_SZ);
 9294    }
 9295    if (ret == 0) {
 9296        /* Step 6. Calculate mu. */
 9297        ret = dilithium_hash256_ctx_msg(&key->shake, tr, DILITHIUM_TR_SZ, 0,
 9298            ctx, ctxLen, msg, msgLen, mu, DILITHIUM_MU_SZ);
 9299    }
 9300    if (ret == 0) {
 9301        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
 9302    }
 9303
 9304    ForceZero(seedMu, sizeof(seedMu));
 9305    return ret;
 9306}
 9307
 9308/* Sign a message with the key and a random number generator.
 9309 *
 9310 * FIPS 204. 5.2: Algorithm 2 ML-DSA.Sign(sk, M, ctx)
 9311 *   ...
 9312 *   5: rnd <- B32  [Randomly generated.]
 9313 *   6: if rnd = NULL then
 9314 *   7:     return falsam
 9315 *   8: end if
 9316 *   9:
 9317 *  10: M' <- BytesToBits(IntegerToBytes(0, 1) || IntegerToBytes(|ctx|, 1) ||
 9318 *                        ctx || M)
 9319 *   ...
 9320 *
 9321 * FIPS 204. 6.2: Algorithm 7 ML-DSA.SignInternal(sk, M', rnd)
 9322 *   ...
 9323 *   6: mu <- H(BytesToBits(tr)||M', 64)
 9324 *   ...
 9325 *
 9326 * @param [in, out] key     Dilithium key.
 9327 * @param [in, out] rng     Random number generator.
 9328 * @param [in]      msg     Message data to sign.
 9329 * @param [in]      msgLen  Length of message data in bytes.
 9330 * @param [out]     sig     Buffer to hold signature.
 9331 * @param [in, out] sigLen  On in, length of buffer in bytes.
 9332 *                          On out, the length of the signature in bytes.
 9333 * @return  0 on success.
 9334 * @return  BUFFER_E when the signature buffer is too small.
 9335 * @return  MEMORY_E when memory allocation fails.
 9336 * @return  Other negative when an error occurs.
 9337 */
 9338#ifdef WOLFSSL_DILITHIUM_NO_CTX
 9339static int dilithium_sign_msg(dilithium_key* key, WC_RNG* rng,
 9340    const byte* msg, word32 msgLen, byte* sig, word32 *sigLen)
 9341{
 9342    int ret = 0;
 9343    const byte* pub_seed = key->k;
 9344    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9345    const byte* tr = k + DILITHIUM_K_SZ;
 9346    byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
 9347    byte* mu = seedMu + DILITHIUM_RND_SZ;
 9348
 9349    /* Must have a random number generator. */
 9350    if (rng == NULL) {
 9351        ret = BAD_FUNC_ARG;
 9352    }
 9353
 9354    if (ret == 0) {
 9355        /* Step 7: Generate random seed. */
 9356        ret = wc_RNG_GenerateBlock(rng, seedMu, DILITHIUM_RND_SZ);
 9357    }
 9358    if (ret == 0) {
 9359        /* Step 6. Calculate mu. */
 9360        ret = dilithium_hash256(&key->shake, tr, DILITHIUM_TR_SZ, msg, msgLen,
 9361            mu, DILITHIUM_MU_SZ);
 9362    }
 9363    if (ret == 0) {
 9364        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
 9365    }
 9366
 9367    ForceZero(seedMu, sizeof(seedMu));
 9368    return ret;
 9369}
 9370#endif /* WOLFSSL_DILITHIUM_NO_CTX */
 9371
 9372/* Sign a pre-hashed message with the key and a seed.
 9373 *
 9374 * FIPS 204. 5.4.1: Algorithm 4 HashML-DSA.Sign(sk, M, ctx, PH)
 9375 *   ...
 9376 *  10: switch PH do
 9377 *  11:    case SHA-256:
 9378 *  12:       OID <-  IntegerToBytes(0x0609608648016503040201, 11)
 9379 *  13:       PHm <- SHA256(M)    (not done here as hash is passed in)
 9380 *   ...
 9381 *  22: end switch
 9382 *  23: M' <- BytesToBits(IntegerToBytes(1, 1) || IntegerToBytes(|ctx|, 1) ||
 9383 *                        ctx || OID || PHm)
 9384 *  24: sigma <- ML-DSA.Sign_internal(sk, M', rnd)
 9385 *  25: return sigma
 9386 *
 9387 * FIPS 204. 6.2: Algorithm 7 ML-DSA.SignInternal(sk, M', rnd)
 9388 *   ...
 9389 *   6: mu <- H(BytesToBits(tr)||M', 64)
 9390 *   ...
 9391 *
 9392 * @param [in, out] key      Dilithium key.
 9393 * @param [in]      seed     Random seed.
 9394 * @param [in]      ctx      Context of signature.
 9395 * @param [in]      ctxLen   Length of context.
 9396 * @param [in]      hashAlg  Hash algorithm used on message.
 9397 * @param [in]      hash     Message hash to sign.
 9398 * @param [in]      hashLen  Length of message hash in bytes.
 9399 * @param [out]     sig      Buffer to hold signature.
 9400 * @param [in, out] sigLen   On in, length of buffer in bytes.
 9401 *                           On out, the length of the signature in bytes.
 9402 * @return  0 on success.
 9403 * @return  BUFFER_E when the signature buffer is too small.
 9404 * @return  MEMORY_E when memory allocation fails.
 9405 * @return  Other negative when an error occurs.
 9406 */
 9407static int dilithium_sign_ctx_hash_with_seed(dilithium_key* key,
 9408    const byte* seed, const byte* ctx, byte ctxLen, int hashAlg,
 9409    const byte* hash, word32 hashLen, byte* sig, word32 *sigLen)
 9410{
 9411    int ret = 0;
 9412    const byte* pub_seed = key->k;
 9413    const byte* k = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9414    const byte* tr = k + DILITHIUM_K_SZ;
 9415    byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
 9416    byte* mu = seedMu + DILITHIUM_RND_SZ;
 9417    byte oidMsgHash[DILITHIUM_HASH_OID_LEN + WC_MAX_DIGEST_SIZE];
 9418    word32 oidMsgHashLen = 0;
 9419
 9420    /* Check that the input hash length is valid. */
 9421    if ((int)hashLen != wc_HashGetDigestSize((enum wc_HashType)hashAlg)) {
 9422        ret = BAD_LENGTH_E;
 9423    }
 9424
 9425    if (ret == 0) {
 9426        XMEMCPY(seedMu, seed, DILITHIUM_RND_SZ);
 9427
 9428        ret = dilithium_get_hash_oid(hashAlg, oidMsgHash, &oidMsgHashLen);
 9429    }
 9430    if (ret == 0) {
 9431        XMEMCPY(oidMsgHash + oidMsgHashLen, hash, hashLen);
 9432        oidMsgHashLen += hashLen;
 9433
 9434        /* Step 6. Calculate mu. */
 9435        ret = dilithium_hash256_ctx_msg(&key->shake, tr, DILITHIUM_TR_SZ, 1,
 9436            ctx, ctxLen, oidMsgHash, oidMsgHashLen, mu, DILITHIUM_MU_SZ);
 9437    }
 9438    if (ret == 0) {
 9439        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
 9440    }
 9441
 9442    ForceZero(seedMu, sizeof(seedMu));
 9443    return ret;
 9444}
 9445
 9446/* Sign a pre-hashed message with the key and a random number generator.
 9447 *
 9448 * FIPS 204. 5.4.1: Algorithm 4 HashML-DSA.Sign(sk, M, ctx, PH)
 9449 *   ...
 9450 *   5: rnd <- B32  [Randomly generated.]
 9451 *   6: if rnd = NULL then
 9452 *   7:     return falsam
 9453 *   8: end if
 9454 *   ...
 9455 *
 9456 * @param [in, out] key      Dilithium key.
 9457 * @param [in, out] rng      Random number generator.
 9458 * @param [in]      ctx      Context of signature.
 9459 * @param [in]      ctxLen   Length of context.
 9460 * @param [in]      hashAlg  Hash algorithm used on message.
 9461 * @param [in]      hash     Message hash to sign.
 9462 * @param [in]      hashLen  Length of message hash in bytes.
 9463 * @param [out]     sig      Buffer to hold signature.
 9464 * @param [in, out] sigLen   On in, length of buffer in bytes.
 9465 *                           On out, the length of the signature in bytes.
 9466 * @return  0 on success.
 9467 * @return  BUFFER_E when the signature buffer is too small.
 9468 * @return  MEMORY_E when memory allocation fails.
 9469 * @return  Other negative when an error occurs.
 9470 */
 9471static int dilithium_sign_ctx_hash(dilithium_key* key, WC_RNG* rng,
 9472    const byte* ctx, byte ctxLen, int hashAlg, const byte* hash, word32 hashLen,
 9473    byte* sig, word32 *sigLen)
 9474{
 9475    int ret = 0;
 9476    byte seed[DILITHIUM_RND_SZ];
 9477
 9478    /* Must have a random number generator. */
 9479    if (rng == NULL) {
 9480        ret = BAD_FUNC_ARG;
 9481    }
 9482    if (ret == 0) {
 9483        /* Step 7: Generate random seed. */
 9484        ret = wc_RNG_GenerateBlock(rng, seed, DILITHIUM_RND_SZ);
 9485    }
 9486
 9487    if (ret == 0) {
 9488        ret = dilithium_sign_ctx_hash_with_seed(key, seed, ctx, ctxLen, hashAlg,
 9489            hash, hashLen, sig, sigLen);
 9490    }
 9491
 9492    ForceZero(seed, sizeof(seed));
 9493    return ret;
 9494}
 9495
 9496#endif /* !WOLFSSL_DILITHIUM_NO_SIGN */
 9497
 9498#ifndef WOLFSSL_DILITHIUM_NO_VERIFY
 9499
 9500#if !defined(WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM) || \
 9501     defined(WC_DILITHIUM_CACHE_PUB_VECTORS)
 9502/* Make public vector from public key data.
 9503 *
 9504 * @param [in, out] key  Key with public key data.
 9505 * @param [out]     t1   Vector in NTT form.
 9506 */
 9507static void dilithium_make_pub_vec(dilithium_key* key, sword32* t1)
 9508{
 9509    const wc_dilithium_params* params = key->params;
 9510    const byte* t1p = key->p + DILITHIUM_PUB_SEED_SZ;
 9511
 9512    dilithium_vec_decode_t1(t1p, params->k, t1);
 9513    dilithium_vec_ntt_full(t1, params->k);
 9514
 9515#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
 9516    key->pubVecSet = 1;
 9517#endif
 9518}
 9519#endif
 9520
 9521/* Verify signature of message using public key.
 9522 *
 9523 * FIPS 204. 6: Algorithm 3 ML-DSA.Verify(pk, M, sigma)
 9524 *  1: (rho, t1) <- pkDecode(pk)
 9525 *  2: (c_tilde, z, h) <- sigDecode(sigma)
 9526 *  3: if h = falsam then return false
 9527 *  4: end if
 9528 *  5: A_circum <- ExpandS(rho)
 9529 *  6: tr <- H(BytesToBits(pk), 512)
 9530 *  7: mu <- H(tr||M, 512)
 9531 *  8: (c1_tilde, c2_tilde) E {0,1}256 x {0,1)2*lambda-256 <- c_tilde
 9532 *  9: c <- SampleInBall(c1_tilde)
 9533 * 10: w'approx <- NTT-1(A_circum o NTT(z) - NTT(c) o NTT(t1.2^d))
 9534 * 11: w1' <- UseHint(h, w'approx)
 9535 * 12: c'_tilde < H(mu||w1Encode(w1'), 2*lambda)
 9536 * 13: return [[ ||z||inf < GAMMA1 - BETA]] and [[c_tilde = c'_tilde]] and
 9537 *             [[number of 1's in h is <= OMEGA
 9538 *
 9539 * @param [in, out] key     Dilithium key.
 9540 * @param [in]      mu      Data to verify.
 9541 * @param [in]      sig     Signature to verify message.
 9542 * @param [in]      sigLen  Length of message in bytes.
 9543 * @param [out]     res     Result of verification.
 9544 * @return  0 on success.
 9545 * @return  SIG_VERIFY_E when hint is malformed.
 9546 * @return  BUFFER_E when the length of the signature does not match
 9547 *          parameters.
 9548 * @return  MEMORY_E when memory allocation fails.
 9549 * @return  Other negative when an error occurs.
 9550 */
 9551static int dilithium_verify_with_mu(dilithium_key* key, const byte* mu,
 9552    const byte* sig, word32 sigLen, int* res)
 9553{
 9554#ifndef WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM
 9555    int ret = 0;
 9556    const wc_dilithium_params* params = key->params;
 9557    const byte* pub_seed = key->p;
 9558    const byte* commit = sig;
 9559    const byte* ze = sig + params->lambda / 4;
 9560    const byte* h = ze + params->zEncSz;
 9561    sword32* a = NULL;
 9562    sword32* t1 = NULL;
 9563    sword32* c = NULL;
 9564    sword32* z = NULL;
 9565    sword32* w = NULL;
 9566    sword32* t1c = NULL;
 9567    byte commit_calc[DILITHIUM_TR_SZ];
 9568    byte* w1e = NULL;
 9569    int valid = 0;
 9570    sword32 hi;
 9571
 9572    /* Ensure the signature is the right size for the parameters. */
 9573    if (sigLen != params->sigSz) {
 9574        ret = BUFFER_E;
 9575    }
 9576    if (ret == 0) {
 9577        /* Step 13: Verify the hint is well-formed. */
 9578        ret = dilithium_check_hint(h, params->k, params->omega);
 9579    }
 9580
 9581    /* Allocate memory for large intermediates. */
 9582#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 9583#ifndef WC_DILITHIUM_FIXED_ARRAY
 9584    if ((ret == 0) && (key->a == NULL)) {
 9585        key->a = (sword32*)XMALLOC(params->aSz, key->heap,
 9586            DYNAMIC_TYPE_DILITHIUM);
 9587        if (key->a == NULL) {
 9588            ret = MEMORY_E;
 9589        }
 9590        else {
 9591            XMEMSET(key->a, 0, params->aSz);
 9592        }
 9593    }
 9594#endif
 9595    if (ret == 0) {
 9596        a = key->a;
 9597    }
 9598#endif
 9599#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
 9600#ifndef WC_DILITHIUM_FIXED_ARRAY
 9601    if ((ret == 0) && (key->t1 == NULL)) {
 9602        key->t1 = (sword32*)XMALLOC(params->s2Sz, key->heap,
 9603            DYNAMIC_TYPE_DILITHIUM);
 9604        if (key->t1 == NULL) {
 9605            ret = MEMORY_E;
 9606        }
 9607        else {
 9608            XMEMSET(key->t1, 0, params->s2Sz);
 9609        }
 9610    }
 9611#endif
 9612    if (ret == 0) {
 9613        t1 = key->t1;
 9614    }
 9615#endif
 9616    if (ret == 0) {
 9617        unsigned int allocSz;
 9618
 9619        /* z, c, w, t1/t1c */
 9620        allocSz = (unsigned int)DILITHIUM_POLY_SIZE + params->s1Sz +
 9621            params->s2Sz + params->s2Sz;
 9622#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 9623        /* a */
 9624        allocSz += params->aSz;
 9625#endif
 9626
 9627        z = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 9628        if (z == NULL) {
 9629            ret = MEMORY_E;
 9630        }
 9631        else {
 9632            XMEMSET(z, 0, allocSz);
 9633            c   = z  + params->s1Sz / sizeof(*z);
 9634            w   = c  + DILITHIUM_N;
 9635#ifndef WC_DILITHIUM_CACHE_PUB_VECTORS
 9636            t1  = w  + params->s2Sz / sizeof(*w);
 9637            t1c = t1;
 9638#else
 9639            t1c = w  + params->s2Sz / sizeof(*w);
 9640#endif
 9641#ifndef WC_DILITHIUM_CACHE_MATRIX_A
 9642            a   = t1 + params->s2Sz / sizeof(*t1);
 9643#endif
 9644            w1e = (byte*)c;
 9645        }
 9646    }
 9647
 9648    if (ret == 0) {
 9649        /* Step 2: Decode z from signature. */
 9650        dilithium_vec_decode_gamma1(ze, params->l, params->gamma1_bits, z);
 9651        /* Step 13: Check z is valid - values are low enough. */
 9652        hi = ((sword32)1 << params->gamma1_bits) - params->beta;
 9653        valid = dilithium_vec_check_low(z, params->l, hi);
 9654    }
 9655    if ((ret == 0) && valid) {
 9656#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
 9657        /* Check that we haven't already cached the public vector. */
 9658        if (!key->pubVecSet)
 9659#endif
 9660        {
 9661            /* Step 1: Decode and NTT vector t1. */
 9662            dilithium_make_pub_vec(key, t1);
 9663        }
 9664
 9665#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 9666        /* Check that we haven't already cached the matrix A. */
 9667        if (!key->aSet)
 9668#endif
 9669        {
 9670            /* Step 5: Expand pub seed to compute matrix A. */
 9671            ret = dilithium_expand_a(&key->shake, pub_seed, params->k,
 9672                params->l, a, key->heap);
 9673#ifdef WC_DILITHIUM_CACHE_MATRIX_A
 9674            /* Whether we have cached A is dependent on success of operation. */
 9675            key->aSet = (ret == 0);
 9676#endif
 9677        }
 9678    }
 9679    if ((ret == 0) && valid) {
 9680        /* Step 9: Compute c from commit. */
 9681        ret = dilithium_sample_in_ball(params->level, &key->shake, commit,
 9682            params->lambda / 4, params->tau, c, key->heap);
 9683    }
 9684    if ((ret == 0) && valid) {
 9685        /* Step 10: w = NTT-1(A o NTT(z) - NTT(c) o NTT(t1)) */
 9686        dilithium_vec_ntt_full(z, params->l);
 9687        dilithium_matrix_mul(w, a, z, params->k, params->l);
 9688    #ifdef WOLFSSL_DILITHIUM_SMALL
 9689        dilithium_vec_red(w, params->k);
 9690    #endif
 9691        dilithium_ntt_small_full(c);
 9692        dilithium_vec_mul(t1c, c, t1, params->k);
 9693        dilithium_vec_sub(w, t1c, params->k);
 9694        dilithium_vec_invntt_full(w, params->k);
 9695        /* Step 11: Use hint to give full w1. */
 9696        dilithium_vec_use_hint(w, params->k, params->gamma2, params->omega, h);
 9697        /* Step 12: Encode w1. */
 9698        dilithium_vec_encode_w1(w, params->k, params->gamma2, w1e);
 9699        /* Step 12: Hash mu and encoded w1. */
 9700        ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ, w1e,
 9701            params->w1EncSz, commit_calc, params->lambda / 4);
 9702    }
 9703    if ((ret == 0) && valid) {
 9704        /* Step 13: Compare commit. */
 9705        valid = (XMEMCMP(commit, commit_calc, params->lambda / 4) == 0);
 9706    }
 9707
 9708    *res = valid;
 9709    XFREE(z, key->heap, DYNAMIC_TYPE_DILITHIUM);
 9710    return ret;
 9711#else
 9712    int ret = 0;
 9713    const wc_dilithium_params* params = key->params;
 9714    const byte* pub_seed = key->p;
 9715    const byte* t1p = pub_seed + DILITHIUM_PUB_SEED_SZ;
 9716    const byte* commit = sig;
 9717    const byte* ze = sig + params->lambda / 4;
 9718    const byte* h = ze + params->zEncSz;
 9719    sword32* t1 = NULL;
 9720    sword32* a = NULL;
 9721    sword32* c = NULL;
 9722    sword32* z = NULL;
 9723    sword32* w = NULL;
 9724#ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9725    sword64* t64 = NULL;
 9726#endif
 9727#ifndef WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC
 9728    byte*    block = NULL;
 9729#endif
 9730    byte* w1e = NULL;
 9731    byte commit_calc[DILITHIUM_TR_SZ];
 9732    int valid = 0;
 9733    sword32 hi;
 9734    unsigned int r;
 9735    byte o;
 9736    byte* encW1;
 9737    byte* seed = commit_calc;
 9738
 9739    /* Ensure the signature is the right size for the parameters. */
 9740    if (sigLen != params->sigSz) {
 9741        ret = BUFFER_E;
 9742    }
 9743    if (ret == 0) {
 9744        /* Step 13: Verify the hint is well-formed. */
 9745        ret = dilithium_check_hint(h, params->k, params->omega);
 9746    }
 9747
 9748#ifndef WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC
 9749    /* Allocate memory for large intermediates. */
 9750    if (ret == 0) {
 9751        /* z, c, w, t1, w1e. */
 9752        unsigned int allocSz;
 9753
 9754        allocSz  = (unsigned int)params->s1Sz + params->w1EncSz +
 9755                   3U * (unsigned int)DILITHIUM_POLY_SIZE +
 9756                   (unsigned int)DILITHIUM_REJ_NTT_POLY_H_SIZE;
 9757    #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9758        allocSz += (unsigned int)DILITHIUM_POLY_SIZE * 2U;
 9759    #endif
 9760        z = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
 9761        if (z == NULL) {
 9762            ret = MEMORY_E;
 9763        }
 9764        else {
 9765            XMEMSET(z, 0, allocSz);
 9766            c     = z + params->s1Sz / sizeof(*t1);
 9767            w     = c + DILITHIUM_N;
 9768            t1    = w + DILITHIUM_N;
 9769            block = (byte*)(t1 + DILITHIUM_N);
 9770            w1e   = block + DILITHIUM_REJ_NTT_POLY_H_SIZE;
 9771            a     = t1;
 9772        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9773            t64   = (sword64*)(w1e + params->w1EncSz);
 9774        #endif
 9775        }
 9776    }
 9777#else
 9778    if (ret == 0) {
 9779        z = key->z;
 9780        c = key->c;
 9781        w = key->w;
 9782        t1 = key->t1;
 9783        w1e = key->w1e;
 9784        a = t1;
 9785    #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9786        t64 = key->t64;
 9787    #endif
 9788    }
 9789#endif
 9790
 9791    if (ret == 0) {
 9792        /* Step 2: Decode z from signature. */
 9793        dilithium_vec_decode_gamma1(ze, params->l, params->gamma1_bits, z);
 9794        /* Step 13: Check z is valid - values are low enough. */
 9795        hi = ((sword32)1 << params->gamma1_bits) - params->beta;
 9796        valid = dilithium_vec_check_low(z, params->l, hi);
 9797    }
 9798    if ((ret == 0) && valid) {
 9799        /* Step 10: NTT(z) */
 9800        dilithium_vec_ntt_full(z, params->l);
 9801
 9802         /* Step 9: Compute c from first 256 bits of commit. */
 9803#ifdef WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC
 9804         ret = dilithium_sample_in_ball_ex(params->level, &key->shake, commit,
 9805             params->lambda / 4, params->tau, c, key->block);
 9806#else
 9807         ret = dilithium_sample_in_ball_ex(params->level, &key->shake, commit,
 9808             params->lambda / 4, params->tau, c, block);
 9809#endif
 9810    }
 9811    if ((ret == 0) && valid) {
 9812        dilithium_ntt_small_full(c);
 9813
 9814        o = 0;
 9815        encW1 = w1e;
 9816
 9817        /* Copy the seed into a buffer that has space for s and r. */
 9818        XMEMCPY(seed, pub_seed, DILITHIUM_PUB_SEED_SZ);
 9819        /* Step 1: Loop over first dimension of matrix. */
 9820        for (r = 0; (ret == 0) && (r < params->k); r++) {
 9821            unsigned int s;
 9822            unsigned int e;
 9823            const sword32* zt = z;
 9824
 9825            /* Step 1: Decode and NTT vector t1. */
 9826            dilithium_decode_t1(t1p, w);
 9827            /* Next polynomial. */
 9828            t1p += DILITHIUM_U * DILITHIUM_N / 8;
 9829
 9830            /* Step 10: - NTT(c) o NTT(t1)) */
 9831            dilithium_ntt_full(w);
 9832    #ifndef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9833        #ifdef WOLFSSL_DILITHIUM_SMALL
 9834            for (e = 0; e < DILITHIUM_N; e++) {
 9835                w[e] = -dilithium_mont_red((sword64)c[e] * w[e]);
 9836            }
 9837        #else
 9838            for (e = 0; e < DILITHIUM_N; e += 8) {
 9839                w[e+0] = -dilithium_mont_red((sword64)c[e+0] * w[e+0]);
 9840                w[e+1] = -dilithium_mont_red((sword64)c[e+1] * w[e+1]);
 9841                w[e+2] = -dilithium_mont_red((sword64)c[e+2] * w[e+2]);
 9842                w[e+3] = -dilithium_mont_red((sword64)c[e+3] * w[e+3]);
 9843                w[e+4] = -dilithium_mont_red((sword64)c[e+4] * w[e+4]);
 9844                w[e+5] = -dilithium_mont_red((sword64)c[e+5] * w[e+5]);
 9845                w[e+6] = -dilithium_mont_red((sword64)c[e+6] * w[e+6]);
 9846                w[e+7] = -dilithium_mont_red((sword64)c[e+7] * w[e+7]);
 9847            }
 9848        #endif
 9849    #else
 9850        #ifdef WOLFSSL_DILITHIUM_SMALL
 9851            for (e = 0; e < DILITHIUM_N; e++) {
 9852                t64[e] = -(sword64)c[e] * w[e];
 9853            }
 9854        #else
 9855            for (e = 0; e < DILITHIUM_N; e += 8) {
 9856                t64[e+0] = -(sword64)c[e+0] * w[e+0];
 9857                t64[e+1] = -(sword64)c[e+1] * w[e+1];
 9858                t64[e+2] = -(sword64)c[e+2] * w[e+2];
 9859                t64[e+3] = -(sword64)c[e+3] * w[e+3];
 9860                t64[e+4] = -(sword64)c[e+4] * w[e+4];
 9861                t64[e+5] = -(sword64)c[e+5] * w[e+5];
 9862                t64[e+6] = -(sword64)c[e+6] * w[e+6];
 9863                t64[e+7] = -(sword64)c[e+7] * w[e+7];
 9864            }
 9865        #endif
 9866    #endif
 9867
 9868            /* Step 5: Expand pub seed to compute matrix A. */
 9869            /* Put r into buffer to be hashed. */
 9870            seed[DILITHIUM_PUB_SEED_SZ + 1] = (byte)r;
 9871            for (s = 0; (ret == 0) && (s < params->l); s++) {
 9872                /* Put s into buffer to be hashed. */
 9873                seed[DILITHIUM_PUB_SEED_SZ + 0] = (byte)s;
 9874                /* Step 3: Create polynomial from hashing seed. */
 9875            #ifdef WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC
 9876                ret = dilithium_rej_ntt_poly_ex(&key->shake, seed, a, key->h);
 9877            #else
 9878                ret = dilithium_rej_ntt_poly_ex(&key->shake, seed, a, block);
 9879            #endif
 9880
 9881                /* Step 10: w = A o NTT(z) - NTT(c) o NTT(t1) */
 9882        #ifndef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9883            #ifdef WOLFSSL_DILITHIUM_SMALL
 9884                for (e = 0; e < DILITHIUM_N; e++) {
 9885                    w[e] += dilithium_mont_red((sword64)a[e] * zt[e]);
 9886                }
 9887            #else
 9888                for (e = 0; e < DILITHIUM_N; e += 8) {
 9889                    w[e+0] += dilithium_mont_red((sword64)a[e+0] * zt[e+0]);
 9890                    w[e+1] += dilithium_mont_red((sword64)a[e+1] * zt[e+1]);
 9891                    w[e+2] += dilithium_mont_red((sword64)a[e+2] * zt[e+2]);
 9892                    w[e+3] += dilithium_mont_red((sword64)a[e+3] * zt[e+3]);
 9893                    w[e+4] += dilithium_mont_red((sword64)a[e+4] * zt[e+4]);
 9894                    w[e+5] += dilithium_mont_red((sword64)a[e+5] * zt[e+5]);
 9895                    w[e+6] += dilithium_mont_red((sword64)a[e+6] * zt[e+6]);
 9896                    w[e+7] += dilithium_mont_red((sword64)a[e+7] * zt[e+7]);
 9897                }
 9898            #endif
 9899        #else
 9900            #ifdef WOLFSSL_DILITHIUM_SMALL
 9901                for (e = 0; e < DILITHIUM_N; e++) {
 9902                    t64[e] += (sword64)a[e] * zt[e];
 9903                }
 9904            #else
 9905                for (e = 0; e < DILITHIUM_N; e += 8) {
 9906                    t64[e+0] += (sword64)a[e+0] * zt[e+0];
 9907                    t64[e+1] += (sword64)a[e+1] * zt[e+1];
 9908                    t64[e+2] += (sword64)a[e+2] * zt[e+2];
 9909                    t64[e+3] += (sword64)a[e+3] * zt[e+3];
 9910                    t64[e+4] += (sword64)a[e+4] * zt[e+4];
 9911                    t64[e+5] += (sword64)a[e+5] * zt[e+5];
 9912                    t64[e+6] += (sword64)a[e+6] * zt[e+6];
 9913                    t64[e+7] += (sword64)a[e+7] * zt[e+7];
 9914                }
 9915            #endif
 9916        #endif
 9917                /* Next polynomial. */
 9918                zt += DILITHIUM_N;
 9919            }
 9920        #ifdef WOLFSSL_DILITHIUM_SMALL_MEM_POLY64
 9921            for (e = 0; e < DILITHIUM_N; e++) {
 9922                w[e] = dilithium_mont_red(t64[e]);
 9923            }
 9924        #endif
 9925
 9926            /* Step 10: w = NTT-1(A o NTT(z) - NTT(c) o NTT(t1)) */
 9927            dilithium_invntt_full(w);
 9928
 9929        #ifndef WOLFSSL_NO_ML_DSA_44
 9930            if (params->gamma2 == DILITHIUM_Q_LOW_88) {
 9931                /* Step 11: Use hint to give full w1. */
 9932                dilithium_use_hint_88(w, h, r, &o);
 9933                /* Step 12: Encode w1. */
 9934                dilithium_encode_w1_88(w, encW1);
 9935                encW1 += DILITHIUM_Q_HI_88_ENC_BITS * 2 * DILITHIUM_N / 16;
 9936            }
 9937            else
 9938        #endif
 9939        #if !defined(WOLFSSL_NO_ML_DSA_65) || !defined(WOLFSSL_NO_ML_DSA_87)
 9940            if (params->gamma2 == DILITHIUM_Q_LOW_32) {
 9941                /* Step 11: Use hint to give full w1. */
 9942                dilithium_use_hint_32(w, h, params->omega, r, &o);
 9943                /* Step 12: Encode w1. */
 9944                dilithium_encode_w1_32(w, encW1);
 9945                encW1 += DILITHIUM_Q_HI_32_ENC_BITS * 2 * DILITHIUM_N / 16;
 9946            }
 9947            else
 9948        #endif
 9949            {
 9950            }
 9951        }
 9952    }
 9953    if ((ret == 0) && valid) {
 9954        /* Step 12: Hash mu and encoded w1. */
 9955        ret = dilithium_hash256(&key->shake, mu, DILITHIUM_MU_SZ, w1e,
 9956            params->w1EncSz, commit_calc, params->lambda / 4);
 9957    }
 9958    if ((ret == 0) && valid) {
 9959        /* Step 13: Compare commit. */
 9960        valid = (XMEMCMP(commit, commit_calc, params->lambda / 4) == 0);
 9961    }
 9962
 9963    *res = valid;
 9964#ifndef WOLFSSL_DILITHIUM_VERIFY_NO_MALLOC
 9965    XFREE(z, key->heap, DYNAMIC_TYPE_DILITHIUM);
 9966#endif
 9967    return ret;
 9968#endif /* !WOLFSSL_DILITHIUM_VERIFY_SMALL_MEM */
 9969}
 9970
 9971/* Verify signature of message using public key.
 9972 *
 9973 * @param [in, out] key     Dilithium key.
 9974 * @param [in]      ctx     Context of verification.
 9975 * @param [in]      ctxLen  Length of context in bytes.
 9976 * @param [in]      msg     Message to verify.
 9977 * @param [in]      msgLen  Length of message in bytes.
 9978 * @param [in]      sig     Signature to verify message.
 9979 * @param [in]      sigLen  Length of message in bytes.
 9980 * @param [out]     res     Result of verification.
 9981 * @return  0 on success.
 9982 * @return  SIG_VERIFY_E when hint is malformed.
 9983 * @return  BUFFER_E when the length of the signature does not match
 9984 *          parameters.
 9985 * @return  MEMORY_E when memory allocation fails.
 9986 * @return  Other negative when an error occurs.
 9987 */
 9988static int dilithium_verify_ctx_msg(dilithium_key* key, const byte* ctx,
 9989    byte ctxLen, const byte* msg, word32 msgLen, const byte* sig, word32 sigLen,
 9990    int* res)
 9991{
 9992    int ret = 0;
 9993    byte tr[DILITHIUM_TR_SZ];
 9994    byte* mu = tr;
 9995
 9996    if (key == NULL) {
 9997        ret = BAD_FUNC_ARG;
 9998    }
 9999
10000    if (ret == 0) {
10001        /* Step 6: Hash public key. */
10002        ret = dilithium_shake256(&key->shake, key->p, key->params->pkSz, tr,
10003            DILITHIUM_TR_SZ);
10004    }
10005    if (ret == 0) {
10006        /* Step 6. Calculate mu. */
10007        ret = dilithium_hash256_ctx_msg(&key->shake, tr, DILITHIUM_TR_SZ, 0,
10008            ctx, (byte)ctxLen, msg, msgLen, mu, DILITHIUM_MU_SZ);
10009    }
10010    if (ret == 0) {
10011        ret = dilithium_verify_with_mu(key, mu, sig, sigLen, res);
10012    }
10013
10014    return ret;
10015}
10016
10017#ifdef WOLFSSL_DILITHIUM_NO_CTX
10018/* Verify signature of message using public key.
10019 *
10020 * @param [in, out] key     Dilithium key.
10021 * @param [in]      msg     Message to verify.
10022 * @param [in]      msgLen  Length of message in bytes.
10023 * @param [in]      sig     Signature to verify message.
10024 * @param [in]      sigLen  Length of message in bytes.
10025 * @param [out]     res     Result of verification.
10026 * @return  0 on success.
10027 * @return  SIG_VERIFY_E when hint is malformed.
10028 * @return  BUFFER_E when the length of the signature does not match
10029 *          parameters.
10030 * @return  MEMORY_E when memory allocation fails.
10031 * @return  Other negative when an error occurs.
10032 */
10033static int dilithium_verify_msg(dilithium_key* key, const byte* msg,
10034    word32 msgLen, const byte* sig, word32 sigLen, int* res)
10035{
10036    int ret = 0;
10037    byte tr[DILITHIUM_TR_SZ];
10038    byte* mu = tr;
10039
10040    if (key == NULL) {
10041        ret = BAD_FUNC_ARG;
10042    }
10043
10044    if (ret == 0) {
10045        /* Step 6: Hash public key. */
10046        ret = dilithium_shake256(&key->shake, key->p, key->params->pkSz, tr,
10047            DILITHIUM_TR_SZ);
10048    }
10049    if (ret == 0) {
10050        /* Step 6. Calculate mu. */
10051        ret = dilithium_hash256(&key->shake, tr, DILITHIUM_TR_SZ, msg, msgLen,
10052            mu, DILITHIUM_MU_SZ);
10053    }
10054    if (ret == 0) {
10055        ret = dilithium_verify_with_mu(key, mu, sig, sigLen, res);
10056    }
10057
10058    return ret;
10059}
10060#endif /* WOLFSSL_DILITHIUM_NO_CTX */
10061
10062/* Verify signature of message using public key.
10063 *
10064 * @param [in, out] key       Dilithium key.
10065 * @param [in]      ctx       Context of verification.
10066 * @param [in]      ctxLen    Length of context in bytes.
10067 * @param [iu]      hashAlg   Hash algorithm used on message.
10068 * @param [in]      hash      Hash of message to verify.
10069 * @param [in]      hashLen   Length of message hash in bytes.
10070 * @param [in]      sig       Signature to verify message.
10071 * @param [in]      sigLen    Length of message in bytes.
10072 * @param [out]     res       Result of verification.
10073 * @return  0 on success.
10074 * @return  SIG_VERIFY_E when hint is malformed.
10075 * @return  BUFFER_E when the length of the signature does not match
10076 *          parameters.
10077 * @return  MEMORY_E when memory allocation fails.
10078 * @return  Other negative when an error occurs.
10079 */
10080static int dilithium_verify_ctx_hash(dilithium_key* key, const byte* ctx,
10081    byte ctxLen, int hashAlg, const byte* hash, word32 hashLen, const byte* sig,
10082    word32 sigLen, int* res)
10083{
10084    int ret = 0;
10085    byte tr[DILITHIUM_TR_SZ];
10086    byte* mu = tr;
10087    byte oidMsgHash[DILITHIUM_HASH_OID_LEN + WC_MAX_DIGEST_SIZE];
10088    word32 oidMsgHashLen = 0;
10089
10090    if (key == NULL) {
10091        ret = BAD_FUNC_ARG;
10092    }
10093    /* Check that the input hash length is valid. */
10094    if ((ret == 0) &&
10095        ((int)hashLen != wc_HashGetDigestSize((enum wc_HashType)hashAlg)))
10096    {
10097        ret = BAD_LENGTH_E;
10098    }
10099
10100    if (ret == 0) {
10101        /* Step 6: Hash public key. */
10102        ret = dilithium_shake256(&key->shake, key->p, key->params->pkSz, tr,
10103            DILITHIUM_TR_SZ);
10104    }
10105    if (ret == 0) {
10106        ret = dilithium_get_hash_oid(hashAlg, oidMsgHash, &oidMsgHashLen);
10107    }
10108    if (ret == 0) {
10109        XMEMCPY(oidMsgHash + oidMsgHashLen, hash, hashLen);
10110        oidMsgHashLen += hashLen;
10111
10112        /* Step 6. Calculate mu. */
10113        ret = dilithium_hash256_ctx_msg(&key->shake, tr, DILITHIUM_TR_SZ, 1,
10114            ctx, (byte)ctxLen, oidMsgHash, oidMsgHashLen, mu, DILITHIUM_MU_SZ);
10115    }
10116    if (ret == 0) {
10117        ret = dilithium_verify_with_mu(key, mu, sig, sigLen, res);
10118    }
10119
10120    return ret;
10121}
10122#endif /* WOLFSSL_DILITHIUM_NO_VERIFY */
10123
10124#ifndef WOLFSSL_DILITHIUM_NO_MAKE_KEY
10125int wc_dilithium_make_key(dilithium_key* key, WC_RNG* rng)
10126{
10127    int ret = 0;
10128
10129    /* Validate parameters. */
10130    if ((key == NULL) || (rng == NULL)) {
10131        ret = BAD_FUNC_ARG;
10132    }
10133
10134#ifdef WOLF_CRYPTO_CB
10135    if (ret == 0) {
10136    #ifndef WOLF_CRYPTO_CB_FIND
10137        if (key->devId != INVALID_DEVID)
10138    #endif
10139        {
10140            ret = wc_CryptoCb_MakePqcSignatureKey(rng,
10141                WC_PQC_SIG_TYPE_DILITHIUM, key->level, key);
10142            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10143                return ret;
10144            /* fall-through when unavailable */
10145            ret = 0;
10146        }
10147    }
10148#endif
10149
10150    if (ret == 0) {
10151        /* Check the level or parameters have been set. */
10152        if (key->params == NULL) {
10153            ret = BAD_STATE_E;
10154        }
10155        else {
10156            /* Make the key. */
10157            ret = dilithium_make_key(key, rng);
10158        }
10159    }
10160
10161#ifdef HAVE_FIPS
10162    /* Pairwise Consistency Test (PCT) per FIPS 140-3 / ISO 19790:2012
10163     * Section 7.10.3.3 (TE10.35.02): sign with new sk, verify with pk.
10164     * Runs on every key generation. */
10165    if (ret == 0) {
10166        static const byte pct_msg[] = "wolfSSL ML-DSA PCT";
10167        WC_DECLARE_VAR(pct_sig, byte, DILITHIUM_MAX_SIG_SIZE, key->heap);
10168        word32 pct_sigSz = DILITHIUM_MAX_SIG_SIZE;
10169        int pct_res = 0;
10170
10171        WC_ALLOC_VAR_EX(pct_sig, byte, DILITHIUM_MAX_SIG_SIZE, key->heap,
10172            DYNAMIC_TYPE_DILITHIUM, ret = MEMORY_E);
10173
10174        if (ret == 0) {
10175            ret = wc_dilithium_sign_ctx_msg(NULL, 0, pct_msg, sizeof(pct_msg),
10176                pct_sig, &pct_sigSz, key, rng);
10177        }
10178
10179        if (ret == 0)
10180            ret = wc_dilithium_verify_ctx_msg(pct_sig, pct_sigSz,
10181                NULL, 0, pct_msg, sizeof(pct_msg), &pct_res, key);
10182
10183        if (ret == 0 && pct_res != 1)
10184            ret = ML_DSA_PCT_E;
10185
10186        if (WC_VAR_OK(pct_sig))
10187            ForceZero(pct_sig, DILITHIUM_MAX_SIG_SIZE);
10188
10189        WC_FREE_VAR_EX(pct_sig, key->heap, DYNAMIC_TYPE_DILITHIUM);
10190
10191        /* FIPS 140-3 IG 10.3.A (TE10.35.02): a key pair that fails the PCT
10192         * must be rendered unusable.  Zeroize the generated key material so
10193         * a caller that ignores the return value cannot use it. */
10194        if (ret != 0) {
10195            wc_dilithium_free(key);
10196        }
10197    }
10198#endif /* HAVE_FIPS */
10199
10200    return ret;
10201}
10202
10203int wc_dilithium_make_key_from_seed(dilithium_key* key, const byte* seed)
10204{
10205    int ret = 0;
10206
10207    /* Validate parameters. */
10208    if ((key == NULL) || (seed == NULL)) {
10209        ret = BAD_FUNC_ARG;
10210    }
10211
10212    if (ret == 0) {
10213        /* Check the level or parameters have been set. */
10214        if (key->params == NULL) {
10215            ret = BAD_STATE_E;
10216        }
10217        else {
10218            /* Make the key. */
10219            ret = dilithium_make_key_from_seed(key, seed);
10220        }
10221    }
10222
10223    /* Note: PCT is performed in wc_dilithium_make_key() which calls this
10224     * function and has the RNG parameter needed for signing. */
10225
10226    return ret;
10227}
10228#endif
10229
10230#ifndef WOLFSSL_DILITHIUM_NO_SIGN
10231/* Sign the message using the dilithium private key.
10232 *
10233 *  ctx         [in]      Context of signature.
10234 *  ctxLen      [in]      Length of context in bytes.
10235 *  msg         [in]      Message to sign.
10236 *  msgLen      [in]      Length of the message in bytes.
10237 *  sig         [out]     Buffer to write signature into.
10238 *  sigLen      [in/out]  On in, size of buffer.
10239 *                        On out, the length of the signature in bytes.
10240 *  key         [in]      Dilithium key to use when signing
10241 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10242 *          or ctx is NULL and ctxLen is not 0,
10243 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10244 *          0 otherwise.
10245 */
10246int wc_dilithium_sign_ctx_msg(const byte* ctx, byte ctxLen, const byte* msg,
10247    word32 msgLen, byte* sig, word32 *sigLen, dilithium_key* key, WC_RNG* rng)
10248{
10249    int ret = 0;
10250
10251    /* Validate parameters. */
10252    if ((msg == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL)) {
10253        ret = BAD_FUNC_ARG;
10254    }
10255    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10256        ret = BAD_FUNC_ARG;
10257    }
10258    if ((ret == 0) && (!key->prvKeySet)) {
10259        ret = BAD_FUNC_ARG;
10260    }
10261
10262#ifdef WOLF_CRYPTO_CB
10263    if (ret == 0) {
10264    #ifndef WOLF_CRYPTO_CB_FIND
10265        if (key->devId != INVALID_DEVID)
10266    #endif
10267        {
10268            ret = wc_CryptoCb_PqcSign(msg, msgLen, sig, sigLen, ctx, ctxLen,
10269                    WC_HASH_TYPE_NONE, rng, WC_PQC_SIG_TYPE_DILITHIUM, key);
10270            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10271                return ret;
10272            /* fall-through when unavailable */
10273            ret = 0;
10274        }
10275    }
10276#endif
10277
10278    if (ret == 0) {
10279        /* Sign message. */
10280        ret = dilithium_sign_ctx_msg(key, rng, ctx, ctxLen, msg, msgLen, sig,
10281            sigLen);
10282    }
10283
10284    return ret;
10285}
10286
10287#ifdef WOLFSSL_DILITHIUM_NO_CTX
10288/* Sign the message using the dilithium private key.
10289 *
10290 *  msg         [in]      Message to sign.
10291 *  msgLen      [in]      Length of the message in bytes.
10292 *  sig         [out]     Buffer to write signature into.
10293 *  sigLen      [in/out]  On in, size of buffer.
10294 *                        On out, the length of the signature in bytes.
10295 *  key         [in]      Dilithium key to use when signing
10296 *  returns BAD_FUNC_ARG when a parameter is NULL or public key not set,
10297 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10298 *          0 otherwise.
10299 * NOTE: This is a pre-FIPS 204 API without context support. New code should
10300 *       use wc_dilithium_sign_ctx_msg() with ctx=NULL/ctxLen=0 instead.
10301 */
10302int wc_dilithium_sign_msg(const byte* msg, word32 msgLen, byte* sig,
10303    word32 *sigLen, dilithium_key* key, WC_RNG* rng)
10304{
10305    int ret = 0;
10306
10307    /* Validate parameters. */
10308    if ((msg == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL)) {
10309        ret = BAD_FUNC_ARG;
10310    }
10311
10312#ifdef WOLF_CRYPTO_CB
10313    if (ret == 0) {
10314    #ifndef WOLF_CRYPTO_CB_FIND
10315        if (key->devId != INVALID_DEVID)
10316    #endif
10317        {
10318            ret = wc_CryptoCb_PqcSign(msg, msgLen, sig, sigLen, NULL, 0,
10319                    WC_HASH_TYPE_NONE, rng, WC_PQC_SIG_TYPE_DILITHIUM, key);
10320            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10321                return ret;
10322            /* fall-through when unavailable */
10323            ret = 0;
10324        }
10325    }
10326#endif
10327
10328    if (ret == 0) {
10329        /* Sign message. */
10330        ret = dilithium_sign_msg(key, rng, msg, msgLen, sig, sigLen);
10331    }
10332
10333    return ret;
10334}
10335#endif /* WOLFSSL_DILITHIUM_NO_CTX */
10336
10337/* Sign the message hash using the dilithium private key.
10338 *
10339 *  ctx         [in]      Context of signature.
10340 *  ctxLen      [in]      Length of context in bytes.
10341 *  hashAlg     [in]      Hash algorithm used on message.
10342 *  hash        [in]      Hash of message to sign.
10343 *  hashLen     [in]      Length of the message hash in bytes.
10344 *  sig         [out]     Buffer to write signature into.
10345 *  sigLen      [in/out]  On in, size of buffer.
10346 *                        On out, the length of the signature in bytes.
10347 *  key         [in]      Dilithium key to use when signing
10348 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10349 *          or ctx is NULL and ctxLen is not 0,
10350 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10351 *          0 otherwise.
10352 */
10353int wc_dilithium_sign_ctx_hash(const byte* ctx, byte ctxLen, int hashAlg,
10354   const byte* hash, word32 hashLen, byte* sig, word32 *sigLen,
10355   dilithium_key* key, WC_RNG* rng)
10356{
10357    int ret = 0;
10358
10359    /* Validate parameters. */
10360    if ((hash == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL)) {
10361        ret = BAD_FUNC_ARG;
10362    }
10363    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10364        ret = BAD_FUNC_ARG;
10365    }
10366
10367#ifdef WOLF_CRYPTO_CB
10368    if (ret == 0) {
10369    #ifndef WOLF_CRYPTO_CB_FIND
10370        if (key->devId != INVALID_DEVID)
10371    #endif
10372        {
10373            ret = wc_CryptoCb_PqcSign(hash, hashLen, sig, sigLen, ctx, ctxLen,
10374                    (word32)hashAlg, rng, WC_PQC_SIG_TYPE_DILITHIUM, key);
10375            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10376                return ret;
10377            /* fall-through when unavailable */
10378            ret = 0;
10379        }
10380    }
10381#endif
10382
10383    if (ret == 0) {
10384        /* Sign message. */
10385        ret = dilithium_sign_ctx_hash(key, rng, ctx, ctxLen, hashAlg, hash,
10386            hashLen, sig, sigLen);
10387    }
10388
10389    return ret;
10390}
10391
10392/* Sign the message using the dilithium private key.
10393 *
10394 *  ctx         [in]      Context of signature.
10395 *  ctxLen      [in]      Length of context in bytes.
10396 *  msg         [in]      Message to sign.
10397 *  msgLen      [in]      Length of the message in bytes.
10398 *  sig         [out]     Buffer to write signature into.
10399 *  sigLen      [in/out]  On in, size of buffer.
10400 *                        On out, the length of the signature in bytes.
10401 *  key         [in]      Dilithium key to use when signing
10402 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10403 *          or ctx is NULL and ctxLen is not 0,
10404 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10405 *          0 otherwise.
10406 */
10407int wc_dilithium_sign_ctx_msg_with_seed(const byte* ctx, byte ctxLen,
10408    const byte* msg, word32 msgLen, byte* sig, word32 *sigLen,
10409    dilithium_key* key, const byte* seed)
10410{
10411    int ret = 0;
10412
10413    /* Validate parameters. */
10414    if ((msg == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL)) {
10415        ret = BAD_FUNC_ARG;
10416    }
10417    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10418        ret = BAD_FUNC_ARG;
10419    }
10420
10421    if (ret == 0) {
10422        /* Sign message. */
10423        ret = dilithium_sign_ctx_msg_with_seed(key, seed, ctx, ctxLen, msg,
10424            msgLen, sig, sigLen);
10425    }
10426
10427    return ret;
10428}
10429
10430#ifdef WOLFSSL_DILITHIUM_NO_CTX
10431/* Sign the message using the dilithium private key.
10432 *
10433 *  msg         [in]      Message to sign.
10434 *  msgLen      [in]      Length of the message in bytes.
10435 *  sig         [out]     Buffer to write signature into.
10436 *  sigLen      [in/out]  On in, size of buffer.
10437 *                        On out, the length of the signature in bytes.
10438 *  key         [in]      Dilithium key to use when signing
10439 *  returns BAD_FUNC_ARG when a parameter is NULL or public key not set,
10440 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10441 *          0 otherwise.
10442 * NOTE: This is a pre-FIPS 204 API without context support. New code should
10443 *       use wc_dilithium_sign_ctx_msg_with_seed() instead.
10444 */
10445int wc_dilithium_sign_msg_with_seed(const byte* msg, word32 msgLen, byte* sig,
10446    word32 *sigLen, dilithium_key* key, const byte* seed)
10447{
10448    int ret = 0;
10449
10450    /* Validate parameters. */
10451    if ((msg == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL)) {
10452        ret = BAD_FUNC_ARG;
10453    }
10454
10455    if (ret == 0) {
10456        /* Sign message. */
10457        ret = dilithium_sign_msg_with_seed(key, seed, msg, msgLen, sig, sigLen);
10458    }
10459
10460    return ret;
10461}
10462#endif /* WOLFSSL_DILITHIUM_NO_CTX */
10463
10464/* Sign the message using the dilithium private key.
10465 *
10466 *  ctx         [in]      Context of signature.
10467 *  ctxLen      [in]      Length of context in bytes.
10468 *  hashAlg     [in]      Hash algorithm used on message.
10469 *  hash        [in]      Hash of message to sign.
10470 *  hashLen     [in]      Length of the message hash in bytes.
10471 *  sig         [out]     Buffer to write signature into.
10472 *  sigLen      [in/out]  On in, size of buffer.
10473 *                        On out, the length of the signature in bytes.
10474 *  key         [in]      Dilithium key to use when signing
10475 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10476 *          or ctx is NULL and ctxLen is not 0,
10477 *          BUFFER_E when outLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10478 *          0 otherwise.
10479 */
10480int wc_dilithium_sign_ctx_hash_with_seed(const byte* ctx, byte ctxLen,
10481    int hashAlg, const byte* hash, word32 hashLen, byte* sig, word32 *sigLen,
10482    dilithium_key* key, const byte* seed)
10483{
10484    int ret = 0;
10485
10486    /* Validate parameters. */
10487    if ((hash == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL) ||
10488            (seed == NULL)) {
10489        ret = BAD_FUNC_ARG;
10490    }
10491    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10492        ret = BAD_FUNC_ARG;
10493    }
10494
10495    if (ret == 0) {
10496        /* Sign message. */
10497        ret = dilithium_sign_ctx_hash_with_seed(key, seed, ctx, ctxLen,
10498            hashAlg, hash, hashLen, sig, sigLen);
10499    }
10500
10501    return ret;
10502}
10503
10504/* Sign using the ML-DSA internal interface with a pre-computed mu value.
10505 *
10506 * This implements ML-DSA.Sign_internal from FIPS 204 Section 6.2.
10507 * The caller provides mu directly (already computed from tr||M'), bypassing
10508 * the external message hashing step. Used by ACVP internal interface tests.
10509 *
10510 *  mu          [in]      Pre-computed mu value (64 bytes).
10511 *  muLen       [in]      Length of mu in bytes (must be 64).
10512 *  sig         [out]     Buffer to write signature into.
10513 *  sigLen      [in/out]  On in, size of buffer.
10514 *                        On out, the length of the signature in bytes.
10515 *  key         [in]      Dilithium key to use when signing.
10516 *  seed        [in]      32-byte random seed (rnd).
10517 *  returns BAD_FUNC_ARG when a parameter is NULL or muLen is not 64,
10518 *          BUFFER_E when sigLen is too small,
10519 *          0 otherwise.
10520 */
10521int wc_dilithium_sign_mu_with_seed(const byte* mu, word32 muLen,
10522    byte* sig, word32 *sigLen, dilithium_key* key, const byte* seed)
10523{
10524    int ret = 0;
10525
10526    /* Validate parameters. */
10527    if ((mu == NULL) || (sig == NULL) || (sigLen == NULL) || (key == NULL) ||
10528            (seed == NULL)) {
10529        ret = BAD_FUNC_ARG;
10530    }
10531    if ((ret == 0) && (muLen != DILITHIUM_MU_SZ)) {
10532        ret = BAD_FUNC_ARG;
10533    }
10534
10535    if (ret == 0) {
10536        /* Build [seed||mu] buffer and call internal sign function. */
10537        byte seedMu[DILITHIUM_RND_SZ + DILITHIUM_MU_SZ];
10538        XMEMCPY(seedMu, seed, DILITHIUM_RND_SZ);
10539        XMEMCPY(seedMu + DILITHIUM_RND_SZ, mu, DILITHIUM_MU_SZ);
10540        ret = dilithium_sign_with_seed_mu(key, seedMu, sig, sigLen);
10541        ForceZero(seedMu, sizeof(seedMu));
10542    }
10543
10544    return ret;
10545}
10546#endif /* !WOLFSSL_DILITHIUM_NO_SIGN */
10547
10548#ifndef WOLFSSL_DILITHIUM_NO_VERIFY
10549/* Verify the message using the dilithium public key.
10550 *
10551 *  sig         [in]  Signature to verify.
10552 *  sigLen      [in]  Size of signature in bytes.
10553 *  ctx         [in]  Context of signature.
10554 *  ctxLen      [in]  Length of context in bytes.
10555 *  msg         [in]  Message to verify.
10556 *  msgLen      [in]  Length of the message in bytes.
10557 *  res         [out] *res is set to 1 on successful verification.
10558 *  key         [in]  Dilithium key to use to verify.
10559 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10560 *          or ctx is NULL and ctxLen is not 0,
10561 *          BUFFER_E when sigLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10562 *          0 otherwise.
10563 */
10564int wc_dilithium_verify_ctx_msg(const byte* sig, word32 sigLen, const byte* ctx,
10565    byte ctxLen, const byte* msg, word32 msgLen, int* res, dilithium_key* key)
10566{
10567    int ret = 0;
10568
10569    /* Validate parameters. */
10570    if ((key == NULL) || (sig == NULL) || (msg == NULL) || (res == NULL)) {
10571        ret = BAD_FUNC_ARG;
10572    }
10573    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10574        ret = BAD_FUNC_ARG;
10575    }
10576    /* Reject msgLen that would cause integer overflow in hash computations */
10577    if ((ret == 0) && (msgLen > WOLFSSL_MAX_32BIT / 2)) {
10578        ret = BAD_FUNC_ARG;
10579    }
10580
10581#ifdef WOLF_CRYPTO_CB
10582    if (ret == 0) {
10583    #ifndef WOLF_CRYPTO_CB_FIND
10584        if (key->devId != INVALID_DEVID)
10585    #endif
10586        {
10587            ret = wc_CryptoCb_PqcVerify(sig, sigLen, msg, msgLen, ctx, ctxLen,
10588                    WC_HASH_TYPE_NONE, res, WC_PQC_SIG_TYPE_DILITHIUM, key);
10589            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10590                return ret;
10591            /* fall-through when unavailable */
10592            ret = 0;
10593        }
10594    }
10595#endif
10596
10597    if (ret == 0) {
10598        /* Verify message with signature. */
10599        ret = dilithium_verify_ctx_msg(key, ctx, ctxLen, msg, msgLen, sig,
10600            sigLen, res);
10601    }
10602
10603    return ret;
10604}
10605
10606#ifdef WOLFSSL_DILITHIUM_NO_CTX
10607/* Verify the message using the dilithium public key.
10608 *
10609 *  sig         [in]  Signature to verify.
10610 *  sigLen      [in]  Size of signature in bytes.
10611 *  msg         [in]  Message to verify.
10612 *  msgLen      [in]  Length of the message in bytes.
10613 *  res         [out] *res is set to 1 on successful verification.
10614 *  key         [in]  Dilithium key to use to verify.
10615 *  returns BAD_FUNC_ARG when a parameter is NULL or contextLen is zero when and
10616 *          BUFFER_E when sigLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10617 *          0 otherwise.
10618 * NOTE: This is a pre-FIPS 204 API without context support. New code should
10619 *       use wc_dilithium_verify_ctx_msg() with ctx=NULL/ctxLen=0 instead.
10620 */
10621int wc_dilithium_verify_msg(const byte* sig, word32 sigLen, const byte* msg,
10622    word32 msgLen, int* res, dilithium_key* key)
10623{
10624    int ret = 0;
10625
10626    /* Validate parameters. */
10627    if ((key == NULL) || (sig == NULL) || (msg == NULL) || (res == NULL)) {
10628        ret = BAD_FUNC_ARG;
10629    }
10630
10631#ifdef WOLF_CRYPTO_CB
10632    if (ret == 0) {
10633    #ifndef WOLF_CRYPTO_CB_FIND
10634        if (key->devId != INVALID_DEVID)
10635    #endif
10636        {
10637            ret = wc_CryptoCb_PqcVerify(sig, sigLen, msg, msgLen, NULL, 0,
10638                    WC_HASH_TYPE_NONE, res, WC_PQC_SIG_TYPE_DILITHIUM, key);
10639            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10640                return ret;
10641            /* fall-through when unavailable */
10642            ret = 0;
10643        }
10644    }
10645#endif
10646
10647    if (ret == 0) {
10648        /* Verify message with signature. */
10649        ret = dilithium_verify_msg(key, msg, msgLen, sig, sigLen, res);
10650    }
10651
10652    return ret;
10653}
10654#endif /* WOLFSSL_DILITHIUM_NO_CTX */
10655
10656/* Verify the message using the dilithium public key.
10657 *
10658 *  sig         [in]  Signature to verify.
10659 *  sigLen      [in]  Size of signature in bytes.
10660 *  ctx         [in]  Context of signature.
10661 *  ctxLen      [in]  Length of context in bytes.
10662 *  hashAlg     [in]  Hash algorithm used on message.
10663 *  hash        [in]  Hash of message to verify.
10664 *  hashLen     [in]  Length of the message hash in bytes.
10665 *  res         [out] *res is set to 1 on successful verification.
10666 *  key         [in]  Dilithium key to use to verify.
10667 *  returns BAD_FUNC_ARG when a parameter is NULL, public key not set
10668 *          or ctx is NULL and ctxLen is not 0,
10669 *          BUFFER_E when sigLen is less than DILITHIUM_LEVEL2_SIG_SIZE,
10670 *          0 otherwise.
10671 */
10672int wc_dilithium_verify_ctx_hash(const byte* sig, word32 sigLen,
10673    const byte* ctx, byte ctxLen, int hashAlg, const byte* hash, word32 hashLen,
10674    int* res, dilithium_key* key)
10675{
10676    int ret = 0;
10677
10678    /* Validate parameters. */
10679    if ((key == NULL) || (sig == NULL) || (hash == NULL) || (res == NULL)) {
10680        ret = BAD_FUNC_ARG;
10681    }
10682    if ((ret == 0) && (ctx == NULL) && (ctxLen > 0)) {
10683        ret = BAD_FUNC_ARG;
10684    }
10685
10686#ifdef WOLF_CRYPTO_CB
10687    if (ret == 0) {
10688    #ifndef WOLF_CRYPTO_CB_FIND
10689        if (key->devId != INVALID_DEVID)
10690    #endif
10691        {
10692            ret = wc_CryptoCb_PqcVerify(sig, sigLen, hash, hashLen, ctx, ctxLen,
10693                    (word32)hashAlg, res, WC_PQC_SIG_TYPE_DILITHIUM, key);
10694            if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
10695                return ret;
10696            /* fall-through when unavailable */
10697            ret = 0;
10698        }
10699    }
10700#endif
10701
10702    if (ret == 0) {
10703        /* Verify message with signature. */
10704        ret = dilithium_verify_ctx_hash(key, ctx, ctxLen, hashAlg, hash,
10705            hashLen, sig, sigLen, res);
10706    }
10707
10708    return ret;
10709}
10710
10711/* Verify using the ML-DSA internal interface with a pre-computed mu value.
10712 *
10713 * This implements ML-DSA.Verify_internal from FIPS 204 Section 6.3.
10714 * The caller provides mu directly (already computed from tr||M'), bypassing
10715 * the external message hashing step. Used by ACVP internal interface tests.
10716 *
10717 *  sig         [in]  Signature to verify.
10718 *  sigLen      [in]  Size of signature in bytes.
10719 *  mu          [in]  Pre-computed mu value (64 bytes).
10720 *  muLen       [in]  Length of mu in bytes (must be 64).
10721 *  res         [out] *res is set to 1 on successful verification.
10722 *  key         [in]  Dilithium key to use to verify.
10723 *  returns BAD_FUNC_ARG when a parameter is NULL or muLen is not 64,
10724 *          0 otherwise.
10725 */
10726int wc_dilithium_verify_mu(const byte* sig, word32 sigLen, const byte* mu,
10727    word32 muLen, int* res, dilithium_key* key)
10728{
10729    int ret = 0;
10730
10731    /* Validate parameters. */
10732    if ((key == NULL) || (sig == NULL) || (mu == NULL) || (res == NULL)) {
10733        ret = BAD_FUNC_ARG;
10734    }
10735    if ((ret == 0) && (muLen != DILITHIUM_MU_SZ)) {
10736        ret = BAD_FUNC_ARG;
10737    }
10738
10739    if (ret == 0) {
10740        ret = dilithium_verify_with_mu(key, mu, sig, sigLen, res);
10741    }
10742
10743    return ret;
10744}
10745#endif /* WOLFSSL_DILITHIUM_NO_VERIFY */
10746
10747#ifndef WC_NO_CONSTRUCTORS
10748/**
10749 * Create a new dilithium key object.
10750 *
10751 * heap  [in]  Dynamic memory hint.
10752 * devId [in]  Device Id.
10753 * returns MEMORY_E when dynamic memory allocation fails
10754 */
10755
10756dilithium_key* wc_dilithium_new(void* heap, int devId)
10757{
10758    int ret;
10759    dilithium_key* key = (dilithium_key*)XMALLOC(sizeof(dilithium_key), heap,
10760        DYNAMIC_TYPE_DILITHIUM);
10761    if (key != NULL) {
10762        ret = wc_dilithium_init_ex(key, heap, devId);
10763        if (ret != 0) {
10764            XFREE(key, heap, DYNAMIC_TYPE_DILITHIUM);
10765            key = NULL;
10766        }
10767    }
10768
10769    return key;
10770}
10771
10772/**
10773 * Delete and free a dilithium key object.
10774 *
10775 * key   [in]       dilithium key object to delete.
10776 * key_p [in, out]  Pointer to key pointer to set to NULL.
10777 * returns BAD_FUNC_ARG when key is NULL
10778 */
10779
10780int wc_dilithium_delete(dilithium_key* key, dilithium_key** key_p)
10781{
10782    void* heap;
10783    if (key == NULL)
10784        return BAD_FUNC_ARG;
10785    heap = key->heap;
10786    wc_dilithium_free(key);
10787    XFREE(key, heap, DYNAMIC_TYPE_DILITHIUM);
10788    if (key_p != NULL)
10789        *key_p = NULL;
10790
10791    return 0;
10792}
10793#endif /* !WC_NO_CONSTRUCTORS */
10794
10795/* Initialize the dilithium private/public key.
10796 *
10797 * key  [in]  Dilithium key.
10798 * returns BAD_FUNC_ARG when key is NULL
10799 */
10800int wc_dilithium_init(dilithium_key* key)
10801{
10802    return wc_dilithium_init_ex(key, NULL, INVALID_DEVID);
10803}
10804
10805/* Initialize the dilithium private/public key.
10806 *
10807 * key  [in]  Dilithium key.
10808 * heap [in]  Heap hint.
10809 * devId[in]  Device ID.
10810 * returns BAD_FUNC_ARG when key is NULL
10811 */
10812int wc_dilithium_init_ex(dilithium_key* key, void* heap, int devId)
10813{
10814    int ret = 0;
10815
10816    (void)devId;
10817
10818    /* Validate parameters. */
10819    if (key == NULL) {
10820        ret = BAD_FUNC_ARG;
10821    }
10822
10823    if (ret == 0) {
10824        /* Ensure all fields reset. */
10825        XMEMSET(key, 0, sizeof(*key));
10826
10827    #ifdef WOLF_CRYPTO_CB
10828        key->devCtx = NULL;
10829        key->devId = devId;
10830    #endif
10831    #ifdef WOLF_PRIVATE_KEY_ID
10832        key->idLen = 0;
10833        key->labelLen = 0;
10834    #endif
10835        key->heap = heap;
10836    }
10837
10838#if defined(USE_INTEL_SPEEDUP)
10839    cpuid_get_flags_ex(&cpuid_flags);
10840#endif
10841
10842    return ret;
10843}
10844
10845#ifdef WOLF_PRIVATE_KEY_ID
10846int wc_dilithium_init_id(dilithium_key* key, const unsigned char* id, int len,
10847    void* heap, int devId)
10848{
10849    int ret = 0;
10850
10851    if (key == NULL) {
10852        ret = BAD_FUNC_ARG;
10853    }
10854    if ((ret == 0) && ((len < 0) || (len > DILITHIUM_MAX_ID_LEN))) {
10855        ret = BUFFER_E;
10856    }
10857
10858    if (ret == 0) {
10859        ret = wc_dilithium_init_ex(key, heap, devId);
10860    }
10861    if ((ret == 0) && (id != NULL) && (len != 0)) {
10862        XMEMCPY(key->id, id, (size_t)len);
10863        key->idLen = len;
10864    }
10865
10866    /* Set the maximum level here */
10867    wc_dilithium_set_level(key, WC_ML_DSA_87);
10868
10869    return ret;
10870}
10871
10872int wc_dilithium_init_label(dilithium_key* key, const char* label, void* heap,
10873    int devId)
10874{
10875    int ret = 0;
10876    int labelLen = 0;
10877
10878    if ((key == NULL) || (label == NULL)) {
10879        ret = BAD_FUNC_ARG;
10880    }
10881    if (ret == 0) {
10882        labelLen = (int)XSTRLEN(label);
10883        if ((labelLen == 0) || (labelLen > DILITHIUM_MAX_LABEL_LEN)) {
10884            ret = BUFFER_E;
10885        }
10886    }
10887
10888    if (ret == 0) {
10889        ret = wc_dilithium_init_ex(key, heap, devId);
10890    }
10891    if (ret == 0) {
10892        XMEMCPY(key->label, label, (size_t)labelLen);
10893        key->labelLen = labelLen;
10894    }
10895
10896    /* Set the maximum level here */
10897    wc_dilithium_set_level(key, WC_ML_DSA_87);
10898
10899    return ret;
10900}
10901#endif
10902
10903/* Set the level of the dilithium private/public key.
10904 *
10905 * key   [out]  Dilithium key.
10906 * level [in]   Either 2,3 or 5.
10907 * returns BAD_FUNC_ARG when key is NULL or level is a bad values.
10908 */
10909int wc_dilithium_set_level(dilithium_key* key, byte level)
10910{
10911    int ret = 0;
10912
10913    /* Validate parameters. */
10914    if (key == NULL) {
10915        ret = BAD_FUNC_ARG;
10916    }
10917    if ((ret == 0) && ((level == WC_ML_DSA_44) || (level == WC_ML_DSA_65) ||
10918            (level == WC_ML_DSA_87))) {
10919        /* Nothing to do. */
10920    }
10921#if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
10922    else if ((ret == 0) && ((level == WC_ML_DSA_44_DRAFT) ||
10923             (level == WC_ML_DSA_65_DRAFT) || (level == WC_ML_DSA_87_DRAFT))) {
10924        /* Nothing to do. */
10925    }
10926#endif
10927    else {
10928        ret = BAD_FUNC_ARG;
10929    }
10930
10931    if (ret == 0) {
10932        /* Get the parameters for level into key. */
10933        ret = dilithium_get_params(level, &key->params);
10934    }
10935    if (ret == 0) {
10936        /* Clear any cached items. */
10937#ifndef WC_DILITHIUM_FIXED_ARRAY
10938    #ifdef WC_DILITHIUM_CACHE_MATRIX_A
10939        XFREE(key->a, key->heap, DYNAMIC_TYPE_DILITHIUM);
10940        key->a = NULL;
10941        key->aSet = 0;
10942    #endif
10943    #ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
10944        XFREE(key->s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
10945        key->s1 = NULL;
10946        key->s2 = NULL;
10947        key->t0 = NULL;
10948        key->privVecsSet = 0;
10949    #endif
10950    #ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
10951        XFREE(key->t1, key->heap, DYNAMIC_TYPE_DILITHIUM);
10952        key->t1 = NULL;
10953        key->pubVecSet = 0;
10954    #endif
10955#endif
10956
10957#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
10958        if (key->k != NULL) {
10959            ForceZero(key->k, key->kSz);
10960            XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
10961            key->k = NULL;
10962            key->kSz = 0;
10963        }
10964        if (key->p != NULL) {
10965            XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
10966            key->p = NULL;
10967        }
10968#endif
10969
10970        /* Store level and indicate public and private key are not set. */
10971        key->level = level % WC_ML_DSA_DRAFT;
10972        key->pubKeySet = 0;
10973        key->prvKeySet = 0;
10974    }
10975
10976    return ret;
10977}
10978
10979/* Get the level of the dilithium private/public key.
10980 *
10981 * key   [in]  Dilithium key.
10982 * level [out] The level.
10983 * returns BAD_FUNC_ARG when key is NULL or level has not been set.
10984 */
10985int wc_dilithium_get_level(dilithium_key* key, byte* level)
10986{
10987    int ret = 0;
10988
10989    /* Validate parameters. */
10990    if ((key == NULL) || (level == NULL)) {
10991        ret = BAD_FUNC_ARG;
10992    }
10993    if ((ret == 0) && (key->level != WC_ML_DSA_44) &&
10994            (key->level != WC_ML_DSA_65) && (key->level != WC_ML_DSA_87)) {
10995        ret = BAD_FUNC_ARG;
10996    }
10997
10998    if (ret == 0) {
10999        /* Return level. */
11000        *level = key->level;
11001    }
11002
11003    return ret;
11004}
11005
11006/* Clears the dilithium key data
11007 *
11008 * key  [in]  Dilithium key.
11009 */
11010void wc_dilithium_free(dilithium_key* key)
11011{
11012    if (key != NULL) {
11013#if defined(WOLF_CRYPTO_CB) && defined(WOLF_CRYPTO_CB_FREE)
11014        if (key->devId != INVALID_DEVID) {
11015            (void)wc_CryptoCb_Free(key->devId, WC_ALGO_TYPE_PK,
11016                             WC_PK_TYPE_PQC_SIG_KEYGEN,
11017                             WC_PQC_SIG_TYPE_DILITHIUM,
11018                             (void*)key);
11019            /* always continue to software cleanup */
11020        }
11021#endif
11022#ifndef WC_DILITHIUM_FIXED_ARRAY
11023        /* Dispose of cached items. */
11024    #ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
11025        XFREE(key->t1, key->heap, DYNAMIC_TYPE_DILITHIUM);
11026    #endif
11027    #ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
11028        XFREE(key->s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
11029    #endif
11030    #ifdef WC_DILITHIUM_CACHE_MATRIX_A
11031        XFREE(key->a, key->heap, DYNAMIC_TYPE_DILITHIUM);
11032    #endif
11033#endif
11034        /* Intel speedup code manually manipulates the state. */
11035#ifndef USE_INTEL_SPEEDUP
11036        /* Free the SHAKE-128/256 object. */
11037        wc_Shake256_Free(&key->shake);
11038#endif
11039#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11040        if (key->k != NULL) {
11041            ForceZero(key->k, key->kSz);
11042            XFREE(key->k, key->heap, DYNAMIC_TYPE_DILITHIUM);
11043        }
11044        if (key->p != NULL) {
11045            XFREE(key->p, key->heap, DYNAMIC_TYPE_DILITHIUM);
11046        }
11047#endif
11048        /* Ensure all private data is zeroized. */
11049        ForceZero(key, sizeof(*key));
11050    }
11051}
11052
11053#ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
11054/* Returns the size of a dilithium private key.
11055 *
11056 * @param [in] key  Dilithium private/public key.
11057 * @return  Private key size on success for set level.
11058 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11059 */
11060int wc_dilithium_size(dilithium_key* key)
11061{
11062    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
11063
11064    if (key != NULL) {
11065    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11066        if (key->params == NULL) {
11067            ret = BAD_FUNC_ARG;
11068        }
11069        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11070            ret = DILITHIUM_LEVEL2_KEY_SIZE;
11071        }
11072        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11073            ret = DILITHIUM_LEVEL3_KEY_SIZE;
11074        }
11075        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11076            ret = DILITHIUM_LEVEL5_KEY_SIZE;
11077        }
11078        else
11079    #endif
11080        if (key->level == WC_ML_DSA_44) {
11081            ret = ML_DSA_LEVEL2_KEY_SIZE;
11082        }
11083        else if (key->level == WC_ML_DSA_65) {
11084            ret = ML_DSA_LEVEL3_KEY_SIZE;
11085        }
11086        else if (key->level == WC_ML_DSA_87) {
11087            ret = ML_DSA_LEVEL5_KEY_SIZE;
11088        }
11089    }
11090
11091    return ret;
11092}
11093
11094#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
11095/* Returns the size of a dilithium private plus public key.
11096 *
11097 * @param [in] key  Dilithium private/public key.
11098 * @return  Private key size on success for set level.
11099 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11100 */
11101int wc_dilithium_priv_size(dilithium_key* key)
11102{
11103    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
11104
11105    if (key != NULL) {
11106    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11107        if (key->params == NULL) {
11108            ret = BAD_FUNC_ARG;
11109        }
11110        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11111            ret = DILITHIUM_LEVEL2_PRV_KEY_SIZE;
11112        }
11113        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11114            ret = DILITHIUM_LEVEL3_PRV_KEY_SIZE;
11115        }
11116        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11117            ret = DILITHIUM_LEVEL5_PRV_KEY_SIZE;
11118        }
11119    #endif
11120        if (key->level == WC_ML_DSA_44) {
11121            ret = ML_DSA_LEVEL2_PRV_KEY_SIZE;
11122        }
11123        else if (key->level == WC_ML_DSA_65) {
11124            ret = ML_DSA_LEVEL3_PRV_KEY_SIZE;
11125        }
11126        else if (key->level == WC_ML_DSA_87) {
11127            ret = ML_DSA_LEVEL5_PRV_KEY_SIZE;
11128        }
11129    }
11130
11131    return ret;
11132}
11133
11134/* Returns the size of a dilithium private plus public key.
11135 *
11136 * @param [in]  key  Dilithium private/public key.
11137 * @param [out] len  Private key size for set level.
11138 * @return  0 on success.
11139 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11140 */
11141int wc_MlDsaKey_GetPrivLen(MlDsaKey* key, int* len)
11142{
11143    int ret = 0;
11144
11145    *len = wc_dilithium_priv_size(key);
11146    if (*len < 0) {
11147        ret = *len;
11148    }
11149
11150    return ret;
11151}
11152#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
11153#endif /* WOLFSSL_DILITHIUM_PRIVATE_KEY */
11154
11155#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
11156/* Returns the size of a dilithium public key.
11157 *
11158 * @param [in] key  Dilithium private/public key.
11159 * @return  Public key size on success for set level.
11160 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11161 */
11162int wc_dilithium_pub_size(dilithium_key* key)
11163{
11164    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
11165
11166    if (key != NULL) {
11167    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11168        if (key->params == NULL) {
11169            ret = BAD_FUNC_ARG;
11170        }
11171        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11172            ret = DILITHIUM_LEVEL2_PUB_KEY_SIZE;
11173        }
11174        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11175            ret = DILITHIUM_LEVEL3_PUB_KEY_SIZE;
11176        }
11177        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11178            ret = DILITHIUM_LEVEL5_PUB_KEY_SIZE;
11179        }
11180        else
11181    #endif
11182        if (key->level == WC_ML_DSA_44) {
11183            ret = ML_DSA_LEVEL2_PUB_KEY_SIZE;
11184        }
11185        else if (key->level == WC_ML_DSA_65) {
11186            ret = ML_DSA_LEVEL3_PUB_KEY_SIZE;
11187        }
11188        else if (key->level == WC_ML_DSA_87) {
11189            ret = ML_DSA_LEVEL5_PUB_KEY_SIZE;
11190        }
11191    }
11192
11193    return ret;
11194}
11195
11196/* Returns the size of a dilithium public key.
11197 *
11198 * @param [in]  key  Dilithium private/public key.
11199 * @param [out] len  Public key size for set level.
11200 * @return  0 on success.
11201 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11202 */
11203int wc_MlDsaKey_GetPubLen(MlDsaKey* key, int* len)
11204{
11205    int ret = 0;
11206
11207    *len = wc_dilithium_pub_size(key);
11208    if (*len < 0) {
11209        ret = *len;
11210    }
11211
11212    return ret;
11213}
11214#endif
11215
11216#if !defined(WOLFSSL_DILITHIUM_NO_SIGN) || !defined(WOLFSSL_DILITHIUM_NO_VERIFY)
11217/* Returns the size of a dilithium signature.
11218 *
11219 * @param [in] key  Dilithium private/public key.
11220 * @return  Signature size on success for set level.
11221 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11222 */
11223int wc_dilithium_sig_size(dilithium_key* key)
11224{
11225    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
11226
11227    if (key != NULL) {
11228    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11229        if (key->params == NULL) {
11230            ret = BAD_FUNC_ARG;
11231        }
11232        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11233            ret = DILITHIUM_LEVEL2_SIG_SIZE;
11234        }
11235        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11236            ret = DILITHIUM_LEVEL3_SIG_SIZE;
11237        }
11238        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11239            ret = DILITHIUM_LEVEL5_SIG_SIZE;
11240        }
11241        else
11242    #endif
11243        if (key->level == WC_ML_DSA_44) {
11244            ret = ML_DSA_LEVEL2_SIG_SIZE;
11245        }
11246        else if (key->level == WC_ML_DSA_65) {
11247            ret = ML_DSA_LEVEL3_SIG_SIZE;
11248        }
11249        else if (key->level == WC_ML_DSA_87) {
11250            ret = ML_DSA_LEVEL5_SIG_SIZE;
11251        }
11252    }
11253
11254    return ret;
11255}
11256
11257/* Returns the size of a dilithium signature.
11258 *
11259 * @param [in]  key  Dilithium private/public key.
11260 * @param [out] len  Signature size for set level.
11261 * @return  0 on success.
11262 * @return  BAD_FUNC_ARG when key is NULL or level not set,
11263 */
11264int wc_MlDsaKey_GetSigLen(MlDsaKey* key, int* len)
11265{
11266    int ret = 0;
11267
11268    *len = wc_dilithium_sig_size(key);
11269    if (*len < 0) {
11270        ret = *len;
11271    }
11272
11273    return ret;
11274}
11275#endif
11276
11277#ifdef WOLFSSL_DILITHIUM_CHECK_KEY
11278/* Check the public key of the dilithium key matches the private key.
11279 *
11280 * @param [in] key  Dilithium private/public key.
11281 * @return  0 on success.
11282 * @return  BAD_FUNC_ARG when key is NULL or no private key available,
11283 * @return  PUBLIC_KEY_E when the public key is not set or doesn't match,
11284 * @return  MEMORY_E when dynamic memory allocation fails.
11285 */
11286int wc_dilithium_check_key(dilithium_key* key)
11287{
11288    int ret = 0;
11289    const wc_dilithium_params* params = NULL;
11290    sword32* a  = NULL;
11291    sword32* s1 = NULL;
11292    sword32* s2 = NULL;
11293    sword32* t  = NULL;
11294    sword32* t0 = NULL;
11295    sword32* t1 = NULL;
11296
11297    /* Validate parameter. */
11298    if (key == NULL) {
11299        ret = BAD_FUNC_ARG;
11300    }
11301    if ((ret == 0) && (!key->prvKeySet)) {
11302        ret = BAD_FUNC_ARG;
11303    }
11304    if ((ret == 0) && (!key->pubKeySet)) {
11305        ret = PUBLIC_KEY_E;
11306    }
11307
11308    /* Any value in public key are valid.
11309     * Public seed is hashed to generate matrix A.
11310     * t1 is the top 10 bits of a number in range of 0..(Q-1).
11311     * Q >> 13 = 0x3ff so all encoded values are valid.
11312     */
11313
11314    if (ret == 0) {
11315        unsigned int allocSz;
11316
11317        params = key->params;
11318
11319        /* s1-L, s2-K, t0-K, t-K, t1-K */
11320        allocSz = (unsigned int)params->s1Sz + 4U * params->s2Sz;
11321#if !defined(WC_DILITHIUM_CACHE_MATRIX_A)
11322        /* A-KxL */
11323        allocSz += params->aSz;
11324#endif
11325
11326        /* Allocate memory for large intermediates. */
11327        s1 = (sword32*)XMALLOC(allocSz, key->heap, DYNAMIC_TYPE_DILITHIUM);
11328        if (s1 == NULL) {
11329            ret = MEMORY_E;
11330        }
11331        else {
11332            XMEMSET(s1, 0, allocSz);
11333            s2 = s1 + params->s1Sz / sizeof(*s1);
11334            t0 = s2 + params->s2Sz / sizeof(*s2);
11335            t  = t0 + params->s2Sz / sizeof(*t0);
11336            t1 = t  + params->s2Sz / sizeof(*t);
11337#if !defined(WC_DILITHIUM_CACHE_MATRIX_A)
11338            a  = t1 + params->s2Sz / sizeof(*t1);
11339#else
11340            a = key->a;
11341#endif
11342        }
11343    }
11344
11345    if (ret == 0) {
11346#ifdef WC_DILITHIUM_CACHE_MATRIX_A
11347        /* Check that we haven't already cached the matrix A. */
11348        if (!key->aSet)
11349#endif
11350        {
11351            const byte* pub_seed = key->p;
11352
11353            ret = dilithium_expand_a(&key->shake, pub_seed, params->k,
11354                params->l, a, key->heap);
11355#ifdef WC_DILITHIUM_CACHE_MATRIX_A
11356            key->aSet = (ret == 0);
11357#endif
11358        }
11359    }
11360    if (ret == 0) {
11361        const byte* s1p = key->k + DILITHIUM_PUB_SEED_SZ + DILITHIUM_K_SZ +
11362                                   DILITHIUM_TR_SZ;
11363        const byte* s2p = s1p + params->s1EncSz;
11364        const byte* t0p = s2p + params->s2EncSz;
11365        const byte* t1p = key->p + DILITHIUM_PUB_SEED_SZ;
11366        sword32* tt = t;
11367        unsigned int i;
11368        unsigned int j;
11369        sword32 x = 0;
11370
11371        /* Get s1, s2 and t0 from private key. */
11372        dilithium_vec_decode_eta_bits(s1p, params->eta, s1, params->l);
11373        dilithium_vec_decode_eta_bits(s2p, params->eta, s2, params->k);
11374        dilithium_vec_decode_t0(t0p, params->k, t0);
11375
11376        /* Get t1 from public key. */
11377        dilithium_vec_decode_t1(t1p, params->k, t1);
11378
11379        /* Calcaluate t = NTT-1(A o NTT(s1)) + s2 */
11380        dilithium_vec_ntt_small_full(s1, params->l);
11381        dilithium_matrix_mul(t, a, s1, params->k, params->l);
11382    #ifdef WOLFSSL_DILITHIUM_SMALL
11383        dilithium_vec_red(t, params->k);
11384    #endif
11385        dilithium_vec_invntt_full(t, params->k);
11386        dilithium_vec_add(t, s2, params->k);
11387        /* Subtract t0 from t. */
11388        dilithium_vec_sub(t, t0, params->k);
11389        /* Make t positive to match t1. */
11390        dilithium_vec_make_pos(t, params->k);
11391
11392        /* Check t - t0 and t1 are the same. */
11393        for (i = 0; i < params->k; i++) {
11394            for (j = 0; j < DILITHIUM_N; j++) {
11395                x |= tt[j] ^ t1[j];
11396            }
11397            tt += DILITHIUM_N;
11398            t1 += DILITHIUM_N;
11399        }
11400        /* Check the public seed is the same in private and public key. */
11401        for (i = 0; i < DILITHIUM_PUB_SEED_SZ; i++) {
11402            x |= key->p[i] ^ key->k[i];
11403        }
11404
11405        if ((ret == 0) && (x != 0)) {
11406            ret = PUBLIC_KEY_E;
11407        }
11408    }
11409
11410    if (key != NULL) {
11411        /* Dispose of allocated memory. */
11412        XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
11413    }
11414    return ret;
11415}
11416#endif /* WOLFSSL_DILITHIUM_CHECK_KEY */
11417
11418#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
11419
11420/* Export the dilithium public key.
11421 *
11422 * @param [in]      key     Dilithium public key.
11423 * @param [out]     out     Array to hold public key.
11424 * @param [in, out] outLen  On in, the number of bytes in array.
11425 *                          On out, the number bytes put into array.
11426 * @return  0 on success.
11427 * @return  BAD_FUNC_ARG when a parameter is NULL.
11428 * @return  BUFFER_E when outLen is less than DILITHIUM_LEVEL2_PUB_KEY_SIZE.
11429 */
11430int wc_dilithium_export_public(dilithium_key* key, byte* out, word32* outLen)
11431{
11432    int ret = 0;
11433    word32 inLen;
11434
11435    /* Validate parameters */
11436    if ((key == NULL) || (out == NULL) || (outLen == NULL)) {
11437        ret = BAD_FUNC_ARG;
11438    }
11439    if (ret == 0) {
11440        /* Get length passed in for checking. */
11441        inLen = *outLen;
11442    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11443        if (key->params == NULL) {
11444            ret = BAD_FUNC_ARG;
11445        }
11446        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11447            /* Set out length. */
11448            *outLen = DILITHIUM_LEVEL2_PUB_KEY_SIZE;
11449            /* Validate length passed in. */
11450            if (inLen < DILITHIUM_LEVEL2_PUB_KEY_SIZE) {
11451                ret = BUFFER_E;
11452            }
11453        }
11454        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11455            /* Set out length. */
11456            *outLen = DILITHIUM_LEVEL3_PUB_KEY_SIZE;
11457            /* Validate length passed in. */
11458            if (inLen < DILITHIUM_LEVEL3_PUB_KEY_SIZE) {
11459                ret = BUFFER_E;
11460            }
11461        }
11462        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11463            /* Set out length. */
11464            *outLen = DILITHIUM_LEVEL5_PUB_KEY_SIZE;
11465            /* Validate length passed in. */
11466            if (inLen < DILITHIUM_LEVEL5_PUB_KEY_SIZE) {
11467                ret = BUFFER_E;
11468            }
11469        }
11470        else
11471    #endif
11472        if (key->level == WC_ML_DSA_44) {
11473            /* Set out length. */
11474            *outLen = ML_DSA_LEVEL2_PUB_KEY_SIZE;
11475            /* Validate length passed in. */
11476            if (inLen < ML_DSA_LEVEL2_PUB_KEY_SIZE) {
11477                ret = BUFFER_E;
11478            }
11479        }
11480        else if (key->level == WC_ML_DSA_65) {
11481            /* Set out length. */
11482            *outLen = ML_DSA_LEVEL3_PUB_KEY_SIZE;
11483            /* Validate length passed in. */
11484            if (inLen < ML_DSA_LEVEL3_PUB_KEY_SIZE) {
11485                ret = BUFFER_E;
11486            }
11487        }
11488        else if (key->level == WC_ML_DSA_87) {
11489            /* Set out length. */
11490            *outLen = ML_DSA_LEVEL5_PUB_KEY_SIZE;
11491            /* Validate length passed in. */
11492            if (inLen < ML_DSA_LEVEL5_PUB_KEY_SIZE) {
11493                ret = BUFFER_E;
11494            }
11495        }
11496        else {
11497            /* Level not set. */
11498            ret = BAD_FUNC_ARG;
11499        }
11500    }
11501
11502    /* Check public key available. */
11503    if ((ret == 0) && (!key->pubKeySet)) {
11504        ret = BAD_FUNC_ARG;
11505    }
11506
11507    if (ret == 0) {
11508        /* Copy public key out. */
11509        XMEMCPY(out, key->p, *outLen);
11510    }
11511
11512    return ret;
11513}
11514
11515/* Import a dilithium public key from a byte array.
11516 *
11517 * Public key encoded in big-endian.
11518 *
11519 * @param [in]      in     Array holding public key.
11520 * @param [in]      inLen  Number of bytes of data in array.
11521 * @param [in, out] key    Dilithium public key.
11522 * @return  0 on success.
11523 * @return  BAD_FUNC_ARG when in or key is NULL or key format is not supported.
11524 */
11525int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key)
11526{
11527    int ret = 0;
11528
11529    /* Validate parameters. */
11530    if ((in == NULL) || (key == NULL)) {
11531        ret = BAD_FUNC_ARG;
11532    }
11533    if (ret == 0) {
11534    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11535        if (key->params == NULL) {
11536            ret = BAD_FUNC_ARG;
11537        }
11538        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11539            /* Check length. */
11540            if (inLen != DILITHIUM_LEVEL2_PUB_KEY_SIZE) {
11541                ret = BAD_FUNC_ARG;
11542            }
11543        }
11544        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11545            /* Check length. */
11546            if (inLen != DILITHIUM_LEVEL3_PUB_KEY_SIZE) {
11547                ret = BAD_FUNC_ARG;
11548            }
11549        }
11550        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11551            /* Check length. */
11552            if (inLen != DILITHIUM_LEVEL5_PUB_KEY_SIZE) {
11553                ret = BAD_FUNC_ARG;
11554            }
11555        }
11556        else
11557    #endif
11558        if (key->level == WC_ML_DSA_44) {
11559            /* Check length. */
11560            if (inLen != ML_DSA_LEVEL2_PUB_KEY_SIZE) {
11561                ret = BAD_FUNC_ARG;
11562            }
11563        }
11564        else if (key->level == WC_ML_DSA_65) {
11565            /* Check length. */
11566            if (inLen != ML_DSA_LEVEL3_PUB_KEY_SIZE) {
11567                ret = BAD_FUNC_ARG;
11568            }
11569        }
11570        else if (key->level == WC_ML_DSA_87) {
11571            /* Check length. */
11572            if (inLen != ML_DSA_LEVEL5_PUB_KEY_SIZE) {
11573                ret = BAD_FUNC_ARG;
11574            }
11575        }
11576        else {
11577            /* Level not set. */
11578            ret = BAD_FUNC_ARG;
11579        }
11580    }
11581
11582
11583#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11584    if (ret == 0) {
11585        ret = dilithium_alloc_pub_buf(key);
11586    }
11587#endif
11588
11589    if (ret == 0) {
11590        /* Copy the private key data in or copy pointer. */
11591    #ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
11592        key->p = in;
11593    #else
11594        XMEMCPY(key->p, in, inLen);
11595    #endif
11596
11597#ifdef WC_DILITHIUM_CACHE_PUB_VECTORS
11598    #ifndef WC_DILITHIUM_FIXED_ARRAY
11599        /* Allocate t1 if required. */
11600        if (key->t1 == NULL) {
11601            key->t1 = (sword32*)XMALLOC(key->params->s2Sz, key->heap,
11602                DYNAMIC_TYPE_DILITHIUM);
11603            if (key->t1 == NULL) {
11604                ret = MEMORY_E;
11605            }
11606            else {
11607                XMEMSET(key->t1, 0, key->params->s2Sz);
11608            }
11609        }
11610    #endif
11611    }
11612    if (ret == 0) {
11613        /* Compute t1 from public key data. */
11614        dilithium_make_pub_vec(key, key->t1);
11615#endif
11616#ifdef WC_DILITHIUM_CACHE_MATRIX_A
11617    #ifndef WC_DILITHIUM_FIXED_ARRAY
11618        /* Allocate matrix a if required. */
11619        if (key->a == NULL) {
11620            key->a = (sword32*)XMALLOC(key->params->aSz, key->heap,
11621                DYNAMIC_TYPE_DILITHIUM);
11622            if (key->a == NULL) {
11623                ret = MEMORY_E;
11624            }
11625            else {
11626                XMEMSET(key->a, 0, key->params->aSz);
11627            }
11628        }
11629    #endif
11630    }
11631    if (ret == 0) {
11632        /* Compute matrix a from public key data. */
11633        ret = dilithium_expand_a(&key->shake, key->p, key->params->k,
11634            key->params->l, key->a, key->heap);
11635        if (ret == 0) {
11636            key->aSet = 1;
11637        }
11638    }
11639    if (ret == 0) {
11640#endif
11641        /* Public key is set. */
11642        key->pubKeySet = 1;
11643    }
11644
11645    return ret;
11646}
11647
11648#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
11649
11650#ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
11651
11652/* Set the private key data into key.
11653 *
11654 * @param [in]     priv    Private key data.
11655 * @param [in]     privSz  Size of private key data in bytes.
11656 * @param in, out] key     Dilithium key to set into.
11657 * @return  0 on success.
11658 * @return  BAD_FUNC_ARG when private key size is invalid.
11659 * @return  MEMORY_E when dynamic memory allocation fails.
11660 * @return  Other negative on hash error.
11661 */
11662static int dilithium_set_priv_key(const byte* priv, word32 privSz,
11663    dilithium_key* key)
11664{
11665    int ret = 0;
11666    int expPrivSz;
11667#ifdef WC_DILITHIUM_CACHE_MATRIX_A
11668    const wc_dilithium_params* params = key->params;
11669#endif
11670
11671    /* Validate parameters. privSz must match the expected size for the
11672     * level set on the key. This is required so that subsequent code
11673     * which reads via key->params stays within the (possibly dynamically
11674     * sized) buffer. */
11675    expPrivSz = wc_dilithium_size(key);
11676    if (expPrivSz < 0) {
11677        ret = BAD_FUNC_ARG;
11678    }
11679    else if (privSz != (word32)expPrivSz) {
11680        ret = BAD_FUNC_ARG;
11681    }
11682
11683#ifdef WOLFSSL_DILITHIUM_DYNAMIC_KEYS
11684    if (ret == 0) {
11685        ret = dilithium_alloc_priv_buf(key);
11686    }
11687#endif
11688
11689    if (ret == 0) {
11690        /* Copy the private key data in or copy pointer. */
11691    #ifdef WOLFSSL_DILITHIUM_ASSIGN_KEY
11692        key->k = priv;
11693    #else
11694        XMEMCPY(key->k, priv, privSz);
11695    #endif
11696    }
11697
11698        /* Allocate and create cached values. */
11699#ifdef WC_DILITHIUM_CACHE_MATRIX_A
11700#ifndef WC_DILITHIUM_FIXED_ARRAY
11701    if (ret == 0) {
11702        /* Allocate matrix a if required. */
11703        if (key->a == NULL) {
11704            key->a = (sword32*)XMALLOC(params->aSz, key->heap,
11705                DYNAMIC_TYPE_DILITHIUM);
11706            if (key->a == NULL) {
11707                ret = MEMORY_E;
11708            }
11709            else {
11710                XMEMSET(key->a, 0, params->aSz);
11711            }
11712        }
11713    }
11714#endif
11715    if (ret == 0) {
11716        /* Compute matrix a from private key data. */
11717        ret = dilithium_expand_a(&key->shake, key->k, params->k, params->l,
11718            key->a, key->heap);
11719        if (ret == 0) {
11720            key->aSet = 1;
11721        }
11722    }
11723#endif
11724#ifdef WC_DILITHIUM_CACHE_PRIV_VECTORS
11725#ifndef WC_DILITHIUM_FIXED_ARRAY
11726    if ((ret == 0) && (key->s1 == NULL)) {
11727        /* Allocate L vector s1, K vector s2 and K vector t0 if required. */
11728        key->s1 = (sword32*)XMALLOC((unsigned int)params->s1Sz + params->s2Sz +
11729            params->s2Sz, key->heap, DYNAMIC_TYPE_DILITHIUM);
11730        if (key->s1 == NULL) {
11731            ret = MEMORY_E;
11732        }
11733        else {
11734            XMEMSET(key->s1, 0, (unsigned int)params->s1Sz + params->s2Sz +
11735                params->s2Sz);
11736        }
11737        if (ret == 0) {
11738            /* Set pointers into allocated memory. */
11739            key->s2 = key->s1 + params->s1Sz / sizeof(*key->s1);
11740            key->t0 = key->s2 + params->s2Sz / sizeof(*key->s2);
11741        }
11742    }
11743#endif
11744    if (ret == 0) {
11745        /* Compute vectors from private key. */
11746        dilithium_make_priv_vecs(key, key->s1, key->s2, key->t0);
11747    }
11748#endif
11749    if (ret == 0) {
11750        /* Private key is set. */
11751        key->prvKeySet = 1;
11752    }
11753
11754    return ret;
11755}
11756
11757/* Import a dilithium private key from a byte array.
11758 *
11759 * @param [in]      priv    Array holding private key.
11760 * @param [in]      privSz  Number of bytes of data in array.
11761 * @param [in, out] key     Dilithium private key.
11762 * @return  0 otherwise.
11763 * @return  BAD_FUNC_ARG when a parameter is NULL or privSz is less than size
11764 *          required for level,
11765 */
11766int wc_dilithium_import_private(const byte* priv, word32 privSz,
11767    dilithium_key* key)
11768{
11769    int ret = 0;
11770
11771    /* Validate parameters. */
11772    if ((priv == NULL) || (key == NULL)) {
11773        ret = BAD_FUNC_ARG;
11774    }
11775    if ((ret == 0) && (key->level != WC_ML_DSA_44) &&
11776            (key->level != WC_ML_DSA_65) && (key->level != WC_ML_DSA_87)) {
11777        ret = BAD_FUNC_ARG;
11778    }
11779
11780    if (ret == 0) {
11781        /* Set the private key data. */
11782        ret = dilithium_set_priv_key(priv, privSz, key);
11783    }
11784
11785    return ret;
11786}
11787
11788#if defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
11789/* Import a dilithium private and public keys from byte array(s).
11790 *
11791 * @param [in] priv    Array holding private key or private+public keys
11792 * @param [in] privSz  Number of bytes of data in private key array.
11793 * @param [in] pub     Array holding public key (or NULL).
11794 * @param [in] pubSz   Number of bytes of data in public key array (or 0).
11795 * @param [in] key     Dilithium private/public key.
11796 * @return  0 on success.
11797 * @return  BAD_FUNC_ARG when a required parameter is NULL an invalid
11798 *          combination of keys/lengths is supplied.
11799 */
11800int wc_dilithium_import_key(const byte* priv, word32 privSz,
11801    const byte* pub, word32 pubSz, dilithium_key* key)
11802{
11803    int ret = 0;
11804
11805    /* Validate parameters. */
11806    if ((priv == NULL) || (key == NULL)) {
11807        ret = BAD_FUNC_ARG;
11808    }
11809    if ((pub == NULL) && (pubSz != 0)) {
11810        ret = BAD_FUNC_ARG;
11811    }
11812    if ((ret == 0) && (key->level != WC_ML_DSA_44) &&
11813            (key->level != WC_ML_DSA_65) && (key->level != WC_ML_DSA_87)) {
11814        ret = BAD_FUNC_ARG;
11815    }
11816
11817    if ((ret == 0) && (pub != NULL)) {
11818        /* Import public key. */
11819        ret = wc_dilithium_import_public(pub, pubSz, key);
11820    }
11821    if (ret == 0) {
11822        ret = dilithium_set_priv_key(priv, privSz, key);
11823    }
11824
11825    return ret;
11826}
11827#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
11828
11829/* Export the dilithium private key.
11830 *
11831 * @param [in]      key     Dilithium private key.
11832 * @param [out]     out     Array to hold private key.
11833 * @param [in, out] outLen  On in, the number of bytes in array.
11834 *                          On out, the number bytes put into array.
11835 * @return  0 on success.
11836 * @return  BAD_FUNC_ARG when a parameter is NULL.
11837 * @return  BUFFER_E when outLen is less than DILITHIUM_LEVEL2_KEY_SIZE.
11838 */
11839int wc_dilithium_export_private(dilithium_key* key, byte* out,
11840    word32* outLen)
11841{
11842    int ret = 0;
11843    word32 inLen = 0;
11844
11845    /* Validate parameters. */
11846    if ((key == NULL) || (out == NULL) || (outLen == NULL)) {
11847        ret = BAD_FUNC_ARG;
11848    }
11849
11850    /* Check private key available. */
11851    if ((ret == 0) && (!key->prvKeySet)) {
11852        ret = BAD_FUNC_ARG;
11853    }
11854
11855    if (ret == 0) {
11856        inLen = *outLen;
11857        /* check and set up out length */
11858    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11859        if (key->params == NULL) {
11860            ret = BAD_FUNC_ARG;
11861        }
11862        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11863            *outLen = DILITHIUM_LEVEL2_KEY_SIZE;
11864        }
11865        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11866            *outLen = DILITHIUM_LEVEL3_KEY_SIZE;
11867        }
11868        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11869            *outLen = DILITHIUM_LEVEL5_KEY_SIZE;
11870        }
11871        else
11872    #endif
11873        if (key->level == WC_ML_DSA_44) {
11874            *outLen = ML_DSA_LEVEL2_KEY_SIZE;
11875        }
11876        else if (key->level == WC_ML_DSA_65) {
11877            *outLen = ML_DSA_LEVEL3_KEY_SIZE;
11878        }
11879        else if (key->level == WC_ML_DSA_87) {
11880            *outLen = ML_DSA_LEVEL5_KEY_SIZE;
11881        }
11882        else {
11883            /* Level not set. */
11884            ret = BAD_FUNC_ARG;
11885        }
11886    }
11887
11888    /* Check array length. */
11889    if ((ret == 0) && (inLen < *outLen)) {
11890        ret = BUFFER_E;
11891    }
11892
11893    if (ret == 0) {
11894        /* Copy private key out key. */
11895        XMEMCPY(out, key->k, *outLen);
11896    }
11897
11898    return ret;
11899}
11900
11901#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
11902/* Export the dilithium private and public key.
11903 *
11904 * @param [in]      key     Dilithium private/public key.
11905 * @param [out]     priv    Array to hold private key.
11906 * @param [in, out] privSz  On in, the number of bytes in private key array.
11907 *                          On out, the number bytes put into private key.
11908 * @param [out]     pub     Array to hold  public key.
11909 * @param [in, out] pubSz   On in, the number of bytes in public key array.
11910 *                          On out, the number bytes put into public key.
11911 * @return  0 on success.
11912 * @return  BAD_FUNC_ARG when a key, priv, privSz, pub or pubSz is NULL.
11913 * @return  BUFFER_E when privSz or pubSz is less than required size.
11914 */
11915int wc_dilithium_export_key(dilithium_key* key, byte* priv, word32 *privSz,
11916    byte* pub, word32 *pubSz)
11917{
11918    int ret;
11919
11920    /* Export private key only. */
11921    ret = wc_dilithium_export_private(key, priv, privSz);
11922    if (ret == 0) {
11923        /* Export public key. */
11924        ret = wc_dilithium_export_public(key, pub, pubSz);
11925    }
11926
11927    return ret;
11928}
11929#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
11930
11931#endif /* WOLFSSL_DILITHIUM_PRIVATE_KEY */
11932
11933#ifndef WOLFSSL_DILITHIUM_NO_ASN1
11934
11935/* Maps ASN.1 OID to wolfCrypt security level macros */
11936static int mapOidToSecLevel(int oid)
11937{
11938    switch (oid) {
11939        case ML_DSA_LEVEL2k:
11940            return WC_ML_DSA_44;
11941        case ML_DSA_LEVEL3k:
11942            return WC_ML_DSA_65;
11943        case ML_DSA_LEVEL5k:
11944            return WC_ML_DSA_87;
11945#ifdef WOLFSSL_DILITHIUM_FIPS204_DRAFT
11946        case DILITHIUM_LEVEL2k:
11947            return WC_ML_DSA_44_DRAFT;
11948        case DILITHIUM_LEVEL3k:
11949            return WC_ML_DSA_65_DRAFT;
11950        case DILITHIUM_LEVEL5k:
11951            return WC_ML_DSA_87_DRAFT;
11952#endif
11953        default:
11954            return ASN_UNKNOWN_OID_E;
11955    }
11956}
11957
11958/* Get OID sum from dilithium key */
11959int dilithium_get_oid_sum(dilithium_key* key, int* keyFormat) {
11960    int ret = 0;
11961
11962    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
11963    if (key->params == NULL) {
11964        ret = BAD_FUNC_ARG;
11965    }
11966    else if (key->params->level == WC_ML_DSA_44_DRAFT) {
11967        *keyFormat = DILITHIUM_LEVEL2k;
11968    }
11969    else if (key->params->level == WC_ML_DSA_65_DRAFT) {
11970        *keyFormat = DILITHIUM_LEVEL3k;
11971    }
11972    else if (key->params->level == WC_ML_DSA_87_DRAFT) {
11973        *keyFormat = DILITHIUM_LEVEL5k;
11974    }
11975    else
11976    #endif /* WOLFSSL_DILITHIUM_FIPS204_DRAFT */
11977    if (key->level == WC_ML_DSA_44) {
11978        *keyFormat = ML_DSA_LEVEL2k;
11979    }
11980    else if (key->level == WC_ML_DSA_65) {
11981        *keyFormat = ML_DSA_LEVEL3k;
11982    }
11983    else if (key->level == WC_ML_DSA_87) {
11984        *keyFormat = ML_DSA_LEVEL5k;
11985    }
11986    else {
11987        /* Level is not set */
11988        ret = ALGO_ID_E;
11989    }
11990
11991    return ret;
11992}
11993
11994#if defined(WOLFSSL_DILITHIUM_PRIVATE_KEY)
11995
11996/* Decode the DER encoded Dilithium key.
11997 *
11998 * @param [in]      input     Array holding DER encoded data.
11999 * @param [in, out] inOutIdx  On in, index into array of start of DER encoding.
12000 *                            On out, index into array after DER encoding.
12001 * @param [in, out] key       Dilithium key structure to hold the decoded key.
12002 *                            If the security level is set in the key structure
12003 *                            on input, the DER key will be decoded as such and
12004 *                            will fail if there is a mismatch. If the level
12005 *                            and parameters are not set in the key structure on
12006 *                            input, the level will be detected from the DER
12007 *                            file based on the algorithm OID, appropriately
12008 *                            decoded, then updated in the key structure on
12009 *                            output. Auto-detection of the security level is
12010 *                            not supported if compiled for FIPS 204 draft mode.
12011 * @param [in]      inSz      Total size of the input DER buffer array.
12012 * @return  0 on success.
12013 * @return  BAD_FUNC_ARG when input, inOutIdx or key is NULL or inSz is 0.
12014 * @return  Other negative on parse error.
12015 */
12016int wc_Dilithium_PrivateKeyDecode(const byte* input, word32* inOutIdx,
12017    dilithium_key* key, word32 inSz)
12018{
12019    int ret = 0;
12020    const byte* seed = NULL;
12021    const byte* privKey = NULL;
12022    const byte* pubKey = NULL;
12023    word32 seedLen = 0;
12024    word32 privKeyLen = 0;
12025    word32 pubKeyLen = 0;
12026    int keyType = 0;
12027    int autoKeyType = ANONk;
12028
12029    /* Validate parameters. */
12030    if ((input == NULL) || (inOutIdx == NULL) || (key == NULL) || (inSz == 0)) {
12031        ret = BAD_FUNC_ARG;
12032    }
12033
12034    if (ret == 0) {
12035        /* Get OID sum for level. */
12036        if (key->level == 0) { /* Check first, because key->params will be NULL
12037                               * when key->level = 0 */
12038            /* Level not set by caller, decode from DER */
12039            keyType = ANONk;
12040        }
12041    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12042        else if (key->params == NULL) {
12043            ret = BAD_FUNC_ARG;
12044        }
12045        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
12046            keyType = DILITHIUM_LEVEL2k;
12047        }
12048        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
12049            keyType = DILITHIUM_LEVEL3k;
12050        }
12051        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
12052            keyType = DILITHIUM_LEVEL5k;
12053        }
12054    #endif
12055        else if (key->level == WC_ML_DSA_44) {
12056            keyType = ML_DSA_LEVEL2k;
12057        }
12058        else if (key->level == WC_ML_DSA_65) {
12059            keyType = ML_DSA_LEVEL3k;
12060        }
12061        else if (key->level == WC_ML_DSA_87) {
12062            keyType = ML_DSA_LEVEL5k;
12063        }
12064        else {
12065            ret = BAD_FUNC_ARG;
12066        }
12067    }
12068
12069    if (ret == 0) {
12070        /* Decode the asymmetric key and get out private and public key data. */
12071#ifndef WOLFSSL_ASN_TEMPLATE
12072        ret = DecodeAsymKey_Assign(input, inOutIdx, inSz,
12073                                   NULL, NULL,
12074                                   &privKey, &privKeyLen,
12075                                   &pubKey, &pubKeyLen, &autoKeyType);
12076#else
12077        ret = DecodeAsymKey_Assign(input, inOutIdx, inSz,
12078                                   &seed, &seedLen,
12079                                   &privKey, &privKeyLen,
12080                                   &pubKey, &pubKeyLen, &autoKeyType);
12081#endif /* WOLFSSL_ASN_TEMPLATE */
12082    }
12083
12084    if (ret == 0) {
12085        if (keyType == ANONk && autoKeyType != ANONk) {
12086            /* Set the security level based on the decoded key. */
12087            ret = mapOidToSecLevel(autoKeyType);
12088            if (ret > 0) {
12089                ret = wc_dilithium_set_level(key, (byte)ret);
12090            }
12091        }
12092        else if (keyType != ANONk && autoKeyType != ANONk) {
12093            if (keyType == autoKeyType)
12094                ret = 0;
12095            else
12096                ret = ASN_PARSE_E;
12097        }
12098        else if (keyType != ANONk && autoKeyType == ANONk) {
12099            ret = 0;
12100        }
12101        else { /* keyType == ANONk && autoKeyType == ANONk */
12102            /*
12103             * When decoding traditional format with not specifying a level will
12104             * cause this error.
12105             */
12106            ret = ASN_PARSE_E;
12107        }
12108    }
12109
12110    if ((ret == 0) && (pubKey == NULL) && (pubKeyLen == 0)) {
12111        /* Check if the public key is included in the private key. */
12112    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12113        if (key->params == NULL) {
12114            ret = BAD_FUNC_ARG;
12115        }
12116        else if ((key->params->level == WC_ML_DSA_44_DRAFT) &&
12117            (privKeyLen == DILITHIUM_LEVEL2_PRV_KEY_SIZE)) {
12118            pubKey = privKey + DILITHIUM_LEVEL2_KEY_SIZE;
12119            pubKeyLen = DILITHIUM_LEVEL2_PUB_KEY_SIZE;
12120            privKeyLen -= DILITHIUM_LEVEL2_PUB_KEY_SIZE;
12121        }
12122        else if ((key->params->level == WC_ML_DSA_65_DRAFT) &&
12123                 (privKeyLen == DILITHIUM_LEVEL3_PRV_KEY_SIZE)) {
12124            pubKey = privKey + DILITHIUM_LEVEL3_KEY_SIZE;
12125            pubKeyLen = DILITHIUM_LEVEL3_PUB_KEY_SIZE;
12126            privKeyLen -= DILITHIUM_LEVEL3_PUB_KEY_SIZE;
12127        }
12128        else if ((key->params->level == WC_ML_DSA_87_DRAFT) &&
12129                 (privKeyLen == DILITHIUM_LEVEL5_PRV_KEY_SIZE)) {
12130            pubKey = privKey + DILITHIUM_LEVEL5_KEY_SIZE;
12131            pubKeyLen = DILITHIUM_LEVEL5_PUB_KEY_SIZE;
12132            privKeyLen -= DILITHIUM_LEVEL5_PUB_KEY_SIZE;
12133        }
12134        else
12135    #endif
12136        if ((key->level == WC_ML_DSA_44) &&
12137            (privKeyLen == ML_DSA_LEVEL2_PRV_KEY_SIZE)) {
12138            pubKey = privKey + ML_DSA_LEVEL2_KEY_SIZE;
12139            pubKeyLen = ML_DSA_LEVEL2_PUB_KEY_SIZE;
12140            privKeyLen -= ML_DSA_LEVEL2_PUB_KEY_SIZE;
12141        }
12142        else if ((key->level == WC_ML_DSA_65) &&
12143                 (privKeyLen == ML_DSA_LEVEL3_PRV_KEY_SIZE)) {
12144            pubKey = privKey + ML_DSA_LEVEL3_KEY_SIZE;
12145            pubKeyLen = ML_DSA_LEVEL3_PUB_KEY_SIZE;
12146            privKeyLen -= ML_DSA_LEVEL3_PUB_KEY_SIZE;
12147        }
12148        else if ((key->level == WC_ML_DSA_87) &&
12149                 (privKeyLen == ML_DSA_LEVEL5_PRV_KEY_SIZE)) {
12150            pubKey = privKey + ML_DSA_LEVEL5_KEY_SIZE;
12151            pubKeyLen = ML_DSA_LEVEL5_PUB_KEY_SIZE;
12152            privKeyLen -= ML_DSA_LEVEL5_PUB_KEY_SIZE;
12153        }
12154    }
12155
12156    if (ret == 0) {
12157        /* Generate a key pair if seed exists and decoded key pair is ignored */
12158        if (seedLen != 0) {
12159#if !defined(WOLFSSL_DILITHIUM_NO_MAKE_KEY)
12160            if (seedLen == DILITHIUM_SEED_SZ) {
12161                ret = wc_dilithium_make_key_from_seed(key, seed);
12162            }
12163            else {
12164                ret = ASN_PARSE_E;
12165            }
12166#else
12167            ret = NOT_COMPILED_IN;
12168#endif
12169        }
12170#if defined(WOLFSSL_DILITHIUM_PUBLIC_KEY)
12171        /* Check whether public key data was found. */
12172        else if (pubKeyLen != 0 && privKeyLen != 0) {
12173            /* Import private and public key data. */
12174            ret = wc_dilithium_import_key(privKey, privKeyLen, pubKey,
12175                pubKeyLen, key);
12176        }
12177#endif
12178        else if (pubKeyLen == 0 && privKeyLen != 0)
12179        {
12180            /* No public key data, only import private key data. */
12181            ret = wc_dilithium_import_private(privKey, privKeyLen, key);
12182        }
12183        else {
12184            /* Not a problem of ASN.1 structure, but the contents is invalid */
12185            ret = ASN_PARSE_E;
12186        }
12187    }
12188
12189    (void)pubKey;
12190    (void)pubKeyLen;
12191
12192    return ret;
12193}
12194
12195#endif /* WOLFSSL_DILITHIUM_PRIVATE_KEY */
12196
12197#endif /* WOLFSSL_DILITHIUM_NO_ASN1 */
12198
12199#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
12200
12201#if defined(WOLFSSL_DILITHIUM_NO_ASN1)
12202#ifndef WOLFSSL_NO_ML_DSA_44
12203static unsigned char ml_dsa_oid_44[] = {
12204    0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x03, 0x11
12205};
12206#if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12207static unsigned char dilithium_oid_44[] = {
12208    0x2b, 0x06, 0x01, 0x04, 0x01, 0x02, 0x82, 0x0b,
12209    0x0c, 0x04, 0x04
12210};
12211#endif
12212#endif
12213#ifndef WOLFSSL_NO_ML_DSA_65
12214static unsigned char ml_dsa_oid_65[] = {
12215    0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x03, 0x12
12216};
12217#if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12218static unsigned char dilithium_oid_65[] = {
12219    0x2b, 0x06, 0x01, 0x04, 0x01, 0x02, 0x82, 0x0b,
12220    0x0c, 0x06, 0x05
12221};
12222#endif
12223#endif
12224#ifndef WOLFSSL_NO_ML_DSA_87
12225static unsigned char ml_dsa_oid_87[] = {
12226    0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x03, 0x13
12227};
12228#if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12229static unsigned char dilithium_oid_87[] = {
12230    0x2b, 0x06, 0x01, 0x04, 0x01, 0x02, 0x82, 0x0b,
12231    0x0c, 0x08, 0x07
12232};
12233#endif
12234#endif
12235
12236static int dilitihium_get_der_length(const byte* input, word32* inOutIdx,
12237    int *length, word32 inSz)
12238{
12239    int ret = 0;
12240    word32 idx = *inOutIdx;
12241    word32 len = 0;
12242
12243    if (idx >= inSz) {
12244        ret = ASN_PARSE_E;
12245    }
12246    else if (input[idx] < 0x80) {
12247        len = input[idx];
12248        idx++;
12249    }
12250    else if ((input[idx] == 0x80) || (input[idx] >= 0x83)) {
12251        ret = ASN_PARSE_E;
12252    }
12253    else if (input[idx] == 0x81) {
12254        if (idx + 1 >= inSz) {
12255            ret = ASN_PARSE_E;
12256        }
12257        else if (input[idx + 1] < 0x80) {
12258            ret = ASN_PARSE_E;
12259        }
12260        else {
12261            len = input[idx + 1];
12262            idx += 2;
12263        }
12264    }
12265    else if (input[idx] == 0x82) {
12266        if (idx + 2 >= inSz) {
12267            ret = ASN_PARSE_E;
12268        }
12269        else {
12270            len = ((word32)input[idx + 1] << 8) + input[idx + 2];
12271            idx += 3;
12272            if (len < 0x100) {
12273                ret = ASN_PARSE_E;
12274            }
12275        }
12276    }
12277
12278    if ((ret == 0) && ((idx + len) > inSz)) {
12279        ret = ASN_PARSE_E;
12280    }
12281
12282    *length = (int)len;
12283    *inOutIdx = idx;
12284    return ret;
12285}
12286
12287static int dilithium_check_type(const byte* input, word32* inOutIdx, byte type,
12288    word32 inSz)
12289{
12290    int ret = 0;
12291    word32 idx = *inOutIdx;
12292
12293    if (idx >= inSz) {
12294        ret = ASN_PARSE_E;
12295    }
12296    else if (input[idx] != type){
12297        ret = ASN_PARSE_E;
12298    }
12299    else {
12300        idx++;
12301    }
12302
12303    *inOutIdx = idx;
12304    return ret;
12305}
12306
12307#endif /* WOLFSSL_DILITHIUM_NO_ASN1 */
12308
12309/* Decode the DER encoded Dilithium public key.
12310 *
12311 * @param [in]      input     Array holding DER encoded data.
12312 * @param [in, out] inOutIdx  On in, index into array of start of DER encoding.
12313 *                            On out, index into array after DER encoding.
12314 * @param [in, out] key       Dilithium key structure to hold the decoded key.
12315 *                            If the security level is set in the key structure
12316 *                            on input, the DER key will be decoded as such
12317 *                            and will fail if there is a mismatch. If the level
12318 *                            and parameters are not set in the key structure on
12319 *                            input, the level will be detected from the DER
12320 *                            file based on the algorithm OID, appropriately
12321 *                            decoded, then updated in the key structure on
12322 *                            output. Auto-detection of the security level is
12323 *                            not supported if compiled for FIPS 204
12324 *                            draft mode.
12325 * @param [in]      inSz      Total size of data in array.
12326 * @return  0 on success.
12327 * @return  BAD_FUNC_ARG when input, inOutIdx or key is NULL or inSz is 0.
12328 * @return  BAD_FUNC_ARG when level not set.
12329 * @return  Other negative on parse error.
12330 */
12331int wc_Dilithium_PublicKeyDecode(const byte* input, word32* inOutIdx,
12332    dilithium_key* key, word32 inSz)
12333{
12334    int ret = 0;
12335    const byte* pubKey = NULL;
12336    word32 pubKeyLen = 0;
12337
12338    /* Validate parameters. */
12339    if ((input == NULL) || (inOutIdx == NULL) || (key == NULL) || (inSz == 0)) {
12340        ret = BAD_FUNC_ARG;
12341    }
12342
12343    if (ret == 0) {
12344        /* Try to import the key directly. */
12345        ret = wc_dilithium_import_public(input, inSz, key);
12346        if (ret != 0) {
12347        #if !defined(WOLFSSL_DILITHIUM_NO_ASN1)
12348            int keyType = 0;
12349        #else
12350            int length = 0;
12351            unsigned char* oid = NULL;
12352            word32 oidLen = 0;
12353            word32 idx = 0;
12354        #endif
12355
12356            /* Start again. */
12357            ret = 0;
12358
12359    #if !defined(WOLFSSL_DILITHIUM_NO_ASN1)
12360            /* Get OID sum for level. */
12361        #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12362            if (key->params == NULL) {
12363                ret = BAD_FUNC_ARG;
12364            }
12365            else if (key->params->level == WC_ML_DSA_44_DRAFT) {
12366                keyType = DILITHIUM_LEVEL2k;
12367            }
12368            else if (key->params->level == WC_ML_DSA_65_DRAFT) {
12369                keyType = DILITHIUM_LEVEL3k;
12370            }
12371            else if (key->params->level == WC_ML_DSA_87_DRAFT) {
12372                keyType = DILITHIUM_LEVEL5k;
12373            }
12374            else
12375        #endif
12376            if (key->level == WC_ML_DSA_44) {
12377                keyType = ML_DSA_LEVEL2k;
12378            }
12379            else if (key->level == WC_ML_DSA_65) {
12380                keyType = ML_DSA_LEVEL3k;
12381            }
12382            else if (key->level == WC_ML_DSA_87) {
12383                keyType = ML_DSA_LEVEL5k;
12384            }
12385            else {
12386                /* Level not set by caller, decode from DER */
12387                keyType = ANONk; /* 0, not a valid key type in this situation*/
12388            }
12389            if (ret == 0) {
12390                /* Decode the asymmetric key and get out public key data. */
12391                ret = DecodeAsymKeyPublic_Assign(input, inOutIdx, inSz,
12392                                                 &pubKey, &pubKeyLen,
12393                                                 &keyType);
12394                if (ret == 0 && key->params == NULL) {
12395                    /* Set the security level based on the decoded key. */
12396                    ret = mapOidToSecLevel(keyType);
12397                    if (ret > 0) {
12398                        ret = wc_dilithium_set_level(key, (byte)ret);
12399                    }
12400                }
12401            }
12402    #else
12403            /* Get OID sum for level. */
12404    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12405            if (key->params == NULL) {
12406                ret = BAD_FUNC_ARG;
12407            }
12408            else
12409        #ifndef WOLFSSL_NO_ML_DSA_44
12410            if (key->params->level == WC_ML_DSA_44_DRAFT) {
12411                oid = dilithium_oid_44;
12412                oidLen = (word32)sizeof(dilithium_oid_44);
12413            }
12414            else
12415        #endif
12416        #ifndef WOLFSSL_NO_ML_DSA_65
12417            if (key->params->level == WC_ML_DSA_65_DRAFT) {
12418                oid = dilithium_oid_65;
12419                oidLen = (word32)sizeof(dilithium_oid_65);
12420            }
12421            else
12422        #endif
12423        #ifndef WOLFSSL_NO_ML_DSA_87
12424            if (key->params->level == WC_ML_DSA_87_DRAFT) {
12425                oid = dilithium_oid_87;
12426                oidLen = (word32)sizeof(dilithium_oid_87);
12427            }
12428            else
12429        #endif
12430    #endif
12431        #ifndef WOLFSSL_NO_ML_DSA_44
12432            if (key->level == WC_ML_DSA_44) {
12433                oid = ml_dsa_oid_44;
12434                oidLen = (word32)sizeof(ml_dsa_oid_44);
12435            }
12436            else
12437        #endif
12438        #ifndef WOLFSSL_NO_ML_DSA_65
12439            if (key->level == WC_ML_DSA_65) {
12440                oid = ml_dsa_oid_65;
12441                oidLen = (word32)sizeof(ml_dsa_oid_65);
12442            }
12443            else
12444        #endif
12445        #ifndef WOLFSSL_NO_ML_DSA_87
12446            if (key->level == WC_ML_DSA_87) {
12447                oid = ml_dsa_oid_87;
12448                oidLen = (word32)sizeof(ml_dsa_oid_87);
12449            }
12450            else
12451        #endif
12452            {
12453                /* Level not set. */
12454                ret = BAD_FUNC_ARG;
12455            }
12456            if (ret == 0) {
12457                ret = dilithium_check_type(input, &idx, 0x30, inSz);
12458            }
12459            if (ret == 0) {
12460                ret = dilitihium_get_der_length(input, &idx, &length, inSz);
12461            }
12462            if (ret == 0) {
12463                ret = dilithium_check_type(input, &idx, 0x30, inSz);
12464            }
12465            if (ret == 0) {
12466                ret = dilitihium_get_der_length(input, &idx, &length, inSz);
12467            }
12468            if (ret == 0) {
12469                ret = dilithium_check_type(input, &idx, 0x06, inSz);
12470            }
12471            if (ret == 0) {
12472                ret = dilitihium_get_der_length(input, &idx, &length, inSz);
12473            }
12474            if (ret == 0) {
12475                if (((word32)length != oidLen) ||
12476                        (XMEMCMP(input + idx, oid, oidLen) != 0)) {
12477                    ret = ASN_PARSE_E;
12478                }
12479                idx += oidLen;
12480            }
12481            if (ret == 0) {
12482                ret = dilithium_check_type(input, &idx, 0x03, inSz);
12483            }
12484            if (ret == 0) {
12485                ret = dilitihium_get_der_length(input, &idx, &length, inSz);
12486            }
12487            if (ret == 0) {
12488                if ((input[idx] != 0) || (length == 0)) {
12489                    ret = ASN_PARSE_E;
12490                }
12491                idx++;
12492                length--;
12493            }
12494            if (ret == 0) {
12495                /* This is the raw point data compressed or uncompressed. */
12496                pubKeyLen = (word32)length;
12497                pubKey = input + idx;
12498
12499                *inOutIdx += idx;
12500            }
12501    #endif
12502            if (ret == 0) {
12503                /* Import public key data. */
12504                ret = wc_dilithium_import_public(pubKey, pubKeyLen, key);
12505            }
12506        }
12507    }
12508    return ret;
12509}
12510
12511#ifndef WOLFSSL_DILITHIUM_NO_ASN1
12512
12513#ifdef WC_ENABLE_ASYM_KEY_EXPORT
12514/* Encode the public part of a Dilithium key in DER.
12515 *
12516 * Pass NULL for output to get the size of the encoding.
12517 *
12518 * @param [in]  key      Dilithium key object.
12519 * @param [out] output   Buffer to put encoded data in.
12520 * @param [in]  len      Size of buffer in bytes.
12521 * @param [in]  withAlg  Whether to use SubjectPublicKeyInfo format.
12522 * @return  Size of encoded data in bytes on success.
12523 * @return  BAD_FUNC_ARG when key is NULL.
12524 * @return  MEMORY_E when dynamic memory allocation failed.
12525 */
12526int wc_Dilithium_PublicKeyToDer(dilithium_key* key, byte* output, word32 len,
12527    int withAlg)
12528{
12529    int ret = 0;
12530    int keyType = 0;
12531    word32 pubKeyLen = 0;
12532
12533    /* Validate parameters. */
12534    if (key == NULL) {
12535        ret = BAD_FUNC_ARG;
12536    }
12537    /* Check we have a public key to encode. */
12538    if ((ret == 0) && (!key->pubKeySet)) {
12539        ret = BAD_FUNC_ARG;
12540    }
12541
12542    if (ret == 0) {
12543        /* Get OID and length for level. */
12544    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12545        if (key->params == NULL) {
12546            ret = BAD_FUNC_ARG;
12547        }
12548        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
12549            keyType = DILITHIUM_LEVEL2k;
12550            pubKeyLen = DILITHIUM_LEVEL2_PUB_KEY_SIZE;
12551        }
12552        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
12553            keyType = DILITHIUM_LEVEL3k;
12554            pubKeyLen = DILITHIUM_LEVEL3_PUB_KEY_SIZE;
12555        }
12556        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
12557            keyType = DILITHIUM_LEVEL5k;
12558            pubKeyLen = DILITHIUM_LEVEL5_PUB_KEY_SIZE;
12559        }
12560        else
12561    #endif
12562        if (key->level == WC_ML_DSA_44) {
12563            keyType = ML_DSA_LEVEL2k;
12564            pubKeyLen = ML_DSA_LEVEL2_PUB_KEY_SIZE;
12565        }
12566        else if (key->level == WC_ML_DSA_65) {
12567            keyType = ML_DSA_LEVEL3k;
12568            pubKeyLen = ML_DSA_LEVEL3_PUB_KEY_SIZE;
12569        }
12570        else if (key->level == WC_ML_DSA_87) {
12571            keyType = ML_DSA_LEVEL5k;
12572            pubKeyLen = ML_DSA_LEVEL5_PUB_KEY_SIZE;
12573        }
12574        else {
12575            /* Level not set. */
12576            ret = BAD_FUNC_ARG;
12577        }
12578    }
12579
12580    if (ret == 0) {
12581        ret = SetAsymKeyDerPublic(key->p, pubKeyLen, output, len, keyType,
12582            withAlg);
12583    }
12584
12585    return ret;
12586}
12587#endif /* WC_ENABLE_ASYM_KEY_EXPORT */
12588
12589#endif /* !WOLFSSL_DILITHIUM_NO_ASN1 */
12590
12591#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
12592
12593#ifdef WOLFSSL_DILITHIUM_PRIVATE_KEY
12594
12595#ifndef WOLFSSL_DILITHIUM_NO_ASN1
12596
12597#ifdef WOLFSSL_DILITHIUM_PUBLIC_KEY
12598/* Encode the private and public data of a Dilithium key in DER.
12599 *
12600 * Pass NULL for output to get the size of the encoding.
12601 *
12602 * @param [in]  key     Dilithium key object.
12603 * @param [out] output  Buffer to put encoded data in.
12604 * @param [in]  len     Size of buffer in bytes.
12605 * @return  Size of encoded data in bytes on success.
12606 * @return  BAD_FUNC_ARG when key is NULL.
12607 * @return  MEMORY_E when dynamic memory allocation failed.
12608 */
12609int wc_Dilithium_KeyToDer(dilithium_key* key, byte* output, word32 len)
12610{
12611    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
12612
12613    /* Validate parameters and check public and private key set. */
12614    if ((key != NULL) && key->prvKeySet && key->pubKeySet) {
12615        /* Create DER for level. */
12616    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12617        if (key->params == NULL) {
12618            ret = BAD_FUNC_ARG;
12619        }
12620        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
12621            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL2_KEY_SIZE, key->p,
12622                DILITHIUM_LEVEL2_PUB_KEY_SIZE, output, len, DILITHIUM_LEVEL2k);
12623        }
12624        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
12625            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL3_KEY_SIZE, key->p,
12626                DILITHIUM_LEVEL3_PUB_KEY_SIZE, output, len, DILITHIUM_LEVEL3k);
12627        }
12628        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
12629            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL5_KEY_SIZE, key->p,
12630                DILITHIUM_LEVEL5_PUB_KEY_SIZE, output, len, DILITHIUM_LEVEL5k);
12631        }
12632        else
12633    #endif
12634        if (key->level == WC_ML_DSA_44) {
12635            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL2_KEY_SIZE, key->p,
12636                ML_DSA_LEVEL2_PUB_KEY_SIZE, output, len, ML_DSA_LEVEL2k);
12637        }
12638        else if (key->level == WC_ML_DSA_65) {
12639            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL3_KEY_SIZE, key->p,
12640                ML_DSA_LEVEL3_PUB_KEY_SIZE, output, len, ML_DSA_LEVEL3k);
12641        }
12642        else if (key->level == WC_ML_DSA_87) {
12643            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL5_KEY_SIZE, key->p,
12644                ML_DSA_LEVEL5_PUB_KEY_SIZE, output, len, ML_DSA_LEVEL5k);
12645        }
12646    }
12647
12648    return ret;
12649}
12650#endif /* WOLFSSL_DILITHIUM_PUBLIC_KEY */
12651
12652/* Encode the private data of a Dilithium key in DER.
12653 *
12654 * Pass NULL for output to get the size of the encoding.
12655 *
12656 * @param [in]  key     Dilithium key object.
12657 * @param [out] output  Buffer to put encoded data in.
12658 * @param [in]  len     Size of buffer in bytes.
12659 * @return  Size of encoded data in bytes on success.
12660 * @return  BAD_FUNC_ARG when key is NULL.
12661 * @return  MEMORY_E when dynamic memory allocation failed.
12662 */
12663int wc_Dilithium_PrivateKeyToDer(dilithium_key* key, byte* output, word32 len)
12664{
12665    int ret = WC_NO_ERR_TRACE(BAD_FUNC_ARG);
12666
12667    /* Validate parameters and check private key set. */
12668    if ((key != NULL) && key->prvKeySet) {
12669        /* Create DER for level. */
12670    #if defined(WOLFSSL_DILITHIUM_FIPS204_DRAFT)
12671        if (key->params == NULL) {
12672            ret = BAD_FUNC_ARG;
12673        }
12674        else if (key->params->level == WC_ML_DSA_44_DRAFT) {
12675            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL2_KEY_SIZE, NULL, 0,
12676                output, len, DILITHIUM_LEVEL2k);
12677        }
12678        else if (key->params->level == WC_ML_DSA_65_DRAFT) {
12679            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL3_KEY_SIZE, NULL, 0,
12680                output, len, DILITHIUM_LEVEL3k);
12681        }
12682        else if (key->params->level == WC_ML_DSA_87_DRAFT) {
12683            ret = SetAsymKeyDer(key->k, DILITHIUM_LEVEL5_KEY_SIZE, NULL, 0,
12684                output, len, DILITHIUM_LEVEL5k);
12685        }
12686        else
12687    #endif
12688        if (key->level == WC_ML_DSA_44) {
12689            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL2_KEY_SIZE, NULL, 0, output,
12690                len, ML_DSA_LEVEL2k);
12691        }
12692        else if (key->level == WC_ML_DSA_65) {
12693            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL3_KEY_SIZE, NULL, 0, output,
12694                len, ML_DSA_LEVEL3k);
12695        }
12696        else if (key->level == WC_ML_DSA_87) {
12697            ret = SetAsymKeyDer(key->k, ML_DSA_LEVEL5_KEY_SIZE, NULL, 0, output,
12698                len, ML_DSA_LEVEL5k);
12699        }
12700    }
12701
12702    return ret;
12703}
12704
12705#endif /* WOLFSSL_DILITHIUM_NO_ASN1 */
12706
12707#endif /* WOLFSSL_DILITHIUM_PRIVATE_KEY */
12708
12709#endif /* HAVE_DILITHIUM */