luna - wolfssl/wolfcrypt/src/wc

   1/* wc_mlkem.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/* Implementation based on FIPS 203:
  23 *   https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf
  24 *
  25 * Original implementation based on NIST 3rd Round submission package.
  26 * See link at:
  27 *   https://csrc.nist.gov/Projects/post-quantum-cryptography/
  28 *   post-quantum-cryptography-standardization/round-3-submissions
  29 */
  30
  31/* Possible ML-KEM options:
  32 *
  33 * WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM                                 Default: OFF
  34 *   Uses less dynamic memory to perform key generation.
  35 *   Has a small performance trade-off.
  36 *   Only usable with C implementation.
  37 *
  38 * WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM                             Default: OFF
  39 *   Uses less dynamic memory to perform encapsulation.
  40 *   Affects decapsulation too as encapsulation called.
  41 *   Has a small performance trade-off.
  42 *   Only usable with C implementation.
  43 *
  44 * WOLFSSL_MLKEM_NO_MAKE_KEY                                       Default: OFF
  45 *   Disable the make key or key generation API.
  46 *   Reduces the code size.
  47 *   Turn on when only doing encapsulation.
  48 *
  49 * WOLFSSL_MLKEM_NO_ENCAPSULATE                                    Default: OFF
  50 *   Disable the encapsulation API.
  51 *   Reduces the code size.
  52 *   Turn on when doing make key/decapsulation.
  53 *
  54 * WOLFSSL_MLKEM_NO_DECAPSULATE                                    Default: OFF
  55 *   Disable the decapsulation API.
  56 *   Reduces the code size.
  57 *   Turn on when only doing encapsulation.
  58 *
  59 * WOLFSSL_MLKEM_CACHE_A                                           Default: OFF
  60 *   Stores the matrix A during key generation for use in encapsulation when
  61 *   performing decapsulation.
  62 *   MlKemKey is 8KB larger but decapsulation is significantly faster.
  63 *   Turn on when performing make key and decapsulation with same object.
  64 *
  65 * WOLFSSL_MLKEM_DYNAMIC_KEYS                                      Default: OFF
  66 *   Dynamically allocates private and public key buffers instead of using
  67 *   static arrays in the MlKemKey struct. Right-sizes buffers to the actual
  68 *   ML-KEM level and only allocates the needed key parts (e.g., no private
  69 *   key buffer for encapsulate-only use).
  70 *   Cannot be used with WOLFSSL_NO_MALLOC.
  71 */
  72
  73#include <wolfssl/wolfcrypt/libwolfssl_sources.h>
  74
  75#ifdef WC_MLKEM_NO_ASM
  76    #undef USE_INTEL_SPEEDUP
  77    #undef WOLFSSL_ARMASM
  78    #undef WOLFSSL_RISCV_ASM
  79#endif
  80
  81#if FIPS_VERSION3_GE(2,0,0)
  82    /* set NO_WRAPPERS before headers, use direct internal f()s not wrappers */
  83    #define FIPS_NO_WRAPPERS
  84#endif
  85
  86#include <wolfssl/wolfcrypt/wc_mlkem.h>
  87#include <wolfssl/wolfcrypt/hash.h>
  88#include <wolfssl/wolfcrypt/memory.h>
  89#ifdef WOLF_CRYPTO_CB
  90    #include <wolfssl/wolfcrypt/cryptocb.h>
  91#endif
  92
  93#ifdef NO_INLINE
  94    #include <wolfssl/wolfcrypt/misc.h>
  95#else
  96    #define WOLFSSL_MISC_INCLUDED
  97    #include <wolfcrypt/src/misc.c>
  98#endif
  99
 100#if defined(USE_INTEL_SPEEDUP) || \
 101    (defined(__aarch64__) && defined(WOLFSSL_ARMASM))
 102    #if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
 103        defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
 104        #error "Can't use small memory with assembly optimized code"
 105    #endif
 106#endif
 107#if defined(WOLFSSL_MLKEM_CACHE_A)
 108    #if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
 109        defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
 110        #error "Can't cache A with small memory code"
 111    #endif
 112#endif
 113
 114#if defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \
 115    defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \
 116    defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
 117    #error "No ML-KEM operations to be built."
 118#endif
 119#if defined(WOLFSSL_MLKEM_DYNAMIC_KEYS) && defined(WOLFSSL_NO_MALLOC)
 120    #error "Cannot use dynamic key buffers without malloc"
 121#endif
 122
 123#ifdef WOLFSSL_HAVE_MLKEM
 124
 125#ifdef DEBUG_MLKEM
 126void print_polys(const char* name, const sword16* a, int d1, int d2);
 127void print_polys(const char* name, const sword16* a, int d1, int d2)
 128{
 129    int i;
 130    int j;
 131    int k;
 132
 133    fprintf(stderr, "%s: %d %d\n", name, d1, d2);
 134    for (i = 0; i < d1; i++) {
 135        for (j = 0; j < d2; j++) {
 136            for (k = 0; k < 256; k++) {
 137                fprintf(stderr, "%9d,", a[(i*d2*256) + (j*256) + k]);
 138                if ((k % 8) == 7) fprintf(stderr, "\n");
 139            }
 140            fprintf(stderr, "\n");
 141        }
 142    }
 143}
 144#endif
 145
 146#ifdef DEBUG_MLKEM
 147void print_data(const char* name, const byte* d, int len);
 148void print_data(const char* name, const byte* d, int len)
 149{
 150    int i;
 151
 152    fprintf(stderr, "%s\n", name);
 153    for (i = 0; i < len; i++) {
 154        fprintf(stderr, "0x%02x,", d[i]);
 155        if ((i % 16) == 15) fprintf(stderr, "\n");
 156    }
 157    fprintf(stderr, "\n");
 158}
 159#endif
 160
 161/******************************************************************************/
 162
 163/* Use SHA3-256 to generate 32-bytes of hash. */
 164#define MLKEM_HASH_H            mlkem_hash256
 165/* Use SHA3-512 to generate 64-bytes of hash. */
 166#define MLKEM_HASH_G            mlkem_hash512
 167/* Use SHAKE-256 as a key derivation function (KDF). */
 168#if defined(USE_INTEL_SPEEDUP) || \
 169        (defined(WOLFSSL_ARMASM) && defined(__aarch64__))
 170    #define MLKEM_KDF               mlkem_kdf
 171#else
 172    #define MLKEM_KDF               wc_Shake256Hash
 173#endif
 174
 175/******************************************************************************/
 176
 177/* Helper function with volatile variable, to force compiler not to optimize
 178 * code in mlkem_from_msg().
 179 */
 180sword16 wc_mlkem_opt_blocker(void);
 181sword16 wc_mlkem_opt_blocker(void) {
 182    static volatile sword16 static_mlkem_opt_blocker = 0;
 183    return static_mlkem_opt_blocker;
 184}
 185
 186/******************************************************************************/
 187
 188#ifndef WOLFSSL_MLKEM_NO_MAKE_KEY
 189/* Get the k value (number of polynomials in a vector) from the key type.
 190 *
 191 * @param  [in]  key  ML-KEM key object.
 192 * @return  k value for the key type, or 0 if not recognized.
 193 */
 194static int mlkemkey_get_k(const MlKemKey* key)
 195{
 196    switch (key->type) {
 197#ifndef WOLFSSL_NO_ML_KEM
 198    #ifdef WOLFSSL_WC_ML_KEM_512
 199        case WC_ML_KEM_512:
 200            return WC_ML_KEM_512_K;
 201    #endif
 202    #ifdef WOLFSSL_WC_ML_KEM_768
 203        case WC_ML_KEM_768:
 204            return WC_ML_KEM_768_K;
 205    #endif
 206    #ifdef WOLFSSL_WC_ML_KEM_1024
 207        case WC_ML_KEM_1024:
 208            return WC_ML_KEM_1024_K;
 209    #endif
 210#endif
 211#ifdef WOLFSSL_MLKEM_KYBER
 212    #ifdef WOLFSSL_KYBER512
 213        case KYBER512:
 214            return KYBER512_K;
 215    #endif
 216    #ifdef WOLFSSL_KYBER768
 217        case KYBER768:
 218            return KYBER768_K;
 219    #endif
 220    #ifdef WOLFSSL_KYBER1024
 221        case KYBER1024:
 222            return KYBER1024_K;
 223    #endif
 224#endif
 225        default:
 226            return 0;
 227    }
 228}
 229#endif
 230
 231#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
 232/* Allocate (or reallocate) the private key buffer, right-sized for k.
 233 *
 234 * @param  [in, out]  key  ML-KEM key object.
 235 * @param  [in]       k    Number of polynomials in a vector.
 236 * @return  0 on success.
 237 * @return  MEMORY_E when dynamic memory allocation fails.
 238 */
 239static int mlkemkey_alloc_priv(MlKemKey* key, unsigned int k)
 240{
 241    word32 sz = (word32)(k * MLKEM_N * sizeof(sword16));
 242    if (key->priv != NULL) {
 243        ForceZero(key->priv, key->privAllocSz);
 244        XFREE(key->priv, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 245        key->priv = NULL;
 246        key->privAllocSz = 0;
 247    }
 248    key->priv = (sword16*)XMALLOC(sz, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 249    if (key->priv == NULL) {
 250        return MEMORY_E;
 251    }
 252    key->privAllocSz = sz;
 253    return 0;
 254}
 255
 256/* Allocate (or reallocate) the public key buffer, right-sized for k.
 257 *
 258 * @param  [in, out]  key  ML-KEM key object.
 259 * @param  [in]       k    Number of polynomials in a vector.
 260 * @return  0 on success.
 261 * @return  MEMORY_E when dynamic memory allocation fails.
 262 */
 263static int mlkemkey_alloc_pub(MlKemKey* key, unsigned int k)
 264{
 265    if (key->pub != NULL) {
 266        XFREE(key->pub, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 267        key->pub = NULL;
 268    }
 269    key->pub = (sword16*)XMALLOC(k * MLKEM_N * sizeof(sword16), key->heap,
 270        DYNAMIC_TYPE_TMP_BUFFER);
 271    if (key->pub == NULL) {
 272        return MEMORY_E;
 273    }
 274    return 0;
 275}
 276
 277#ifdef WOLFSSL_MLKEM_CACHE_A
 278/* Allocate (or reallocate) the A matrix buffer, right-sized for k.
 279 *
 280 * @param  [in, out]  key  ML-KEM key object.
 281 * @param  [in]       k    Number of polynomials in a vector.
 282 * @return  0 on success.
 283 * @return  MEMORY_E when dynamic memory allocation fails.
 284 */
 285static int mlkemkey_alloc_a(MlKemKey* key, unsigned int k)
 286{
 287    int ret = 0;
 288
 289    if (key->a != NULL) {
 290        XFREE(key->a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 291        key->a = NULL;
 292    }
 293    key->a = (sword16*)XMALLOC(k * k * MLKEM_N * sizeof(sword16), key->heap,
 294        DYNAMIC_TYPE_TMP_BUFFER);
 295    if (key->a == NULL) {
 296        ret = MEMORY_E;
 297    }
 298
 299    return ret;
 300}
 301#endif /* WOLFSSL_MLKEM_CACHE_A */
 302#endif /* WOLFSSL_MLKEM_DYNAMIC_KEYS */
 303
 304/******************************************************************************/
 305
 306#ifndef WC_NO_CONSTRUCTORS
 307/**
 308 * Create a new ML-KEM key object.
 309 *
 310 * Allocates and initializes a ML-KEM key object.
 311 *
 312 * @param  [in]  type   Type of key:
 313 *                        WC_ML_KEM_512, WC_ML_KEM_768, WC_ML_KEM_1024,
 314 *                        KYBER512, KYBER768, KYBER1024.
 315 * @param  [in]  heap   Dynamic memory hint.
 316 * @param  [in]  devId  Device Id.
 317 * @return  Pointer to new MlKemKey object on success.
 318 * @return  NULL on failure.
 319 */
 320MlKemKey* wc_MlKemKey_New(int type, void* heap, int devId)
 321{
 322    int ret;
 323    MlKemKey* key;
 324
 325    key = (MlKemKey*)XMALLOC(sizeof(MlKemKey), heap, DYNAMIC_TYPE_TMP_BUFFER);
 326    if (key != NULL) {
 327        ret = wc_MlKemKey_Init(key, type, heap, devId);
 328        if (ret != 0) {
 329            XFREE(key, heap, DYNAMIC_TYPE_TMP_BUFFER);
 330            key = NULL;
 331        }
 332    }
 333
 334    return key;
 335}
 336
 337/**
 338 * Delete and free a ML-KEM key object.
 339 *
 340 * Frees resources associated with a ML-KEM key object and sets pointer to NULL.
 341 *
 342 * @param  [in]       key    ML-KEM key object to delete.
 343 * @param  [in, out]  key_p  Pointer to key pointer to set to NULL.
 344 * @return  0 on success.
 345 * @return  BAD_FUNC_ARG when key is NULL.
 346 */
 347int wc_MlKemKey_Delete(MlKemKey* key, MlKemKey** key_p)
 348{
 349    int ret = 0;
 350
 351    if (key == NULL) {
 352        ret = BAD_FUNC_ARG;
 353    }
 354    else {
 355        void* heap = key->heap;
 356
 357        wc_MlKemKey_Free(key);
 358        XFREE(key, heap, DYNAMIC_TYPE_TMP_BUFFER);
 359        if (key_p != NULL) {
 360            *key_p = NULL;
 361        }
 362    }
 363
 364    return ret;
 365}
 366#endif /* !WC_NO_CONSTRUCTORS */
 367
 368/**
 369 * Initialize the ML-KEM key.
 370 *
 371 * @param  [out]  key    ML-KEM key object to initialize.
 372 * @param  [in]   type   Type of key:
 373 *                         WC_ML_KEM_512, WC_ML_KEM_768, WC_ML_KEM_1024,
 374 *                         KYBER512, KYBER768, KYBER1024.
 375 * @param  [in]   heap   Dynamic memory hint.
 376 * @param  [in]   devId  Device Id.
 377 * @return  0 on success.
 378 * @return  BAD_FUNC_ARG when key is NULL or type is unrecognized.
 379 * @return  NOT_COMPILED_IN when key type is not supported.
 380 */
 381int wc_MlKemKey_Init(MlKemKey* key, int type, void* heap, int devId)
 382{
 383    int ret = 0;
 384
 385    /* Validate key. */
 386    if (key == NULL) {
 387        ret = BAD_FUNC_ARG;
 388    }
 389    if (ret == 0) {
 390        /* Validate type. */
 391        switch (type) {
 392    #ifndef WOLFSSL_NO_ML_KEM
 393        case WC_ML_KEM_512:
 394        #ifndef WOLFSSL_WC_ML_KEM_512
 395            /* Code not compiled in for ML-KEM-512. */
 396            ret = NOT_COMPILED_IN;
 397        #endif
 398            break;
 399        case WC_ML_KEM_768:
 400        #ifndef WOLFSSL_WC_ML_KEM_768
 401            /* Code not compiled in for ML-KEM-768. */
 402            ret = NOT_COMPILED_IN;
 403        #endif
 404            break;
 405        case WC_ML_KEM_1024:
 406        #ifndef WOLFSSL_WC_ML_KEM_1024
 407            /* Code not compiled in for ML-KEM-1024. */
 408            ret = NOT_COMPILED_IN;
 409        #endif
 410            break;
 411    #endif
 412    #ifdef WOLFSSL_MLKEM_KYBER
 413        case KYBER512:
 414        #ifndef WOLFSSL_KYBER512
 415            /* Code not compiled in for Kyber-512. */
 416            ret = NOT_COMPILED_IN;
 417        #endif
 418            break;
 419        case KYBER768:
 420        #ifndef WOLFSSL_KYBER768
 421            /* Code not compiled in for Kyber-768. */
 422            ret = NOT_COMPILED_IN;
 423        #endif
 424            break;
 425        case KYBER1024:
 426        #ifndef WOLFSSL_KYBER1024
 427            /* Code not compiled in for Kyber-1024. */
 428            ret = NOT_COMPILED_IN;
 429        #endif
 430            break;
 431    #endif
 432        default:
 433            /* No other values supported. */
 434            ret = BAD_FUNC_ARG;
 435            break;
 436        }
 437    }
 438    if (ret == 0) {
 439        /* Keep type for parameters. */
 440        key->type = type;
 441        /* Cache heap pointer. */
 442        key->heap = heap;
 443    #ifdef WOLF_CRYPTO_CB
 444        key->devCtx = NULL;
 445        key->devId = devId;
 446    #endif
 447#ifdef WOLF_PRIVATE_KEY_ID
 448        key->idLen = 0;
 449        key->labelLen = 0;
 450#endif
 451        key->flags = 0;
 452
 453    #ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
 454        key->priv = NULL;
 455        key->pub = NULL;
 456        key->privAllocSz = 0;
 457    #ifdef WOLFSSL_MLKEM_CACHE_A
 458        key->a = NULL;
 459    #endif
 460    #endif
 461
 462        /* Zero out all data. */
 463        XMEMSET(&key->prf, 0, sizeof(key->prf));
 464
 465        /* Initialize the hash algorithm object. */
 466        ret = mlkem_hash_new(&key->hash, heap, devId);
 467    }
 468    if (ret == 0) {
 469        /* Initialize the PRF algorithm object. */
 470        ret = mlkem_prf_new(&key->prf, heap, devId);
 471    }
 472    if (ret == 0) {
 473        mlkem_init();
 474    }
 475
 476    (void)devId;
 477
 478    return ret;
 479}
 480
 481#ifdef WOLF_PRIVATE_KEY_ID
 482/**
 483 * Initialize the ML-KEM key with an id.
 484 *
 485 * @param  [out]  key    ML-KEM key object to initialize.
 486 * @param  [in]   type   Type of key:
 487 *                         WC_ML_KEM_512, WC_ML_KEM_768, WC_ML_KEM_1024,
 488 *                         KYBER512, KYBER768, KYBER1024.
 489 * @param  [in]   id     Identifier of key.
 490 * @param  [in]   len    Length of key identifier in bytes.
 491 * @param  [in]   heap   Dynamic memory hint.
 492 * @param  [in]   devId  Device Id.
 493 * @return  0 on success.
 494 * @return  BAD_FUNC_ARG when key is NULL, id is NULL but len is not zero, or
 495 *          type is unrecognized.
 496 * @return  BUFFER_E when len is out of range.
 497 * @return  NOT_COMPILED_IN when key type is not supported.
 498 */
 499int wc_MlKemKey_Init_Id(MlKemKey* key, int type, const unsigned char* id,
 500    int len, void* heap, int devId)
 501{
 502    int ret = 0;
 503
 504    /* Validate parameters. */
 505    if ((key == NULL) || (id == NULL && len != 0)) {
 506        ret = BAD_FUNC_ARG;
 507    }
 508    if ((ret == 0) && ((len < 0) || (len > MLKEM_MAX_ID_LEN))) {
 509        ret = BUFFER_E;
 510    }
 511
 512    if (ret == 0) {
 513        /* Initialize key. */
 514        ret = wc_MlKemKey_Init(key, type, heap, devId);
 515    }
 516    if ((ret == 0) && (id != NULL) && (len != 0)) {
 517        /* Store key identifier. */
 518        XMEMCPY(key->id, id, (size_t)len);
 519        key->idLen = len;
 520    }
 521
 522    return ret;
 523}
 524
 525/**
 526 * Initialize the ML-KEM key with a label.
 527 *
 528 * @param  [out]  key    ML-KEM key object to initialize.
 529 * @param  [in]   type   Type of key:
 530 *                         WC_ML_KEM_512, WC_ML_KEM_768, WC_ML_KEM_1024,
 531 *                         KYBER512, KYBER768, KYBER1024.
 532 * @param  [in]   label  Label of key. Must be a null-terminated string.
 533 * @param  [in]   heap   Dynamic memory hint.
 534 * @param  [in]   devId  Device Id.
 535 * @return  0 on success.
 536 * @return  BAD_FUNC_ARG when key or label is NULL, or type is unrecognized.
 537 * @return  BUFFER_E when label is too small or big.
 538 * @return  NOT_COMPILED_IN when key type is not supported.
 539 */
 540int wc_MlKemKey_Init_Label(MlKemKey* key, int type, const char* label,
 541    void* heap, int devId)
 542{
 543    int ret = 0;
 544    int labelLen = 0;
 545
 546    /* Validate parameters. */
 547    if ((key == NULL) || (label == NULL)) {
 548        ret = BAD_FUNC_ARG;
 549    }
 550    if (ret == 0) {
 551        /* Validate label length. */
 552        labelLen = (int)XSTRLEN(label);
 553        if ((labelLen == 0) || (labelLen > MLKEM_MAX_LABEL_LEN)) {
 554            ret = BUFFER_E;
 555        }
 556    }
 557
 558    if (ret == 0) {
 559        /* Initialize key. */
 560        ret = wc_MlKemKey_Init(key, type, heap, devId);
 561    }
 562    if (ret == 0) {
 563        /* Don't save string in key->label with null terminator.
 564         * Use key->labelLen to get the length if required. */
 565        XMEMCPY(key->label, label, (size_t)labelLen);
 566        key->labelLen = labelLen;
 567    }
 568
 569    return ret;
 570}
 571#endif
 572
 573/**
 574 * Free the ML-KEM key object.
 575 *
 576 * @param  [in, out]  key   ML-KEM key object to dispose of.
 577 * @return  0 on success.
 578 */
 579int wc_MlKemKey_Free(MlKemKey* key)
 580{
 581    if (key != NULL) {
 582#if defined(WOLF_CRYPTO_CB) && defined(WOLF_CRYPTO_CB_FREE)
 583        if (key->devId != INVALID_DEVID) {
 584            (void)wc_CryptoCb_Free(key->devId, WC_ALGO_TYPE_PK,
 585                WC_PK_TYPE_PQC_KEM_KEYGEN, WC_PQC_KEM_TYPE_KYBER, (void*)key);
 586            /* always continue to software cleanup */
 587        }
 588#endif
 589        /* Dispose of PRF object. */
 590        mlkem_prf_free(&key->prf);
 591        /* Dispose of hash object. */
 592        mlkem_hash_free(&key->hash);
 593        /* Ensure all private data is zeroed. */
 594        ForceZero(&key->hash, sizeof(key->hash));
 595        ForceZero(&key->prf, sizeof(key->prf));
 596#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
 597        if (key->priv != NULL) {
 598            ForceZero(key->priv, key->privAllocSz);
 599            XFREE(key->priv, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 600            key->priv = NULL;
 601            key->privAllocSz = 0;
 602        }
 603        if (key->pub != NULL) {
 604            XFREE(key->pub, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 605            key->pub = NULL;
 606        }
 607    #ifdef WOLFSSL_MLKEM_CACHE_A
 608        if (key->a != NULL) {
 609            XFREE(key->a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 610            key->a = NULL;
 611        }
 612    #endif
 613#else
 614        ForceZero(key->priv, sizeof(key->priv));
 615#endif
 616        ForceZero(key->z, sizeof(key->z));
 617
 618        /* Clear flags as values are no longer set. */
 619        key->flags = 0;
 620    }
 621
 622    return 0;
 623}
 624
 625/******************************************************************************/
 626
 627#ifndef WOLFSSL_MLKEM_NO_MAKE_KEY
 628/**
 629 * Make a ML-KEM key object using a random number generator.
 630 *
 631 * FIPS 203 - Algorithm 19: ML-KEM.KeyGen()
 632 * Generates an encapsulation key and a corresponding decapsulation key.
 633 *   1: d <- B_32                                        >  d is 32 random bytes
 634 *   2: z <- B_32                                        >  z is 32 random bytes
 635 *   3: if d == NULL or z == NULL then
 636 *   4:   return falsum
 637 *                  > return an error indication if random bit generation failed
 638 *   5: end if
 639 *   6: (ek,dk) <- ML-KEM.KeyGen_Internal(d, z)
 640 *                                       > run internal key generation algorithm
 641 *   7: return (ek,dk)
 642 *
 643 * @param  [in, out]  key   ML-KEM key object.
 644 * @param  [in]       rng   Random number generator.
 645 * @return  0 on success.
 646 * @return  BAD_FUNC_ARG when key or rng is NULL.
 647 * @return  MEMORY_E when dynamic memory allocation failed.
 648 * @return  RNG_FAILURE_E when generating random numbers failed.
 649 * @return  DRBG_CONT_FAILURE when random number generator health check fails.
 650 * @return  ML_KEM_PCT_E when pairwise consistency test fails. FIPS only.
 651 * @return  BAD_COND_E when fault attack detected.
 652 * @return  NOT_COMPILED_IN when no random number generator is compiled in or
 653 *          key type is not supported.
 654 */
 655int wc_MlKemKey_MakeKey(MlKemKey* key, WC_RNG* rng)
 656{
 657#ifndef WC_NO_RNG
 658    int ret = 0;
 659    unsigned char rand[WC_ML_KEM_MAKEKEY_RAND_SZ];
 660
 661    /* Validate parameters. */
 662    if ((key == NULL) || (rng == NULL)) {
 663        ret = BAD_FUNC_ARG;
 664    }
 665
 666#ifdef WOLF_CRYPTO_CB
 667#ifndef WOLF_CRYPTO_CB_FIND
 668    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
 669#else
 670    if (ret == 0) {
 671#endif
 672        ret = wc_CryptoCb_MakePqcKemKey(rng, WC_PQC_KEM_TYPE_KYBER, key->type,
 673            key);
 674        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
 675            return ret;
 676        /* fall-through when unavailable */
 677        ret = 0;
 678    }
 679#endif
 680
 681    if (ret == 0) {
 682        /* Generate random to use with PRFs.
 683         * Step 1: d is 32 random bytes
 684         * Step 2: z is 32 random bytes
 685         */
 686        ret = wc_RNG_GenerateBlock(rng, rand, WC_ML_KEM_SYM_SZ * 2);
 687        /* Step 3: ret is not zero when d == NULL or z == NULL. */
 688    }
 689    if (ret == 0) {
 690        /* Make a key pair from the random.
 691         * Step 6. run internal key generation algorithm
 692         * Step 7. public and private key are stored in key
 693         */
 694        ret = wc_MlKemKey_MakeKeyWithRandom(key, rand, sizeof(rand));
 695    }
 696
 697#ifdef HAVE_FIPS
 698    /* Pairwise Consistency Test (PCT) per FIPS 140-3 / ISO 19790:2012
 699     * Section 7.10.3.3: encapsulate with ek, decapsulate with dk,
 700     * verify shared secrets match. */
 701    if (ret == 0) {
 702        WC_DECLARE_VAR(pct_ct, byte, WC_ML_KEM_MAX_CIPHER_TEXT_SIZE,
 703            key->heap);
 704        byte pct_ss1[WC_ML_KEM_SS_SZ];
 705        byte pct_ss2[WC_ML_KEM_SS_SZ];
 706        word32 ctSz = 0;
 707
 708        WC_ALLOC_VAR_EX(pct_ct, byte, WC_ML_KEM_MAX_CIPHER_TEXT_SIZE,
 709            key->heap, DYNAMIC_TYPE_TMP_BUFFER, ret = MEMORY_E);
 710
 711        if (ret == 0)
 712            ret = wc_MlKemKey_CipherTextSize(key, &ctSz);
 713
 714        if (ret == 0)
 715            ret = wc_MlKemKey_Encapsulate(key, pct_ct, pct_ss1, rng);
 716
 717        if (ret == 0)
 718            ret = wc_MlKemKey_Decapsulate(key, pct_ss2, pct_ct, ctSz);
 719
 720        if (ret == 0) {
 721            if (XMEMCMP(pct_ss1, pct_ss2, WC_ML_KEM_SS_SZ) != 0)
 722                ret = ML_KEM_PCT_E;
 723        }
 724
 725        ForceZero(pct_ss1, sizeof(pct_ss1));
 726        ForceZero(pct_ss2, sizeof(pct_ss2));
 727        if (WC_VAR_OK(pct_ct))
 728            ForceZero(pct_ct, WC_ML_KEM_MAX_CIPHER_TEXT_SIZE);
 729
 730        WC_FREE_VAR_EX(pct_ct, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 731
 732        /* FIPS 140-3 IG 10.3.A (TE10.35.02): a key pair that fails the PCT
 733         * must be rendered unusable.  Zeroize the generated key material so
 734         * a caller that ignores the return value cannot use it. */
 735        if (ret != 0) {
 736            wc_MlKemKey_Free(key);
 737        }
 738    }
 739#endif /* HAVE_FIPS */
 740
 741    /* Ensure seeds are zeroized. */
 742    ForceZero((void*)rand, (word32)sizeof(rand));
 743
 744    /* Step 4: return ret != 0 on falsum or internal key generation failure. */
 745    return ret;
 746#else
 747    (void)key;
 748    (void)rng;
 749    return NOT_COMPILED_IN;
 750#endif /* WC_NO_RNG */
 751}
 752
 753/**
 754 * Make a ML-KEM key object using random data.
 755 *
 756 * FIPS 203 - Algorithm 16: ML-KEM.KeyGen_internal(d,z)
 757 * Uses randomness to generate an encapsulation key and a corresponding
 758 * decapsulation key.
 759 *   1: (ek_PKE,dk_PKE) <- K-PKE.KeyGen(d)        > run key generation for K-PKE
 760 *   ...
 761 *
 762 * FIPS 203 - Algorithm 13: K-PKE.KeyGen(d)
 763 * Uses randomness to generate an encryption key and a corresponding decryption
 764 * key.
 765 *   1: (rho,sigma) <- G(d||k)
 766 *                         > expand 32+1 bytes to two pseudorandom 32-byte seeds
 767 *   2: N <- 0
 768 *   3-7: generate matrix A_hat
 769 *   8-11: generate s
 770 *   12-15: generate e
 771 *   16-18: calculate t_hat from A_hat, s and e
 772 *   ...
 773 *
 774 * @param  [in, out]  key   ML-KEM key object.
 775 * @param  [in]       rand  Random data.
 776 * @param  [in]       len   Length of random data in bytes.
 777 * @return  0 on success.
 778 * @return  BAD_FUNC_ARG when key or rand is NULL.
 779 * @return  BUFFER_E when length is not WC_ML_KEM_MAKEKEY_RAND_SZ.
 780 * @return  NOT_COMPILED_IN when key type is not supported.
 781 * @return  MEMORY_E when dynamic memory allocation failed.
 782 * @return  BAD_COND_E when fault attack detected.
 783 */
 784int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand,
 785    int len)
 786{
 787    byte buf[2 * WC_ML_KEM_SYM_SZ + 1];
 788    byte* rho = buf;
 789#ifndef WC_MLKEM_FAULT_HARDEN
 790    byte* sigma = buf + WC_ML_KEM_SYM_SZ;
 791#else
 792    byte sigma[WC_ML_KEM_SYM_SZ + 1];
 793#endif
 794#ifndef WOLFSSL_NO_MALLOC
 795    sword16* e = NULL;
 796#else
 797#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM
 798#ifndef WOLFSSL_MLKEM_CACHE_A
 799    sword16 e[(WC_ML_KEM_MAX_K + 1) * WC_ML_KEM_MAX_K * MLKEM_N];
 800#else
 801    sword16 e[WC_ML_KEM_MAX_K * MLKEM_N];
 802#endif
 803#else
 804    sword16 e[WC_ML_KEM_MAX_K * MLKEM_N];
 805#endif
 806#endif
 807#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM
 808    sword16* a = NULL;
 809#endif
 810    sword16* s = NULL;
 811    sword16* t = NULL;
 812    int ret = 0;
 813    int k = 0;
 814
 815    /* Validate parameters. */
 816    if ((key == NULL) || (rand == NULL)) {
 817        ret = BAD_FUNC_ARG;
 818    }
 819    if ((ret == 0) && (len != WC_ML_KEM_MAKEKEY_RAND_SZ)) {
 820        ret = BUFFER_E;
 821    }
 822
 823    if (ret == 0) {
 824        key->flags = 0;
 825
 826        /* Establish parameters based on key type. */
 827        k = mlkemkey_get_k(key);
 828        if (k == 0) {
 829            ret = NOT_COMPILED_IN;
 830        }
 831    }
 832
 833#ifndef WOLFSSL_NO_MALLOC
 834    if (ret == 0) {
 835        /* Allocate dynamic memory for matrix and error vector. */
 836#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM
 837#ifndef WOLFSSL_MLKEM_CACHE_A
 838        /* e (v) | a (m) */
 839        e = (sword16*)XMALLOC((size_t)((k + 1) * k * MLKEM_N) * sizeof(sword16),
 840            key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 841#else
 842        /* e (v) */
 843        e = (sword16*)XMALLOC((size_t)(k * MLKEM_N) * sizeof(sword16),
 844            key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 845#endif
 846#else
 847        /* e (v) */
 848        e = (sword16*)XMALLOC((size_t)(k * MLKEM_N) * sizeof(sword16),
 849            key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 850#endif
 851        if (e == NULL) {
 852            ret = MEMORY_E;
 853        }
 854    }
 855#endif
 856#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
 857    if (ret == 0) {
 858        ret = mlkemkey_alloc_priv(key, (unsigned int)k);
 859    }
 860    if (ret == 0) {
 861        ret = mlkemkey_alloc_pub(key, (unsigned int)k);
 862    }
 863#ifdef WOLFSSL_MLKEM_CACHE_A
 864    if (ret == 0) {
 865        ret = mlkemkey_alloc_a(key, (unsigned int)k);
 866    }
 867#endif
 868#endif
 869    if (ret == 0) {
 870        const byte* d = rand;
 871
 872#ifdef WOLFSSL_MLKEM_CACHE_A
 873        a = key->a;
 874#elif !defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM)
 875        /* Matrix A allocated at end of error vector. */
 876        a = e + (k * MLKEM_N);
 877#endif
 878
 879#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
 880        if (key->type & MLKEM_KYBER)
 881#endif
 882#ifdef WOLFSSL_MLKEM_KYBER
 883        {
 884            /* Expand 32 bytes of random to 64. */
 885            ret = MLKEM_HASH_G(&key->hash, d, WC_ML_KEM_SYM_SZ, NULL, 0, buf);
 886        }
 887#endif
 888#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
 889        else
 890#endif
 891#ifndef WOLFSSL_NO_ML_KEM
 892        {
 893            buf[0] = (byte)k;
 894            /* Expand 33 bytes of random to 64.
 895             * Alg 13: Step 1: (rho,sigma) <- G(d||k)
 896             */
 897            ret = MLKEM_HASH_G(&key->hash, d, WC_ML_KEM_SYM_SZ, buf, 1, buf);
 898        }
 899#endif
 900    }
 901#ifdef WC_MLKEM_FAULT_HARDEN
 902    if (ret == 0) {
 903        XMEMCPY(sigma, buf + WC_ML_KEM_SYM_SZ, WC_ML_KEM_SYM_SZ);
 904        /* Check that correct data was copied and pointer was not faulted. */
 905        if (XMEMCMP(sigma, rho, WC_ML_KEM_SYM_SZ) == 0) {
 906            ret = BAD_COND_E;
 907        }
 908        /* Check that sigma is after rho - rho pointer may have been modified.
 909         */
 910        if (XMEMCMP(sigma, rho + WC_ML_KEM_SYM_SZ, WC_ML_KEM_SYM_SZ) != 0) {
 911            ret = BAD_COND_E;
 912        }
 913    }
 914#endif
 915    if (ret == 0) {
 916        const byte* z = rand + WC_ML_KEM_SYM_SZ;
 917        s = key->priv;
 918        t = key->pub;
 919
 920        /* Cache the public seed for use in encapsulation and encoding public
 921         * key. */
 922        XMEMCPY(key->pubSeed, rho, WC_ML_KEM_SYM_SZ);
 923        /* Cache the z value for decapsulation and encoding private key. */
 924        XMEMCPY(key->z, z, sizeof(key->z));
 925
 926        /* Initialize PRF for use in noise generation. */
 927        mlkem_prf_init(&key->prf);
 928#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM
 929        /* Generate noise using PRF.
 930         * Alg 13: Steps 8-15: generate s and e
 931         */
 932        ret = mlkem_get_noise(&key->prf, k, s, e, NULL, sigma);
 933    }
 934    if (ret == 0) {
 935        /* Generate the matrix A.
 936         * Alg 13: Steps 3-7
 937         */
 938        ret = mlkem_gen_matrix(&key->prf, a, k, rho, 0);
 939    }
 940    if (ret == 0) {
 941        /* Generate key pair from random data.
 942         * Alg 13: Steps 16-18.
 943         */
 944        mlkem_keygen(s, t, e, a, k);
 945#else
 946        /* Generate noise using PRF.
 947         * Alg 13: Steps 8-11: generate s
 948         */
 949        ret = mlkem_get_noise(&key->prf, k, s, NULL, NULL, sigma);
 950    }
 951    if (ret == 0) {
 952        /* Generate key pair from private vector and seeds.
 953         * Alg 13: Steps 3-7: generate matrix A_hat
 954         * Alg 13: Steps 12-15: generate e
 955         * Alg 13: Steps 16-18: calculate t_hat from A_hat, s and e
 956         */
 957        ret = mlkem_keygen_seeds(s, t, &key->prf, e, k, rho, sigma);
 958    }
 959    if (ret == 0) {
 960#endif
 961        /* Private and public key are set/available. */
 962        key->flags |= MLKEM_FLAG_PRIV_SET | MLKEM_FLAG_PUB_SET;
 963#ifdef WOLFSSL_MLKEM_CACHE_A
 964        key->flags |= MLKEM_FLAG_A_SET;
 965#endif
 966    }
 967
 968#ifndef WOLFSSL_NO_MALLOC
 969    /* Free dynamic memory allocated in function. */
 970    if (key != NULL) {
 971        XFREE(e, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
 972    }
 973#endif
 974
 975    /* Note: PCT is performed in wc_MlKemKey_MakeKey() which calls this
 976     * function and has the RNG parameter needed for encapsulation. */
 977
 978    return ret;
 979}
 980#endif /* !WOLFSSL_MLKEM_NO_MAKE_KEY */
 981
 982/******************************************************************************/
 983
 984/**
 985 * Get the size in bytes of cipher text for key.
 986 *
 987 * @param  [in]   key  ML-KEM key object.
 988 * @param  [out]  len  Length of cipher text in bytes.
 989 * @return  0 on success.
 990 * @return  BAD_FUNC_ARG when key or len is NULL.
 991 * @return  NOT_COMPILED_IN when key type is not supported.
 992 */
 993int wc_MlKemKey_CipherTextSize(MlKemKey* key, word32* len)
 994{
 995    int ret = 0;
 996
 997    /* Validate parameters. */
 998    if ((key == NULL) || (len == NULL)) {
 999        ret = BAD_FUNC_ARG;
1000    }
1001
1002    if (ret == 0) {
1003        /* Return in 'len' size of the cipher text for the type of this key. */
1004        switch (key->type) {
1005#ifndef WOLFSSL_NO_ML_KEM
1006    #ifdef WOLFSSL_WC_ML_KEM_512
1007        case WC_ML_KEM_512:
1008            *len = WC_ML_KEM_512_CIPHER_TEXT_SIZE;
1009            break;
1010    #endif
1011    #ifdef WOLFSSL_WC_ML_KEM_768
1012        case WC_ML_KEM_768:
1013            *len = WC_ML_KEM_768_CIPHER_TEXT_SIZE;
1014            break;
1015    #endif
1016    #ifdef WOLFSSL_WC_ML_KEM_1024
1017        case WC_ML_KEM_1024:
1018            *len = WC_ML_KEM_1024_CIPHER_TEXT_SIZE;
1019            break;
1020    #endif
1021#endif
1022#ifdef WOLFSSL_MLKEM_KYBER
1023    #ifdef WOLFSSL_KYBER512
1024        case KYBER512:
1025            *len = KYBER512_CIPHER_TEXT_SIZE;
1026            break;
1027    #endif
1028    #ifdef WOLFSSL_KYBER768
1029        case KYBER768:
1030            *len = KYBER768_CIPHER_TEXT_SIZE;
1031            break;
1032    #endif
1033    #ifdef WOLFSSL_KYBER1024
1034        case KYBER1024:
1035            *len = KYBER1024_CIPHER_TEXT_SIZE;
1036            break;
1037    #endif
1038#endif
1039        default:
1040            /* No other values supported. */
1041            ret = NOT_COMPILED_IN;
1042            break;
1043        }
1044    }
1045
1046    return ret;
1047}
1048
1049/**
1050 * Size of a shared secret in bytes. Always WC_ML_KEM_SS_SZ.
1051 *
1052 * @param  [in]   key  ML-KEM key object. Not used.
1053 * @param  [out]  len  Size of the shared secret created with a ML-KEM key.
1054 * @return  0 on success.
1055 * @return  BAD_FUNC_ARG when len is NULL.
1056 */
1057int wc_MlKemKey_SharedSecretSize(MlKemKey* key, word32* len)
1058{
1059    int ret = 0;
1060
1061    if (len == NULL) {
1062        ret = BAD_FUNC_ARG;
1063    }
1064    else {
1065        *len = WC_ML_KEM_SS_SZ;
1066    }
1067
1068    (void)key;
1069    return ret;
1070}
1071
1072#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
1073    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
1074/* Encapsulate data and derive secret.
1075 *
1076 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE, m, r)
1077 * Uses the encryption key to encrypt a plaintext message using the randomness
1078 * r.
1079 *   1: N <- 0
1080 *   2: t_hat <- ByteDecode_12(ek_PKE[0:384k])
1081 *                                   > run ByteDecode_12 k times to decode t_hat
1082 *   3: rho <- ek_PKE[384k : 384k + 32]
1083 *                                            > extract 32-byte seed from ek_PKE
1084 *   4-8: generate matrix A_hat
1085 *   9-12: generate y
1086 *   13-16: generate e_1
1087 *   17: generate e_2
1088 *   18-19: calculate u
1089 *   20: mu <- Decompress_1(ByteDecode_1(m))
1090 *   21: calculate v
1091 *   22: c_1 <- ByteEncode_d_u(Compress_d_u(u))
1092 *                                 > run ByteEncode_d_u and Compress_d_u k times
1093 *   23: c_2 <- ByteEncode_d_v(Compress_d_v(v))
1094 *   24: return c <- (c_1||c_2)
1095 *
1096 * @param  [in]  key  ML-KEM key object.
1097 * @param  [in]  m    Random bytes.
1098 * @param  [in]  r    Seed to feed to PRF when generating y, e1 and e2.
1099 * @param  [out] c    Calculated cipher text.
1100 * @return  0 on success.
1101 * @return  NOT_COMPILED_IN when key type is not supported.
1102 */
1103static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
1104{
1105    int ret = 0;
1106    sword16* a = NULL;
1107#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
1108    sword16* mu = NULL;
1109    sword16* e1 = NULL;
1110    sword16* e2 = NULL;
1111#endif
1112    unsigned int k = 0;
1113    unsigned int compVecSz = 0;
1114#ifndef WOLFSSL_NO_MALLOC
1115    sword16* y = NULL;
1116#else
1117#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
1118    sword16 y[((WC_ML_KEM_MAX_K + 3) * WC_ML_KEM_MAX_K + 3) * MLKEM_N];
1119#else
1120    sword16 y[3 * WC_ML_KEM_MAX_K * MLKEM_N];
1121#endif
1122#endif
1123    sword16* u = 0;
1124    sword16* v = 0;
1125
1126    /* Establish parameters based on key type. */
1127    switch (key->type) {
1128#ifndef WOLFSSL_NO_ML_KEM
1129#ifdef WOLFSSL_WC_ML_KEM_512
1130    case WC_ML_KEM_512:
1131        k = WC_ML_KEM_512_K;
1132        compVecSz = WC_ML_KEM_512_POLY_VEC_COMPRESSED_SZ;
1133        break;
1134#endif
1135#ifdef WOLFSSL_WC_ML_KEM_768
1136    case WC_ML_KEM_768:
1137        k = WC_ML_KEM_768_K;
1138        compVecSz = WC_ML_KEM_768_POLY_VEC_COMPRESSED_SZ;
1139        break;
1140#endif
1141#ifdef WOLFSSL_WC_ML_KEM_1024
1142    case WC_ML_KEM_1024:
1143        k = WC_ML_KEM_1024_K;
1144        compVecSz = WC_ML_KEM_1024_POLY_VEC_COMPRESSED_SZ;
1145        break;
1146#endif
1147#endif
1148#ifdef WOLFSSL_MLKEM_KYBER
1149#ifdef WOLFSSL_KYBER512
1150    case KYBER512:
1151        k = KYBER512_K;
1152        compVecSz = KYBER512_POLY_VEC_COMPRESSED_SZ;
1153        break;
1154#endif
1155#ifdef WOLFSSL_KYBER768
1156    case KYBER768:
1157        k = KYBER768_K;
1158        compVecSz = KYBER768_POLY_VEC_COMPRESSED_SZ;
1159        break;
1160#endif
1161#ifdef WOLFSSL_KYBER1024
1162    case KYBER1024:
1163        k = KYBER1024_K;
1164        compVecSz = KYBER1024_POLY_VEC_COMPRESSED_SZ;
1165        break;
1166#endif
1167#endif
1168    default:
1169        /* No other values supported. */
1170        ret = NOT_COMPILED_IN;
1171        break;
1172    }
1173
1174#ifndef WOLFSSL_NO_MALLOC
1175    if (ret == 0) {
1176        /* Allocate dynamic memory for all matrices, vectors and polynomials. */
1177#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
1178        y = (sword16*)XMALLOC(((k + 3) * k + 3) * MLKEM_N * sizeof(sword16),
1179            key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1180#else
1181        y = (sword16*)XMALLOC(3 * k * MLKEM_N * sizeof(sword16), key->heap,
1182            DYNAMIC_TYPE_TMP_BUFFER);
1183#endif
1184        if (y == NULL) {
1185            ret = MEMORY_E;
1186        }
1187    }
1188#endif
1189
1190#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
1191    if (ret == 0) {
1192        /* Assign allocated dynamic memory to pointers.
1193         * y (b) | a (m) | mu (p) | e1 (p) | e2 (v) | u (v) | v (p) */
1194        a  = y  + MLKEM_N * k;
1195        mu = a  + MLKEM_N * k * k;
1196        e1 = mu + MLKEM_N;
1197        e2 = e1 + MLKEM_N * k;
1198
1199        /* Convert msg to a polynomial.
1200         * Step 20: mu <- Decompress_1(ByteDecode_1(m)) */
1201        mlkem_from_msg(mu, m);
1202
1203        /* Initialize the PRF for use in the noise generation. */
1204        mlkem_prf_init(&key->prf);
1205        /* Generate noise using PRF.
1206         * Steps 9-17: generate y, e_1, e_2
1207         */
1208        ret = mlkem_get_noise(&key->prf, (int)k, y, e1, e2, r);
1209    }
1210    #ifdef WOLFSSL_MLKEM_CACHE_A
1211    if ((ret == 0) && ((key->flags & MLKEM_FLAG_A_SET) != 0)) {
1212        unsigned int i;
1213        /* Transpose matrix.
1214         *   Steps 4-8: generate matrix A_hat (from original) */
1215        for (i = 0; i < k; i++) {
1216            unsigned int j;
1217            for (j = 0; j < k; j++) {
1218                XMEMCPY(&a[(i * k + j) * MLKEM_N],
1219                        &key->a[(j * k + i) * MLKEM_N],
1220                        MLKEM_N * 2);
1221            }
1222        }
1223    }
1224    else
1225    #endif /* WOLFSSL_MLKEM_CACHE_A */
1226    if (ret == 0) {
1227        /* Generate the transposed matrix.
1228         *   Step 4-8: generate matrix A_hat */
1229        ret = mlkem_gen_matrix(&key->prf, a, (int)k, key->pubSeed, 1);
1230    }
1231    if (ret == 0) {
1232        /* Assign remaining allocated dynamic memory to pointers.
1233         * y (b) | a (m) | mu (p) | e1 (p) | e2 (v) | u (v) | v (p) */
1234        u  = e2 + MLKEM_N;
1235        v  = u  + MLKEM_N * k;
1236
1237        /* Perform encapsulation maths.
1238         *   Steps 18-19, 21: calculate u and v */
1239        mlkem_encapsulate(key->pub, u, v, a, y, e1, e2, mu, (int)k);
1240    }
1241#else /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */
1242    if (ret == 0) {
1243        /* Assign allocated dynamic memory to pointers.
1244         * y (v) | a (v) | u (v) */
1245        a = y + MLKEM_N * k;
1246
1247        /* Initialize the PRF for use in the noise generation. */
1248        mlkem_prf_init(&key->prf);
1249        /* Generate noise using PRF.
1250         * Steps 9-12: generate y */
1251        ret = mlkem_get_noise(&key->prf, (int)k, y, NULL, NULL, r);
1252    }
1253    if (ret == 0) {
1254        /* Assign remaining allocated dynamic memory to pointers.
1255         * y (v) | at (v) | u (v) */
1256        u  = a + MLKEM_N * k;
1257        v  = a;
1258
1259        /* Perform encapsulation maths.
1260         *   Steps 13-17: generate e_1 and e_2
1261         *   Steps 18-19, 21: calculate u and v */
1262        ret = mlkem_encapsulate_seeds(key->pub, &key->prf, u, a, y, (int)k, m,
1263            key->pubSeed, r);
1264    }
1265#endif /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */
1266
1267    if (ret == 0) {
1268        byte* c1 = c;
1269        byte* c2 = c + compVecSz;
1270
1271    #if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
1272        if (k == WC_ML_KEM_512_K) {
1273            /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
1274            mlkem_vec_compress_10(c1, u, k);
1275            /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
1276            mlkem_compress_4(c2, v);
1277            /* Step 24: return c <- (c_1||c_2) */
1278        }
1279    #endif
1280    #if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
1281        if (k == WC_ML_KEM_768_K) {
1282            /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
1283            mlkem_vec_compress_10(c1, u, k);
1284            /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
1285            mlkem_compress_4(c2, v);
1286            /* Step 24: return c <- (c_1||c_2) */
1287        }
1288    #endif
1289    #if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
1290        if (k == WC_ML_KEM_1024_K) {
1291            /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
1292            mlkem_vec_compress_11(c1, u);
1293            /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
1294            mlkem_compress_5(c2, v);
1295            /* Step 24: return c <- (c_1||c_2) */
1296        }
1297    #endif
1298    }
1299
1300#ifndef WOLFSSL_NO_MALLOC
1301    /* Dispose of dynamic memory allocated in function. */
1302    XFREE(y, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1303#endif
1304
1305    return ret;
1306}
1307#endif
1308
1309#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
1310    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
1311static int wc_mlkemkey_check_h(MlKemKey* key)
1312{
1313    int ret = 0;
1314
1315    /* If public hash (h) is not stored against key, calculate it
1316     * (fields set explicitly instead of using decode).
1317     * Step 1: ... H(ek)...
1318     */
1319    if ((key->flags & MLKEM_FLAG_H_SET) == 0) {
1320    #ifndef WOLFSSL_NO_MALLOC
1321        byte* pubKey = NULL;
1322        word32 pubKeyLen;
1323    #else
1324        byte pubKey[WC_ML_KEM_MAX_PUBLIC_KEY_SIZE];
1325        word32 pubKeyLen;
1326    #endif
1327
1328        /* Determine how big an encoded public key will be. */
1329        ret = wc_MlKemKey_PublicKeySize(key, &pubKeyLen);
1330        if (ret == 0) {
1331    #ifndef WOLFSSL_NO_MALLOC
1332            /* Allocate dynamic memory for encoded public key. */
1333            pubKey = (byte*)XMALLOC(pubKeyLen, key->heap,
1334                DYNAMIC_TYPE_TMP_BUFFER);
1335            if (pubKey == NULL) {
1336                ret = MEMORY_E;
1337            }
1338        }
1339        if (ret == 0) {
1340    #endif
1341            /* Encode public key - h is hash of encoded public key. */
1342            ret = wc_MlKemKey_EncodePublicKey(key, pubKey, pubKeyLen);
1343        }
1344    #ifndef WOLFSSL_NO_MALLOC
1345        /* Dispose of encoded public key. */
1346        XFREE(pubKey, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1347    #endif
1348    }
1349    if ((ret == 0) && ((key->flags & MLKEM_FLAG_H_SET) == 0)) {
1350        /* Implementation issue if h not cached and flag not set. */
1351        ret = BAD_STATE_E;
1352    }
1353
1354    return ret;
1355}
1356#endif
1357
1358#ifndef WOLFSSL_MLKEM_NO_ENCAPSULATE
1359/**
1360 * Encapsulate with random number generator and derive secret.
1361 *
1362 * FIPS 203, Algorithm 20: ML-KEM.Encaps(ek)
1363 * Uses the encapsulation key to generate a shared secret key and an associated
1364 * ciphertext.
1365 *   1: m <- B_32                                         > m is 32 random bytes
1366 *   2: if m == NULL then
1367 *   3:     return falsum
1368 *   4: end if
1369 *   5: (K,c) <- ML-KEM.Encaps_internal(ek,m)
1370 *                                        > run internal encapsulation algorithm
1371 *   6: return (K,c)
1372 *
1373 * @param  [in]   key  ML-KEM key object.
1374 * @param  [out]  ct   Cipher text.
1375 * @param  [out]  ss   Shared secret generated.
1376 * @param  [in]   rng  Random number generator.
1377 * @return  0 on success.
1378 * @return  BAD_FUNC_ARG when key, ct, ss or rng is NULL.
1379 * @return  BAD_STATE_E when public key not set.
1380 * @return  NOT_COMPILED_IN when key type is not supported.
1381 * @return  MEMORY_E when dynamic memory allocation failed.
1382 */
1383int wc_MlKemKey_Encapsulate(MlKemKey* key, unsigned char* ct, unsigned char* ss,
1384    WC_RNG* rng)
1385{
1386#ifndef WC_NO_RNG
1387    int ret = 0;
1388    unsigned char m[WC_ML_KEM_ENC_RAND_SZ];
1389#ifdef WOLF_CRYPTO_CB
1390    word32 ctlen = 0;
1391#endif
1392
1393    /* Validate parameters. */
1394    if ((key == NULL) || (ct == NULL) || (ss == NULL) || (rng == NULL)) {
1395        ret = BAD_FUNC_ARG;
1396    }
1397    /* Check the public key has been set. */
1398    else if ((key->flags & MLKEM_FLAG_PUB_SET) == 0) {
1399        ret = BAD_STATE_E;
1400    }
1401
1402#ifdef WOLF_CRYPTO_CB
1403    if (ret == 0) {
1404        ret = wc_MlKemKey_CipherTextSize(key, &ctlen);
1405    }
1406#ifndef WOLF_CRYPTO_CB_FIND
1407    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
1408#else
1409    if (ret == 0) {
1410#endif
1411        ret = wc_CryptoCb_PqcEncapsulate(ct, ctlen, ss, WC_ML_KEM_SS_SZ, rng,
1412            WC_PQC_KEM_TYPE_KYBER, key);
1413        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
1414            return ret;
1415        /* fall-through when unavailable */
1416        ret = 0;
1417    }
1418#endif
1419
1420    if (ret == 0) {
1421        /* Generate seed for use with PRFs.
1422         * Step 1: m is 32 random bytes
1423         */
1424        ret = wc_RNG_GenerateBlock(rng, m, sizeof(m));
1425        /* Step 2: ret is not zero when m == NULL. */
1426    }
1427    if (ret == 0) {
1428        /* Encapsulate with the random.
1429         * Step 5: run internal encapsulation algorithm
1430         */
1431        ret = wc_MlKemKey_EncapsulateWithRandom(key, ct, ss, m, sizeof(m));
1432    }
1433
1434    /* Step 3: return ret != 0 on falsum or internal key generation failure. */
1435    return ret;
1436#else
1437    (void)key;
1438    (void)ct;
1439    (void)ss;
1440    (void)rng;
1441    return NOT_COMPILED_IN;
1442#endif /* WC_NO_RNG */
1443}
1444
1445/**
1446 * Encapsulate with random data and derive secret.
1447 *
1448 * FIPS 203, Algorithm 17: ML-KEM.Encaps_internal(ek, m)
1449 * Uses the encapsulation key and randomness to generate a key and an associated
1450 * ciphertext.
1451 *   Step 1: (K,r) <- G(m||H(ek))
1452 *                                 > derive shared secret key K and randomness r
1453 *   Step 2: c <- K-PKE.Encrypt(ek, m, r)
1454 *                                     > encrypt m using K-PKE with randomness r
1455 *   Step 3: return (K,c)
1456 *
1457 * @param  [in]   key   ML-KEM key object.
1458 * @param  [out]  ct    Cipher text.
1459 * @param  [out]  ss    Shared secret generated.
1460 * @param  [in]   rand  Random bytes.
1461 * @param  [in]   len   Length of random bytes.
1462 * @return  0 on success.
1463 * @return  BAD_FUNC_ARG when key, ct, ss or rand is NULL.
1464 * @return  BUFFER_E when len is not WC_ML_KEM_ENC_RAND_SZ.
1465 * @return  BAD_STATE_E when public key not set.
1466 * @return  NOT_COMPILED_IN when key type is not supported.
1467 * @return  MEMORY_E when dynamic memory allocation failed.
1468 */
1469int wc_MlKemKey_EncapsulateWithRandom(MlKemKey* key, unsigned char* ct,
1470    unsigned char* ss, const unsigned char* rand, int len)
1471{
1472#ifdef WOLFSSL_MLKEM_KYBER
1473    byte msg[WC_ML_KEM_SYM_SZ];
1474#endif
1475    byte kr[2 * WC_ML_KEM_SYM_SZ + 1];
1476    int ret = 0;
1477#ifdef WOLFSSL_MLKEM_KYBER
1478    unsigned int cSz = 0;
1479#endif
1480
1481    /* Validate parameters. */
1482    if ((key == NULL) || (ct == NULL) || (ss == NULL) || (rand == NULL)) {
1483        ret = BAD_FUNC_ARG;
1484    }
1485    if ((ret == 0) && (len != WC_ML_KEM_ENC_RAND_SZ)) {
1486        ret = BUFFER_E;
1487    }
1488    /* Check the public key has been set. */
1489    if ((ret == 0) && ((key->flags & MLKEM_FLAG_PUB_SET) == 0)) {
1490        ret = BAD_STATE_E;
1491    }
1492
1493#ifdef WOLFSSL_MLKEM_KYBER
1494    if (ret == 0) {
1495        /* Establish parameters based on key type. */
1496        switch (key->type) {
1497#ifndef WOLFSSL_NO_ML_KEM
1498    #ifdef WOLFSSL_WC_ML_KEM_512
1499        case WC_ML_KEM_512:
1500    #endif
1501    #ifdef WOLFSSL_WC_ML_KEM_768
1502        case WC_ML_KEM_768:
1503    #endif
1504    #ifdef WOLFSSL_WC_ML_KEM_1024
1505        case WC_ML_KEM_1024:
1506    #endif
1507            break;
1508#endif
1509    #ifdef WOLFSSL_KYBER512
1510        case KYBER512:
1511            cSz = KYBER512_CIPHER_TEXT_SIZE;
1512            break;
1513    #endif
1514    #ifdef WOLFSSL_KYBER768
1515        case KYBER768:
1516            cSz = KYBER768_CIPHER_TEXT_SIZE;
1517            break;
1518    #endif
1519    #ifdef WOLFSSL_KYBER1024
1520        case KYBER1024:
1521            cSz = KYBER1024_CIPHER_TEXT_SIZE;
1522            break;
1523    #endif
1524        default:
1525            /* No other values supported. */
1526            ret = NOT_COMPILED_IN;
1527            break;
1528        }
1529    }
1530#endif
1531
1532    if (ret == 0) {
1533        ret = wc_mlkemkey_check_h(key);
1534    }
1535
1536#ifdef WOLFSSL_MLKEM_KYBER
1537    if (ret == 0) {
1538#ifndef WOLFSSL_NO_ML_KEM
1539        if (key->type & MLKEM_KYBER)
1540#endif
1541        {
1542            /* Hash random to anonymize as seed data. */
1543            ret = MLKEM_HASH_H(&key->hash, rand, WC_ML_KEM_SYM_SZ, msg);
1544        }
1545    }
1546#endif
1547    if (ret == 0) {
1548        /* Hash message into seed buffer. */
1549#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1550        if (key->type & MLKEM_KYBER)
1551#endif
1552#ifdef WOLFSSL_MLKEM_KYBER
1553        {
1554            ret = MLKEM_HASH_G(&key->hash, msg, WC_ML_KEM_SYM_SZ, key->h,
1555                WC_ML_KEM_SYM_SZ, kr);
1556        }
1557#endif
1558#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1559        else
1560#endif
1561#ifndef WOLFSSL_NO_ML_KEM
1562        {
1563            /* Step 1: (K,r) <- G(m||H(ek)) */
1564            ret = MLKEM_HASH_G(&key->hash, rand, WC_ML_KEM_SYM_SZ, key->h,
1565                WC_ML_KEM_SYM_SZ, kr);
1566        }
1567#endif
1568    }
1569
1570    if (ret == 0) {
1571        /* Encapsulate the message using the key and the seed. */
1572#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1573        if (key->type & MLKEM_KYBER)
1574#endif
1575#ifdef WOLFSSL_MLKEM_KYBER
1576        {
1577            ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, ct);
1578        }
1579#endif
1580#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1581        else
1582#endif
1583#ifndef WOLFSSL_NO_ML_KEM
1584        {
1585            /* Step 2: c <- K-PKE.Encrypt(ek,m,r) */
1586            ret = mlkemkey_encapsulate(key, rand, kr + WC_ML_KEM_SYM_SZ, ct);
1587        }
1588#endif
1589    }
1590
1591#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1592    if (key->type & MLKEM_KYBER)
1593#endif
1594#ifdef WOLFSSL_MLKEM_KYBER
1595    {
1596        if (ret == 0) {
1597            /* Hash the cipher text after the seed. */
1598            ret = MLKEM_HASH_H(&key->hash, ct, cSz, kr + WC_ML_KEM_SYM_SZ);
1599        }
1600        if (ret == 0) {
1601            /* Derive the secret from the seed and hash of cipher text. */
1602            ret = MLKEM_KDF(kr, 2 * WC_ML_KEM_SYM_SZ, ss, WC_ML_KEM_SS_SZ);
1603        }
1604    }
1605#endif
1606#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1607    else
1608#endif
1609#ifndef WOLFSSL_NO_ML_KEM
1610    {
1611        if (ret == 0) {
1612            /* return (K,c) */
1613            XMEMCPY(ss, kr, WC_ML_KEM_SS_SZ);
1614        }
1615    }
1616#endif
1617
1618    ForceZero(kr, sizeof(kr));
1619
1620    return ret;
1621}
1622#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE */
1623
1624/******************************************************************************/
1625
1626#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
1627/* Decapsulate cipher text to the message using key.
1628 *
1629 * FIPS 203, Algorithm 15: K-PKE.Decrypt(dk_PKE,c)
1630 * Uses the decryption key to decrypt a ciphertext.
1631 *   1: c1 <- c[0 : 32.d_u.k]
1632 *   2: c2 <- c[32.d_u.k : 32(d_u.k + d_v)]
1633 *   3: u' <- Decompress_d_u(ByteDecode_d_u(c1))
1634 *   4: v' <- Decompress_d_v(ByteDecode_d_v(c2))
1635 *   ...
1636 *   6: w <- v' - InvNTT(s_hat_trans o NTT(u'))
1637 *   7: m <- ByteEncode_1(Compress_1(w))
1638 *   8: return m
1639 *
1640 * @param  [in]   key  ML-KEM key object.
1641 * @param  [out]  m    Message that was encapsulated.
1642 * @param  [in]   c    Cipher text.
1643 * @return  0 on success.
1644 * @return  NOT_COMPILED_IN when key type is not supported.
1645 * @return  MEMORY_E when dynamic memory allocation failed.
1646 */
1647static MLKEM_NOINLINE int mlkemkey_decapsulate(MlKemKey* key, byte* m,
1648    const byte* c)
1649{
1650    int ret = 0;
1651    sword16* v;
1652    sword16* w;
1653    unsigned int k = 0;
1654    unsigned int compVecSz;
1655#if defined(WOLFSSL_SMALL_STACK) || \
1656    (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC))
1657    sword16* u = NULL;
1658#else
1659    sword16 u[(WC_ML_KEM_MAX_K + 1) * MLKEM_N];
1660#endif
1661
1662    /* Establish parameters based on key type. */
1663    switch (key->type) {
1664#ifndef WOLFSSL_NO_ML_KEM
1665#ifdef WOLFSSL_WC_ML_KEM_512
1666    case WC_ML_KEM_512:
1667        k = WC_ML_KEM_512_K;
1668        compVecSz = WC_ML_KEM_512_POLY_VEC_COMPRESSED_SZ;
1669        break;
1670#endif
1671#ifdef WOLFSSL_WC_ML_KEM_768
1672    case WC_ML_KEM_768:
1673        k = WC_ML_KEM_768_K;
1674        compVecSz = WC_ML_KEM_768_POLY_VEC_COMPRESSED_SZ;
1675        break;
1676#endif
1677#ifdef WOLFSSL_WC_ML_KEM_1024
1678    case WC_ML_KEM_1024:
1679        k = WC_ML_KEM_1024_K;
1680        compVecSz = WC_ML_KEM_1024_POLY_VEC_COMPRESSED_SZ;
1681        break;
1682#endif
1683#endif
1684#ifdef WOLFSSL_MLKEM_KYBER
1685#ifdef WOLFSSL_KYBER512
1686    case KYBER512:
1687        k = KYBER512_K;
1688        compVecSz = KYBER512_POLY_VEC_COMPRESSED_SZ;
1689        break;
1690#endif
1691#ifdef WOLFSSL_KYBER768
1692    case KYBER768:
1693        k = KYBER768_K;
1694        compVecSz = KYBER768_POLY_VEC_COMPRESSED_SZ;
1695        break;
1696#endif
1697#ifdef WOLFSSL_KYBER1024
1698    case KYBER1024:
1699        k = KYBER1024_K;
1700        compVecSz = KYBER1024_POLY_VEC_COMPRESSED_SZ;
1701        break;
1702#endif
1703#endif
1704    default:
1705        /* No other values supported. */
1706        ret = NOT_COMPILED_IN;
1707        break;
1708    }
1709
1710#if defined(WOLFSSL_SMALL_STACK) || \
1711    (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC))
1712    if (ret == 0) {
1713        /* Allocate dynamic memory for a vector and a polynomial. */
1714        u = (sword16*)XMALLOC((k + 1) * MLKEM_N * sizeof(sword16), key->heap,
1715            DYNAMIC_TYPE_TMP_BUFFER);
1716        if (u == NULL) {
1717            ret = MEMORY_E;
1718        }
1719    }
1720#endif
1721    if (ret == 0) {
1722        /* Step 1: c1 <- c[0 : 32.d_u.k] */
1723        const byte* c1 = c;
1724        /* Step 2: c2 <- c[32.d_u.k : 32(d_u.k + d_v)] */
1725        const byte* c2 = c + compVecSz;
1726
1727        /* Assign allocated dynamic memory to pointers.
1728         * u (v) | v (p) */
1729        v = u + k * MLKEM_N;
1730        w = u;
1731
1732    #if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
1733        if (k == WC_ML_KEM_512_K) {
1734            /* Step 3: u' <- Decompress_d_u(ByteDecode_d_u(c1)) */
1735            mlkem_vec_decompress_10(u, c1, k);
1736            /* Step 4: v' <- Decompress_d_v(ByteDecode_d_v(c2)) */
1737            mlkem_decompress_4(v, c2);
1738        }
1739    #endif
1740    #if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
1741        if (k == WC_ML_KEM_768_K) {
1742            /* Step 3: u' <- Decompress_d_u(ByteDecode_d_u(c1)) */
1743            mlkem_vec_decompress_10(u, c1, k);
1744            /* Step 4: v' <- Decompress_d_v(ByteDecode_d_v(c2)) */
1745            mlkem_decompress_4(v, c2);
1746        }
1747    #endif
1748    #if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
1749        if (k == WC_ML_KEM_1024_K) {
1750            /* Step 3: u' <- Decompress_d_u(ByteDecode_d_u(c1)) */
1751            mlkem_vec_decompress_11(u, c1);
1752            /* Step 4: v' <- Decompress_d_v(ByteDecode_d_v(c2)) */
1753            mlkem_decompress_5(v, c2);
1754        }
1755    #endif
1756
1757        /* Decapsulate the cipher text into polynomial.
1758         * Step 6: w <- v' - InvNTT(s_hat_trans o NTT(u')) */
1759        mlkem_decapsulate(key->priv, w, u, v, (int)k);
1760
1761        /* Convert the polynomial into a array of bytes (message).
1762         * Step 7: m <- ByteEncode_1(Compress_1(w)) */
1763        mlkem_to_msg(m, w);
1764        /* Step 8: return m */
1765    }
1766
1767#if defined(WOLFSSL_SMALL_STACK) || \
1768    (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC))
1769    /* Dispose of dynamically memory allocated in function. */
1770    XFREE(u, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1771#endif
1772
1773    return ret;
1774}
1775
1776/**
1777 * Decapsulate the cipher text to calculate the shared secret.
1778 *
1779 * Validates the cipher text by encapsulating and comparing with data passed in.
1780 *
1781 * FIPS 203, Algorithm 21: ML-KEM.Decaps(dk, c)
1782 * Uses the decapsulation key to produce a shared secret key from a ciphertext.
1783 *   1: K' <- ML-KEM.Decaps_internal(dk,c)
1784 *                                        > run internal decapsulation algorithm
1785 *   2: return K'
1786 *
1787 * FIPS 203, Algorithm 18: ML-KEM.Decaps_internal(dk, c)
1788 * Uses the decapsulation key to produce a shared secret key from a ciphertext.
1789 *   ...
1790 *   1: dk_PKE <- dk[0 : 384k]
1791 *                        > extract (from KEM decaps key) the PKE decryption key
1792 *   2: ek_PKE <- dk[384k : 768k + 32]
1793 *                                                  > extract PKE encryption key
1794 *   3: h <- dk[768k + 32 : 768k + 64]
1795 *                                          > extract hash of PKE encryption key
1796 *   4: z <- dk[768k + 64 : 768k + 96]
1797 *                                            > extract implicit rejection value
1798 *   5: m' <- K-PKE.Decrypt(dk_PKE, c)                      > decrypt ciphertext
1799 *   6: (K', r') <- G(m'||h)
1800 *   7: K_bar <- J(z||c)
1801 *   8: c' <- K-PKE.Encrypt(ek_PKE, m', r')
1802 *                                  > re-encrypt using the derived randomness r'
1803 *   9: if c != c' then
1804 *  10:      K' <- K_bar
1805 *                            > if ciphertexts do not match, "implicitly reject"
1806 *  11: end if
1807 *  12: return K'
1808 *
1809 * @param  [in]   key  ML-KEM key object.
1810 * @param  [out]  ss   Shared secret.
1811 * @param  [in]   ct   Cipher text.
1812 * @param  [in]   len  Length of cipher text.
1813 * @return  0 on success.
1814 * @return  BAD_FUNC_ARG when key, ss or ct are NULL.
1815 * @return  BAD_STATE_E when private key is not set.
1816 * @return  NOT_COMPILED_IN when key type is not supported.
1817 * @return  BUFFER_E when len is not the length of cipher text for the key type.
1818 * @return  MEMORY_E when dynamic memory allocation failed.
1819 */
1820int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss,
1821    const unsigned char* ct, word32 len)
1822{
1823    byte msg[WC_ML_KEM_SYM_SZ];
1824    byte kr[2 * WC_ML_KEM_SYM_SZ + 1];
1825    int ret = 0;
1826    unsigned int ctSz = 0;
1827    unsigned int i = 0;
1828    int fail = 0;
1829#if !defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)
1830    byte* cmp = NULL;
1831#else
1832    byte cmp[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE];
1833#endif
1834
1835    /* Validate parameters. */
1836    if ((key == NULL) || (ss == NULL) || (ct == NULL)) {
1837        ret = BAD_FUNC_ARG;
1838    }
1839    if ((ret == 0) && ((key->flags & MLKEM_FLAG_PRIV_SET) == 0)) {
1840        ret = BAD_STATE_E;
1841    }
1842
1843    if (ret == 0) {
1844        /* Establish cipher text size based on key type. */
1845        switch (key->type) {
1846#ifndef WOLFSSL_NO_ML_KEM
1847    #ifdef WOLFSSL_WC_ML_KEM_512
1848        case WC_ML_KEM_512:
1849            ctSz = WC_ML_KEM_512_CIPHER_TEXT_SIZE;
1850            break;
1851    #endif
1852    #ifdef WOLFSSL_WC_ML_KEM_768
1853        case WC_ML_KEM_768:
1854            ctSz = WC_ML_KEM_768_CIPHER_TEXT_SIZE;
1855            break;
1856    #endif
1857    #ifdef WOLFSSL_WC_ML_KEM_1024
1858        case WC_ML_KEM_1024:
1859            ctSz = WC_ML_KEM_1024_CIPHER_TEXT_SIZE;
1860            break;
1861    #endif
1862#endif
1863#ifdef WOLFSSL_MLKEM_KYBER
1864    #ifdef WOLFSSL_KYBER512
1865        case KYBER512:
1866            ctSz = KYBER512_CIPHER_TEXT_SIZE;
1867            break;
1868    #endif
1869    #ifdef WOLFSSL_KYBER768
1870        case KYBER768:
1871            ctSz = KYBER768_CIPHER_TEXT_SIZE;
1872            break;
1873    #endif
1874    #ifdef WOLFSSL_KYBER1024
1875        case KYBER1024:
1876            ctSz = KYBER1024_CIPHER_TEXT_SIZE;
1877            break;
1878    #endif
1879#endif
1880        default:
1881            /* No other values supported. */
1882            ret = NOT_COMPILED_IN;
1883            break;
1884        }
1885    }
1886
1887    /* Ensure the cipher text passed in is the correct size. */
1888    if ((ret == 0) && (len != ctSz)) {
1889        ret = BUFFER_E;
1890    }
1891
1892#ifdef WOLF_CRYPTO_CB
1893#ifndef WOLF_CRYPTO_CB_FIND
1894    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
1895#else
1896    if (ret == 0) {
1897#endif
1898        ret = wc_CryptoCb_PqcDecapsulate(ct, ctSz, ss, WC_ML_KEM_SS_SZ,
1899            WC_PQC_KEM_TYPE_KYBER, key);
1900        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE))
1901            return ret;
1902        /* fall-through when unavailable */
1903        ret = 0;
1904    }
1905#endif
1906
1907#if !defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)
1908    if (ret == 0) {
1909        /* Allocate memory for cipher text that is generated. */
1910        cmp = (byte*)XMALLOC(ctSz, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1911        if (cmp == NULL) {
1912            ret = MEMORY_E;
1913        }
1914    }
1915#endif
1916
1917    if (ret == 0) {
1918        /* Decapsulate the cipher text. */
1919        ret = mlkemkey_decapsulate(key, msg, ct);
1920    }
1921    if (ret == 0) {
1922        /* Check we have H, hash of public, set. */
1923        ret = wc_mlkemkey_check_h(key);
1924    }
1925    if (ret == 0) {
1926        /* Hash message into seed buffer. */
1927        ret = MLKEM_HASH_G(&key->hash, msg, WC_ML_KEM_SYM_SZ, key->h,
1928            WC_ML_KEM_SYM_SZ, kr);
1929    }
1930    if (ret == 0) {
1931        /* Encapsulate the message. */
1932        ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, cmp);
1933    }
1934    if (ret == 0) {
1935        /* Compare generated cipher text with that passed in. */
1936        fail = mlkem_cmp(ct, cmp, (int)ctSz);
1937
1938#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1939        if (key->type & MLKEM_KYBER)
1940#endif
1941#ifdef WOLFSSL_MLKEM_KYBER
1942        {
1943            /* Hash the cipher text after the seed. */
1944            ret = MLKEM_HASH_H(&key->hash, ct, ctSz, kr + WC_ML_KEM_SYM_SZ);
1945            if (ret == 0) {
1946                /* Change seed to z on comparison failure. */
1947                for (i = 0; i < WC_ML_KEM_SYM_SZ; i++) {
1948                    kr[i] ^= (kr[i] ^ key->z[i]) & fail;
1949                }
1950
1951                /* Derive the secret from the seed and hash of cipher text. */
1952                ret = MLKEM_KDF(kr, 2 * WC_ML_KEM_SYM_SZ, ss, WC_ML_KEM_SS_SZ);
1953            }
1954        }
1955#endif
1956#if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM)
1957        else
1958#endif
1959#ifndef WOLFSSL_NO_ML_KEM
1960        {
1961            ret = mlkem_derive_secret(&key->prf, key->z, ct, ctSz, msg);
1962            if (ret == 0) {
1963               /* Set secret to kr or fake secret on comparison failure. */
1964               for (i = 0; i < WC_ML_KEM_SYM_SZ; i++) {
1965                   ss[i] = (byte)(kr[i] ^ ((kr[i] ^ msg[i]) & fail));
1966               }
1967            }
1968        }
1969#endif
1970    }
1971
1972#if !defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)
1973    /* Dispose of dynamic memory allocated in function. */
1974    if (key != NULL) {
1975        XFREE(cmp, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
1976    }
1977#endif
1978
1979    ForceZero(msg, sizeof(msg));
1980    ForceZero(kr, sizeof(kr));
1981
1982    return ret;
1983}
1984#endif /* WOLFSSL_MLKEM_NO_DECAPSULATE */
1985
1986/******************************************************************************/
1987
1988/**
1989 * Get the public key and public seed from bytes.
1990 *
1991 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE, m, r)
1992 *   ...
1993 *   2: t <- ByteDecode_12(ek_PKE[0 : 384k])
1994 *   3: rho <- ek_PKE[384k :  384k + 32]
1995 *   ...
1996 *
1997 * @param [out] pub      Public key - vector.
1998 * @param [out] pubSeed  Public seed.
1999 * @param [in]  p        Public key data.
2000 * @param [in]  k        Number of polynomials in vector.
2001 */
2002static void mlkemkey_decode_public(sword16* pub, byte* pubSeed, const byte* p,
2003    unsigned int k)
2004{
2005    unsigned int i;
2006
2007    /* Decode public key that is vector of polynomials.
2008     * Step 2: t <- ByteDecode_12(ek_PKE[0 : 384k]) */
2009    mlkem_from_bytes(pub, p, (int)k);
2010    p += k * WC_ML_KEM_POLY_SIZE;
2011
2012    /* Read public key seed.
2013     * Step 3: rho <- ek_PKE[384k :  384k + 32] */
2014    for (i = 0; i < WC_ML_KEM_SYM_SZ; i++) {
2015        pubSeed[i] = p[i];
2016    }
2017}
2018
2019/**
2020 * Decode the private key.
2021 *
2022 * Private Vector | Public Key | Public Hash | Randomizer
2023 *
2024 * FIPS 203, Algorithm 18: ML-KEM.Decaps_internal(dk, c)
2025 *   1: dk_PKE <- dk[0 : 384k]
2026 *                        > extract (from KEM decaps key) the PKE decryption key
2027 *   2: ek_PKE <- dk[384k : 768k + 32]
2028 *                                                  > extract PKE encryption key
2029 *   3: h <- dk[768k + 32 : 768k + 64]
2030 *                                          > extract hash of PKE encryption key
2031 *   4: z <- dk[768k + 64 : 768k + 96]
2032 *                                            > extract implicit rejection value
2033 *
2034 * FIPS 203, Algorithm 15: K-PKE.Decrypt(dk_PKE, c)
2035 *   ...
2036 *   5: s_hat <- ByteDecode_12(dk_PKE)
2037 *   ...
2038 *
2039 * @param  [in, out]  key  ML-KEM key object.
2040 * @param  [in]       in   Buffer holding encoded key.
2041 * @param  [in]       len  Length of data in buffer.
2042 * @return  0 on success.
2043 * @return  BAD_FUNC_ARG when key or in is NULL.
2044 * @return  NOT_COMPILED_IN when key type is not supported.
2045 * @return  BUFFER_E when len is not the correct size.
2046 * @return  PUBLIC_KEY_E when public key data doesn't match parameters.
2047 * @return  MLKEM_PUB_HASH_E when public key hash doesn't match stored hash.
2048 * @return  MEMORY_E when dynamic memory allocation failed.
2049 */
2050int wc_MlKemKey_DecodePrivateKey(MlKemKey* key, const unsigned char* in,
2051    word32 len)
2052{
2053    int ret = 0;
2054    word32 privLen = 0;
2055    word32 pubLen = 0;
2056    unsigned int k = 0;
2057    const unsigned char* p = in;
2058
2059    /* Validate parameters. */
2060    if ((key == NULL) || (in == NULL)) {
2061        ret = BAD_FUNC_ARG;
2062    }
2063
2064    if (ret == 0) {
2065        /* Establish parameters based on key type. */
2066        switch (key->type) {
2067#ifndef WOLFSSL_NO_ML_KEM
2068    #ifdef WOLFSSL_WC_ML_KEM_512
2069        case WC_ML_KEM_512:
2070            k = WC_ML_KEM_512_K;
2071            privLen = WC_ML_KEM_512_PRIVATE_KEY_SIZE;
2072            pubLen = WC_ML_KEM_512_PUBLIC_KEY_SIZE;
2073            break;
2074    #endif
2075    #ifdef WOLFSSL_WC_ML_KEM_768
2076        case WC_ML_KEM_768:
2077            k = WC_ML_KEM_768_K;
2078            privLen = WC_ML_KEM_768_PRIVATE_KEY_SIZE;
2079            pubLen = WC_ML_KEM_768_PUBLIC_KEY_SIZE;
2080            break;
2081    #endif
2082    #ifdef WOLFSSL_WC_ML_KEM_1024
2083        case WC_ML_KEM_1024:
2084            k = WC_ML_KEM_1024_K;
2085            privLen = WC_ML_KEM_1024_PRIVATE_KEY_SIZE;
2086            pubLen = WC_ML_KEM_1024_PUBLIC_KEY_SIZE;
2087            break;
2088    #endif
2089#endif
2090#ifdef WOLFSSL_MLKEM_KYBER
2091    #ifdef WOLFSSL_KYBER512
2092        case KYBER512:
2093            k = KYBER512_K;
2094            privLen = KYBER512_PRIVATE_KEY_SIZE;
2095            pubLen = KYBER512_PUBLIC_KEY_SIZE;
2096            break;
2097    #endif
2098    #ifdef WOLFSSL_KYBER768
2099        case KYBER768:
2100            k = KYBER768_K;
2101            privLen = KYBER768_PRIVATE_KEY_SIZE;
2102            pubLen = KYBER768_PUBLIC_KEY_SIZE;
2103            break;
2104    #endif
2105    #ifdef WOLFSSL_KYBER1024
2106        case KYBER1024:
2107            k = KYBER1024_K;
2108            privLen = KYBER1024_PRIVATE_KEY_SIZE;
2109            pubLen = KYBER1024_PUBLIC_KEY_SIZE;
2110            break;
2111    #endif
2112#endif
2113        default:
2114            /* No other values supported. */
2115            ret = NOT_COMPILED_IN;
2116            break;
2117        }
2118    }
2119    /* Ensure the data is the correct length for the key type. */
2120    if ((ret == 0) && (len != privLen)) {
2121        ret = BUFFER_E;
2122    }
2123
2124#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
2125    if (ret == 0) {
2126        ret = mlkemkey_alloc_priv(key, k);
2127    }
2128    if (ret == 0) {
2129        ret = mlkemkey_alloc_pub(key, k);
2130    }
2131#endif
2132    if (ret == 0) {
2133        /* Decode private key that is vector of polynomials.
2134         * Alg 18 Step 1: dk_PKE <- dk[0 : 384k]
2135         * Alg 15 Step 5: s_hat <- ByteDecode_12(dk_PKE) */
2136        mlkem_from_bytes(key->priv, p, (int)k);
2137        p += k * WC_ML_KEM_POLY_SIZE;
2138
2139        /* Decode the public key that is after the private key. */
2140        mlkemkey_decode_public(key->pub, key->pubSeed, p, k);
2141        ret = mlkem_check_public(key->pub, (int)k);
2142        if (ret != 0) {
2143            ForceZero(key->priv, k * MLKEM_N * sizeof(sword16));
2144        }
2145    }
2146    if (ret == 0) {
2147        /* Compute the hash of the public key. */
2148        ret = MLKEM_HASH_H(&key->hash, p, pubLen, key->h);
2149        if (ret != 0) {
2150            ForceZero(key->priv, k * MLKEM_N * sizeof(sword16));
2151        }
2152    }
2153
2154    if (ret == 0) {
2155        p += pubLen;
2156        /* Compare computed public key hash with stored hash */
2157        if (XMEMCMP(key->h, p, WC_ML_KEM_SYM_SZ) != 0) {
2158            ForceZero(key->priv, k * MLKEM_N * sizeof(sword16));
2159            ret = MLKEM_PUB_HASH_E;
2160        }
2161    }
2162
2163    if (ret == 0) {
2164        /* Copy the hash of the encoded public key that is after public key. */
2165        XMEMCPY(key->h, p, sizeof(key->h));
2166        p += WC_ML_KEM_SYM_SZ;
2167        /* Copy the z (randomizer) that is after hash. */
2168        XMEMCPY(key->z, p, sizeof(key->z));
2169
2170        /* Set flags */
2171        key->flags |= MLKEM_FLAG_H_SET | MLKEM_FLAG_BOTH_SET;
2172    }
2173
2174    return ret;
2175}
2176
2177/**
2178 * Decode public key.
2179 *
2180 * Public vector | Public Seed
2181 *
2182 * @param  [in, out]  key  ML-KEM key object.
2183 * @param  [in]       in   Buffer holding encoded key.
2184 * @param  [in]       len  Length of data in buffer.
2185 * @return  0 on success.
2186 * @return  BAD_FUNC_ARG when key or in is NULL.
2187 * @return  NOT_COMPILED_IN when key type is not supported.
2188 * @return  BUFFER_E when len is not the correct size.
2189 * @return  PUBLIC_KEY_E when public key data doesn't match parameters.
2190 * @return  MEMORY_E when dynamic memory allocation failed.
2191 */
2192int wc_MlKemKey_DecodePublicKey(MlKemKey* key, const unsigned char* in,
2193    word32 len)
2194{
2195    int ret = 0;
2196    word32 pubLen = 0;
2197    unsigned int k = 0;
2198    const unsigned char* p = in;
2199
2200    if ((key == NULL) || (in == NULL)) {
2201        ret = BAD_FUNC_ARG;
2202    }
2203
2204    if (ret == 0) {
2205        /* Establish parameters based on key type. */
2206        switch (key->type) {
2207#ifndef WOLFSSL_NO_ML_KEM
2208    #ifdef WOLFSSL_WC_ML_KEM_512
2209        case WC_ML_KEM_512:
2210            k = WC_ML_KEM_512_K;
2211            pubLen = WC_ML_KEM_512_PUBLIC_KEY_SIZE;
2212            break;
2213    #endif
2214    #ifdef WOLFSSL_WC_ML_KEM_768
2215        case WC_ML_KEM_768:
2216            k = WC_ML_KEM_768_K;
2217            pubLen = WC_ML_KEM_768_PUBLIC_KEY_SIZE;
2218            break;
2219    #endif
2220    #ifdef WOLFSSL_WC_ML_KEM_1024
2221        case WC_ML_KEM_1024:
2222            k = WC_ML_KEM_1024_K;
2223            pubLen = WC_ML_KEM_1024_PUBLIC_KEY_SIZE;
2224            break;
2225    #endif
2226#endif
2227#ifdef WOLFSSL_MLKEM_KYBER
2228    #ifdef WOLFSSL_KYBER512
2229        case KYBER512:
2230            k = KYBER512_K;
2231            pubLen = KYBER512_PUBLIC_KEY_SIZE;
2232            break;
2233    #endif
2234    #ifdef WOLFSSL_KYBER768
2235        case KYBER768:
2236            k = KYBER768_K;
2237            pubLen = KYBER768_PUBLIC_KEY_SIZE;
2238            break;
2239    #endif
2240    #ifdef WOLFSSL_KYBER1024
2241        case KYBER1024:
2242            k = KYBER1024_K;
2243            pubLen = KYBER1024_PUBLIC_KEY_SIZE;
2244            break;
2245    #endif
2246#endif
2247        default:
2248            /* No other values supported. */
2249            ret = NOT_COMPILED_IN;
2250            break;
2251        }
2252    }
2253    /* Ensure the data is the correct length for the key type. */
2254    if ((ret == 0) && (len != pubLen)) {
2255        ret = BUFFER_E;
2256    }
2257
2258#ifdef WOLFSSL_MLKEM_DYNAMIC_KEYS
2259    if (ret == 0) {
2260        ret = mlkemkey_alloc_pub(key, k);
2261    }
2262#endif
2263    if (ret == 0) {
2264        /* Decode public key and check public key matches parameters. */
2265        mlkemkey_decode_public(key->pub, key->pubSeed, p, k);
2266        ret = mlkem_check_public(key->pub, (int)k);
2267    }
2268    if (ret == 0) {
2269        /* Calculate public hash. */
2270        ret = MLKEM_HASH_H(&key->hash, in, len, key->h);
2271    }
2272    if (ret == 0) {
2273        /* Record public key and public hash set. */
2274        key->flags |= MLKEM_FLAG_PUB_SET | MLKEM_FLAG_H_SET;
2275    }
2276
2277    return ret;
2278}
2279
2280/**
2281 * Get the size in bytes of encoded private key for the key.
2282 *
2283 * @param  [in]   key  ML-KEM key object.
2284 * @param  [out]  len  Length of encoded private key in bytes.
2285 * @return  0 on success.
2286 * @return  BAD_FUNC_ARG when key or len is NULL.
2287 * @return  NOT_COMPILED_IN when key type is not supported.
2288 */
2289int wc_MlKemKey_PrivateKeySize(MlKemKey* key, word32* len)
2290{
2291    int ret = 0;
2292
2293    /* Validate parameters. */
2294    if ((key == NULL) || (len == NULL)) {
2295        ret = BAD_FUNC_ARG;
2296    }
2297
2298    if (ret == 0) {
2299        /* Return in 'len' size of the encoded private key for the type of this
2300         * key. */
2301        switch (key->type) {
2302#ifndef WOLFSSL_NO_ML_KEM
2303    #ifdef WOLFSSL_WC_ML_KEM_512
2304        case WC_ML_KEM_512:
2305            *len = WC_ML_KEM_512_PRIVATE_KEY_SIZE;
2306            break;
2307    #endif
2308    #ifdef WOLFSSL_WC_ML_KEM_768
2309        case WC_ML_KEM_768:
2310            *len = WC_ML_KEM_768_PRIVATE_KEY_SIZE;
2311            break;
2312    #endif
2313    #ifdef WOLFSSL_WC_ML_KEM_1024
2314        case WC_ML_KEM_1024:
2315            *len = WC_ML_KEM_1024_PRIVATE_KEY_SIZE;
2316            break;
2317    #endif
2318#endif
2319#ifdef WOLFSSL_MLKEM_KYBER
2320    #ifdef WOLFSSL_KYBER512
2321        case KYBER512:
2322            *len = KYBER512_PRIVATE_KEY_SIZE;
2323            break;
2324    #endif
2325    #ifdef WOLFSSL_KYBER768
2326        case KYBER768:
2327            *len = KYBER768_PRIVATE_KEY_SIZE;
2328            break;
2329    #endif
2330    #ifdef WOLFSSL_KYBER1024
2331        case KYBER1024:
2332            *len = KYBER1024_PRIVATE_KEY_SIZE;
2333            break;
2334    #endif
2335#endif
2336        default:
2337            /* No other values supported. */
2338            ret = NOT_COMPILED_IN;
2339            break;
2340        }
2341    }
2342
2343    return ret;
2344}
2345
2346/**
2347 * Get the size in bytes of encoded public key for the key.
2348 *
2349 * @param  [in]   key  ML-KEM key object.
2350 * @param  [out]  len  Length of encoded public key in bytes.
2351 * @return  0 on success.
2352 * @return  BAD_FUNC_ARG when key or len is NULL.
2353 * @return  NOT_COMPILED_IN when key type is not supported.
2354 */
2355int wc_MlKemKey_PublicKeySize(MlKemKey* key, word32* len)
2356{
2357    int ret = 0;
2358
2359    /* Validate parameters. */
2360    if ((key == NULL) || (len == NULL)) {
2361        ret = BAD_FUNC_ARG;
2362    }
2363
2364    if (ret == 0) {
2365        /* Return in 'len' size of the encoded public key for the type of this
2366         * key. */
2367        switch (key->type) {
2368#ifndef WOLFSSL_NO_ML_KEM
2369    #ifdef WOLFSSL_WC_ML_KEM_512
2370        case WC_ML_KEM_512:
2371            *len = WC_ML_KEM_512_PUBLIC_KEY_SIZE;
2372            break;
2373    #endif
2374    #ifdef WOLFSSL_WC_ML_KEM_768
2375        case WC_ML_KEM_768:
2376            *len = WC_ML_KEM_768_PUBLIC_KEY_SIZE;
2377            break;
2378    #endif
2379    #ifdef WOLFSSL_WC_ML_KEM_1024
2380        case WC_ML_KEM_1024:
2381            *len = WC_ML_KEM_1024_PUBLIC_KEY_SIZE;
2382            break;
2383    #endif
2384#endif
2385#ifdef WOLFSSL_MLKEM_KYBER
2386    #ifdef WOLFSSL_KYBER512
2387        case KYBER512:
2388            *len = KYBER512_PUBLIC_KEY_SIZE;
2389            break;
2390    #endif
2391    #ifdef WOLFSSL_KYBER768
2392        case KYBER768:
2393            *len = KYBER768_PUBLIC_KEY_SIZE;
2394            break;
2395    #endif
2396    #ifdef WOLFSSL_KYBER1024
2397        case KYBER1024:
2398            *len = KYBER1024_PUBLIC_KEY_SIZE;
2399            break;
2400    #endif
2401#endif
2402        default:
2403            /* No other values supported. */
2404            ret = NOT_COMPILED_IN;
2405            break;
2406        }
2407    }
2408
2409    return ret;
2410}
2411
2412/**
2413 * Encode the private key.
2414 *
2415 * Private Vector | Public Key | Public Hash | Randomizer
2416 *
2417 * FIPS 203, Algorithm 16: ML-KEM.KeyGen_internal(d,z)
2418 *   ...
2419 *   3: dk <- (dk_PKE||ek||H(ek)||z)
2420 *   ...
2421 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
2422 *   ...
2423 *   20: dk_PKE  <- ByteEncode_12(s_hat)
2424 *   ...
2425 *
2426 * @param  [in]   key  ML-KEM key object.
2427 * @param  [out]  out  Buffer to hold data.
2428 * @param  [in]   len  Size of buffer in bytes.
2429 * @return  0 on success.
2430 * @return  BAD_FUNC_ARG when key or out is NULL.
2431 * @return  BAD_STATE_E when private/public key not available.
2432 * @return  NOT_COMPILED_IN when key type is not supported.
2433 */
2434int wc_MlKemKey_EncodePrivateKey(MlKemKey* key, unsigned char* out, word32 len)
2435{
2436    int ret = 0;
2437    unsigned int k = 0;
2438    unsigned int pubLen = 0;
2439    unsigned int privLen = 0;
2440    unsigned char* p = out;
2441
2442    if ((key == NULL) || (out == NULL)) {
2443        ret = BAD_FUNC_ARG;
2444    }
2445    if ((ret == 0) &&
2446            ((key->flags & MLKEM_FLAG_BOTH_SET) != MLKEM_FLAG_BOTH_SET)) {
2447        ret = BAD_STATE_E;
2448    }
2449
2450    if (ret == 0) {
2451        switch (key->type) {
2452#ifndef WOLFSSL_NO_ML_KEM
2453    #ifdef WOLFSSL_WC_ML_KEM_512
2454        case WC_ML_KEM_512:
2455            k = WC_ML_KEM_512_K;
2456            pubLen = WC_ML_KEM_512_PUBLIC_KEY_SIZE;
2457            privLen = WC_ML_KEM_512_PRIVATE_KEY_SIZE;
2458            break;
2459    #endif
2460    #ifdef WOLFSSL_WC_ML_KEM_768
2461        case WC_ML_KEM_768:
2462            k = WC_ML_KEM_768_K;
2463            pubLen = WC_ML_KEM_768_PUBLIC_KEY_SIZE;
2464            privLen = WC_ML_KEM_768_PRIVATE_KEY_SIZE;
2465            break;
2466    #endif
2467    #ifdef WOLFSSL_WC_ML_KEM_1024
2468        case WC_ML_KEM_1024:
2469            k = WC_ML_KEM_1024_K;
2470            pubLen = WC_ML_KEM_1024_PUBLIC_KEY_SIZE;
2471            privLen = WC_ML_KEM_1024_PRIVATE_KEY_SIZE;
2472            break;
2473    #endif
2474#endif
2475#ifdef WOLFSSL_MLKEM_KYBER
2476    #ifdef WOLFSSL_KYBER512
2477        case KYBER512:
2478            k = KYBER512_K;
2479            pubLen = KYBER512_PUBLIC_KEY_SIZE;
2480            privLen = KYBER512_PRIVATE_KEY_SIZE;
2481            break;
2482    #endif
2483    #ifdef WOLFSSL_KYBER768
2484        case KYBER768:
2485            k = KYBER768_K;
2486            pubLen = KYBER768_PUBLIC_KEY_SIZE;
2487            privLen = KYBER768_PRIVATE_KEY_SIZE;
2488            break;
2489    #endif
2490    #ifdef WOLFSSL_KYBER1024
2491        case KYBER1024:
2492            k = KYBER1024_K;
2493            pubLen = KYBER1024_PUBLIC_KEY_SIZE;
2494            privLen = KYBER1024_PRIVATE_KEY_SIZE;
2495            break;
2496    #endif
2497#endif
2498        default:
2499            /* No other values supported. */
2500            ret = NOT_COMPILED_IN;
2501            break;
2502        }
2503    }
2504    /* Check buffer is big enough for encoding. */
2505    if ((ret == 0) && (len != privLen)) {
2506        ret = BUFFER_E;
2507    }
2508
2509    if (ret == 0) {
2510        /* Encode private key that is vector of polynomials. */
2511        mlkem_to_bytes(p, key->priv, (int)k);
2512        p += WC_ML_KEM_POLY_SIZE * k;
2513
2514        /* Encode public key - calculates hash of public key. */
2515        ret = wc_MlKemKey_EncodePublicKey(key, p, pubLen);
2516        p += pubLen;
2517    }
2518    if (ret == 0) {
2519        /* Append public hash. */
2520        XMEMCPY(p, key->h, sizeof(key->h));
2521        p += WC_ML_KEM_SYM_SZ;
2522        /* Append z (randomizer). */
2523        XMEMCPY(p, key->z, sizeof(key->z));
2524    }
2525
2526    return ret;
2527}
2528
2529/**
2530 * Encode the public key.
2531 *
2532 * Public vector | Public Seed
2533 *
2534 * FIPS 203, Algorithm 16: ML-KEM.KeyGen_internal(d,z)
2535 *   ...
2536 *   2: ek <- ek_PKE
2537 *   ...
2538 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
2539 *   ...
2540 *   19: ek_PKE  <- ByteEncode_12(t_hat)||rho
2541 *   ...
2542 *
2543 * @param  [in]   key  ML-KEM key object.
2544 * @param  [out]  out  Buffer to hold data.
2545 * @param  [in]   len  Size of buffer in bytes.
2546 * @return  0 on success.
2547 * @return  BAD_FUNC_ARG when key or out is NULL.
2548 * @return  BAD_STATE_E when public key not available.
2549 * @return  NOT_COMPILED_IN when key type is not supported.
2550 */
2551int wc_MlKemKey_EncodePublicKey(MlKemKey* key, unsigned char* out, word32 len)
2552{
2553    int ret = 0;
2554    unsigned int k = 0;
2555    unsigned int pubLen = 0;
2556    unsigned char* p = out;
2557
2558    if ((key == NULL) || (out == NULL)) {
2559        ret = BAD_FUNC_ARG;
2560    }
2561    if ((ret == 0) &&
2562            ((key->flags & MLKEM_FLAG_PUB_SET) != MLKEM_FLAG_PUB_SET)) {
2563        ret = BAD_STATE_E;
2564    }
2565
2566    if (ret == 0) {
2567        switch (key->type) {
2568#ifndef WOLFSSL_NO_ML_KEM
2569    #ifdef WOLFSSL_WC_ML_KEM_512
2570        case WC_ML_KEM_512:
2571            k = WC_ML_KEM_512_K;
2572            pubLen = WC_ML_KEM_512_PUBLIC_KEY_SIZE;
2573            break;
2574    #endif
2575    #ifdef WOLFSSL_WC_ML_KEM_768
2576        case WC_ML_KEM_768:
2577            k = WC_ML_KEM_768_K;
2578            pubLen = WC_ML_KEM_768_PUBLIC_KEY_SIZE;
2579            break;
2580    #endif
2581    #ifdef WOLFSSL_WC_ML_KEM_1024
2582        case WC_ML_KEM_1024:
2583            k = WC_ML_KEM_1024_K;
2584            pubLen = WC_ML_KEM_1024_PUBLIC_KEY_SIZE;
2585            break;
2586    #endif
2587#endif
2588#ifdef WOLFSSL_MLKEM_KYBER
2589    #ifdef WOLFSSL_KYBER512
2590        case KYBER512:
2591            k = KYBER512_K;
2592            pubLen = KYBER512_PUBLIC_KEY_SIZE;
2593            break;
2594    #endif
2595    #ifdef WOLFSSL_KYBER768
2596        case KYBER768:
2597            k = KYBER768_K;
2598            pubLen = KYBER768_PUBLIC_KEY_SIZE;
2599            break;
2600    #endif
2601    #ifdef WOLFSSL_KYBER1024
2602        case KYBER1024:
2603            k = KYBER1024_K;
2604            pubLen = KYBER1024_PUBLIC_KEY_SIZE;
2605            break;
2606    #endif
2607#endif
2608        default:
2609            /* No other values supported. */
2610            ret = NOT_COMPILED_IN;
2611            break;
2612        }
2613    }
2614    /* Check buffer is big enough for encoding. */
2615    if ((ret == 0) && (len != pubLen)) {
2616        ret = BUFFER_E;
2617    }
2618
2619    if (ret == 0) {
2620        int i;
2621
2622        /* Encode public key polynomial by polynomial. */
2623        mlkem_to_bytes(p, key->pub, (int)k);
2624        p += k * WC_ML_KEM_POLY_SIZE;
2625
2626        /* Append public seed. */
2627        for (i = 0; i < WC_ML_KEM_SYM_SZ; i++) {
2628            p[i] = key->pubSeed[i];
2629        }
2630
2631        /* Make sure public hash is set. */
2632        if ((key->flags & MLKEM_FLAG_H_SET) == 0) {
2633            ret = MLKEM_HASH_H(&key->hash, out, len, key->h);
2634        }
2635    }
2636    if (ret == 0) {
2637        /* Public hash is set. */
2638        key->flags |= MLKEM_FLAG_H_SET;
2639    }
2640
2641    return ret;
2642}
2643
2644#endif /* WOLFSSL_HAVE_MLKEM */