/*
 * 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/crypto.h>
#include "slh_dsa_local.h"
#include "slh_dsa_key.h"

/* k = 14, 17, 22, 33, 35 (number of trees) */
#define SLH_MAX_K 35
/* a = 6, 8, 9, 12 or 14  - There are (2^a) merkle trees */
#define SLH_MAX_A 9

#define SLH_MAX_K_TIMES_A (SLH_MAX_A * SLH_MAX_K)
#define SLH_MAX_ROOTS (SLH_MAX_K_TIMES_A * SLH_MAX_N)

static void slh_base_2b(const uint8_t *in, uint32_t b, uint32_t *out, size_t out_len);

/**
 * @brief Generate FORS secret values
 * See FIPS 205 Section 8.1 Algorithm 14.
 *
 * @param ctx Contains SLH_DSA algorithm functions and constants.
 * @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 the layer address of zero, with the
 *             tree address and key pair address set to the index of the WOTS+
 *             key within the XMSS tree that signs the FORS key.
 * @param id The index of the FORS secret value within the sets of FORS trees.
 *               (which must be < 2^(hm - height)
 * @param pk_out The generated FORS secret value of size |n|
 * @param pk_out_len The maximum size of |pk_out|
 * @returns 1 on success, or 0 on error.
 */
static int slh_fors_sk_gen(SLH_DSA_HASH_CTX *ctx, const uint8_t *sk_seed,
    const uint8_t *pk_seed, uint8_t *adrs, uint32_t id,
    uint8_t *pk_out, size_t pk_out_len)
{
    const SLH_DSA_KEY *key = ctx->key;
    SLH_ADRS_DECLARE(sk_adrs);
    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    adrsf->copy(sk_adrs, adrs);
    adrsf->set_type_and_clear(sk_adrs, SLH_ADRS_TYPE_FORS_PRF);
    adrsf->copy_keypair_address(sk_adrs, adrs);
    adrsf->set_tree_index(sk_adrs, id);
    return key->hash_func->PRF(ctx, pk_seed, sk_seed, sk_adrs, pk_out, pk_out_len);
}

/**
 * @brief Computes the nodes of a Merkle tree.
 * See FIPS 205 Section 8.2 Algorithm 18
 *
 * The leaf nodes are hashes of FORS secret values.
 * Each parent node is a hash of its 2 children.
 * Note this is a recursive function.
 *
 * @param ctx Contains SLH_DSA algorithm functions and constants.
 * @param sk_seed A SLH_DSA private key seed of size |n|
 * @param pk_seed A SLH_DSA public key seed of size |n|
 * @param adrs The ADRS object must have a layer address of zero, and the
 *             tree address set to the XMSS tree that signs the FORS key,
 *             the type set to FORS_TREE, and the keypair address set to the
 *             index of the WOTS+ key that signs the FORS key.
 * @param node_id The target node index
 * @param height The target node height
 * @param node The returned hash for a node of size|n|
 * @param node_len The maximum size of |node|
 * @returns 1 on success, or 0 on error.
 */
static int slh_fors_node(SLH_DSA_HASH_CTX *ctx, const uint8_t *sk_seed,
    const uint8_t *pk_seed, uint8_t *adrs, uint32_t node_id,
    uint32_t height, uint8_t *node, size_t node_len)
{
    int ret = 0;
    const SLH_DSA_KEY *key = ctx->key;
    uint8_t sk[SLH_MAX_N], lnode[SLH_MAX_N], rnode[SLH_MAX_N];
    uint32_t n = key->params->n;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    if (height == 0) {
        /* Gets here for leaf nodes */
        if (!slh_fors_sk_gen(ctx, sk_seed, pk_seed, adrs, node_id, sk, sizeof(sk)))
            return 0;
        adrsf->set_tree_height(adrs, 0);
        adrsf->set_tree_index(adrs, node_id);
        ret = key->hash_func->F(ctx, pk_seed, adrs, sk, n, node, node_len);
        OPENSSL_cleanse(sk, n);
        return ret;
    } else {
        if (!slh_fors_node(ctx, sk_seed, pk_seed, adrs, 2 * node_id, height - 1,
                lnode, sizeof(rnode))
            || !slh_fors_node(ctx, sk_seed, pk_seed, adrs, 2 * node_id + 1,
                height - 1, rnode, sizeof(rnode)))
            return 0;
        adrsf->set_tree_height(adrs, height);
        adrsf->set_tree_index(adrs, node_id);
        if (!key->hash_func->H(ctx, pk_seed, adrs, lnode, rnode, node, node_len))
            return 0;
    }
    return 1;
}

/**
 * @brief Generate an FORS signature
 * See FIPS 205 Section 8.3 Algorithm 16
 *
 * A FORS signature has a size of (k * (1 + a) * n) bytes
 * There are k trees, each of which have a private key value of size |n| followed
 * by an authentication path of size |a| (where each path is size |n|)
 *
 * @param ctx Contains SLH_DSA algorithm functions and constants.
 * @param md A message digest of size |(k * a + 7) / 8| bytes to sign
 * @param sk_seed A private key seed of size |n|
 * @param pk_seed A public key seed of size |n|
 * @param adrs The ADRS object must have a layer address of zero, and the
 *             tree address set to the XMSS tree that signs the FORS key,
 *             the type set to FORS_TREE, and the keypair address set to the
 *             index of the WOTS+ key that signs the FORS key.
 * @param sig_wpkt A WPACKET object to write the generated XMSS signature to
 * @param sig_len  The size of |sig| which is (2 * n + 3) * n + tree_height * n.
 * @returns 1 on success, or 0 on error.
 */
int ossl_slh_fors_sign(SLH_DSA_HASH_CTX *ctx, const uint8_t *md,
    const uint8_t *sk_seed, const uint8_t *pk_seed,
    uint8_t *adrs, WPACKET *sig_wpkt)
{
    const SLH_DSA_KEY *key = ctx->key;
    uint32_t tree_id, layer, s, tree_offset;
    uint32_t ids[SLH_MAX_K];
    const SLH_DSA_PARAMS *params = key->params;
    uint32_t n = params->n;
    uint32_t k = params->k; /* number of trees */
    uint32_t a = params->a;
    uint32_t two_power_a = (1 << a); /* this is t in FIPS 205 */
    uint32_t tree_id_times_two_power_a = 0;
    uint8_t out[SLH_MAX_N];

    /*
     * Split md into k a-bit values e.g with k = 14, a = 12
     * ids[0..13] = 12 bits each of md
     */
    slh_base_2b(md, a, ids, k);

    for (tree_id = 0; tree_id < k; ++tree_id) {
        /* Get the tree[i] leaf id */
        uint32_t node_id = ids[tree_id]; /* |id| = |a| bits */

        /*
         * Give each of the k trees a unique range at each level.
         * e.g. If we have 4096 leaf nodes (2^a = 2^12) for each tree
         * tree i will use indexes from 4096 * i + (0..4095) for its bottom level.
         * For the next level up from the bottom there would be 2048 nodes
         * (so tree i uses indexes 2048 * i + (0...2047) for this level)
         */
        tree_offset = tree_id_times_two_power_a;

        if (!slh_fors_sk_gen(ctx, sk_seed, pk_seed, adrs,
                node_id + tree_id_times_two_power_a, out, sizeof(out))
            || !WPACKET_memcpy(sig_wpkt, out, n))
            return 0;

        /*
         * Traverse from the bottom of the tree (layer = 0)
         * up to the root (layer = a - 1).
         * NOTE: This is a really inefficient way of doing this, since at
         * layer a - 1 it calculates most of the hashes of the entire tree as
         * well as all the leaf nodes. So it is calculating nodes multiple times.
         */
        for (layer = 0; layer < a; ++layer) {
            s = node_id ^ 1; /* XOR gets the index of the other child in a binary tree */
            if (!slh_fors_node(ctx, sk_seed, pk_seed, adrs,
                    s + tree_offset, layer, out, sizeof(out)))
                return 0;
            node_id >>= 1; /* Get the parent node id */
            tree_offset >>= 1; /* Each layer up has half as many nodes */
            if (!WPACKET_memcpy(sig_wpkt, out, n))
                return 0;
        }
        tree_id_times_two_power_a += two_power_a;
    }
    return 1;
}

/**
 * @brief Compute a candidate FORS public key from a message and signature.
 * See FIPS 205 Section 8.4 Algorithm 17.
 *
 * A FORS signature has a size of (k * (a + 1) * n) bytes
 *
 * @param ctx Contains SLH_DSA algorithm functions and constants.
 * @param fors_sig_rpkt A PACKET object to read a FORS signature from
 * @param md A message digest of size (k * a / 8) bytes
 * @param pk_seed A public key seed of size |n|
 * @param adrs The ADRS object must have a layer address of zero, and the
 *             tree address set to the XMSS tree that signs the FORS key,
 *             the type set to FORS_TREE, and the keypair address set to the
 *             index of the WOTS+ key that signs the FORS key.
 * @param pk_out The returned candidate FORS 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_fors_pk_from_sig(SLH_DSA_HASH_CTX *ctx, PACKET *fors_sig_rpkt,
    const uint8_t *md, const uint8_t *pk_seed,
    uint8_t *adrs, uint8_t *pk_out, size_t pk_out_len)
{
    const SLH_DSA_KEY *key = ctx->key;
    int ret = 0;
    uint32_t i, j, aoff = 0;
    uint32_t ids[SLH_MAX_K];
    const SLH_DSA_PARAMS *params = key->params;
    uint32_t a = params->a;
    uint32_t k = params->k;
    uint32_t n = params->n;
    uint32_t two_power_a = (1 << a);
    const uint8_t *sk, *authj; /* Pointers to |sig| buffer inside fors_sig_rpkt */
    uint8_t roots[SLH_MAX_ROOTS];
    size_t roots_len = 0; /* The size of |roots| */
    uint8_t *node0, *node1; /* Pointers into roots[] */
    WPACKET root_pkt, *wroot_pkt = &root_pkt; /* Points to |roots| buffer */

    SLH_ADRS_DECLARE(pk_adrs);
    SLH_ADRS_FUNC_DECLARE(key, adrsf);
    SLH_ADRS_FN_DECLARE(adrsf, set_tree_index);
    SLH_ADRS_FN_DECLARE(adrsf, set_tree_height);
    SLH_HASH_FUNC_DECLARE(key, hashf);
    SLH_HASH_FN_DECLARE(hashf, F);
    SLH_HASH_FN_DECLARE(hashf, H);

    if (!WPACKET_init_static_len(wroot_pkt, roots, sizeof(roots), 0))
        return 0;

    /* Split md into k a-bit values e.g ids[0..k-1] = 12 bits each of md */
    slh_base_2b(md, a, ids, k);

    /* Compute the roots of k Merkle trees */
    for (i = 0; i < k; ++i) {
        uint32_t id = ids[i];
        uint32_t node_id = id + aoff;

        set_tree_height(adrs, 0);
        set_tree_index(adrs, node_id);

        /* Regenerate the public key of the leaf */
        if (!PACKET_get_bytes(fors_sig_rpkt, &sk, n)
            || !WPACKET_allocate_bytes(wroot_pkt, n, &node0)
            || !F(ctx, pk_seed, adrs, sk, n, node0, n))
            goto err;

        /* This omits the copying of the nodes that the FIPS 205 code does */
        node1 = node0;
        for (j = 0; j < a; ++j) {
            /* Get this layers other child public key */
            if (!PACKET_get_bytes(fors_sig_rpkt, &authj, n))
                goto err;
            /* Hash the children together to get the parent nodes public key */
            set_tree_height(adrs, j + 1);
            if ((id & 1) == 0) {
                node_id >>= 1;
                set_tree_index(adrs, node_id);
                if (!H(ctx, pk_seed, adrs, node0, authj, node1, n))
                    goto err;
            } else {
                node_id = (node_id - 1) >> 1;
                set_tree_index(adrs, node_id);
                if (!H(ctx, pk_seed, adrs, authj, node0, node1, n))
                    goto err;
            }
            id >>= 1;
        }
        aoff += two_power_a;
    }
    if (!WPACKET_get_total_written(wroot_pkt, &roots_len))
        goto err;

    /* The public key is the hash of all the roots of the k trees */
    adrsf->copy(pk_adrs, adrs);
    adrsf->set_type_and_clear(pk_adrs, SLH_ADRS_TYPE_FORS_ROOTS);
    adrsf->copy_keypair_address(pk_adrs, adrs);
    ret = hashf->T(ctx, pk_seed, pk_adrs, roots, roots_len, pk_out, pk_out_len);
err:
    if (!WPACKET_finish(wroot_pkt))
        ret = 0;
    return ret;
}

/**
 * @brief Convert a byte string into a base 2^b representation
 * See FIPS 205 Algorithm 4
 *
 * @param in An input byte stream with a size >= |outlen * b / 8|
 * @param b The bit size to divide |in| into
 *          This is one of 6, 8, 9, 12 or 14 for FORS.
 * @param out The array of returned base 2^b integers that represents the first
 *            |outlen|*|b| bits of |in|
 * @param out_len The size of |out|
 */
static void slh_base_2b(const uint8_t *in, uint32_t b,
    uint32_t *out, size_t out_len)
{
    size_t consumed = 0;
    uint32_t bits = 0;
    uint32_t total = 0;
    uint32_t mask = (1 << b) - 1;

    for (consumed = 0; consumed < out_len; consumed++) {
        while (bits < b) {
            total <<= 8;
            total += *in++;
            bits += 8;
        }
        bits -= b;
        *out++ = (total >> bits) & mask;
    }
}