/*

 * Copyright 2024-2025 The OpenSSL Project Authors. All Rights Reserved.

 *

 * Licensed under the Apache License 2.0 (the "License").  You may not use

 * this file except in compliance with the License.  You can obtain a copy

 * in the file LICENSE in the source distribution or at

 * https://www.openssl.org/source/license.html

 */

#include <string.h>

#include <openssl/err.h>

#include <openssl/core_dispatch.h>

#include <openssl/core_names.h>

#include <openssl/params.h>

#include <openssl/rand.h>

#include <openssl/proverr.h>

#include "slh_dsa_local.h"

#include "slh_dsa_key.h"

#include "internal/encoder.h"

static int slh_dsa_compute_pk_root(SLH_DSA_HASH_CTX *ctx, SLH_DSA_KEY *out, int verify);

static void slh_dsa_key_hash_cleanup(SLH_DSA_KEY *key)

{

    OPENSSL_free(key->propq);

    if (key->md_big != key->md)

        EVP_MD_free(key->md_big);

    key->md_big = NULL;

    EVP_MD_free(key->md);

    EVP_MAC_free(key->hmac);

    key->md = NULL;

}

static int slh_dsa_key_hash_init(SLH_DSA_KEY *key)

{

    int is_shake = key->params->is_shake;

    int security_category = key->params->security_category;

    const char *digest_alg = is_shake ? "SHAKE-256" : "SHA2-256";

    key->md = EVP_MD_fetch(key->libctx, digest_alg, key->propq);

    if (key->md == NULL)

        return 0;

    /*

     * SHA2 algorithm(s) require SHA256 + HMAC_SHA(X) & MGF1(SHAX)

     * SHAKE algorithm(s) use SHAKE for all functions.

     */

    if (is_shake == 0) {

        if (security_category == 1) {

            /* For category 1 SHA2-256 is used for all hash operations */

            key->md_big = key->md;

        } else {

            /* Security categories 3 & 5 also need SHA-512 */

            key->md_big = EVP_MD_fetch(key->libctx, "SHA2-512", key->propq);

            if (key->md_big == NULL)

                goto err;

        }

        key->hmac = EVP_MAC_fetch(key->libctx, "HMAC", key->propq);

        if (key->hmac == NULL)

            goto err;

    }

    key->adrs_func = ossl_slh_get_adrs_fn(is_shake == 0);

    key->hash_func = ossl_slh_get_hash_fn(is_shake);

    return 1;

 err:

    slh_dsa_key_hash_cleanup(key);

    return 0;

}

static void slh_dsa_key_hash_dup(SLH_DSA_KEY *dst, const SLH_DSA_KEY *src)

{

    if (src->md_big != NULL && src->md_big != src->md)

        EVP_MD_up_ref(src->md_big);

    if (src->md != NULL)

        EVP_MD_up_ref(src->md);

    if (src->hmac != NULL)

        EVP_MAC_up_ref(src->hmac);

}

/**

 * @brief Return the libctx associated with a SLH_DSA_KEY object

 *

 * @param key A SLH_DSA_KEY to extract the libctx from.

 * @returns The new OSSL_LIB_CTX object on success, or NULL failure

 */

OSSL_LIB_CTX *ossl_slh_dsa_key_get0_libctx(const SLH_DSA_KEY *key)

{

    return key != NULL ? key->libctx : NULL;

}

/**

 * @brief Create a new SLH_DSA_KEY object

 *

 * @param libctx A OSSL_LIB_CTX object used for fetching algorithms.

 * @param propq The property query used for fetching algorithms

 * @param alg The algorithm name associated with the key type

 * @returns The new SLH_DSA_KEY object on success, or NULL on malloc failure

 */

SLH_DSA_KEY *ossl_slh_dsa_key_new(OSSL_LIB_CTX *libctx, const char *propq,

                                  const char *alg)

{

    SLH_DSA_KEY *ret;

    const SLH_DSA_PARAMS *params = ossl_slh_dsa_params_get(alg);

    if (params == NULL)

        return NULL;

    ret = OPENSSL_zalloc(sizeof(*ret));

    if (ret != NULL) {

        ret->libctx = libctx;

        ret->params = params;

        if (propq != NULL) {

            ret->propq = OPENSSL_strdup(propq);

            if (ret->propq == NULL)

                goto err;

        }

        if (!slh_dsa_key_hash_init(ret))

            goto err;

    }

    return ret;

 err:

    ossl_slh_dsa_key_free(ret);

    return NULL;

}

/**

 * @brief Destroy a SLH_DSA_KEY object

 */

void ossl_slh_dsa_key_free(SLH_DSA_KEY *key)

{

    if (key == NULL)

        return;

    slh_dsa_key_hash_cleanup(key);

    OPENSSL_cleanse(&key->priv, sizeof(key->priv) >> 1);

    OPENSSL_free(key);

}

/**

 * @brief Duplicate a key

 *

 * @param src A SLH_DSA_KEY object to copy

 * @param selection to select public and/or private components. Selecting the

 *                  private key will also select the public key

 * @returns The duplicated key, or NULL on failure.

 */

SLH_DSA_KEY *ossl_slh_dsa_key_dup(const SLH_DSA_KEY *src, int selection)

{

    SLH_DSA_KEY *ret = NULL;

    if (src == NULL)

        return NULL;

    ret = OPENSSL_zalloc(sizeof(*ret));

    if (ret != NULL) {

        *ret = *src; /* this copies everything including the keydata in priv[] */

        ret->propq = NULL;

        ret->pub = NULL;

        ret->has_priv = 0;

        slh_dsa_key_hash_dup(ret, src);

        if (src->propq != NULL) {

            ret->propq = OPENSSL_strdup(src->propq);

            if (ret->propq == NULL)

                goto err;

        }

        if ((selection & OSSL_KEYMGMT_SELECT_KEYPAIR) != 0) {

            /* The public components are present if the private key is present */

            if (src->pub != NULL)

                ret->pub = SLH_DSA_PUB(ret);

            if ((selection & OSSL_KEYMGMT_SELECT_PRIVATE_KEY) != 0)

                ret->has_priv = src->has_priv;

        }

    }

    return ret;

 err:

    ossl_slh_dsa_key_free(ret);

    return NULL;

}

/**

 * @brief Are 2 keys equal?

 *

 * To be equal the keys must have the same key data and algorithm name.

 *

 * @param key1 A SLH_DSA_KEY object

 * @param key2 A SLH_DSA_KEY object

 * @param selection to select public and/or private component comparison.

 * @returns 1 if the keys are equal otherwise it returns 0.

 */

int ossl_slh_dsa_key_equal(const SLH_DSA_KEY *key1, const SLH_DSA_KEY *key2,

                           int selection)

{

    int key_checked = 0;

    /* The parameter sets must match - i.e. The same algorithm name */

    if (key1->params != key2->params)

        return 0;

    if ((selection & OSSL_KEYMGMT_SELECT_KEYPAIR) != 0) {

        if ((selection & OSSL_KEYMGMT_SELECT_PUBLIC_KEY) != 0) {

            if (key1->pub != NULL && key2->pub != NULL) {

                if (memcmp(key1->pub, key2->pub, key1->params->pk_len) != 0)

                    return 0;

                key_checked = 1;

            }

        }

        if (!key_checked

                && (selection & OSSL_KEYMGMT_SELECT_PRIVATE_KEY) != 0) {

            if (key1->has_priv && key2->has_priv) {

                if (memcmp(key1->priv, key2->priv,

                           key1->params->pk_len) != 0)

                    return 0;

                key_checked = 1;

            }

        }

        return key_checked;

    }

    return 1;

}

int ossl_slh_dsa_key_has(const SLH_DSA_KEY *key, int selection)

{

    if ((selection & OSSL_KEYMGMT_SELECT_KEYPAIR) != 0) {

        if (key->pub == NULL)

            return 0; /* No public key */

        if ((selection & OSSL_KEYMGMT_SELECT_PRIVATE_KEY) != 0

                && key->has_priv == 0)

            return 0; /* No private key */

        return 1;

    }

    return 0;

}

int ossl_slh_dsa_key_pairwise_check(const SLH_DSA_KEY *key)

{

    int ret;

    SLH_DSA_HASH_CTX *ctx = NULL;

    if (key->pub == NULL || key->has_priv == 0)

        return 0;

    ctx = ossl_slh_dsa_hash_ctx_new(key);

    if (ctx == NULL)

        return 0;

    ret = slh_dsa_compute_pk_root(ctx, (SLH_DSA_KEY *)key, 1);

    ossl_slh_dsa_hash_ctx_free(ctx);

    return ret;

}

void ossl_slh_dsa_key_reset(SLH_DSA_KEY *key)

{

    key->pub = NULL;

    if (key->has_priv) {

        key->has_priv = 0;

        OPENSSL_cleanse(key->priv, sizeof(key->priv));

    }

}

/**

 * @brief Load a SLH_DSA key from raw data.

 *

 * @param key An SLH_DSA key to load into

 * @param params An array of parameters containing key data.

 * @param include_private Set to 1 to optionally include the private key data

 *                        if it exists.

 * @returns 1 on success, or 0 on failure.

 */

int ossl_slh_dsa_key_fromdata(SLH_DSA_KEY *key, const OSSL_PARAM *param_pub,

                              const OSSL_PARAM *param_priv,

                              int include_private)

{

    size_t priv_len, key_len, data_len = 0;

    void *p;

    if (key == NULL)

        return 0;

    /* The private key consists of 4 elements SK_SEED, SK_PRF, PK_SEED and PK_ROOT */

    priv_len = ossl_slh_dsa_key_get_priv_len(key);

    /* The size of either SK_SEED + SK_PRF OR PK_SEED + PK_ROOT */

    key_len = priv_len >> 1;

    /* Private key is optional */

    if (include_private) {

        if (param_priv != NULL) {

            p = key->priv;

            if (!OSSL_PARAM_get_octet_string(param_priv, &p, priv_len, &data_len))

                return 0;

            /* If the data read includes all 4 elements then we are finished */

            if (data_len == priv_len) {

                key->has_priv = 1;

                key->pub = SLH_DSA_PUB(key);

                return 1;

            }

            /* Otherwise it must be just SK_SEED + SK_PRF */

            if (data_len != key_len)

                goto err;

            key->has_priv = 1;

        }

    }

    /*

     * In the case where the passed in private key does not contain the public key

     * there MUST be a separate public key, since the private key cannot exist

     * without the public key elements. NOTE that this does not accept half of

     * the public key, (Keygen must be used for this case currently).

     */

    p = SLH_DSA_PUB(key);

    if (param_pub == NULL

            || !OSSL_PARAM_get_octet_string(param_pub, &p, key_len, &data_len)

            || data_len != key_len)

        goto err;

    key->pub = p;

    return 1;

 err:

    ossl_slh_dsa_key_reset(key);

    return 0;

}

/**

 * Generate the public key root from private key (seed and prf) and public key seed.

 * See FIPS 205 Section 9.1 Algorithm 18

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

 * @param out An SLH_DSA key containing the private key (seed and prf) and public key seed.

 *            The public root key is written to this key.

 * @param validate If set to 1 the computed public key is not written to the key,

 *                 but will be compared to the existing value.

 * @returns 1 if the root key is generated or compared successfully, or 0 on error.

 */

static int slh_dsa_compute_pk_root(SLH_DSA_HASH_CTX *ctx, SLH_DSA_KEY *out,

                                   int validate)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    SLH_ADRS_DECLARE(adrs);

    const SLH_DSA_PARAMS *params = key->params;

    size_t n = params->n;

    uint8_t pk_root[SLH_DSA_MAX_N], *dst;

    adrsf->zero(adrs);

    adrsf->set_layer_address(adrs, params->d - 1);

    dst = validate ? pk_root : SLH_DSA_PK_ROOT(out);

    /* Generate the ROOT public key */

    return ossl_slh_xmss_node(ctx, SLH_DSA_SK_SEED(key), 0, params->hm,

                              SLH_DSA_PK_SEED(key), adrs, dst, n)

        && (validate == 0 || memcmp(dst, SLH_DSA_PK_ROOT(out), n) == 0);

}

/**

 * @brief Generate a SLH_DSA keypair. The private key seed and prf as well as the

 * public key seed are generated using an approved DRBG's. The public key root is

 * calculated using these generated values.

 * See FIPS 205 Section 10.1 Algorithm 21

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants

 * @param out An SLH_DSA key to write key pair data to.

 * @param lib_ctx A library context for fetching RAND algorithms

 * @param entropy Optional entropy to use instead of using a DRBG.

 *        Required for ACVP testing. It may be NULL.

 * @param entropy_len the size of |entropy|. If set it must be at least 3 * |n|.

 * @returns 1 if the key is generated or 0 otherwise.

 */

int ossl_slh_dsa_generate_key(SLH_DSA_HASH_CTX *ctx, SLH_DSA_KEY *out,

                              OSSL_LIB_CTX *lib_ctx,

                              const uint8_t *entropy, size_t entropy_len)

{

    size_t n = out->params->n;

    size_t secret_key_len = 2 * n; /* The length of SK_SEED + SK_PRF */

    size_t pk_seed_len = n;        /* The length of PK_SEED */

    size_t entropy_len_expected = secret_key_len + pk_seed_len;

    uint8_t *priv = SLH_DSA_PRIV(out);

    uint8_t *pub = SLH_DSA_PUB(out);

    if (entropy != NULL && entropy_len != 0) {

        if (entropy_len != entropy_len_expected)

            goto err;

        memcpy(priv, entropy, entropy_len_expected);

    } else {

        if (RAND_priv_bytes_ex(lib_ctx, priv, secret_key_len, 0) <= 0

                || RAND_bytes_ex(lib_ctx, pub, pk_seed_len, 0) <= 0)

            goto err;

    }

    if (!slh_dsa_compute_pk_root(ctx, out, 0))

        goto err;

    out->pub = pub;

    out->has_priv = 1;

    return 1;

err:

    out->pub = NULL;

    out->has_priv = 0;

    OPENSSL_cleanse(priv, secret_key_len);

    return 0;

}

/**

 * @brief This is used when a SLH key is used for an operation.

 * This checks that the algorithm is the same (i.e. uses the same parameters)

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants to be used for

 *            an operation.

 * @param key A SLH_DSA key to use for an operation.

 *

 * @returns 1 if the algorithm matches, or 0 otherwise.

 */

int ossl_slh_dsa_key_type_matches(const SLH_DSA_KEY *key, const char *alg)

{

    return (OPENSSL_strcasecmp(key->params->alg, alg) == 0);

}

/* Returns the public key data or NULL if there is no public key */

const uint8_t *ossl_slh_dsa_key_get_pub(const SLH_DSA_KEY *key)

{

    return key->pub;

}

/* Returns the constant 2 * |n| which is the size of PK_SEED + PK_ROOT */

size_t ossl_slh_dsa_key_get_pub_len(const SLH_DSA_KEY *key)

{

    return 2 * key->params->n;

}

/* Returns the private key data or NULL if there is no private key */

const uint8_t *ossl_slh_dsa_key_get_priv(const SLH_DSA_KEY *key)

{

    return key->has_priv ? key->priv : NULL;

}

/*

 * Returns the constant 4 * |n| which is the size of both

 * the private and public key components.

 * SK_SEED + SK_ROOT + PK_SEED + PK_ROOT

 */

size_t ossl_slh_dsa_key_get_priv_len(const SLH_DSA_KEY *key)

{

    return 4 * key->params->n;

}

size_t ossl_slh_dsa_key_get_n(const SLH_DSA_KEY *key)

{

    return key->params->n;

}

int ossl_slh_dsa_key_get_security_category(const SLH_DSA_KEY *key)

{

    return key->params->security_category;

}

size_t ossl_slh_dsa_key_get_sig_len(const SLH_DSA_KEY *key)

{

    return key->params->sig_len;

}

const char *ossl_slh_dsa_key_get_name(const SLH_DSA_KEY *key)

{

    return key->params->alg;

}

int ossl_slh_dsa_key_get_type(const SLH_DSA_KEY *key)

{

    return key->params->type;

}

int ossl_slh_dsa_set_priv(SLH_DSA_KEY *key, const uint8_t *priv, size_t priv_len)

{

    if (ossl_slh_dsa_key_get_priv_len(key) != priv_len)

        return 0;

    memcpy(key->priv, priv, priv_len);

    key->has_priv = 1;

    key->pub = SLH_DSA_PUB(key);

    return 1;

}

int ossl_slh_dsa_set_pub(SLH_DSA_KEY *key, const uint8_t *pub, size_t pub_len)

{

    if (ossl_slh_dsa_key_get_pub_len(key) != pub_len)

        return 0;

    key->pub = SLH_DSA_PUB(key);

    memcpy(key->pub, pub, pub_len);

    key->has_priv = 0;

    return 1;

}

#ifndef FIPS_MODULE

int ossl_slh_dsa_key_to_text(BIO *out, const SLH_DSA_KEY *key, int selection)

{

    const char *name;

    if (out == NULL || key == NULL) {

        ERR_raise(ERR_LIB_PROV, ERR_R_PASSED_NULL_PARAMETER);

        return 0;

    }

    name = ossl_slh_dsa_key_get_name(key);

    if (ossl_slh_dsa_key_get_pub(key) == NULL) {

        /* Regardless of the |selection|, there must be a public key */

        ERR_raise_data(ERR_LIB_PROV, PROV_R_MISSING_KEY,

                       "no %s key material available", name);

        return 0;

    }

    if ((selection & OSSL_KEYMGMT_SELECT_PRIVATE_KEY) != 0) {

        if (ossl_slh_dsa_key_get_priv(key) == NULL) {

            ERR_raise_data(ERR_LIB_PROV, PROV_R_MISSING_KEY,

                           "no %s key material available", name);

            return 0;

        }

        if (BIO_printf(out, "%s Private-Key:\n", name) <= 0)

            return 0;

        if (!ossl_bio_print_labeled_buf(out, "priv:", ossl_slh_dsa_key_get_priv(key),

                                        ossl_slh_dsa_key_get_priv_len(key)))

            return 0;

    } else if ((selection & OSSL_KEYMGMT_SELECT_PUBLIC_KEY) != 0) {

        if (BIO_printf(out, "%s Public-Key:\n", name) <= 0)

            return 0;

    }

    if (!ossl_bio_print_labeled_buf(out, "pub:", ossl_slh_dsa_key_get_pub(key),

                                    ossl_slh_dsa_key_get_pub_len(key)))

        return 0;

    return 1;

}

#endif /* FIPS_MODULE */