/*
 * Copyright (C) 2024 - 2026, Stephan Mueller <smueller@chronox.de>
 *
 * License: see LICENSE file in root directory
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, ALL OF
 * WHICH ARE HEREBY DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
 * OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF NOT ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */
/*
 * This code is derived in parts from the code distribution provided with
 * https://github.com/sphincs/sphincsplus
 *
 * That code is released under Public Domain
 * (https://creativecommons.org/share-your-work/public-domain/cc0/).
 */

#include "alignment.h"
#include "build_bug_on.h"
#include "compare.h"
#include "cpufeatures.h"
#include "helper.h"
#include "lc_rng.h"
#include "lc_memcmp_secure.h"
#include "lc_memcpy_secure.h"
#include "signature_domain_separation.h"
#include "small_stack_support.h"
#include "sphincs_type.h"
#include "sphincs_address.h"
#include "sphincs_fors.h"
#include "sphincs_hash.h"
#include "sphincs_internal.h"
#include "sphincs_merkle.h"
#include "sphincs_pct.h"
#include "sphincs_selftest.h"
#include "sphincs_thash.h"
#include "sphincs_utils.h"
#include "sphincs_wots.h"
#include "timecop.h"
#include "ret_checkers.h"
#include "visibility.h"

#include "avx2/sphincs_fors_avx2.h"
#include "avx2/sphincs_merkle_avx2.h"
#include "avx2/sphincs_wots_avx2.h"

#if defined(LC_HOST_AARCH64)
#include "armv8/sphincs_fors_armv8.h"
#include "armv8/sphincs_merkle_armv8.h"
#include "armv8/sphincs_wots_armv8.h"
#endif

struct lc_sphincs_sigver_func_ctx {
	fors_pk_from_sig_f fors_pk_from_sig;
	wots_pk_from_sig_f wots_pk_from_sig;
};

static const struct lc_sphincs_sigver_func_ctx f_ctx_c = {
	.fors_pk_from_sig = fors_pk_from_sig_c,
	.wots_pk_from_sig = wots_pk_from_sig_c,
};

static const struct lc_sphincs_sigver_func_ctx f_ctx_avx2 __maybe_unused = {
	.fors_pk_from_sig = fors_pk_from_sig_avx2,
	.wots_pk_from_sig = wots_pk_from_sig_avx2,
};

#if defined(LC_HOST_AARCH64)
static const struct lc_sphincs_sigver_func_ctx f_ctx_armv8 __maybe_unused = {
	.fors_pk_from_sig = fors_pk_from_sig_armv8,
	.wots_pk_from_sig = wots_pk_from_sig_armv8,
};
#endif

static const struct lc_sphincs_sigver_func_ctx *lc_sphincs_sigver_get_ctx(void)
{
	enum lc_cpu_features feat __maybe_unused = lc_cpu_feature_available();

#if (defined(LC_HOST_X86_64) && !defined(LC_SPHINCS_TYPE_128F_ASCON) &&        \
     !defined(LC_SPHINCS_TYPE_128S_ASCON))
	if (feat & LC_CPU_FEATURE_INTEL_AVX2) {
		return &f_ctx_avx2;
	} else
#endif /* LC_HOST_X86_64 */
#if (defined(LC_HOST_AARCH64) && !defined(LINUX_KERNEL) &&                     \
     !defined(LC_SPHINCS_TYPE_128F_ASCON) &&                                   \
     !defined(LC_SPHINCS_TYPE_128S_ASCON))
		/*
		 * TODO See issue in Kbuild.slh-dsa - enable NEON intrinsics
		 * for the Linux kernel.
		 */
		if (feat & LC_CPU_FEATURE_ARM_NEON) {
			return &f_ctx_armv8;
		}
#endif /* LC_HOST_AARCH64 */

	return &f_ctx_c;
}

/**
 * Verifies a detached signature and message under a given public key.
 */
int lc_sphincs_verify_ctx_nocheck(const struct lc_sphincs_sig *sig,
				  struct lc_sphincs_ctx *ctx, const uint8_t *m,
				  size_t mlen, const struct lc_sphincs_pk *pk)
{
	struct workspace {
		uint8_t pk_aligned[sizeof(pk->pk)] __align(sizeof(uint64_t));
		uint64_t tree;
		uint32_t idx_leaf;
		uint32_t wots_addr[8];
		uint32_t tree_addr[8];
		uint32_t wots_pk_addr[8];
		uint8_t root[LC_SPX_N];
		uint8_t leaf[LC_SPX_N];
		uint8_t wots_pk[LC_SPX_WOTS_BYTES] __align(sizeof(uint64_t));
		uint8_t mhash[LC_SPX_FORS_MSG_BYTES];
	};
	LC_HASH_CTX_ON_STACK(hash_ctx, LC_SPHINCS_HASH_TYPE);
	unsigned int i;
	const struct lc_sphincs_sigver_func_ctx *f_ctx =
		lc_sphincs_sigver_get_ctx();
	spx_ctx ctx_int;
	const uint8_t *pub_root = pk->pk + LC_SPX_N;
	const uint8_t *wots_sig = sig->sight;
	int ret = 0;
	LC_DECLARE_MEM(ws, struct workspace, sizeof(uint64_t));

	CKNULL(sig, -EINVAL);
	CKNULL(pk, -EINVAL);

	ctx_int.pub_seed = pk->pk;
	if (aligned(pk->pk, sizeof(uint64_t) - 1)) {
		ctx_int.pub_seed = pk->pk;
	} else {
		lc_memcpy_secure(ws->pk_aligned, sizeof(ws->pk_aligned), pk,
				 sizeof(pk->pk));
		ctx_int.pub_seed = ws->pk_aligned;
	}

	set_type(ws->wots_addr, LC_SPX_ADDR_TYPE_WOTS);
	set_type(ws->tree_addr, LC_SPX_ADDR_TYPE_HASHTREE);
	set_type(ws->wots_pk_addr, LC_SPX_ADDR_TYPE_WOTSPK);

	/* Derive the message digest and leaf index from R || PK || M. */
	/* The additional LC_SPX_N is a result of the hash domain separator. */
	CKINT(hash_message(ws->mhash, &ws->tree, &ws->idx_leaf, sig->r, pk->pk,
			   m, mlen, ctx));

	/* Layer correctly defaults to 0, so no need to set_layer_addr */
	set_tree_addr(ws->wots_addr, ws->tree);
	set_keypair_addr(ws->wots_addr, ws->idx_leaf);

	CKINT(f_ctx->fors_pk_from_sig(ws->root, sig->sigfors, ws->mhash,
				      &ctx_int, ws->wots_addr));

	/* For each subtree.. */
	for (i = 0; i < LC_SPX_D; i++) {
		set_layer_addr(ws->tree_addr, i);
		set_tree_addr(ws->tree_addr, ws->tree);

		copy_subtree_addr(ws->wots_addr, ws->tree_addr);
		set_keypair_addr(ws->wots_addr, ws->idx_leaf);

		copy_keypair_addr(ws->wots_pk_addr, ws->wots_addr);

		/*
		 * The WOTS public key is only correct if the signature was
		 * correct. Initially, root is the FORS pk, but on subsequent
		 * iterations it is the root of the subtree below the currently
		 * processed subtree.
		 */
		CKINT(f_ctx->wots_pk_from_sig(ws->wots_pk, wots_sig, ws->root,
					      &ctx_int, ws->wots_addr));
		wots_sig += LC_SPX_WOTS_BYTES;

		/* Compute the leaf node using the WOTS public key. */
		CKINT(thash(hash_ctx, ws->leaf, ws->wots_pk, LC_SPX_WOTS_LEN,
			    pk->pk, ws->wots_pk_addr));

		/* Compute the root node of this subtree. */
		CKINT(compute_root(ws->root, ws->leaf, ws->idx_leaf, 0,
				   wots_sig, LC_SPX_TREE_HEIGHT, pk->pk,
				   ws->tree_addr));
		wots_sig += LC_SPX_TREE_HEIGHT * LC_SPX_N;

		/* Update the indices for the next layer. */
		ws->idx_leaf = (ws->tree & ((1 << LC_SPX_TREE_HEIGHT) - 1));
		ws->tree = ws->tree >> LC_SPX_TREE_HEIGHT;
	}

	/* Check if the root node equals the root node in the public key. */
	CKRET_HARDENED(lc_memcmp_secure(ws->root, sizeof(ws->root), pub_root,
					LC_SPX_N),
		       -EBADMSG);

out:
	LC_RELEASE_MEM(ws);
	lc_hash_zero(hash_ctx);
	return ret;
}

LC_INTERFACE_FUNCTION(int, lc_sphincs_verify_ctx,
		      const struct lc_sphincs_sig *sig,
		      struct lc_sphincs_ctx *ctx, const uint8_t *m, size_t mlen,
		      const struct lc_sphincs_pk *pk)
{
	sphincs_selftest_sigver();
	LC_SELFTEST_COMPLETED(LC_ALG_STATUS_SLHDSA_SIGVER);

	return lc_sphincs_verify_ctx_nocheck(sig, ctx, m, mlen, pk);
}

LC_INTERFACE_FUNCTION(int, lc_sphincs_verify, const struct lc_sphincs_sig *sig,
		      const uint8_t *m, size_t mlen,
		      const struct lc_sphincs_pk *pk)
{
	LC_SPHINCS_CTX_ON_STACK(sphincs_ctx);
	int ret = lc_sphincs_verify_ctx(sig, sphincs_ctx, m, mlen, pk);

	lc_sphincs_ctx_zero(sphincs_ctx);
	return ret;
}

LC_INTERFACE_FUNCTION(int, lc_sphincs_verify_init, struct lc_sphincs_ctx *ctx,
		      const struct lc_sphincs_pk *pk)
{
	int ret = 0;

	/*
	 * We do not need the PK here - but we leave it to have a consistent
	 * API with ML-DSA.
	 */
	(void)pk;

	CKNULL(ctx, -EINVAL);

	sphincs_selftest_sigver();
	LC_SELFTEST_COMPLETED(LC_ALG_STATUS_SLHDSA_SIGVER);

	if (!ctx->sphincs_prehash_type) {
#if (LC_SPHINCS_NIST_CATEGORY == 1)
		ctx->sphincs_prehash_type = lc_sha3_256;
#elif (LC_SPHINCS_NIST_CATEGORY == 3)
		ctx->sphincs_prehash_type = lc_sha3_384;
#elif (LC_SPHINCS_NIST_CATEGORY == 5)
		ctx->sphincs_prehash_type = lc_sha3_512;
#else
#error "Unknown NIST category"
#endif
	}

	/*
	 * If the prehash type was set, we do not check it here, but the
	 * signature_domain_separation function will force the proper hash.
	 */

	/* Initialize the hash */
	LC_HASH_SET_CTX((&ctx->sphincs_hash_ctx), ctx->sphincs_prehash_type);

	CKINT(lc_hash_init(&ctx->sphincs_hash_ctx));

out:
	return ret;
}

LC_INTERFACE_FUNCTION(int, lc_sphincs_verify_update, struct lc_sphincs_ctx *ctx,
		      const uint8_t *m, size_t mlen)
{
	int ret = 0;

	CKNULL(ctx, -EINVAL);

	lc_hash_update(&ctx->sphincs_hash_ctx, m, mlen);

out:
	return ret;
}

LC_INTERFACE_FUNCTION(int, lc_sphincs_verify_final,
		      const struct lc_sphincs_sig *sig,
		      struct lc_sphincs_ctx *ctx,
		      const struct lc_sphincs_pk *pk)
{
	uint8_t digest[LC_SHA3_512_SIZE_DIGEST];
	int ret;

	CKNULL(ctx, -EINVAL);

	/*
	 * Catch SHAKE algos and set the digestsize following the
	 * specification.
	 */
	if (!lc_hash_digestsize(&ctx->sphincs_hash_ctx)) {
#if (LC_SPHINCS_NIST_CATEGORY == 1)
		lc_hash_set_digestsize(&ctx->sphincs_hash_ctx, 32);
#elif (LC_SPHINCS_NIST_CATEGORY == 3)
		lc_hash_set_digestsize(&ctx->sphincs_hash_ctx, 48);
#elif (LC_SPHINCS_NIST_CATEGORY == 5)
		lc_hash_set_digestsize(&ctx->sphincs_hash_ctx, 64);
#else
#error "Unknown NIST category"
#endif
	}

	/* Should never happen */
	if (lc_hash_digestsize(&ctx->sphincs_hash_ctx) > sizeof(digest))
		return -EFAULT;

	lc_hash_final(&ctx->sphincs_hash_ctx, digest);

	ret = lc_sphincs_verify_ctx(sig, ctx, digest,
				    lc_hash_digestsize(&ctx->sphincs_hash_ctx),
				    pk);

	/* Zeroize hash context in case of successful signature operation */
	if (!ret || ret == -EBADMSG)
		lc_hash_zero(&ctx->sphincs_hash_ctx);

out:
	return ret;
}
