/*

 * 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 "slh_dsa_local.h"

#include "slh_dsa_key.h"

/**

 * @brief Compute the root Public key of a XMSS tree.

 * See FIPS 205 Section 6.1 Algorithm 9.

 * This is a recursive function that starts at an leaf index, that calculates

 * the hash of each parent using 2 child nodes.

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

 * @param sk_seed A SLH-DSA private key seed of size |n|

 * @param nodeid The index of the target node being computed

 *               (which must be < 2^(hm - height)

 * @param h The height within the tree of the node being computed.

 *          (which must be <= hm) (hm is one of 3, 4, 8 or 9)

 *          At height=0 There are 2^hm leaf nodes,

 *          and the root node is at height = hm)

 * @param pk_seed A SLH-DSA public key seed of size |n|

 * @param adrs An ADRS object containing the layer address and tree address set

 *             to the XMSS tree within which the XMSS tree is being computed.

 * @param pk_out The generated public key of size |n|

 * @param pk_out_len The maximum size of |pk_out|

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

 */

int ossl_slh_xmss_node(SLH_DSA_HASH_CTX *ctx, const uint8_t *sk_seed,

                       uint32_t node_id, uint32_t h,

                       const uint8_t *pk_seed, uint8_t *adrs,

                       uint8_t *pk_out, size_t pk_out_len)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    if (h == 0) {

        /* For leaf nodes generate the public key */

        adrsf->set_type_and_clear(adrs, SLH_ADRS_TYPE_WOTS_HASH);

        adrsf->set_keypair_address(adrs, node_id);

        if (!ossl_slh_wots_pk_gen(ctx, sk_seed, pk_seed, adrs,

                                  pk_out, pk_out_len))

            return 0;

    } else {

        uint8_t lnode[SLH_MAX_N], rnode[SLH_MAX_N];

        if (!ossl_slh_xmss_node(ctx, sk_seed, 2 * node_id, h - 1, pk_seed, adrs,

                                lnode, sizeof(lnode))

                || !ossl_slh_xmss_node(ctx, sk_seed, 2 * node_id + 1, h - 1,

                                       pk_seed, adrs, rnode, sizeof(rnode)))

            return 0;

        adrsf->set_type_and_clear(adrs, SLH_ADRS_TYPE_TREE);

        adrsf->set_tree_height(adrs, h);

        adrsf->set_tree_index(adrs, node_id);

        if (!key->hash_func->H(ctx, pk_seed, adrs, lnode, rnode, pk_out, pk_out_len))

            return 0;

    }

    return 1;

}

/**

 * @brief Generate an XMSS signature using a message and key.

 * See FIPS 205 Section 6.2 Algorithm 10

 *

 * The generated signature consists of:

 *  - A WOTS+ signature of size (2 * n + 3) * n

 *  - An array of authentication paths of size (XMSS tree_height) * n.

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

 * @param msg A message of size |n| bytes to sign

 * @param sk_seed A private key seed of size |n|

 * @param node_id The index of a WOTS+ key within the XMSS tree to use for signing.

 * @param pk_seed A public key seed f size |n|

 * @param adrs An ADRS object containing the layer address and tree address set

 *              to the XMSS key being used to sign the message.

 * @param sig_wpkt A WPACKET object to write the generated XMSS signature to.

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

 */

int ossl_slh_xmss_sign(SLH_DSA_HASH_CTX *ctx, const uint8_t *msg,

                       const uint8_t *sk_seed, uint32_t node_id,

                       const uint8_t *pk_seed, uint8_t *adrs, WPACKET *sig_wpkt)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    SLH_ADRS_DECLARE(tmp_adrs);

    size_t n = key->params->n;

    uint32_t h, hm = key->params->hm;

    uint32_t id = node_id;

    uint8_t *auth_path; /* Pointer to a buffer offset inside |sig_wpkt| */

    size_t auth_path_len = n;

    /*

     * This code reverses the order of the FIPS 205 code so that it does the

     * sign first. This simplifies the WPACKET writing.

     */

    adrsf->copy(tmp_adrs, adrs);

    adrsf->set_type_and_clear(adrs, SLH_ADRS_TYPE_WOTS_HASH);

    adrsf->set_keypair_address(adrs, node_id);

    if (!ossl_slh_wots_sign(ctx, msg, sk_seed, pk_seed, adrs, sig_wpkt))

        return 0;

    adrsf->copy(adrs, tmp_adrs);

    for (h = 0; h < hm; ++h) {

        if (!WPACKET_allocate_bytes(sig_wpkt, auth_path_len, &auth_path)

                || !ossl_slh_xmss_node(ctx, sk_seed, id ^ 1, h, pk_seed, adrs,

                                       auth_path, auth_path_len))

            return 0;

        id >>= 1;

    }

    return 1;

}

/**

 * @brief Compute a candidate XMSS public key from a message and XMSS signature

 * See FIPS 205 Section 6.3 Algorithm 11

 *

 * * The signature consists of:

 *  - A WOTS+ signature of size (2 * n + 3) * n

 *  - An array of authentication paths of size (XMSS tree height) * n.

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

 * @param node_id Must be set to the |node_id| used in xmss_sign().

 * @param sig_rpkt A Packet to read a XMSS signature from.

 * @param msg A message of size |n| bytes

 * @param sk_seed A private key seed of size |n|

 * @param pk_seed A public key seed of size |n|

 * @param adrs An ADRS object containing a layer address and tree address of an

 *             XMSS key used for signing the message.

 * @param pk_out The returned candidate XMSS public key of size |n|.

 * @param pk_out_len The maximum size of |pk_out|.

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

 */

int ossl_slh_xmss_pk_from_sig(SLH_DSA_HASH_CTX *ctx, uint32_t node_id,

                              PACKET *sig_rpkt, const uint8_t *msg,

                              const uint8_t *pk_seed, uint8_t *adrs,

                              uint8_t *pk_out, size_t pk_out_len)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_HASH_FUNC_DECLARE(key, hashf);

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    SLH_HASH_FN_DECLARE(hashf, H);

    SLH_ADRS_FN_DECLARE(adrsf, set_tree_index);

    SLH_ADRS_FN_DECLARE(adrsf, set_tree_height);

    uint32_t k;

    size_t n = key->params->n;

    uint32_t hm = key->params->hm;

    uint8_t *node = pk_out;

    const uint8_t *auth_path; /* Pointer to buffer offset in |pkt_sig| */

    adrsf->set_type_and_clear(adrs, SLH_ADRS_TYPE_WOTS_HASH);

    adrsf->set_keypair_address(adrs, node_id);

    if (!ossl_slh_wots_pk_from_sig(ctx, sig_rpkt, msg, pk_seed, adrs,

                                   node, pk_out_len))

        return 0;

    adrsf->set_type_and_clear(adrs, SLH_ADRS_TYPE_TREE);

    for (k = 0; k < hm; ++k) {

        if (!PACKET_get_bytes(sig_rpkt, &auth_path, n))

            return 0;

        set_tree_height(adrs, k + 1);

        if ((node_id & 1) == 0) { /* even */

            node_id >>= 1;

            set_tree_index(adrs, node_id);

            if (!H(ctx, pk_seed, adrs, node, auth_path, node, pk_out_len))

                return 0;

        } else { /* odd */

            node_id = (node_id - 1) >> 1;

            set_tree_index(adrs, node_id);

            if (!H(ctx, pk_seed, adrs, auth_path, node, node, pk_out_len))

                return 0;

        }

    }

    return 1;

}