/*
 * 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 "internal/deprecated.h" /* PKCS1_MGF1() */

#include <string.h>
#include <openssl/evp.h>
#include <openssl/core_names.h>
#include <openssl/rsa.h> /* PKCS1_MGF1() */
#include "slh_dsa_local.h"
#include "slh_dsa_key.h"

#define MAX_DIGEST_SIZE 64 /* SHA-512 is used for security category 3 & 5 */

static OSSL_SLH_HASHFUNC_H_MSG slh_hmsg_sha2;
static OSSL_SLH_HASHFUNC_PRF slh_prf_sha2;
static OSSL_SLH_HASHFUNC_PRF_MSG slh_prf_msg_sha2;
static OSSL_SLH_HASHFUNC_F slh_f_sha2;
static OSSL_SLH_HASHFUNC_H slh_h_sha2;
static OSSL_SLH_HASHFUNC_T slh_t_sha2;

static OSSL_SLH_HASHFUNC_H_MSG slh_hmsg_shake;
static OSSL_SLH_HASHFUNC_PRF slh_prf_shake;
static OSSL_SLH_HASHFUNC_PRF_MSG slh_prf_msg_shake;
static OSSL_SLH_HASHFUNC_F slh_f_shake;
static OSSL_SLH_HASHFUNC_H slh_h_shake;
static OSSL_SLH_HASHFUNC_T slh_t_shake;

static ossl_inline int xof_digest_3(EVP_MD_CTX *ctx,
    const uint8_t *in1, size_t in1_len,
    const uint8_t *in2, size_t in2_len,
    const uint8_t *in3, size_t in3_len,
    uint8_t *out, size_t out_len)
{
    return (EVP_DigestInit_ex2(ctx, NULL, NULL) == 1
        && EVP_DigestUpdate(ctx, in1, in1_len) == 1
        && EVP_DigestUpdate(ctx, in2, in2_len) == 1
        && EVP_DigestUpdate(ctx, in3, in3_len) == 1
        && EVP_DigestFinalXOF(ctx, out, out_len) == 1);
}

static ossl_inline int xof_digest_4(EVP_MD_CTX *ctx,
    const uint8_t *in1, size_t in1_len,
    const uint8_t *in2, size_t in2_len,
    const uint8_t *in3, size_t in3_len,
    const uint8_t *in4, size_t in4_len,
    uint8_t *out, size_t out_len)
{
    return (EVP_DigestInit_ex2(ctx, NULL, NULL) == 1
        && EVP_DigestUpdate(ctx, in1, in1_len) == 1
        && EVP_DigestUpdate(ctx, in2, in2_len) == 1
        && EVP_DigestUpdate(ctx, in3, in3_len) == 1
        && EVP_DigestUpdate(ctx, in4, in4_len) == 1
        && EVP_DigestFinalXOF(ctx, out, out_len) == 1);
}

/* See FIPS 205 Section 11.1 */
static int
slh_hmsg_shake(SLH_DSA_HASH_CTX *ctx, const uint8_t *r,
    const uint8_t *pk_seed, const uint8_t *pk_root,
    const uint8_t *msg, size_t msg_len,
    uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t m = params->m;
    size_t n = params->n;

    return xof_digest_4(ctx->md_ctx, r, n, pk_seed, n, pk_root, n,
        msg, msg_len, out, m);
}

static int
slh_prf_shake(SLH_DSA_HASH_CTX *ctx,
    const uint8_t *pk_seed, const uint8_t *sk_seed,
    const uint8_t *adrs, uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t n = params->n;

    return xof_digest_3(ctx->md_ctx, pk_seed, n, adrs, SLH_ADRS_SIZE,
        sk_seed, n, out, n);
}

static int
slh_prf_msg_shake(SLH_DSA_HASH_CTX *ctx, const uint8_t *sk_prf,
    const uint8_t *opt_rand, const uint8_t *msg, size_t msg_len,
    WPACKET *pkt)
{
    unsigned char out[SLH_MAX_N];
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t n = params->n;

    return xof_digest_3(ctx->md_ctx, sk_prf, n, opt_rand, n, msg, msg_len, out, n)
        && WPACKET_memcpy(pkt, out, n);
}

static int
slh_f_shake(SLH_DSA_HASH_CTX *ctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *m1, size_t m1_len, uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t n = params->n;

    return xof_digest_3(ctx->md_ctx, pk_seed, n, adrs, SLH_ADRS_SIZE, m1, m1_len, out, n);
}

static int
slh_h_shake(SLH_DSA_HASH_CTX *ctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *m1, const uint8_t *m2, uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t n = params->n;

    return xof_digest_4(ctx->md_ctx, pk_seed, n, adrs, SLH_ADRS_SIZE, m1, n, m2, n, out, n);
}

static int
slh_t_shake(SLH_DSA_HASH_CTX *ctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *ml, size_t ml_len, uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = ctx->key->params;
    size_t n = params->n;

    return xof_digest_3(ctx->md_ctx, pk_seed, n, adrs, SLH_ADRS_SIZE, ml, ml_len, out, n);
}

static ossl_inline int
digest_4(EVP_MD_CTX *ctx,
    const uint8_t *in1, size_t in1_len, const uint8_t *in2, size_t in2_len,
    const uint8_t *in3, size_t in3_len, const uint8_t *in4, size_t in4_len,
    uint8_t *out)
{
    return (EVP_DigestInit_ex2(ctx, NULL, NULL) == 1
        && EVP_DigestUpdate(ctx, in1, in1_len) == 1
        && EVP_DigestUpdate(ctx, in2, in2_len) == 1
        && EVP_DigestUpdate(ctx, in3, in3_len) == 1
        && EVP_DigestUpdate(ctx, in4, in4_len) == 1
        && EVP_DigestFinal_ex(ctx, out, NULL) == 1);
}

/* FIPS 205 Section 11.2.1 and 11.2.2 */

static int
slh_hmsg_sha2(SLH_DSA_HASH_CTX *hctx, const uint8_t *r, const uint8_t *pk_seed,
    const uint8_t *pk_root, const uint8_t *msg, size_t msg_len,
    uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *params = hctx->key->params;
    size_t m = params->m;
    size_t n = params->n;
    uint8_t seed[2 * SLH_MAX_N + MAX_DIGEST_SIZE];
    int sz = EVP_MD_get_size(hctx->key->md_big);
    size_t seed_len = (size_t)sz + 2 * n;

    if (sz <= 0)
        return 0;

    memcpy(seed, r, n);
    memcpy(seed + n, pk_seed, n);
    return digest_4(hctx->md_big_ctx, r, n, pk_seed, n, pk_root, n, msg, msg_len,
               seed + 2 * n)
        && (PKCS1_MGF1(out, (long)m, seed, (long)seed_len, hctx->key->md_big) == 0);
}

static int
slh_prf_msg_sha2(SLH_DSA_HASH_CTX *hctx,
    const uint8_t *sk_prf, const uint8_t *opt_rand,
    const uint8_t *msg, size_t msg_len, WPACKET *pkt)
{
    int ret;
    const SLH_DSA_KEY *key = hctx->key;
    EVP_MAC_CTX *mctx = hctx->hmac_ctx;
    const SLH_DSA_PARAMS *prms = key->params;
    size_t n = prms->n;
    uint8_t mac[MAX_DIGEST_SIZE] = { 0 };
    OSSL_PARAM *p = NULL;
    OSSL_PARAM params[3];

    /*
     * Due to the way HMAC works, it is not possible to do this code early
     * in hmac_ctx_new() since it requires a key in order to set the digest.
     * So we do a lazy update here on the first call.
     */
    if (hctx->hmac_digest_used == 0) {
        p = params;
        /* The underlying digest to be used */
        *p++ = OSSL_PARAM_construct_utf8_string(OSSL_MAC_PARAM_DIGEST,
            (char *)EVP_MD_get0_name(key->md_big), 0);
        if (key->propq != NULL)
            *p++ = OSSL_PARAM_construct_utf8_string(OSSL_MAC_PARAM_PROPERTIES,
                (char *)key->propq, 0);
        *p = OSSL_PARAM_construct_end();
        p = params;
        hctx->hmac_digest_used = 1;
    }

    ret = EVP_MAC_init(mctx, sk_prf, n, p) == 1
        && EVP_MAC_update(mctx, opt_rand, n) == 1
        && EVP_MAC_update(mctx, msg, msg_len) == 1
        && EVP_MAC_final(mctx, mac, NULL, sizeof(mac)) == 1
        && WPACKET_memcpy(pkt, mac, n); /* Truncate output to n bytes */
    return ret;
}

static ossl_inline int
do_hash(EVP_MD_CTX *ctx, size_t n, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *m, size_t m_len, size_t b, uint8_t *out, size_t out_len)
{
    int ret;
    uint8_t zeros[128] = { 0 };
    uint8_t digest[MAX_DIGEST_SIZE];

    ret = digest_4(ctx, pk_seed, n, zeros, b - n, adrs, SLH_ADRSC_SIZE,
        m, m_len, digest);
    /* Truncated returned value is n = 16 bytes */
    memcpy(out, digest, n);
    return ret;
}

static int
slh_prf_sha2(SLH_DSA_HASH_CTX *hctx, const uint8_t *pk_seed,
    const uint8_t *sk_seed, const uint8_t *adrs,
    uint8_t *out, size_t out_len)
{
    size_t n = hctx->key->params->n;

    return do_hash(hctx->md_ctx, n, pk_seed, adrs, sk_seed, n,
        OSSL_SLH_DSA_SHA2_NUM_ZEROS_H_AND_T_BOUND1, out, out_len);
}

static int
slh_f_sha2(SLH_DSA_HASH_CTX *hctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *m1, size_t m1_len, uint8_t *out, size_t out_len)
{
    return do_hash(hctx->md_ctx, hctx->key->params->n, pk_seed, adrs, m1, m1_len,
        OSSL_SLH_DSA_SHA2_NUM_ZEROS_H_AND_T_BOUND1, out, out_len);
}

static int
slh_h_sha2(SLH_DSA_HASH_CTX *hctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *m1, const uint8_t *m2, uint8_t *out, size_t out_len)
{
    uint8_t m[SLH_MAX_N * 2];
    const SLH_DSA_PARAMS *prms = hctx->key->params;
    size_t n = prms->n;

    memcpy(m, m1, n);
    memcpy(m + n, m2, n);
    return do_hash(hctx->md_big_ctx, n, pk_seed, adrs, m, 2 * n,
        prms->sha2_h_and_t_bound, out, out_len);
}

static int
slh_t_sha2(SLH_DSA_HASH_CTX *hctx, const uint8_t *pk_seed, const uint8_t *adrs,
    const uint8_t *ml, size_t ml_len, uint8_t *out, size_t out_len)
{
    const SLH_DSA_PARAMS *prms = hctx->key->params;

    return do_hash(hctx->md_big_ctx, prms->n, pk_seed, adrs, ml, ml_len,
        prms->sha2_h_and_t_bound, out, out_len);
}

const SLH_HASH_FUNC *ossl_slh_get_hash_fn(int is_shake)
{
    static const SLH_HASH_FUNC methods[] = {
        { slh_hmsg_shake,
            slh_prf_shake,
            slh_prf_msg_shake,
            slh_f_shake,
            slh_h_shake,
            slh_t_shake },
        { slh_hmsg_sha2,
            slh_prf_sha2,
            slh_prf_msg_sha2,
            slh_f_sha2,
            slh_h_sha2,
            slh_t_sha2 }
    };
    return &methods[is_shake ? 0 : 1];
}