// SPDX-License-Identifier: GPL-2.0 OR BSD-2-Clause
/*
 * Copyright (C) 2024 - 2025, 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.
 */

#include <crypto/internal/sig.h>
#include <crypto/scatterwalk.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/scatterlist.h>
#include <linux/types.h>

#include "lc_dilithium.h"
#include "lc_sha3.h"

#include "leancrypto_kernel.h"

enum lc_kernel_dilithium_ed25519_key_type {
	lc_kernel_dilithium_ed25519_key_unset = 0,
	lc_kernel_dilithium_ed25519_key_sk = 1,
	lc_kernel_dilithium_ed25519_key_pk = 2,
};

struct lc_kernel_dilithium_ed25519_ctx {
	union {
		struct lc_dilithium_ed25519_sk sk;
		struct lc_dilithium_ed25519_pk pk;
	};
	enum lc_kernel_dilithium_ed25519_key_type key_type;
};

/* src -> message */
/* dst -> signature */
static int lc_kernel_dilithium_ed25519_sign(struct crypto_sig *tfm,
					    const void *src, unsigned int slen,
					    void *dst, unsigned int dlen)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);
	struct lc_dilithium_ed25519_sig *sig;
	uint8_t *sig_ptr, *sig_ed25519_ptr;
	size_t sig_len, sig_ed25519_len;
	enum lc_dilithium_type type;
	int ret;

	if (unlikely(ctx->key_type != lc_kernel_dilithium_ed25519_key_sk))
		return -EINVAL;

	type = lc_dilithium_ed25519_sk_type(&ctx->sk);
	if (dlen != lc_dilithium_ed25519_sig_size(type))
		return -EINVAL;

	sig = kmalloc(sizeof(struct lc_dilithium_ed25519_sig), GFP_KERNEL);
	if (!sig)
		return -ENOMEM;

	ret = lc_dilithium_ed25519_sign(sig, src, slen, &ctx->sk,
					lc_seeded_rng);
	if (ret)
		goto out;

	ret = lc_dilithium_ed25519_sig_ptr(&sig_ptr, &sig_len, &sig_ed25519_ptr,
					   &sig_ed25519_len, sig);
	if (ret)
		goto out;

	if (dlen < sig_len + sig_ed25519_len) {
		ret = -EOVERFLOW;
		goto out;
	}

	/* Copy the Dilithium signature */
	memcpy(dst, sig_ptr, sig_len);

	/*
	 * Copy the ED25519 signature which is simply concatenated to
	 * the Dilithium signature.
	 */
	memcpy((uint8_t *)dst + sig_len, sig_ed25519_ptr, sig_ed25519_len);

out:
	free_zero(sig);
	return ret;
}

/* src -> Dilithium signature || ED25519 signature */
/* msg -> message */
static int lc_kernel_dilithium_ed25519_verify(struct crypto_sig *tfm,
					      const void *src,
					      unsigned int slen,
					      const void *msg,
					      unsigned int msg_len)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);
	struct lc_dilithium_ed25519_sig *sig;
	size_t sig_len, sig_ed25519_len;
	enum lc_dilithium_type type;
	uint8_t *sig_ptr, *sig_ed25519_ptr;
	int ret;

	if (unlikely(ctx->key_type != lc_kernel_dilithium_ed25519_key_pk))
		return -EINVAL;

	type = lc_dilithium_ed25519_pk_type(&ctx->pk);
	if (slen < lc_dilithium_ed25519_sig_size(type))
		return -EINVAL;

	sig = kmalloc(sizeof(struct lc_dilithium_ed25519_sig), GFP_KERNEL);
	if (!sig)
		return -ENOMEM;

	/*
	 * Obtain the empty pointers to fill it with a signature. Thus, we
	 * need to set the signature type here as the signature struct is
	 * currently unset.
	 */
	sig->dilithium_type = type;
	ret = lc_dilithium_ed25519_sig_ptr(&sig_ptr, &sig_len, &sig_ed25519_ptr,
					   &sig_ed25519_len, sig);
	if (ret)
		goto out;

	ret = lc_dilithium_ed25519_sig_load(
		sig, src, sig_len, (uint8_t *)src + sig_len, sig_ed25519_len);
	if (ret)
		goto out;

	ret = lc_dilithium_ed25519_verify(sig, msg, msg_len, &ctx->pk);

out:
	free_zero(sig);
	return ret;
}

static int lc_kernel_dilithium_ed25519_set_pub_key_int(
	struct crypto_sig *tfm, const void *key, unsigned int keylen,
	enum lc_dilithium_type type)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);
	int ret;

	if (keylen < LC_ED25519_PUBLICKEYBYTES)
		return -EINVAL;

	ctx->key_type = lc_kernel_dilithium_ed25519_key_unset;

	/*
	 * Load the Dilithium and the ED25519 keys - they are expected to be
	 * concatenated in the linear buffer of key.
	 */
	ret = lc_dilithium_ed25519_pk_load(
		&ctx->pk, key, keylen - LC_ED25519_PUBLICKEYBYTES,
		key + keylen - LC_ED25519_PUBLICKEYBYTES,
		LC_ED25519_PUBLICKEYBYTES);

	if (!ret) {
		if (lc_dilithium_ed25519_pk_type(&ctx->pk) != type)
			ret = -EOPNOTSUPP;
		else
			ctx->key_type = lc_kernel_dilithium_ed25519_key_pk;
	}

	return ret;
}

static unsigned int
lc_kernel_dilithium_ed25519_87_key_size(struct crypto_sig *tfm)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);

	switch (ctx->key_type) {
	case lc_kernel_dilithium_ed25519_key_sk:
		return sizeof(struct lc_dilithium_87_ed25519_sk);

	case lc_kernel_dilithium_ed25519_key_unset:
	case lc_kernel_dilithium_ed25519_key_pk:
	default:
		return sizeof(struct lc_dilithium_87_ed25519_pk);
	}
}

static unsigned int
lc_kernel_dilithium_ed25519_65_key_size(struct crypto_sig *tfm)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);

	switch (ctx->key_type) {
	case lc_kernel_dilithium_ed25519_key_sk:
		return sizeof(struct lc_dilithium_65_ed25519_sk);

	case lc_kernel_dilithium_ed25519_key_unset:
	case lc_kernel_dilithium_ed25519_key_pk:
	default:
		return sizeof(struct lc_dilithium_65_ed25519_pk);
	}
}

static unsigned int
lc_kernel_dilithium_ed25519_44_key_size(struct crypto_sig *tfm)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);

	switch (ctx->key_type) {
	case lc_kernel_dilithium_ed25519_key_sk:
		return sizeof(struct lc_dilithium_44_ed25519_sk);

	case lc_kernel_dilithium_ed25519_key_unset:
	case lc_kernel_dilithium_ed25519_key_pk:
	default:
		return sizeof(struct lc_dilithium_44_ed25519_pk);
	}
}

static int lc_kernel_dilithium_ed25519_44_set_pub_key(struct crypto_sig *tfm,
						      const void *key,
						      unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_pub_key_int(tfm, key, keylen,
							   LC_DILITHIUM_44);
}

static int lc_kernel_dilithium_ed25519_65_set_pub_key(struct crypto_sig *tfm,
						      const void *key,
						      unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_pub_key_int(tfm, key, keylen,
							   LC_DILITHIUM_65);
}

static int lc_kernel_dilithium_ed25519_87_set_pub_key(struct crypto_sig *tfm,
						      const void *key,
						      unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_pub_key_int(tfm, key, keylen,
							   LC_DILITHIUM_87);
}

static int lc_kernel_dilithium_ed25519_set_priv_key_int(
	struct crypto_sig *tfm, const void *key, unsigned int keylen,
	enum lc_dilithium_type type)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);
	int ret;

	if (keylen < LC_ED25519_SECRETKEYBYTES)
		return -EINVAL;

	ctx->key_type = lc_kernel_dilithium_ed25519_key_unset;

	/*
	 * Load the Dilithium and the ED25519 keys - they are expected to be
	 * concatenated in the linear buffer of key.
	 */
	ret = lc_dilithium_ed25519_sk_load(
		&ctx->sk, key, keylen - LC_ED25519_SECRETKEYBYTES,
		key + keylen - LC_ED25519_SECRETKEYBYTES,
		LC_ED25519_SECRETKEYBYTES);

	if (!ret) {
		if (lc_dilithium_ed25519_sk_type(&ctx->sk) != type)
			ret = -EOPNOTSUPP;
		else
			ctx->key_type = lc_kernel_dilithium_ed25519_key_sk;
	}

	return ret;
}

static int lc_kernel_dilithium_ed25519_44_set_priv_key(struct crypto_sig *tfm,
						       const void *key,
						       unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_priv_key_int(tfm, key, keylen,
							    LC_DILITHIUM_44);
}

static int lc_kernel_dilithium_ed25519_65_set_priv_key(struct crypto_sig *tfm,
						       const void *key,
						       unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_priv_key_int(tfm, key, keylen,
							    LC_DILITHIUM_65);
}

static int lc_kernel_dilithium_ed25519_87_set_priv_key(struct crypto_sig *tfm,
						       const void *key,
						       unsigned int keylen)
{
	return lc_kernel_dilithium_ed25519_set_priv_key_int(tfm, key, keylen,
							    LC_DILITHIUM_87);
}

static unsigned int lc_kernel_dilithium_ed25519_max_size(struct crypto_sig *tfm)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);
	enum lc_dilithium_type type;

	switch (ctx->key_type) {
	case lc_kernel_dilithium_ed25519_key_sk:
		type = lc_dilithium_ed25519_sk_type(&ctx->sk);
		/* When SK is set -> generate a signature */
		return lc_dilithium_ed25519_sig_size(type);
	case lc_kernel_dilithium_ed25519_key_pk:
		type = lc_dilithium_ed25519_pk_type(&ctx->pk);
		/* When PK is set, this is a safety valve, result is boolean */
		return lc_dilithium_ed25519_sig_size(type);
	default:
		return 0;
	}
}

static int lc_kernel_dilithium_ed25519_alg_init(struct crypto_sig *tfm)
{
	return 0;
}

static void lc_kernel_dilithium_ed25519_alg_exit(struct crypto_sig *tfm)
{
	struct lc_kernel_dilithium_ed25519_ctx *ctx = crypto_sig_ctx(tfm);

	ctx->key_type = lc_kernel_dilithium_ed25519_key_unset;
}

static struct sig_alg lc_kernel_dilithium_87_ed25519 = {
	.sign = lc_kernel_dilithium_ed25519_sign,
	.verify = lc_kernel_dilithium_ed25519_verify,
	.set_pub_key = lc_kernel_dilithium_ed25519_87_set_pub_key,
	.set_priv_key = lc_kernel_dilithium_ed25519_87_set_priv_key,
	.key_size = lc_kernel_dilithium_ed25519_87_key_size,
	.max_size = lc_kernel_dilithium_ed25519_max_size,
	.init = lc_kernel_dilithium_ed25519_alg_init,
	.exit = lc_kernel_dilithium_ed25519_alg_exit,
	.base.cra_name = "dilithium87-ed25519",
	.base.cra_driver_name = "dilithium87-ed25519-leancrypto",
	.base.cra_ctxsize = sizeof(struct lc_kernel_dilithium_ed25519_ctx),
	.base.cra_module = THIS_MODULE,
	.base.cra_priority = LC_KERNEL_DEFAULT_PRIO,
};

static struct sig_alg lc_kernel_dilithium_65_ed25519 = {
	.sign = lc_kernel_dilithium_ed25519_sign,
	.verify = lc_kernel_dilithium_ed25519_verify,
	.set_pub_key = lc_kernel_dilithium_ed25519_65_set_pub_key,
	.set_priv_key = lc_kernel_dilithium_ed25519_65_set_priv_key,
	.key_size = lc_kernel_dilithium_ed25519_65_key_size,
	.max_size = lc_kernel_dilithium_ed25519_max_size,
	.init = lc_kernel_dilithium_ed25519_alg_init,
	.exit = lc_kernel_dilithium_ed25519_alg_exit,
	.base.cra_name = "dilithium65_ed25519",
	.base.cra_driver_name = "dilithium65-ed25519-leancrypto",
	.base.cra_ctxsize = sizeof(struct lc_kernel_dilithium_ed25519_ctx),
	.base.cra_module = THIS_MODULE,
	.base.cra_priority = LC_KERNEL_DEFAULT_PRIO,
};

static struct sig_alg lc_kernel_dilithium_44_ed25519 = {
	.sign = lc_kernel_dilithium_ed25519_sign,
	.verify = lc_kernel_dilithium_ed25519_verify,
	.set_pub_key = lc_kernel_dilithium_ed25519_44_set_pub_key,
	.set_priv_key = lc_kernel_dilithium_ed25519_44_set_priv_key,
	.key_size = lc_kernel_dilithium_ed25519_44_key_size,
	.max_size = lc_kernel_dilithium_ed25519_max_size,
	.init = lc_kernel_dilithium_ed25519_alg_init,
	.exit = lc_kernel_dilithium_ed25519_alg_exit,
	.base.cra_name = "dilithium44_ed25519",
	.base.cra_driver_name = "dilithium44-ed25519-leancrypto",
	.base.cra_ctxsize = sizeof(struct lc_kernel_dilithium_ed25519_ctx),
	.base.cra_module = THIS_MODULE,
	.base.cra_priority = LC_KERNEL_DEFAULT_PRIO,
};

int __init lc_kernel_dilithium_44_ed25519_init(void)
{
	return crypto_register_sig(&lc_kernel_dilithium_44_ed25519);
}

void lc_kernel_dilithium_44_ed25519_exit(void)
{
	crypto_unregister_sig(&lc_kernel_dilithium_44_ed25519);
}

int __init lc_kernel_dilithium_65_ed25519_init(void)
{
	return crypto_register_sig(&lc_kernel_dilithium_65_ed25519);
}

void lc_kernel_dilithium_65_ed25519_exit(void)
{
	crypto_unregister_sig(&lc_kernel_dilithium_65_ed25519);
}

int __init lc_kernel_dilithium_ed25519_init(void)
{
	return crypto_register_sig(&lc_kernel_dilithium_87_ed25519);
}

void lc_kernel_dilithium_ed25519_exit(void)
{
	crypto_unregister_sig(&lc_kernel_dilithium_87_ed25519);
}
