/* wc_lms.c
 *
 * Copyright (C) 2006-2026 wolfSSL Inc.
 *
 * This file is part of wolfSSL.
 *
 * wolfSSL is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * wolfSSL is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
 */

#include <wolfssl/wolfcrypt/libwolfssl_sources.h>

#if defined(WOLFSSL_HAVE_LMS)

#if FIPS_VERSION3_GE(2,0,0)
    /* set NO_WRAPPERS before headers, use direct internal f()s not wrappers */
    #define FIPS_NO_WRAPPERS
#endif
#include <wolfssl/wolfcrypt/wc_lms.h>
#include <wolfssl/wolfcrypt/hash.h>

#ifdef NO_INLINE
    #include <wolfssl/wolfcrypt/misc.h>
#else
    #define WOLFSSL_MISC_INCLUDED
    #include <wolfcrypt/src/misc.c>
#endif

#ifdef WOLF_CRYPTO_CB
    #include <wolfssl/wolfcrypt/cryptocb.h>
#endif


/* Calculate u. Appendix B. Works for w of 1, 2, 4, or 8.
 *
 * @param [in] w  Winternitz width.
 */
#define LMS_U(w, hLen)              \
    (8 * (hLen) / (w))
/* Calculate u. Appendix B. Works for w of 1, 2, 4, or 8.
 *
 * @param [in] w   Winternitz width.
 * @param [in] wb  Winternitz width length in bits.
 */
#define LMS_V(w, wb)                \
    (2 + (8 - (wb)) / (w))
/* Calculate ls. Appendix B. Works for w of 1, 2, 4, or 8.
 *
 * @param [in] w   Winternitz width.
 * @param [in] wb  Winternitz width length in bits.
 */
#define LMS_LS(w, wb)               \
    (16 - LMS_V(w, wb) * (w))
/* Calculate p. Appendix B. Works for w of 1, 2, 4, or 8.
 *
 * @param [in] w   Winternitz width.
 * @param [in] wb  Winternitz width length in bits.
 */
#define LMS_P(w, wb, hLen)          \
    (LMS_U(w, hLen) + LMS_V(w, wb))
/* Calculate signature length.
 *
 * @param [in] l  Number of levels.
 * @param [in] h  Height of the trees.
 * @param [in] p  Number of n-byte string elements in signature for a tree.
 */
#define LMS_PARAMS_SIG_LEN(l, h, p, hLen)               \
    (4 + (l) * (4 + 4 + 4 + (hLen) * (1 + (p) + (h))) + \
     ((l) - 1) * LMS_PUBKEY_LEN(hLen))

#ifndef WOLFSSL_WC_LMS_SMALL
    /* Root levels and leaf cache bits. */
    #define LMS_PARAMS_CACHE(h)                             \
        (((h) < LMS_ROOT_LEVELS) ? (h) : LMS_ROOT_LEVELS),  \
        (((h) < LMS_CACHE_BITS ) ? (h) : LMS_CACHE_BITS )
#else
    /* Root levels and leaf cache bits aren't in structure. */
    #define LMS_PARAMS_CACHE(h) /* null expansion */
#endif

/* Define parameters entry for LMS.
 *
 * @param [in] l   Number of levels.
 * @param [in] h   Height of the trees.
 * @param [in] w   Winternitz width.
 * @param [in] wb  Winternitz width length in bits.
 * @param [in] t   LMS type.
 * @param [in] t2  LM-OTS type.
 */
#define LMS_PARAMS(l, h, w, wb, t, t2, hLen)                        \
    { l, h, w, LMS_LS(w, wb), LMS_P(w, wb, hLen), t, t2,            \
      (hLen), LMS_PARAMS_SIG_LEN(l, h, LMS_P(w, wb, hLen), hLen),   \
      LMS_PARAMS_CACHE(h) }


/* Initialize the working state for LMS operations.
 *
 * @param [in, out] state   LMS state.
 * @param [in]      params  LMS parameters.
 */
static int wc_lmskey_state_init(LmsState* state, const LmsParams* params)
{
    int ret;

    /* Zero out every field. */
    XMEMSET(state, 0, sizeof(LmsState));

    /* Keep a reference to the parameters for use in operations. */
    state->params = params;

#ifdef WOLFSSL_LMS_SHAKE256
    if (LMS_IS_SHAKE(params->lmOtsType)) {
        ret = wc_InitShake256(LMS_STATE_SHAKE(state), NULL, INVALID_DEVID);
        if (ret == 0) {
            ret = wc_InitShake256(LMS_STATE_SHAKE_K(state), NULL, INVALID_DEVID);
            if (ret != 0) {
                wc_Shake256_Free(LMS_STATE_SHAKE(state));
            }
        }
        return ret;
    }
#endif

    /* Initialize the two hash algorithms. */
    ret = wc_InitSha256(LMS_STATE_HASH(state));
    if (ret == 0) {
        ret = wc_InitSha256(LMS_STATE_HASH_K(state));
        if (ret != 0) {
            wc_Sha256Free(LMS_STATE_HASH(state));
        }
    }

    return ret;
}

/* Free the working state for LMS operations.
 *
 * @param [in] state  LMS state.
 */
static void wc_lmskey_state_free(LmsState* state)
{
#ifdef WOLFSSL_LMS_SHAKE256
    if (LMS_IS_SHAKE(state->params->lmOtsType)) {
        wc_Shake256_Free(LMS_STATE_SHAKE_K(state));
        wc_Shake256_Free(LMS_STATE_SHAKE(state));
        return;
    }
#endif
    wc_Sha256Free(LMS_STATE_HASH_K(state));
    wc_Sha256Free(LMS_STATE_HASH(state));
}

/* Supported LMS parameters. */
static const wc_LmsParamsMap wc_lms_map[] = {
#ifndef WOLFSSL_NO_LMS_SHA256_256
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_L1_H15_W2, "LMS/HSS L1_H15_W2",
      LMS_PARAMS(1, 15, 2, 1, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H15_W4, "LMS/HSS L1_H15_W4",
      LMS_PARAMS(1, 15, 4, 2, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_LEVELS >= 2
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L2_H10_W2, "LMS/HSS L2_H10_W2",
      LMS_PARAMS(2, 10, 2, 1, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H10_W4, "LMS/HSS L2_H10_W4",
      LMS_PARAMS(2, 10, 4, 2, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H10_W8, "LMS/HSS L2_H10_W8",
      LMS_PARAMS(2, 10, 8, 3, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 3
    { WC_LMS_PARM_L3_H5_W2 , "LMS/HSS L3_H5_W2" ,
      LMS_PARAMS(3,  5, 2, 1, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L3_H5_W4 , "LMS/HSS L3_H5_W4" ,
      LMS_PARAMS(3,  5, 4, 2, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L3_H5_W8 , "LMS/HSS L3_H5_W8" ,
      LMS_PARAMS(3,  5, 8, 3, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L3_H10_W4, "LMS/HSS L3_H10_W4",
      LMS_PARAMS(3, 10, 4, 2, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 4
    { WC_LMS_PARM_L4_H5_W8 , "LMS/HSS L4_H5_W8" ,
      LMS_PARAMS(4,  5, 8, 3, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif

    /* For when user sets L, H, W explicitly. */
    { WC_LMS_PARM_L1_H5_W1 , "LMS/HSS_L1_H5_W1" ,
      LMS_PARAMS(1,  5, 1, 1, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H5_W2 , "LMS/HSS_L1_H5_W2" ,
      LMS_PARAMS(1,  5, 2, 1, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H5_W4 , "LMS/HSS_L1_H5_W4" ,
      LMS_PARAMS(1,  5, 4, 2, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H5_W8 , "LMS/HSS_L1_H5_W8" ,
      LMS_PARAMS(1,  5, 8, 3, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L1_H10_W2 , "LMS/HSS_L1_H10_W2",
      LMS_PARAMS(1, 10, 2, 1, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H10_W4 , "LMS/HSS_L1_H10_W4",
      LMS_PARAMS(1, 10, 4, 2, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H10_W8 , "LMS/HSS_L1_H10_W8",
      LMS_PARAMS(1, 10, 8, 3, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_L1_H15_W8 , "LMS/HSS L1_H15_W8",
      LMS_PARAMS(1, 15, 8, 3, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    { WC_LMS_PARM_L1_H20_W2 , "LMS/HSS_L1_H20_W2",
      LMS_PARAMS(1, 20, 2, 1, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H20_W4 , "LMS/HSS_L1_H20_W4",
      LMS_PARAMS(1, 20, 4, 2, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H20_W8 , "LMS/HSS_L1_H20_W8",
      LMS_PARAMS(1, 20, 8, 3, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_LEVELS >= 2
    { WC_LMS_PARM_L2_H5_W2 , "LMS/HSS_L2_H5_W2" ,
      LMS_PARAMS(2,  5, 2, 1, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H5_W4 , "LMS/HSS_L2_H5_W4" ,
      LMS_PARAMS(2,  5, 4, 2, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H5_W8 , "LMS/HSS_L2_H5_W8" ,
      LMS_PARAMS(2,  5, 8, 3, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_L2_H15_W2 , "LMS/HSS_L2_H15_W2",
      LMS_PARAMS(2, 15, 2, 1, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H15_W4 , "LMS/HSS_L2_H15_W4",
      LMS_PARAMS(2, 15, 4, 2, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H15_W8 , "LMS/HSS_L2_H15_W8",
      LMS_PARAMS(2, 15, 8, 3, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    { WC_LMS_PARM_L2_H20_W2 , "LMS/HSS_L2_H20_W2",
      LMS_PARAMS(2, 20, 2, 1, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H20_W4 , "LMS/HSS_L2_H20_W4",
      LMS_PARAMS(2, 20, 4, 2, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L2_H20_W8 , "LMS/HSS_L2_H20_W8",
      LMS_PARAMS(2, 20, 8, 3, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 3
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L3_H10_W8 , "LMS/HSS L3_H10_W8",
      LMS_PARAMS(3, 10, 8, 3, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 4
    { WC_LMS_PARM_L4_H5_W2 , "LMS/HSS L4_H5_W2" ,
      LMS_PARAMS(4,  5, 2, 1, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L4_H5_W4 , "LMS/HSS L4_H5_W4" ,
      LMS_PARAMS(4,  5, 4, 2, LMS_SHA256_M32_H5 , LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L4_H10_W4 , "LMS/HSS L4_H10_W4",
      LMS_PARAMS(4, 10, 4, 2, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L4_H10_W8 , "LMS/HSS L4_H10_W8",
      LMS_PARAMS(4, 10, 8, 3, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_HEIGHT >= 25
    { WC_LMS_PARM_L1_H25_W1 , "LMS/HSS_L1_H25_W1",
      LMS_PARAMS(1, 25, 1, 1, LMS_SHA256_M32_H25, LMOTS_SHA256_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H25_W2 , "LMS/HSS_L1_H25_W2",
      LMS_PARAMS(1, 25, 2, 1, LMS_SHA256_M32_H25, LMOTS_SHA256_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H25_W4 , "LMS/HSS_L1_H25_W4",
      LMS_PARAMS(1, 25, 4, 2, LMS_SHA256_M32_H25, LMOTS_SHA256_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_L1_H25_W8 , "LMS/HSS_L1_H25_W8",
      LMS_PARAMS(1, 25, 8, 3, LMS_SHA256_M32_H25, LMOTS_SHA256_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_L1_H10_W1 , "LMS/HSS_L1_H10_W1",
      LMS_PARAMS(1, 10, 1, 1, LMS_SHA256_M32_H10, LMOTS_SHA256_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_L1_H15_W1 , "LMS/HSS_L1_H15_W1",
      LMS_PARAMS(1, 15, 1, 1, LMS_SHA256_M32_H15, LMOTS_SHA256_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    { WC_LMS_PARM_L1_H20_W1 , "LMS/HSS_L1_H20_W1",
      LMS_PARAMS(1, 20, 1, 1, LMS_SHA256_M32_H20, LMOTS_SHA256_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif /* !WOLFSSL_NO_LMS_SHA256_256 */

#ifdef WOLFSSL_LMS_SHA256_192
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_SHA256_192_L1_H15_W2, "LMS/HSS_SHA256/192 L1_H15_W2",
      LMS_PARAMS(1, 15, 2, 2, LMS_SHA256_M24_H15, LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H15_W4, "LMS/HSS_SHA256/192 L1_H15_W4",
      LMS_PARAMS(1, 15, 4, 3, LMS_SHA256_M24_H15, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_LEVELS >= 2
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_SHA256_192_L2_H10_W2, "LMS/HSS SHA256/192 L2_H10_W2",
      LMS_PARAMS(2, 10, 2, 2, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L2_H10_W4, "LMS/HSS SHA256/192 L2_H10_W4",
      LMS_PARAMS(2, 10, 4, 3, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L2_H10_W8, "LMS/HSS SHA256/192 L2_H10_W8",
      LMS_PARAMS(2, 10, 8, 4, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 3
    { WC_LMS_PARM_SHA256_192_L3_H5_W2 , "LMS/HSS_SHA256/192 L3_H5_W2" ,
      LMS_PARAMS(3,  5, 2, 2, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L3_H5_W4 , "LMS/HSS_SHA256/192 L3_H5_W4" ,
      LMS_PARAMS(3,  5, 4, 3, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L3_H5_W8 , "LMS/HSS_SHA256/192 L3_H5_W8" ,
      LMS_PARAMS(3,  5, 8, 4, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_SHA256_192_L3_H10_W4, "LMS/HSS_SHA256/192 L3_H10_W4",
      LMS_PARAMS(3, 10, 4, 3, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#endif
#if LMS_MAX_LEVELS >= 4
    { WC_LMS_PARM_SHA256_192_L4_H5_W8 , "LMS/HSS_SHA256/192 L4_H5_W8" ,
      LMS_PARAMS(4,  5, 8, 4, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif

    { WC_LMS_PARM_SHA256_192_L1_H5_W1 , "LMS/HSS_SHA256/192_L1_H5_W1" ,
      LMS_PARAMS(1,  5, 1, 2, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H5_W2 , "LMS/HSS_SHA256/192_L1_H5_W2" ,
      LMS_PARAMS(1,  5, 2, 2, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H5_W4 , "LMS/HSS_SHA256/192_L1_H5_W4" ,
      LMS_PARAMS(1,  5, 4, 3, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H5_W8 , "LMS/HSS_SHA256/192_L1_H5_W8" ,
      LMS_PARAMS(1,  5, 8, 4, LMS_SHA256_M24_H5 , LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_SHA256_192_L1_H10_W2 , "LMS/HSS_SHA256/192_L1_H10_W2",
      LMS_PARAMS(1, 10, 2, 2, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H10_W4 , "LMS/HSS_SHA256/192_L1_H10_W4",
      LMS_PARAMS(1, 10, 4, 3, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H10_W8 , "LMS/HSS_SHA256/192_L1_H10_W8",
      LMS_PARAMS(1, 10, 8, 4, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    { WC_LMS_PARM_SHA256_192_L1_H20_W2 , "LMS/HSS_SHA256/192_L1_H20_W2",
      LMS_PARAMS(1, 20, 2, 2, LMS_SHA256_M24_H20, LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H20_W4 , "LMS/HSS_SHA256/192_L1_H20_W4",
      LMS_PARAMS(1, 20, 4, 3, LMS_SHA256_M24_H20, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H20_W8 , "LMS/HSS_SHA256/192_L1_H20_W8",
      LMS_PARAMS(1, 20, 8, 4, LMS_SHA256_M24_H20, LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 25
    { WC_LMS_PARM_SHA256_192_L1_H25_W1 , "LMS/HSS_SHA256/192_L1_H25_W1",
      LMS_PARAMS(1, 25, 1, 2, LMS_SHA256_M24_H25, LMOTS_SHA256_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H25_W2 , "LMS/HSS_SHA256/192_L1_H25_W2",
      LMS_PARAMS(1, 25, 2, 2, LMS_SHA256_M24_H25, LMOTS_SHA256_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H25_W4 , "LMS/HSS_SHA256/192_L1_H25_W4",
      LMS_PARAMS(1, 25, 4, 3, LMS_SHA256_M24_H25, LMOTS_SHA256_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H25_W8 , "LMS/HSS_SHA256/192_L1_H25_W8",
      LMS_PARAMS(1, 25, 8, 4, LMS_SHA256_M24_H25, LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 10
    { WC_LMS_PARM_SHA256_192_L1_H10_W1 , "LMS/HSS_SHA256/192_L1_H10_W1",
      LMS_PARAMS(1, 10, 1, 2, LMS_SHA256_M24_H10, LMOTS_SHA256_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 15
    { WC_LMS_PARM_SHA256_192_L1_H15_W1 , "LMS/HSS_SHA256/192_L1_H15_W1",
      LMS_PARAMS(1, 15, 1, 2, LMS_SHA256_M24_H15, LMOTS_SHA256_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHA256_192_L1_H15_W8 , "LMS/HSS_SHA256/192_L1_H15_W8",
      LMS_PARAMS(1, 15, 8, 4, LMS_SHA256_M24_H15, LMOTS_SHA256_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    { WC_LMS_PARM_SHA256_192_L1_H20_W1 , "LMS/HSS_SHA256/192_L1_H20_W1",
      LMS_PARAMS(1, 20, 1, 2, LMS_SHA256_M24_H20, LMOTS_SHA256_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#endif /* WOLFSSL_LMS_SHA256_192 */

#ifdef WOLFSSL_LMS_SHAKE256
#ifndef WOLFSSL_NO_LMS_SHAKE256_256
    /* SHAKE256/256 L1 H5 */
    { WC_LMS_PARM_SHAKE_L1_H5_W1 , "LMS/HSS_SHAKE256/256_L1_H5_W1",
      LMS_PARAMS(1,  5, 1, 1, LMS_SHAKE_M32_H5 , LMOTS_SHAKE_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H5_W2 , "LMS/HSS_SHAKE256/256_L1_H5_W2",
      LMS_PARAMS(1,  5, 2, 1, LMS_SHAKE_M32_H5 , LMOTS_SHAKE_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H5_W4 , "LMS/HSS_SHAKE256/256_L1_H5_W4",
      LMS_PARAMS(1,  5, 4, 2, LMS_SHAKE_M32_H5 , LMOTS_SHAKE_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H5_W8 , "LMS/HSS_SHAKE256/256_L1_H5_W8",
      LMS_PARAMS(1,  5, 8, 3, LMS_SHAKE_M32_H5 , LMOTS_SHAKE_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    /* SHAKE256/256 L1 H10 */
    { WC_LMS_PARM_SHAKE_L1_H10_W1 , "LMS/HSS_SHAKE256/256_L1_H10_W1",
      LMS_PARAMS(1, 10, 1, 1, LMS_SHAKE_M32_H10, LMOTS_SHAKE_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H10_W2 , "LMS/HSS_SHAKE256/256_L1_H10_W2",
      LMS_PARAMS(1, 10, 2, 1, LMS_SHAKE_M32_H10, LMOTS_SHAKE_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H10_W4 , "LMS/HSS_SHAKE256/256_L1_H10_W4",
      LMS_PARAMS(1, 10, 4, 2, LMS_SHAKE_M32_H10, LMOTS_SHAKE_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H10_W8 , "LMS/HSS_SHAKE256/256_L1_H10_W8",
      LMS_PARAMS(1, 10, 8, 3, LMS_SHAKE_M32_H10, LMOTS_SHAKE_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 15
    /* SHAKE256/256 L1 H15 */
    { WC_LMS_PARM_SHAKE_L1_H15_W1 , "LMS/HSS_SHAKE256/256_L1_H15_W1",
      LMS_PARAMS(1, 15, 1, 1, LMS_SHAKE_M32_H15, LMOTS_SHAKE_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H15_W2 , "LMS/HSS_SHAKE256/256_L1_H15_W2",
      LMS_PARAMS(1, 15, 2, 1, LMS_SHAKE_M32_H15, LMOTS_SHAKE_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H15_W4 , "LMS/HSS_SHAKE256/256_L1_H15_W4",
      LMS_PARAMS(1, 15, 4, 2, LMS_SHAKE_M32_H15, LMOTS_SHAKE_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H15_W8 , "LMS/HSS_SHAKE256/256_L1_H15_W8",
      LMS_PARAMS(1, 15, 8, 3, LMS_SHAKE_M32_H15, LMOTS_SHAKE_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    /* SHAKE256/256 L1 H20 */
    { WC_LMS_PARM_SHAKE_L1_H20_W1 , "LMS/HSS_SHAKE256/256_L1_H20_W1",
      LMS_PARAMS(1, 20, 1, 1, LMS_SHAKE_M32_H20, LMOTS_SHAKE_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H20_W2 , "LMS/HSS_SHAKE256/256_L1_H20_W2",
      LMS_PARAMS(1, 20, 2, 1, LMS_SHAKE_M32_H20, LMOTS_SHAKE_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H20_W4 , "LMS/HSS_SHAKE256/256_L1_H20_W4",
      LMS_PARAMS(1, 20, 4, 2, LMS_SHAKE_M32_H20, LMOTS_SHAKE_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H20_W8 , "LMS/HSS_SHAKE256/256_L1_H20_W8",
      LMS_PARAMS(1, 20, 8, 3, LMS_SHAKE_M32_H20, LMOTS_SHAKE_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 25
    /* SHAKE256/256 L1 H25 */
    { WC_LMS_PARM_SHAKE_L1_H25_W1 , "LMS/HSS_SHAKE256/256_L1_H25_W1",
      LMS_PARAMS(1, 25, 1, 1, LMS_SHAKE_M32_H25, LMOTS_SHAKE_N32_W1,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H25_W2 , "LMS/HSS_SHAKE256/256_L1_H25_W2",
      LMS_PARAMS(1, 25, 2, 1, LMS_SHAKE_M32_H25, LMOTS_SHAKE_N32_W2,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H25_W4 , "LMS/HSS_SHAKE256/256_L1_H25_W4",
      LMS_PARAMS(1, 25, 4, 2, LMS_SHAKE_M32_H25, LMOTS_SHAKE_N32_W4,
                 WC_SHA256_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE_L1_H25_W8 , "LMS/HSS_SHAKE256/256_L1_H25_W8",
      LMS_PARAMS(1, 25, 8, 3, LMS_SHAKE_M32_H25, LMOTS_SHAKE_N32_W8,
                 WC_SHA256_DIGEST_SIZE) },
#endif
#endif /* !WOLFSSL_NO_LMS_SHAKE256_256 */

    /* SHAKE256/192 L1 H5 */
    { WC_LMS_PARM_SHAKE192_L1_H5_W1 , "LMS/HSS_SHAKE256/192_L1_H5_W1",
      LMS_PARAMS(1,  5, 1, 2, LMS_SHAKE_M24_H5 , LMOTS_SHAKE_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H5_W2 , "LMS/HSS_SHAKE256/192_L1_H5_W2",
      LMS_PARAMS(1,  5, 2, 2, LMS_SHAKE_M24_H5 , LMOTS_SHAKE_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H5_W4 , "LMS/HSS_SHAKE256/192_L1_H5_W4",
      LMS_PARAMS(1,  5, 4, 3, LMS_SHAKE_M24_H5 , LMOTS_SHAKE_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H5_W8 , "LMS/HSS_SHAKE256/192_L1_H5_W8",
      LMS_PARAMS(1,  5, 8, 4, LMS_SHAKE_M24_H5 , LMOTS_SHAKE_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#if LMS_MAX_HEIGHT >= 10
    /* SHAKE256/192 L1 H10 */
    { WC_LMS_PARM_SHAKE192_L1_H10_W1 , "LMS/HSS_SHAKE256/192_L1_H10_W1",
      LMS_PARAMS(1, 10, 1, 2, LMS_SHAKE_M24_H10, LMOTS_SHAKE_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H10_W2 , "LMS/HSS_SHAKE256/192_L1_H10_W2",
      LMS_PARAMS(1, 10, 2, 2, LMS_SHAKE_M24_H10, LMOTS_SHAKE_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H10_W4 , "LMS/HSS_SHAKE256/192_L1_H10_W4",
      LMS_PARAMS(1, 10, 4, 3, LMS_SHAKE_M24_H10, LMOTS_SHAKE_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H10_W8 , "LMS/HSS_SHAKE256/192_L1_H10_W8",
      LMS_PARAMS(1, 10, 8, 4, LMS_SHAKE_M24_H10, LMOTS_SHAKE_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 15
    /* SHAKE256/192 L1 H15 */
    { WC_LMS_PARM_SHAKE192_L1_H15_W1 , "LMS/HSS_SHAKE256/192_L1_H15_W1",
      LMS_PARAMS(1, 15, 1, 2, LMS_SHAKE_M24_H15, LMOTS_SHAKE_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H15_W2 , "LMS/HSS_SHAKE256/192_L1_H15_W2",
      LMS_PARAMS(1, 15, 2, 2, LMS_SHAKE_M24_H15, LMOTS_SHAKE_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H15_W4 , "LMS/HSS_SHAKE256/192_L1_H15_W4",
      LMS_PARAMS(1, 15, 4, 3, LMS_SHAKE_M24_H15, LMOTS_SHAKE_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H15_W8 , "LMS/HSS_SHAKE256/192_L1_H15_W8",
      LMS_PARAMS(1, 15, 8, 4, LMS_SHAKE_M24_H15, LMOTS_SHAKE_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 20
    /* SHAKE256/192 L1 H20 */
    { WC_LMS_PARM_SHAKE192_L1_H20_W1 , "LMS/HSS_SHAKE256/192_L1_H20_W1",
      LMS_PARAMS(1, 20, 1, 2, LMS_SHAKE_M24_H20, LMOTS_SHAKE_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H20_W2 , "LMS/HSS_SHAKE256/192_L1_H20_W2",
      LMS_PARAMS(1, 20, 2, 2, LMS_SHAKE_M24_H20, LMOTS_SHAKE_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H20_W4 , "LMS/HSS_SHAKE256/192_L1_H20_W4",
      LMS_PARAMS(1, 20, 4, 3, LMS_SHAKE_M24_H20, LMOTS_SHAKE_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H20_W8 , "LMS/HSS_SHAKE256/192_L1_H20_W8",
      LMS_PARAMS(1, 20, 8, 4, LMS_SHAKE_M24_H20, LMOTS_SHAKE_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#if LMS_MAX_HEIGHT >= 25
    /* SHAKE256/192 L1 H25 */
    { WC_LMS_PARM_SHAKE192_L1_H25_W1 , "LMS/HSS_SHAKE256/192_L1_H25_W1",
      LMS_PARAMS(1, 25, 1, 2, LMS_SHAKE_M24_H25, LMOTS_SHAKE_N24_W1,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H25_W2 , "LMS/HSS_SHAKE256/192_L1_H25_W2",
      LMS_PARAMS(1, 25, 2, 2, LMS_SHAKE_M24_H25, LMOTS_SHAKE_N24_W2,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H25_W4 , "LMS/HSS_SHAKE256/192_L1_H25_W4",
      LMS_PARAMS(1, 25, 4, 3, LMS_SHAKE_M24_H25, LMOTS_SHAKE_N24_W4,
                 WC_SHA256_192_DIGEST_SIZE) },
    { WC_LMS_PARM_SHAKE192_L1_H25_W8 , "LMS/HSS_SHAKE256/192_L1_H25_W8",
      LMS_PARAMS(1, 25, 8, 4, LMS_SHAKE_M24_H25, LMOTS_SHAKE_N24_W8,
                 WC_SHA256_192_DIGEST_SIZE) },
#endif
#endif /* WOLFSSL_LMS_SHAKE256 */
};
/* Number of parameter sets supported. */
#define WC_LMS_MAP_LEN      ((int)(sizeof(wc_lms_map) / sizeof(*wc_lms_map)))

/* Initialize LMS key.
 *
 * Call this before setting the params of an LMS key.
 * Must call wc_LmsKey_Free before calling this function again.
 *
 * @param [out] key    LMS key to initialize.
 * @param [in]  heap   Heap hint.
 * @param [in]  devId  Device identifier.
 *                     Use INVALID_DEVID when not using a device.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key is NULL.
 */
int wc_LmsKey_Init(LmsKey* key, void* heap, int devId)
{
    int ret = 0;

    (void)heap;
    (void)devId;

    /* Validate parameters. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }
    if (ret == 0) {
        /* Clear the key data. */
        XMEMSET(key, 0, sizeof(LmsKey));

    #ifndef WOLFSSL_LMS_VERIFY_ONLY
        /* Initialize other fields. */
        key->heap = heap;
    #endif
    #ifdef WOLF_CRYPTO_CB
        key->devId = devId;
    #endif
        /* Start in initialized state. */
        key->state = WC_LMS_STATE_INITED;
    }

    return ret;
}

#ifdef WOLF_PRIVATE_KEY_ID
/* Initialize an LmsKey and bind it to a device-side key identifier.
 *
 * @param [in,out] key    LmsKey to initialize.
 * @param [in]     id     Identifier bytes (may be NULL when len is 0).
 * @param [in]     len    Length of id; must be in [0, LMS_MAX_ID_LEN].
 * @param [in]     heap   Heap hint forwarded to wc_LmsKey_Init.
 * @param [in]     devId  Device identifier.
 *
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key is NULL, or id is NULL and len is not 0.
 * @return  BUFFER_E when len is negative or exceeds LMS_MAX_ID_LEN.
 */
int wc_LmsKey_InitId(LmsKey* key, const unsigned char* id, int len, void* heap,
    int devId)
{
    int ret = 0;

    if ((key == NULL) || ((id == NULL) && (len != 0))) {
        ret = BAD_FUNC_ARG;
    }
    if ((ret == 0) && ((len < 0) || (len > LMS_MAX_ID_LEN))) {
        ret = BUFFER_E;
    }
    if (ret == 0) {
        ret = wc_LmsKey_Init(key, heap, devId);
    }
    if ((ret == 0) && (id != NULL) && (len != 0)) {
        XMEMCPY(key->id, id, (size_t)len);
        key->idLen = len;
    }

    return ret;
}

/* Initialize an LmsKey and bind it to a device-side key label.
 *
 * @param [in,out] key    LmsKey to initialize.
 * @param [in]     label  NUL-terminated label string (must be non-empty).
 * @param [in]     heap   Heap hint forwarded to wc_LmsKey_Init.
 * @param [in]     devId  Device identifier.
 *
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or label is NULL.
 * @return  BUFFER_E when label is empty or longer than LMS_MAX_LABEL_LEN.
 */
int wc_LmsKey_InitLabel(LmsKey* key, const char* label, void* heap, int devId)
{
    int ret = 0;
    size_t labelLen = 0;

    if ((key == NULL) || (label == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    if (ret == 0) {
        labelLen = XSTRLEN(label);
        if ((labelLen == 0) || (labelLen > LMS_MAX_LABEL_LEN)) {
            ret = BUFFER_E;
        }
    }
    if (ret == 0) {
        ret = wc_LmsKey_Init(key, heap, devId);
    }
    if (ret == 0) {
        XMEMCPY(key->label, label, labelLen);
        key->labelLen = (int)labelLen;
    }

    return ret;
}
#endif /* WOLF_PRIVATE_KEY_ID */

/* Get the string representation of the LMS parameter set.
 *
 * @param [in] lmsParm  LMS parameter set identifier.
 * @return  String representing LMS parameter set on success.
 * @return  NULL when parameter set not supported.
 */
const char* wc_LmsKey_ParmToStr(enum wc_LmsParm lmsParm)
{
    const char* str = NULL;
    int i;

    /* Search through table for matching numeric identifier. */
    for (i = 0; i < WC_LMS_MAP_LEN; i++) {
        if (lmsParm == wc_lms_map[i].id) {
            /* Get string corresponding to numeric identifier. */
            str = wc_lms_map[i].str;
            break;
        }
    }

    /* Return the string or NULL. */
    return str;
}

/* Set the wc_LmsParm of an LMS key.
 *
 * Use this if you wish to set a key with a predefined parameter set,
 * such as WC_LMS_PARM_L2_H10_W8.
 *
 * Key must be inited before calling this.
 *
 * @param [in, out] key      LMS key to set parameters on.
 * @param [in]      lmsParm  Identifier of parameters.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key is NULL.
 * @return  BAD_FUNC_ARG when parameters not supported.
 */
int wc_LmsKey_SetLmsParm(LmsKey* key, enum wc_LmsParm lmsParm)
{
    int ret = 0;

    /* Validate parameters. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }

    /* Check state is valid. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_INITED)) {
        WOLFSSL_MSG("error: LmsKey needs init");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        int i;

        ret = BAD_FUNC_ARG;
        /* Search through table for matching numeric identifier. */
        for (i = 0; i < WC_LMS_MAP_LEN; i++) {
            if (lmsParm == wc_lms_map[i].id) {
                /* Set the parameters into the key. */
                key->params = &wc_lms_map[i].params;
                ret = 0;
                break;
            }
        }
    }

    if (ret == 0) {
        /* Move the state to params set.
         * Key is ready for MakeKey or Reload. */
        key->state = WC_LMS_STATE_PARMSET;
    }

    return ret;
}

/* Set the parameters of an LMS key.
 *
 * Key must be inited before calling this.
 *
 * @param [in, out] key         LMS key to set parameters on.
 * @param [in]      levels      Number of tree levels.
 * @param [in]      height      Height of each tree.
 * @param [in]      winternitz  Width or Winternitz coefficient.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key is NULL.
 * @return  BAD_FUNC_ARG when parameters not supported.
 * @return  BAD_STATE_E when key state is not initialized.
 * */
int wc_LmsKey_SetParameters(LmsKey* key, int levels, int height,
    int winternitz)
{
    int ret = 0;

    /* Validate parameters. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }

    /* Check state is valid. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_INITED)) {
        WOLFSSL_MSG("error: LmsKey needs init");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        int i;

        ret = BAD_FUNC_ARG;
        /* Search through table for matching levels, height and width. */
        for (i = 0; i < WC_LMS_MAP_LEN; i++) {
            if ((levels == wc_lms_map[i].params.levels) &&
                    (height == wc_lms_map[i].params.height) &&
                    (winternitz == wc_lms_map[i].params.width)) {
                /* Set the parameters into the key. */
                key->params = &wc_lms_map[i].params;
                ret = 0;
                break;
            }
        }
    }

    if (ret == 0) {
        /* Move the state to params set.
         * Key is ready for MakeKey or Reload. */
        key->state = WC_LMS_STATE_PARMSET;
    }

    return ret;
}

/* Set the parameters of an LMS key including hash length.
 *
 * Key must be inited before calling this.
 *
 * @param [in, out] key         LMS key to set parameters on.
 * @param [in]      levels      Number of tree levels.
 * @param [in]      height      Height of each tree.
 * @param [in]      winternitz  Width or Winternitz coefficient.
 * @param [in]      hash        Hash algorithm to use. Valid values:
 *                               - LMS_SHA256
 *                               - LMS_SHA256_192
 *                               - LMS_SHAKE256
 *                               - LMS_SHAKE256_192
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key is NULL.
 * @return  BAD_FUNC_ARG when parameters not supported.
 * @return  BAD_STATE_E when key state is not initialized.
 * */
int wc_LmsKey_SetParameters_ex(LmsKey* key, int levels, int height,
    int winternitz, int hash)
{
    int ret = 0;

    /* Validate parameters. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }

    /* Check state is valid. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_INITED)) {
        WOLFSSL_MSG("error: LmsKey needs init");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        int i;

        ret = BAD_FUNC_ARG;
        /* Search through table for matching levels, height and width. */
        for (i = 0; i < WC_LMS_MAP_LEN; i++) {
            if ((levels == wc_lms_map[i].params.levels) &&
                    (height == wc_lms_map[i].params.height) &&
                    (winternitz == wc_lms_map[i].params.width) &&
                    (hash == (wc_lms_map[i].params.lmsType & LMS_HASH_MASK))) {
                /* Set the parameters into the key. */
                key->params = &wc_lms_map[i].params;
                ret = 0;
                break;
            }
        }
    }

    if (ret == 0) {
        /* Move the state to params set.
         * Key is ready for MakeKey or Reload. */
        key->state = WC_LMS_STATE_PARMSET;
    }

    return ret;
}

/* Get the parameters of an LMS key.
 *
 * Key must be inited and parameters set before calling this.
 *
 * @param [in]  key         LMS key.
 * @param [out] levels      Number of levels of trees.
 * @param [out] height      Height of the trees.
 * @param [out] winternitz  Winternitz width.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key, key->params, levels, height or winternitz is
 *          NULL.
 */
int wc_LmsKey_GetParameters(const LmsKey* key, int* levels, int* height,
    int* winternitz)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (levels == NULL) || (height == NULL) ||
            (winternitz == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    /* Validate the parameters are available. */
    if ((ret == 0) && (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        /* Set the levels, height and Winternitz width from parameters. */
        *levels = key->params->levels;
        *height = key->params->height;
        *winternitz = key->params->width;
    }

    return ret;
}

/* Get the parameters of an LMS key.
 *
 * Key must be inited and parameters set before calling this.
 *
 * @param [in]  key         LMS key.
 * @param [out] levels      Number of levels of trees.
 * @param [out] height      Height of the trees.
 * @param [out] winternitz  Winternitz width.
 * @param [out] hash        Hash algorithm.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key, key->params, levels, height, winternitz or
 *          hash is NULL.
 */
int wc_LmsKey_GetParameters_ex(const LmsKey* key, int* levels, int* height,
    int* winternitz, int* hash)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (levels == NULL) || (height == NULL) ||
            (winternitz == NULL) || (hash == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    /* Validate the parameters are available. */
    if ((ret == 0) && (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        /* Set the levels, height and Winternitz width from parameters. */
        *levels = key->params->levels;
        *height = key->params->height;
        *winternitz = key->params->width;
        *hash = key->params->lmsType & LMS_HASH_MASK;
    }

    return ret;
}

/* Frees the LMS key from memory.
 *
 * This does not affect the private key saved to non-volatile storage.
 *
 * @param [in, out] key  LMS key to free.
 */
void wc_LmsKey_Free(LmsKey* key)
{
    if (key != NULL) {
    #ifndef WOLFSSL_LMS_VERIFY_ONLY
        if (key->priv_data != NULL) {
            const LmsParams* params = key->params;
            word32 priv_data_len = LMS_PRIV_DATA_LEN(params->levels,
                params->height, params->p, params->rootLevels,
                params->cacheBits, params->hash_len);

        #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
            priv_data_len += HSS_PRIVATE_KEY_LEN(key->params->hash_len);
        #endif
            ForceZero(key->priv_data, priv_data_len);
            XFREE(key->priv_data, key->heap, DYNAMIC_TYPE_LMS);
            key->priv_data = NULL;
        }

        ForceZero(key->priv_raw, HSS_MAX_PRIVATE_KEY_LEN);
        ForceZero(&key->priv, sizeof(HssPrivKey));

        key->write_private_key = NULL;
        key->read_private_key = NULL;
        key->context = NULL;
        key->heap = NULL;
    #endif
        XMEMSET(key->pub, 0, sizeof(key->pub));
        key->params = NULL;
    #ifdef WOLF_CRYPTO_CB
        key->devId = INVALID_DEVID;
        key->devCtx = NULL;
    #endif
    #ifdef WOLF_PRIVATE_KEY_ID
        XMEMSET(key->id, 0, sizeof(key->id));
        key->idLen = 0;
        XMEMSET(key->label, 0, sizeof(key->label));
        key->labelLen = 0;
    #endif

        key->state = WC_LMS_STATE_FREED;
    }
}

#ifndef WOLFSSL_LMS_VERIFY_ONLY
/* Set the write private key callback to the LMS key structure.
 *
 * The callback must be able to write/update the private key to
 * non-volatile storage.
 *
 * @param [in, out] key       LMS key.
 * @param [in]      write_cb  Callback function that stores private key.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or write_cb is NULL.
 * @return  BAD_STATE_E when key state is invalid.
 */
int wc_LmsKey_SetWriteCb(LmsKey* key, wc_lms_write_private_key_cb write_cb)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (write_cb == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    /* Changing the write callback of an already working key is forbidden. */
    if ((ret == 0) && (key->state == WC_LMS_STATE_OK)) {
        WOLFSSL_MSG("error: wc_LmsKey_SetWriteCb: key in use");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        /* Set the callback into the key. */
        key->write_private_key = write_cb;
    }

    return ret;
}

/* Set the read private key callback to the LMS key structure.
 *
 * The callback must be able to read the private key from
 * non-volatile storage.
 *
 * @param [in, out] key      LMS key.
 * @param [in]      read_cb  Callback function that loads private key.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or read_cb is NULL.
 * @return  BAD_STATE_E when key state is invalid.
 * */
int wc_LmsKey_SetReadCb(LmsKey* key, wc_lms_read_private_key_cb read_cb)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (read_cb == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    /* Changing the read callback of an already working key is forbidden. */
    if ((ret == 0) && (key->state == WC_LMS_STATE_OK)) {
        WOLFSSL_MSG("error: wc_LmsKey_SetReadCb: key in use");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        /* Set the callback into the key. */
        key->read_private_key = read_cb;
    }

    return ret;
}

/* Sets the context to be used by write and read callbacks.
 *
 * E.g. this could be a filename if the callbacks write/read to file.
 *
 * @param [in, out] key      LMS key.
 * @param [in]      context  Pointer to data for read/write callbacks.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or context is NULL.
 * @return  BAD_STATE_E when key state is invalid.
 * */
int wc_LmsKey_SetContext(LmsKey* key, void* context)
{
    int ret = 0;

    /* Validate parameters. NULL context is allowed: callers with stub
     * read/write callbacks (e.g. HSM-backed keys whose private state lives
     * in the device) have no meaningful context to pass. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }
    /* Setting context of an already working key is forbidden. */
    if ((ret == 0) && (key->state == WC_LMS_STATE_OK)) {
        WOLFSSL_MSG("error: wc_LmsKey_SetContext: key in use");
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        /* Set the callback context into the key. */
        key->context = context;
    }

    return ret;
}

/* Make the LMS private/public key pair. The key must have its parameters
 * set before calling this.
 *
 * Write/read callbacks, and context data, must be set prior.
 * Key must have parameters set.
 *
 * @param [in, out] key   LMS key.
 * @param [in]      rng   Random number generator.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or rng is NULL.
 * @return  BAD_STATE_E when key is in an invalid state.
 * @return  BAD_FUNC_ARG when write callback or callback context not set.
 * @return  BAD_STATE_E when no more signatures can be created.
 */
int wc_LmsKey_MakeKey(LmsKey* key, WC_RNG* rng)
{
    int ret = 0;
    word32 priv_data_len = 0;

    /* Validate parameters. */
    if ((key == NULL) || (rng == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    /* Check state. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_PARMSET)) {
        WOLFSSL_MSG("error: LmsKey not ready for generation");
        ret = BAD_STATE_E;
    }

#ifdef WOLF_CRYPTO_CB
    /* HSM-backed keys skip the software write/context callbacks because the
     * device owns the private state. On CRYPTOCB_UNAVAILABLE fall-through the
     * software checks below still run. */
    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
        ret = wc_CryptoCb_PqcStatefulSigKeyGen(WC_PQC_STATEFUL_SIG_TYPE_LMS,
            key, rng);
        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
            /* This should not happen, but check whether signatures can be
             * created. */
            if ((ret == 0) && (wc_LmsKey_SigsLeft(key) == 0)) {
                WOLFSSL_MSG("error: generated LMS key signatures exhausted");
                key->state = WC_LMS_STATE_NOSIGS;
                ret = BAD_STATE_E;
            }
            /* On success, mirror the software path's terminal state so
             * subsequent Sign/Verify calls don't fail with BAD_STATE_E. */
            if (ret == 0) {
                key->state = WC_LMS_STATE_OK;
            }
            return ret;
        }
        ret = 0; /* fall through to software path */
    }
#endif

    /* Check write callback set. */
    if ((ret == 0) && (key->write_private_key == NULL)) {
        WOLFSSL_MSG("error: LmsKey write callback is not set");
        ret = BAD_FUNC_ARG;
    }
    if (ret == 0) {
        const LmsParams* params = key->params;
        priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
            params->p, params->rootLevels, params->cacheBits, params->hash_len);

    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        priv_data_len += HSS_PRIVATE_KEY_LEN(key->params->hash_len);
    #endif
    }
    if ((ret == 0) && (key->priv_data == NULL)) {
        /* Allocate memory for the private key data. */
        key->priv_data = (byte *)XMALLOC(priv_data_len, key->heap,
            DYNAMIC_TYPE_LMS);
        /* Check pointer is valid. */
        if (key->priv_data == NULL) {
            ret = MEMORY_E;
        }
    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        else {
            XMEMSET(key->priv_data, 0, priv_data_len);
        }
    #endif
    }
    if (ret == 0) {
        WC_DECLARE_VAR(state, LmsState, 1, 0);

        /* Allocate memory for working state. */
        WC_ALLOC_VAR_EX(state, LmsState, 1, NULL, DYNAMIC_TYPE_TMP_BUFFER,
            ret=MEMORY_E);
        if (WC_VAR_OK(state))
        {
            /* Initialize working state for use. */
            ret = wc_lmskey_state_init(state, key->params);
            if (ret == 0) {
                /* Make the HSS key. */
                ret = wc_hss_make_key(state, rng, key->priv_raw, &key->priv,
                    key->priv_data, key->pub);
                wc_lmskey_state_free(state);
            }
            ForceZero(state, sizeof(LmsState));
            WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
        }
    }
    if (ret == 0) {
        int rv;
        /* Write private key to storage. */
    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        XMEMCPY(key->priv_data + priv_data_len -
            HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->priv_raw,
            HSS_PRIVATE_KEY_LEN(key->params->hash_len));
        rv = key->write_private_key(key->priv_data, priv_data_len,
            key->context);
    #else
        rv = key->write_private_key(key->priv_raw,
            HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
    #endif
        if (rv != WC_LMS_RC_SAVED_TO_NV_MEMORY) {
            ret = IO_FAILED_E;
        }
    }

    /* This should not happen, but check whether signatures can be created. */
    if ((ret == 0) && (wc_LmsKey_SigsLeft(key) == 0)) {
        WOLFSSL_MSG("error: generated LMS key signatures exhausted");
        key->state = WC_LMS_STATE_NOSIGS;
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        /* Update state. */
        key->state = WC_LMS_STATE_OK;
    }

    return ret;
}

/* Reload a key that has been prepared with the appropriate params and data.
 *
 * Use this if you wish to resume signing with an existing key.
 * Call this function after initializing and setting parameters.
 *
 * Write/read callbacks, and context data, must be set prior.
 * Key must have parameters set.
 *
 * @param [in, out] key  LMS key.
 *
 * Returns 0 on success. */
int wc_LmsKey_Reload(LmsKey* key)
{
    int ret = 0;
    word32 priv_data_len = 0;

    /* Validate parameter. */
    if (key == NULL) {
        ret = BAD_FUNC_ARG;
    }
    /* Check state. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_PARMSET)) {
        WOLFSSL_MSG("error: LmsKey not ready for reload");
        ret = BAD_STATE_E;
    }

#ifdef WOLF_CRYPTO_CB
    /* State for HSM-backed keys lives in the device; no software reload. */
    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
        WOLFSSL_MSG("wc_LmsKey_Reload is a no-op for HSM-backed keys");
        key->state = WC_LMS_STATE_OK;
        return 0;
    }
#endif

    /* Check read callback present. */
    if ((ret == 0) && (key->read_private_key == NULL)) {
        WOLFSSL_MSG("error: LmsKey read callback is not set");
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        const LmsParams* params = key->params;
        priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
            params->p, params->rootLevels, params->cacheBits, params->hash_len);

    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        priv_data_len += HSS_PRIVATE_KEY_LEN(params->hash_len);
    #endif
    }
    if ((ret == 0) && (key->priv_data == NULL)) {
        /* Allocate memory for the private key data. */
        key->priv_data = (byte *)XMALLOC(priv_data_len, key->heap,
            DYNAMIC_TYPE_LMS);
        /* Check pointer is valid. */
        if (key->priv_data == NULL) {
            ret = MEMORY_E;
        }
    }
    if (ret == 0) {
        int rv;

        /* Load private key. */
    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        const LmsParams* params = key->params;

        rv = key->read_private_key(key->priv_data, priv_data_len, key->context);
    #else
        rv = key->read_private_key(key->priv_raw,
            HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
    #endif
        if (rv != WC_LMS_RC_READ_TO_MEMORY) {
            ret = IO_FAILED_E;
        }
    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        if (ret == 0) {
            XMEMCPY(key->priv_raw, key->priv_data + priv_data_len -
                HSS_PRIVATE_KEY_LEN(params->hash_len),
                HSS_PRIVATE_KEY_LEN(params->hash_len));
        }
    #endif
    }

    /* Double check the key actually has signatures left. */
    if ((ret == 0) && (wc_LmsKey_SigsLeft(key) == 0)) {
        WOLFSSL_MSG("error: reloaded LMS key signatures exhausted");
        key->state = WC_LMS_STATE_NOSIGS;
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        WC_DECLARE_VAR(state, LmsState, 1, 0);

        /* Allocate memory for working state. */
        WC_ALLOC_VAR_EX(state, LmsState, 1, NULL, DYNAMIC_TYPE_TMP_BUFFER,
            ret=MEMORY_E);
        if (WC_VAR_OK(state))
        {
            /* Initialize working state for use. */
            ret = wc_lmskey_state_init(state, key->params);
            if (ret == 0) {
                /* Reload the key ready for signing. */
                ret = wc_hss_reload_key(state, key->priv_raw, &key->priv,
                    key->priv_data, NULL);
                wc_lmskey_state_free(state);
            }
            ForceZero(state, sizeof(LmsState));
            WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
        }
    }

    if (ret == 0) {
        /* Update state. */
        key->state = WC_LMS_STATE_OK;
    }

    return ret;
}

/* Get the raw private key length based on parameter set of key.
 *
 * @param [in]  key  LMS key.
 * @param [out] len  Length of private key.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or len is NULL or parameters not set.
 */
int wc_LmsKey_GetPrivLen(const LmsKey* key, word32* len)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (len == NULL) || (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        /* Return private key length from parameter set. */
        *len = HSS_PRIVATE_KEY_LEN(key->params->hash_len);
    }

    return ret;
}

/* Sign a message.
 *
 * @param [in, out] key    LMS key to sign with.
 * @param [out]     sig    Signature data. Buffer must be big enough to hold
 *                         signature data.
 * @param [out]     sigSz  Length of signature data.
 * @param [in]      msg    Message to sign.
 * @param [in]      msgSz  Length of message in bytes.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key, sig, sigSz or msg is NULL.
 * @return  BAD_FUNC_ARG when msgSz is less than 0.
 * @return  BAD_FUNC_ARG when a write private key is not set.
 * @return  BAD_FUNC_ARG when a read/write private key context is not set.
 * @return  BUFFER_E when sigSz is too small.
 * @return  BAD_STATE_E when wrong state for operation.
 * @return  IO_FAILED_E when reading or writing private key failed.
 */
int wc_LmsKey_Sign(LmsKey* key, byte* sig, word32* sigSz, const byte* msg,
    int msgSz)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (key->params == NULL) || (sig == NULL) ||
            (sigSz == NULL) || (msg == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    if ((ret == 0) && (msgSz < 0)) {
        ret = BAD_FUNC_ARG;
    }
    /* Check state. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_OK)) {
        if (key->state == WC_LMS_STATE_NOSIGS) {
            WOLFSSL_MSG("error: LMS signatures exhausted");
        }
        else {
            /* The key state is not ready for signing. */
            WOLFSSL_MSG("error: can't sign, LMS key not in good state");
        }
        ret = BAD_STATE_E;
    }
    /* Check signature buffer size. */
    if ((ret == 0) && (*sigSz < key->params->sig_len)) {
        /* Signature buffer too small. */
        WOLFSSL_MSG("error: LMS sig buffer too small");
        ret = BUFFER_E;
    }

#ifdef WOLF_CRYPTO_CB
    /* HSM-backed keys skip the software write/context callbacks because the
     * device owns the private state. On CRYPTOCB_UNAVAILABLE fall-through the
     * software checks below still run. */
    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
        ret = wc_CryptoCb_PqcStatefulSigSign(msg, (word32)msgSz, sig, sigSz,
            WC_PQC_STATEFUL_SIG_TYPE_LMS, key);
        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
            return ret;
        }
        ret = 0; /* fall through to software path */
    }
#endif

    /* Check read and write callbacks available. */
    if ((ret == 0) && (key->write_private_key == NULL)) {
        WOLFSSL_MSG("error: LmsKey write/read callbacks are not set");
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        WC_DECLARE_VAR(state, LmsState, 1, 0);

        /* Allocate memory for working state. */
        WC_ALLOC_VAR_EX(state, LmsState, 1, NULL, DYNAMIC_TYPE_TMP_BUFFER,
            ret=MEMORY_E);
        if (WC_VAR_OK(state))
        {
            /* Initialize working state for use. */
            ret = wc_lmskey_state_init(state, key->params);
            if (ret == 0) {
                /* Sign message. */
                ret = wc_hss_sign(state, key->priv_raw, &key->priv,
                    key->priv_data, msg, (word32)msgSz, sig);
                wc_lmskey_state_free(state);
            }
            ForceZero(state, sizeof(LmsState));
            WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
        }
    }
    if (ret == 0) {
        *sigSz = (word32)key->params->sig_len;
    }
    if (ret == 0) {
        int rv;

        /* Write private key to storage. */
    #ifdef WOLFSSL_WC_LMS_SERIALIZE_STATE
        const LmsParams* params = key->params;
        word32 priv_data_len = LMS_PRIV_DATA_LEN(params->levels, params->height,
            params->p, params->rootLevels, params->cacheBits,
            params->hash_len) + HSS_PRIVATE_KEY_LEN(key->params->hash_len);

        XMEMCPY(key->priv_data + priv_data_len -
            HSS_PRIVATE_KEY_LEN(params->hash_len), key->priv_raw,
            HSS_PRIVATE_KEY_LEN(params->hash_len));
        rv = key->write_private_key(key->priv_data, priv_data_len,
            key->context);
    #else
        rv = key->write_private_key(key->priv_raw,
            HSS_PRIVATE_KEY_LEN(key->params->hash_len), key->context);
    #endif
        if (rv != WC_LMS_RC_SAVED_TO_NV_MEMORY) {
            /* Write to NV storage failed. Erase the signature from
             * memory to prevent OTS key reuse if state is rolled back. */
            ForceZero(sig, key->params->sig_len);
            ret = IO_FAILED_E;
        }
    }

    return ret;
}

/* Returns whether signatures can be created with key.
 *
 * @param [in]  key  LMS key.
 *
 * @return  1 if there are signatures remaining.
 * @return  0 if available signatures are exhausted.
 */
int wc_LmsKey_SigsLeft(LmsKey* key)
{
    int ret = 0;

    /* NULL keys have no signatures remaining. */
    if (key != NULL) {
    #ifdef WOLF_CRYPTO_CB
        if (key->devId != INVALID_DEVID) {
            word32 sigsLeft = 0;
            int cbRet = wc_CryptoCb_PqcStatefulSigSigsLeft(
                WC_PQC_STATEFUL_SIG_TYPE_LMS, key, &sigsLeft);
            if (cbRet == 0) {
                return (sigsLeft != 0) ? 1 : 0;
            }
            if (cbRet != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
                /* Device returned an actual error; the device owns the
                 * private state so no safe software fallback exists. */
                WOLFSSL_MSG("PqcStatefulSigSigsLeft returned an error");
                return 0;
            }
            /* Cryptocb declined. priv_raw reflects software state from the
             * CRYPTOCB_UNAVAILABLE fall-through in MakeKey/Reload, so the
             * software check below is valid. */
            WOLFSSL_MSG("LMS SigsLeft not supported by device, using software");
        }
    #endif
        ret = wc_hss_sigsleft(key->params, key->priv_raw);
    }

    return ret;
}

#endif /* ifndef WOLFSSL_LMS_VERIFY_ONLY*/

/* Get the public key length based on parameter set of key.
 *
 * @param [in]  key  LMS key.
 * @param [out] len  Length of public key.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or len is NULL or parameters not set.
 */
int wc_LmsKey_GetPubLen(const LmsKey* key, word32* len)
{
    int ret = 0;

    /* Validate parameters */
    if ((key == NULL) || (len == NULL) || (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        *len = HSS_PUBLIC_KEY_LEN(key->params->hash_len);
    }

    return ret;
}

/* Export a generated public key and parameter set from one LmsKey
 * to another, with explicit heap and device bindings.
 *
 * The destination key must be unused or have been freed.
 * The destination is fully initialized as a verify-only key.
 *
 * @param [out] keyDst  LMS key to copy into.
 * @param [in]  keySrc  LMS key to copy.
 * @param [in]  heap    Heap hint for keyDst.
 * @param [in]  devId   Device identifier for keyDst.
 *                      Use INVALID_DEVID when not using a device.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when keyDst or keySrc is NULL.
 * @return  BAD_STATE_E when in the wrong state for the operation.
 */
int wc_LmsKey_ExportPub_ex(LmsKey* keyDst, const LmsKey* keySrc,
                           void* heap, int devId)
{
    int ret = 0;

    if ((keyDst == NULL) || (keySrc == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    if ((ret == 0) && (keySrc->state != WC_LMS_STATE_OK) &&
            (keySrc->state != WC_LMS_STATE_VERIFYONLY) &&
            (keySrc->state != WC_LMS_STATE_NOSIGS)) {
        ret = BAD_STATE_E;
    }

    if (ret == 0) {
        ret = wc_LmsKey_Init(keyDst, heap, devId);
    }
    if (ret == 0) {
        keyDst->params = keySrc->params;
        XMEMCPY(keyDst->pub, keySrc->pub, sizeof(keySrc->pub));

        /* Mark this key as verify only, to prevent misuse. */
        keyDst->state = WC_LMS_STATE_VERIFYONLY;
    }

    return ret;
}

/* Export a generated public key and parameter set from one LmsKey
 * to another. Use this to prepare a signature verification LmsKey
 * that is pub only.
 *
 * Though the public key is all that is used to verify signatures,
 * the parameter set is needed to calculate the signature length
 * before hand.
 *
 * The destination key is left with no heap hint and no device
 * binding. Callers that need either should use
 * wc_LmsKey_ExportPub_ex.
 *
 * @param [out] keyDst  LMS key to copy into.
 * @param [in]  keySrc  LMS key to copy.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when keyDst or keySrc is NULL.
 */
int wc_LmsKey_ExportPub(LmsKey* keyDst, const LmsKey* keySrc)
{
    return wc_LmsKey_ExportPub_ex(keyDst, keySrc,
    #ifndef WOLFSSL_LMS_VERIFY_ONLY
        (keySrc != NULL) ? keySrc->heap : NULL,
    #else
        NULL,
    #endif
        INVALID_DEVID);
}

/* Exports the raw LMS public key buffer from key to out buffer.
 * The out buffer should be large enough to hold the public key, and
 * outLen should indicate the size of the buffer.
 *
 * Call wc_LmsKey_GetPubLen beforehand to determine pubLen.
 *
 * @param [in]      key     LMS key.
 * @param [out]     out     Buffer to hold encoded public key.
 * @param [in, out] outLen  On in, length of out in bytes.
 *                          On out, the length of the public key in bytes.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key, out or outLen is NULL.
 * @return  BUFFER_E when outLen is too small to hold encoded public key.
 */
int wc_LmsKey_ExportPubRaw(const LmsKey* key, byte* out, word32* outLen)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (out == NULL) || (outLen == NULL) ||
            (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    /* Check size of out is sufficient. */
    if ((ret == 0) &&
            (*outLen < (word32)HSS_PUBLIC_KEY_LEN(key->params->hash_len))) {
        ret = BUFFER_E;
    }

    if (ret == 0) {
        /* Return encoded public key. */
        XMEMCPY(out, key->pub, HSS_PUBLIC_KEY_LEN(key->params->hash_len));
        *outLen = HSS_PUBLIC_KEY_LEN(key->params->hash_len);
    }

    return ret;
}

/* Imports a raw public key buffer from in array to LmsKey key.
 *
 * If the LMS parameters have already been configured (via
 * wc_LmsKey_SetLmsParm or wc_LmsKey_SetParameters), the levels /
 * lms_algorithm_type / lmots_algorithm_type encoded in the raw key are
 * checked for consistency and inLen must match wc_LmsKey_GetPubLen.
 *
 * If the parameters have not yet been set (key->params == NULL), they
 * are derived from the raw public key prefix (RFC 8554 sec 3.3 / sec
 * 6.1: u32str(L) || lms_algorithm_type || lmots_algorithm_type) and
 * matched against the static parameter map. The candidate is held in
 * a local until the length check passes, so a length mismatch leaves
 * key->params NULL.
 *
 * Accepts a key in INITED, PARMSET or VERIFYONLY state. WC_LMS_STATE_OK
 * is rejected because the key already has private material loaded and
 * silently overwriting key->pub would create an inconsistent priv/pub
 * pair.
 *
 * @param [in, out] key    LMS key to put public key in.
 * @param [in]      in     Buffer holding encoded public key.
 * @param [in]      inLen  Length of encoded public key in bytes.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or in is NULL, or when the raw key's
 *          levels / lmsType / lmOtsType disagree with pre-set params.
 * @return  BAD_STATE_E when wrong state for operation.
 * @return  BUFFER_E when inLen is too small to contain the LMS type
 *          fields, or doesn't match the public key length determined
 *          by parameters.
 * @return  NOT_COMPILED_IN when the derived parameter set isn't built in.
 */
int wc_LmsKey_ImportPubRaw(LmsKey* key, const byte* in, word32 inLen)
{
    int              ret = 0;
    const LmsParams* matched = NULL;

    /* Validate parameters. */
    if ((key == NULL) || (in == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    /* Reject states where re-importing the public bytes would desync
     * the key. INITED (params unset, will derive), PARMSET (params set
     * but no private material) and VERIFYONLY (already a pub-only key)
     * are all safe. OK means a private key is loaded; silently
     * overwriting key->pub would create a priv/pub mismatch. Mirrors
     * the wc_XmssKey_ImportPubRaw_ex post-condition added in the same
     * RFC 9802 series. */
    if ((ret == 0) &&
            (key->state != WC_LMS_STATE_INITED) &&
            (key->state != WC_LMS_STATE_PARMSET) &&
            (key->state != WC_LMS_STATE_VERIFYONLY)) {
        WOLFSSL_MSG("error: LMS key not ready for import");
        ret = BAD_STATE_E;
    }
    /* Need at least L || lmsType || lmOtsType to derive or validate. */
    if ((ret == 0) && (inLen < (word32)(LMS_L_LEN + 2 * LMS_TYPE_LEN))) {
        ret = BUFFER_E;
    }

    if (ret == 0) {
        word32 levels = 0;
        word32 lmsType = 0;
        word32 lmOtsType = 0;

        /* RFC 8554 sec 3.3 / sec 6.1: HSS public key = u32str(L) || pub[0],
         * where pub[0] starts with lms_algorithm_type || lmots_algorithm_type.
         */
        ato32(in + 0,                   &levels);
        ato32(in + LMS_L_LEN,           &lmsType);
        ato32(in + LMS_L_LEN + LMS_TYPE_LEN, &lmOtsType);

        /* The wire format carries only the RFC type code (low 12 bits);
         * params->lmsType / lmOtsType also pack a wolfSSL-internal hash
         * family flag in the high 4 bits (LMS_HASH_MASK). Compare on the
         * RFC code only -- safe as long as low-12-bit codes stay globally
         * distinct across hash families (see wc_lms_impl.c step 3.d-e
         * note). */
        if (key->params == NULL) {
            /* Auto-derive: find matching entry in the static map. Hold
             * the candidate in a local until the length check passes to
             * avoid leaving key->params half-set on failure. */
            int i;
            ret = WC_NO_ERR_TRACE(NOT_COMPILED_IN);
            for (i = 0; i < WC_LMS_MAP_LEN; i++) {
                if (((word32)wc_lms_map[i].params.levels == levels) &&
                    ((word32)(wc_lms_map[i].params.lmsType & LMS_H_W_MASK) ==
                                                                     lmsType) &&
                    ((word32)(wc_lms_map[i].params.lmOtsType& LMS_H_W_MASK) ==
                                                                   lmOtsType)) {
                    matched = &wc_lms_map[i].params;
                    ret = 0;
                    break;
                }
            }
            if (ret != 0) {
                WOLFSSL_MSG("error: LMS params from pub key not supported");
            }
        }
        else {
            /* Validate against pre-set params. */
            if (((word32)key->params->levels != levels) ||
                ((word32)(key->params->lmsType & LMS_H_W_MASK) != lmsType) ||
                ((word32)(key->params->lmOtsType & LMS_H_W_MASK) != lmOtsType)){
                WOLFSSL_MSG("error: LMS pub key doesn't match set params");
                ret = BAD_FUNC_ARG;
            }
            else {
                matched = key->params;
            }
        }
    }
    if ((ret == 0) &&
            (inLen != (word32)HSS_PUBLIC_KEY_LEN(matched->hash_len))) {
        ret = BUFFER_E;
    }

    if (ret == 0) {
        /* Commit params (no-op when already set) and copy the key.
         * State is INITED/PARMSET/VERIFYONLY here (OK is rejected
         * above), so promoting to VERIFYONLY is always correct. */
        key->params = matched;
        XMEMCPY(key->pub, in, inLen);
        key->state = WC_LMS_STATE_VERIFYONLY;
    }

    return ret;
}

/* Given a levels, height, winternitz parameter set, determine
 * the signature length.
 *
 * Call this before wc_LmsKey_Sign so you know the length of
 * the required signature buffer.
 *
 * @param [in]  key  LMS key.
 * @param [out] len  Length of a signature in bytes.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when key or len is NULL, or when the LMS
 *          parameters have not been configured on the key.
 */
int wc_LmsKey_GetSigLen(const LmsKey* key, word32* len)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (len == NULL) || (key->params == NULL)) {
        ret = BAD_FUNC_ARG;
    }

    if (ret == 0) {
        *len = key->params->sig_len;
    }

    return ret;
}

/* Verify the signature of the message with public key.
 *
 * @param [in] key    LMS key.
 * @param [in] sig    Signature to verify.
 * @param [in] sigSz  Size of signature in bytes.
 * @param [in] msg    Message to verify.
 * @param [in] msgSz  Length of the message in bytes.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when a key, sig or msg is NULL.
 * @return  BAD_FUNC_ARG when msgSz is negative.
 * @return  SIG_VERIFY_E when signature did not verify message.
 * @return  BAD_STATE_E when wrong state for operation.
 * @return  BUFFER_E when sigSz is invalid for parameters.
 */
int wc_LmsKey_Verify(LmsKey* key, const byte* sig, word32 sigSz,
    const byte* msg, int msgSz)
{
    int ret = 0;

    /* Validate parameters. */
    if ((key == NULL) || (key->params == NULL) || (sig == NULL) ||
            (msg == NULL)) {
        ret = BAD_FUNC_ARG;
    }
    if ((ret == 0) && (msgSz < 0)) {
        ret = BAD_FUNC_ARG;
    }
    /* Check state. */
    if ((ret == 0) && (key->state != WC_LMS_STATE_OK) &&
            (key->state != WC_LMS_STATE_VERIFYONLY)) {
        /* LMS key not ready for verification. Param str must be
         * set first, and Reload() called. */
        WOLFSSL_MSG("error: LMS key not ready for verification");
        ret = BAD_STATE_E;
    }
    /* Check signature length. */
    if ((ret == 0) && (sigSz != key->params->sig_len)) {
        ret = BUFFER_E;
    }

#ifdef WOLF_CRYPTO_CB
    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
        int res = 0;
        ret = wc_CryptoCb_PqcStatefulSigVerify(sig, sigSz, msg, (word32)msgSz,
            &res, WC_PQC_STATEFUL_SIG_TYPE_LMS, key);
        if (ret != WC_NO_ERR_TRACE(CRYPTOCB_UNAVAILABLE)) {
            if (ret == 0 && res != 1)
                ret = SIG_VERIFY_E;
            return ret;
        }
        ret = 0; /* fall through to software path */
    }
#endif

    if (ret == 0) {
        WC_DECLARE_VAR(state, LmsState, 1, 0);

        /* Allocate memory for working state. */
        WC_ALLOC_VAR_EX(state, LmsState, 1, NULL, DYNAMIC_TYPE_TMP_BUFFER,
            ret=MEMORY_E);
        if (WC_VAR_OK(state))
        {
            /* Initialize working state for use. */
            ret = wc_lmskey_state_init(state, key->params);
            if (ret == 0) {
                /* Verify signature of message with public key. */
                ret = wc_hss_verify(state, key->pub, msg, (word32)msgSz, sig,
                    sigSz);
                wc_lmskey_state_free(state);
            }
            ForceZero(state, sizeof(LmsState));
            WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
        }
    }

    return ret;
}

#ifndef WOLFSSL_LMS_VERIFY_ONLY

/* Get the Key ID from the LMS key.
 *
 * PRIV = Q | PARAMS | SEED | I
 * where I is the Key ID.
 *
 * @param [in]  key    LMS key.
 * @param [out] kid    Key ID data.
 * @param [out] kidSz  Size of key ID.
 * @return  0 on success.
 * @return  BAD_FUNC_ARG when a key, kid or kidSz is NULL.
 * @return  NOT_COMPILED_IN when key is on a device.
 */
int wc_LmsKey_GetKid(LmsKey* key, const byte** kid, word32* kidSz)
{
    int ret = 0;
    word32 offset;

    if ((key == NULL) || (key->params == NULL) || (kid == NULL) ||
            (kidSz == NULL)) {
        ret = BAD_FUNC_ARG;
    }

#ifdef WOLF_CRYPTO_CB
    /* priv_raw is not populated for HSM-backed keys where the device owns
     * the private state. Extend the CryptoCb surface if device-side KID
     * retrieval becomes a requirement.
     */
    if ((ret == 0) && (key->devId != INVALID_DEVID)) {
        WOLFSSL_MSG(
                "wc_LmsKey_GetKid: priv_raw may be uninitialised for HSM keys");
    }
#endif
    if (ret == 0) {
        /* SEED length is hash length. */
        offset = HSS_Q_LEN + HSS_PRIV_KEY_PARAM_SET_LEN + key->params->hash_len;
        *kid = key->priv_raw + offset;
        *kidSz = HSS_PRIVATE_KEY_LEN(key->params->hash_len) - offset;
    }

    return ret;
}


/* Get the Key ID from the raw private key data.
 *
 * PRIV = Q | PARAMS | SEED | I
 * where I is the Key ID.
 *
 * @param [in] priv    Private key data.
 * @param [in] privSz  Size of private key data.
 * @param  Pointer to 16 byte Key ID in the private key.
 * @return  NULL on failure.
 */
const byte* wc_LmsKey_GetKidFromPrivRaw(const byte* priv, word32 privSz)
{
    const byte* ret;

    if ((priv == NULL) ||
            (privSz < HSS_Q_LEN + HSS_PRIV_KEY_PARAM_SET_LEN + LMS_I_LEN)) {
        ret = NULL;
    }
    else {
        word32 seedSz = privSz - HSS_Q_LEN - HSS_PRIV_KEY_PARAM_SET_LEN -
                        LMS_I_LEN;
        if ((seedSz != WC_SHA256_192_DIGEST_SIZE) &&
                (seedSz != WC_SHA256_DIGEST_SIZE)) {
            ret = NULL;
        }
        else {
            ret = priv + privSz - LMS_I_LEN;
        }
    }

    return ret;
}

#endif

#endif /* WOLFSSL_HAVE_LMS */
