// SPDX-License-Identifier: GPL-2.0-or-later
/*
 * crypto_sig wrapper around ML-DSA library.
 */
#include <linux/init.h>
#include <linux/module.h>
#include <crypto/internal/sig.h>
#include <crypto/mldsa.h>

struct crypto_mldsa_ctx {
	u8 pk[MAX(MAX(MLDSA44_PUBLIC_KEY_SIZE,
		      MLDSA65_PUBLIC_KEY_SIZE),
		  MLDSA87_PUBLIC_KEY_SIZE)];
	unsigned int pk_len;
	enum mldsa_alg strength;
	bool key_set;
};

static int crypto_mldsa_sign(struct crypto_sig *tfm,
			     const void *msg, unsigned int msg_len,
			     void *sig, unsigned int sig_len)
{
	return -EOPNOTSUPP;
}

static int crypto_mldsa_verify(struct crypto_sig *tfm,
			       const void *sig, unsigned int sig_len,
			       const void *msg, unsigned int msg_len)
{
	const struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	if (unlikely(!ctx->key_set))
		return -EINVAL;

	return mldsa_verify(ctx->strength, sig, sig_len, msg, msg_len,
			    ctx->pk, ctx->pk_len);
}

static unsigned int crypto_mldsa_key_size(struct crypto_sig *tfm)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	switch (ctx->strength) {
	case MLDSA44:
		return MLDSA44_PUBLIC_KEY_SIZE;
	case MLDSA65:
		return MLDSA65_PUBLIC_KEY_SIZE;
	case MLDSA87:
		return MLDSA87_PUBLIC_KEY_SIZE;
	default:
		WARN_ON_ONCE(1);
		return 0;
	}
}

static int crypto_mldsa_set_pub_key(struct crypto_sig *tfm,
				    const void *key, unsigned int keylen)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);
	unsigned int expected_len = crypto_mldsa_key_size(tfm);

	if (keylen != expected_len)
		return -EINVAL;

	ctx->pk_len = keylen;
	memcpy(ctx->pk, key, keylen);
	ctx->key_set = true;
	return 0;
}

static int crypto_mldsa_set_priv_key(struct crypto_sig *tfm,
				     const void *key, unsigned int keylen)
{
	return -EOPNOTSUPP;
}

static unsigned int crypto_mldsa_max_size(struct crypto_sig *tfm)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	switch (ctx->strength) {
	case MLDSA44:
		return MLDSA44_SIGNATURE_SIZE;
	case MLDSA65:
		return MLDSA65_SIGNATURE_SIZE;
	case MLDSA87:
		return MLDSA87_SIGNATURE_SIZE;
	default:
		WARN_ON_ONCE(1);
		return 0;
	}
}

static int crypto_mldsa44_alg_init(struct crypto_sig *tfm)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	ctx->strength = MLDSA44;
	ctx->key_set = false;
	return 0;
}

static int crypto_mldsa65_alg_init(struct crypto_sig *tfm)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	ctx->strength = MLDSA65;
	ctx->key_set = false;
	return 0;
}

static int crypto_mldsa87_alg_init(struct crypto_sig *tfm)
{
	struct crypto_mldsa_ctx *ctx = crypto_sig_ctx(tfm);

	ctx->strength = MLDSA87;
	ctx->key_set = false;
	return 0;
}

static void crypto_mldsa_alg_exit(struct crypto_sig *tfm)
{
}

static struct sig_alg crypto_mldsa_algs[] = {
	{
		.sign			= crypto_mldsa_sign,
		.verify			= crypto_mldsa_verify,
		.set_pub_key		= crypto_mldsa_set_pub_key,
		.set_priv_key		= crypto_mldsa_set_priv_key,
		.key_size		= crypto_mldsa_key_size,
		.max_size		= crypto_mldsa_max_size,
		.init			= crypto_mldsa44_alg_init,
		.exit			= crypto_mldsa_alg_exit,
		.base.cra_name		= "mldsa44",
		.base.cra_driver_name	= "mldsa44-lib",
		.base.cra_ctxsize	= sizeof(struct crypto_mldsa_ctx),
		.base.cra_module	= THIS_MODULE,
		.base.cra_priority	= 5000,
	}, {
		.sign			= crypto_mldsa_sign,
		.verify			= crypto_mldsa_verify,
		.set_pub_key		= crypto_mldsa_set_pub_key,
		.set_priv_key		= crypto_mldsa_set_priv_key,
		.key_size		= crypto_mldsa_key_size,
		.max_size		= crypto_mldsa_max_size,
		.init			= crypto_mldsa65_alg_init,
		.exit			= crypto_mldsa_alg_exit,
		.base.cra_name		= "mldsa65",
		.base.cra_driver_name	= "mldsa65-lib",
		.base.cra_ctxsize	= sizeof(struct crypto_mldsa_ctx),
		.base.cra_module	= THIS_MODULE,
		.base.cra_priority	= 5000,
	}, {
		.sign			= crypto_mldsa_sign,
		.verify			= crypto_mldsa_verify,
		.set_pub_key		= crypto_mldsa_set_pub_key,
		.set_priv_key		= crypto_mldsa_set_priv_key,
		.key_size		= crypto_mldsa_key_size,
		.max_size		= crypto_mldsa_max_size,
		.init			= crypto_mldsa87_alg_init,
		.exit			= crypto_mldsa_alg_exit,
		.base.cra_name		= "mldsa87",
		.base.cra_driver_name	= "mldsa87-lib",
		.base.cra_ctxsize	= sizeof(struct crypto_mldsa_ctx),
		.base.cra_module	= THIS_MODULE,
		.base.cra_priority	= 5000,
	},
};

static int __init mldsa_init(void)
{
	int ret, i;

	for (i = 0; i < ARRAY_SIZE(crypto_mldsa_algs); i++) {
		ret = crypto_register_sig(&crypto_mldsa_algs[i]);
		if (ret < 0)
			goto error;
	}
	return 0;

error:
	pr_err("Failed to register (%d)\n", ret);
	for (i--; i >= 0; i--)
		crypto_unregister_sig(&crypto_mldsa_algs[i]);
	return ret;
}
module_init(mldsa_init);

static void mldsa_exit(void)
{
	for (int i = 0; i < ARRAY_SIZE(crypto_mldsa_algs); i++)
		crypto_unregister_sig(&crypto_mldsa_algs[i]);
}
module_exit(mldsa_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Crypto API support for ML-DSA signature verification");
MODULE_ALIAS_CRYPTO("mldsa44");
MODULE_ALIAS_CRYPTO("mldsa65");
MODULE_ALIAS_CRYPTO("mldsa87");