// SPDX-License-Identifier: GPL-2.0-or-later
/*
 * Crypto API support for AES block cipher
 *
 * Copyright 2026 Google LLC
 */

#include <crypto/aes-cbc-macs.h>
#include <crypto/aes.h>
#include <crypto/algapi.h>
#include <crypto/internal/hash.h>
#include <linux/module.h>

static_assert(__alignof__(struct aes_key) <= CRYPTO_MINALIGN);

static int crypto_aes_setkey(struct crypto_tfm *tfm, const u8 *in_key,
                             unsigned int key_len)
{
        struct aes_key *key = crypto_tfm_ctx(tfm);

        return aes_preparekey(key, in_key, key_len);
}

static void crypto_aes_encrypt(struct crypto_tfm *tfm, u8 *out, const u8 *in)
{
        const struct aes_key *key = crypto_tfm_ctx(tfm);

        aes_encrypt(key, out, in);
}

static void crypto_aes_decrypt(struct crypto_tfm *tfm, u8 *out, const u8 *in)
{
        const struct aes_key *key = crypto_tfm_ctx(tfm);

        aes_decrypt(key, out, in);
}

static_assert(__alignof__(struct aes_cmac_key) <= CRYPTO_MINALIGN);
#define AES_CMAC_KEY(tfm) ((struct aes_cmac_key *)crypto_shash_ctx(tfm))
#define AES_CMAC_CTX(desc) ((struct aes_cmac_ctx *)shash_desc_ctx(desc))

static int __maybe_unused crypto_aes_cmac_setkey(struct crypto_shash *tfm,
                                                 const u8 *in_key,
                                                 unsigned int key_len)
{
        return aes_cmac_preparekey(AES_CMAC_KEY(tfm), in_key, key_len);
}

static int __maybe_unused crypto_aes_xcbc_setkey(struct crypto_shash *tfm,
                                                 const u8 *in_key,
                                                 unsigned int key_len)
{
        if (key_len != AES_KEYSIZE_128)
                return -EINVAL;
        aes_xcbcmac_preparekey(AES_CMAC_KEY(tfm), in_key);
        return 0;
}

static int __maybe_unused crypto_aes_cmac_init(struct shash_desc *desc)
{
        aes_cmac_init(AES_CMAC_CTX(desc), AES_CMAC_KEY(desc->tfm));
        return 0;
}

static int __maybe_unused crypto_aes_cmac_update(struct shash_desc *desc,
                                                 const u8 *data,
                                                 unsigned int len)
{
        aes_cmac_update(AES_CMAC_CTX(desc), data, len);
        return 0;
}

static int __maybe_unused crypto_aes_cmac_final(struct shash_desc *desc,
                                                u8 *out)
{
        aes_cmac_final(AES_CMAC_CTX(desc), out);
        return 0;
}

static int __maybe_unused crypto_aes_cmac_digest(struct shash_desc *desc,
                                                 const u8 *data,
                                                 unsigned int len, u8 *out)
{
        aes_cmac(AES_CMAC_KEY(desc->tfm), data, len, out);
        return 0;
}

static_assert(__alignof__(struct aes_enckey) <= CRYPTO_MINALIGN);
#define AES_CBCMAC_KEY(tfm) ((struct aes_enckey *)crypto_shash_ctx(tfm))
#define AES_CBCMAC_CTX(desc) ((struct aes_cbcmac_ctx *)shash_desc_ctx(desc))

static int __maybe_unused crypto_aes_cbcmac_setkey(struct crypto_shash *tfm,
                                                   const u8 *in_key,
                                                   unsigned int key_len)
{
        return aes_prepareenckey(AES_CBCMAC_KEY(tfm), in_key, key_len);
}

static int __maybe_unused crypto_aes_cbcmac_init(struct shash_desc *desc)
{
        aes_cbcmac_init(AES_CBCMAC_CTX(desc), AES_CBCMAC_KEY(desc->tfm));
        return 0;
}

static int __maybe_unused crypto_aes_cbcmac_update(struct shash_desc *desc,
                                                   const u8 *data,
                                                   unsigned int len)
{
        aes_cbcmac_update(AES_CBCMAC_CTX(desc), data, len);
        return 0;
}

static int __maybe_unused crypto_aes_cbcmac_final(struct shash_desc *desc,
                                                  u8 *out)
{
        aes_cbcmac_final(AES_CBCMAC_CTX(desc), out);
        return 0;
}

static int __maybe_unused crypto_aes_cbcmac_digest(struct shash_desc *desc,
                                                   const u8 *data,
                                                   unsigned int len, u8 *out)
{
        aes_cbcmac_init(AES_CBCMAC_CTX(desc), AES_CBCMAC_KEY(desc->tfm));
        aes_cbcmac_update(AES_CBCMAC_CTX(desc), data, len);
        aes_cbcmac_final(AES_CBCMAC_CTX(desc), out);
        return 0;
}

static struct crypto_alg alg = {
        .cra_name = "aes",
        .cra_driver_name = "aes-lib",
        .cra_priority = 100,
        .cra_flags = CRYPTO_ALG_TYPE_CIPHER,
        .cra_blocksize = AES_BLOCK_SIZE,
        .cra_ctxsize = sizeof(struct aes_key),
        .cra_module = THIS_MODULE,
        .cra_u = { .cipher = { .cia_min_keysize = AES_MIN_KEY_SIZE,
                               .cia_max_keysize = AES_MAX_KEY_SIZE,
                               .cia_setkey = crypto_aes_setkey,
                               .cia_encrypt = crypto_aes_encrypt,
                               .cia_decrypt = crypto_aes_decrypt } }
};

static struct shash_alg mac_algs[] = {
#if IS_ENABLED(CONFIG_CRYPTO_CMAC)
        {
                .base.cra_name = "cmac(aes)",
                .base.cra_driver_name = "cmac-aes-lib",
                .base.cra_priority = 300,
                .base.cra_blocksize = AES_BLOCK_SIZE,
                .base.cra_ctxsize = sizeof(struct aes_cmac_key),
                .base.cra_module = THIS_MODULE,
                .digestsize = AES_BLOCK_SIZE,
                .setkey = crypto_aes_cmac_setkey,
                .init = crypto_aes_cmac_init,
                .update = crypto_aes_cmac_update,
                .final = crypto_aes_cmac_final,
                .digest = crypto_aes_cmac_digest,
                .descsize = sizeof(struct aes_cmac_ctx),
        },
#endif
#if IS_ENABLED(CONFIG_CRYPTO_XCBC)
        {
                /*
                 * Note that the only difference between xcbc(aes) and cmac(aes)
                 * is the preparekey function.
                 */
                .base.cra_name = "xcbc(aes)",
                .base.cra_driver_name = "xcbc-aes-lib",
                .base.cra_priority = 300,
                .base.cra_blocksize = AES_BLOCK_SIZE,
                .base.cra_ctxsize = sizeof(struct aes_cmac_key),
                .base.cra_module = THIS_MODULE,
                .digestsize = AES_BLOCK_SIZE,
                .setkey = crypto_aes_xcbc_setkey,
                .init = crypto_aes_cmac_init,
                .update = crypto_aes_cmac_update,
                .final = crypto_aes_cmac_final,
                .digest = crypto_aes_cmac_digest,
                .descsize = sizeof(struct aes_cmac_ctx),
        },
#endif
#if IS_ENABLED(CONFIG_CRYPTO_CCM)
        {
                .base.cra_name = "cbcmac(aes)",
                .base.cra_driver_name = "cbcmac-aes-lib",
                .base.cra_priority = 300,
                .base.cra_blocksize = AES_BLOCK_SIZE,
                .base.cra_ctxsize = sizeof(struct aes_enckey),
                .base.cra_module = THIS_MODULE,
                .digestsize = AES_BLOCK_SIZE,
                .setkey = crypto_aes_cbcmac_setkey,
                .init = crypto_aes_cbcmac_init,
                .update = crypto_aes_cbcmac_update,
                .final = crypto_aes_cbcmac_final,
                .digest = crypto_aes_cbcmac_digest,
                .descsize = sizeof(struct aes_cbcmac_ctx),
        },
#endif
};

static int __init crypto_aes_mod_init(void)
{
        int err = crypto_register_alg(&alg);

        if (err)
                return err;

        if (ARRAY_SIZE(mac_algs) > 0) {
                err = crypto_register_shashes(mac_algs, ARRAY_SIZE(mac_algs));
                if (err)
                        goto err_unregister_alg;
        } /* Else, CONFIG_CRYPTO_HASH might not be enabled. */
        return 0;

err_unregister_alg:
        crypto_unregister_alg(&alg);
        return err;
}
module_init(crypto_aes_mod_init);

static void __exit crypto_aes_mod_exit(void)
{
        if (ARRAY_SIZE(mac_algs) > 0)
                crypto_unregister_shashes(mac_algs, ARRAY_SIZE(mac_algs));
        crypto_unregister_alg(&alg);
}
module_exit(crypto_aes_mod_exit);

MODULE_DESCRIPTION("Crypto API support for AES block cipher");
MODULE_IMPORT_NS("CRYPTO_INTERNAL");
MODULE_LICENSE("GPL");
MODULE_ALIAS_CRYPTO("aes");
MODULE_ALIAS_CRYPTO("aes-lib");
#if IS_ENABLED(CONFIG_CRYPTO_CMAC)
MODULE_ALIAS_CRYPTO("cmac(aes)");
MODULE_ALIAS_CRYPTO("cmac-aes-lib");
#endif
#if IS_ENABLED(CONFIG_CRYPTO_XCBC)
MODULE_ALIAS_CRYPTO("xcbc(aes)");
MODULE_ALIAS_CRYPTO("xcbc-aes-lib");
#endif
#if IS_ENABLED(CONFIG_CRYPTO_CCM)
MODULE_ALIAS_CRYPTO("cbcmac(aes)");
MODULE_ALIAS_CRYPTO("cbcmac-aes-lib");
#endif