/*

 * 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 Generate a Hypertree Signature

 * See FIPS 205 Section 7.1 Algorithm 12

 *

 * This writes |d| XMSS signatures i.e. ((|h| + |d| * |len|) * |n|)

 * where the first signature uses the XMSS key at the lowest layer, and the last

 * signature uses the XMSS key at the top layer.

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

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

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

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

 * @param tree_id Index of the XMSS tree that will sign the message

 * @param leaf_id Index of the WOTS+ key within the XMSS tree that will sign the message

 * @param sig_wpkt A WPACKET object to write the Hypertree Signature to.

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

 */

int ossl_slh_ht_sign(SLH_DSA_HASH_CTX *ctx,

                     const uint8_t *msg, const uint8_t *sk_seed,

                     const uint8_t *pk_seed,

                     uint64_t tree_id, uint32_t leaf_id, WPACKET *sig_wpkt)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    SLH_ADRS_DECLARE(adrs);

    uint8_t root[SLH_MAX_N];

    uint32_t layer, mask;

    const SLH_DSA_PARAMS *params = key->params;

    uint32_t n = params->n;

    uint32_t d = params->d;

    uint32_t hm = params->hm;

    uint8_t *psig;

    PACKET rpkt, *xmss_sig_rpkt = &rpkt;

    mask = (1 << hm) - 1; /* A mod 2^h = A & ((2^h - 1))) */

    adrsf->zero(adrs);

    /*

     * For each XMSS tree there is a current leaf node that is used for signing.

     * The first iteration of the loop signs the input message using the bottom

     * tree. Subsequent passes use the parent trees leaf node to sign the current

     * trees public key.

     * Each node in an XMSS tree has a sibling (except for the root node),

     * so starting at the leaf node it traverses up the tree calculating

     * hashes for all the siblings in the path to the root node,

     * which are then stored in the XMSS signature. The verify then just needs

     * the hash of the leaf node which is can then combine with the signature

     * path hashes to work all the way up to the root node to calculate the

     * public key.

     */

    memcpy(root, msg, n);

    for (layer = 0; layer < d; ++layer) {

        /* type = SLH_ADRS_TYPE_WOTS_HASH */

        adrsf->set_layer_address(adrs, layer);

        adrsf->set_tree_address(adrs, tree_id);

        psig = WPACKET_get_curr(sig_wpkt);

        if (!ossl_slh_xmss_sign(ctx, root, sk_seed, leaf_id, pk_seed, adrs,

                                sig_wpkt))

            return 0;

        /*

         * On the last loop it skips getting the public key since it is not needed

         * to calculate another signature. If this was called it should equal

         * the PK_ROOT (i.e. the public key of the top level tree).

         */

        if (layer < d - 1) {

            if (!PACKET_buf_init(xmss_sig_rpkt, psig,

                                 WPACKET_get_curr(sig_wpkt) - psig))

                return 0;

            if (!ossl_slh_xmss_pk_from_sig(ctx, leaf_id, xmss_sig_rpkt, root,

                                           pk_seed, adrs, root, sizeof(root)))

                return 0;

            leaf_id = tree_id & mask;

            tree_id >>= hm;

        }

    }

    return 1;

}

/**

 * @brief Verify a Hypertree Signature

 * See FIPS 205 Section 7.2 Algorithm 13

 *

 * @param ctx Contains SLH_DSA algorithm functions and constants.

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

 * @param sig A HT signature of size (|h| + |d| * |len|) * |n| bytes

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

 * @param tree_id Index of the XMSS tree that signed the message

 * @param leaf_id Index of the WOTS+ key within the XMSS tree that signed the message

 * @param pk_root The known Hypertree public key of size |n|

 *

 * @returns 1 if the computed XMSS public key matches pk_root, or 0 otherwise.

 */

int ossl_slh_ht_verify(SLH_DSA_HASH_CTX *ctx, const uint8_t *msg, PACKET *sig_pkt,

                       const uint8_t *pk_seed, uint64_t tree_id, uint32_t leaf_id,

                       const uint8_t *pk_root)

{

    const SLH_DSA_KEY *key = ctx->key;

    SLH_ADRS_FUNC_DECLARE(key, adrsf);

    SLH_ADRS_DECLARE(adrs);

    uint8_t node[SLH_MAX_N];

    const SLH_DSA_PARAMS *params = key->params;

    uint32_t tree_height = params->hm;

    uint32_t n = params->n;

    uint32_t d = params->d;

    uint32_t mask = (1 << tree_height) - 1;

    uint32_t layer;

    adrsf->zero(adrs);

    memcpy(node, msg, n);

    for (layer = 0; layer < d; ++layer) {

        adrsf->set_layer_address(adrs, layer);

        adrsf->set_tree_address(adrs, tree_id);

        if (!ossl_slh_xmss_pk_from_sig(ctx, leaf_id, sig_pkt, node,

                                       pk_seed, adrs, node, sizeof(node)))

            return 0;

        leaf_id = tree_id & mask;

        tree_id >>= tree_height;

    }

    return (memcmp(node, pk_root, n) == 0);

}