/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * AES block cipher, optimized for ARM64
 *
 * Copyright (C) 2013 - 2017 Linaro Ltd <ard.biesheuvel@linaro.org>
 * Copyright 2026 Google LLC
 */

#include <asm/neon.h>
#include <asm/simd.h>
#include <linux/unaligned.h>
#include <linux/cpufeature.h>

static __ro_after_init DEFINE_STATIC_KEY_FALSE(have_aes);

struct aes_block {
        u8 b[AES_BLOCK_SIZE];
};

asmlinkage void __aes_arm64_encrypt(const u32 rk[], u8 out[AES_BLOCK_SIZE],
                                    const u8 in[AES_BLOCK_SIZE], int rounds);
asmlinkage void __aes_arm64_decrypt(const u32 inv_rk[], u8 out[AES_BLOCK_SIZE],
                                    const u8 in[AES_BLOCK_SIZE], int rounds);
asmlinkage void __aes_ce_encrypt(const u32 rk[], u8 out[AES_BLOCK_SIZE],
                                 const u8 in[AES_BLOCK_SIZE], int rounds);
asmlinkage void __aes_ce_decrypt(const u32 inv_rk[], u8 out[AES_BLOCK_SIZE],
                                 const u8 in[AES_BLOCK_SIZE], int rounds);
asmlinkage u32 __aes_ce_sub(u32 l);
asmlinkage void __aes_ce_invert(struct aes_block *out,
                                const struct aes_block *in);

/*
 * Expand an AES key using the crypto extensions if supported and usable or
 * generic code otherwise.  The expanded key format is compatible between the
 * two cases.  The outputs are @rndkeys (required) and @inv_rndkeys (optional).
 */
static void aes_expandkey_arm64(u32 rndkeys[], u32 *inv_rndkeys,
                                const u8 *in_key, int key_len, int nrounds)
{
        /*
         * The AES key schedule round constants
         */
        static u8 const rcon[] = {
                0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36,
        };

        u32 kwords = key_len / sizeof(u32);
        struct aes_block *key_enc, *key_dec;
        int i, j;

        if (!IS_ENABLED(CONFIG_KERNEL_MODE_NEON) ||
            !static_branch_likely(&have_aes) || unlikely(!may_use_simd())) {
                aes_expandkey_generic(rndkeys, inv_rndkeys, in_key, key_len);
                return;
        }

        for (i = 0; i < kwords; i++)
                rndkeys[i] = get_unaligned_le32(&in_key[i * sizeof(u32)]);

        scoped_ksimd() {
                for (i = 0; i < sizeof(rcon); i++) {
                        u32 *rki = &rndkeys[i * kwords];
                        u32 *rko = rki + kwords;

                        rko[0] = ror32(__aes_ce_sub(rki[kwords - 1]), 8) ^
                                 rcon[i] ^ rki[0];
                        rko[1] = rko[0] ^ rki[1];
                        rko[2] = rko[1] ^ rki[2];
                        rko[3] = rko[2] ^ rki[3];

                        if (key_len == AES_KEYSIZE_192) {
                                if (i >= 7)
                                        break;
                                rko[4] = rko[3] ^ rki[4];
                                rko[5] = rko[4] ^ rki[5];
                        } else if (key_len == AES_KEYSIZE_256) {
                                if (i >= 6)
                                        break;
                                rko[4] = __aes_ce_sub(rko[3]) ^ rki[4];
                                rko[5] = rko[4] ^ rki[5];
                                rko[6] = rko[5] ^ rki[6];
                                rko[7] = rko[6] ^ rki[7];
                        }
                }

                /*
                 * Generate the decryption keys for the Equivalent Inverse
                 * Cipher.  This involves reversing the order of the round
                 * keys, and applying the Inverse Mix Columns transformation on
                 * all but the first and the last one.
                 */
                if (inv_rndkeys) {
                        key_enc = (struct aes_block *)rndkeys;
                        key_dec = (struct aes_block *)inv_rndkeys;
                        j = nrounds;

                        key_dec[0] = key_enc[j];
                        for (i = 1, j--; j > 0; i++, j--)
                                __aes_ce_invert(key_dec + i, key_enc + j);
                        key_dec[i] = key_enc[0];
                }
        }
}

static void aes_preparekey_arch(union aes_enckey_arch *k,
                                union aes_invkey_arch *inv_k,
                                const u8 *in_key, int key_len, int nrounds)
{
        aes_expandkey_arm64(k->rndkeys, inv_k ? inv_k->inv_rndkeys : NULL,
                            in_key, key_len, nrounds);
}

/*
 * This is here temporarily until the remaining AES mode implementations are
 * migrated from arch/arm64/crypto/ to lib/crypto/arm64/.
 */
int ce_aes_expandkey(struct crypto_aes_ctx *ctx, const u8 *in_key,
                     unsigned int key_len)
{
        if (aes_check_keylen(key_len) != 0)
                return -EINVAL;
        ctx->key_length = key_len;
        aes_expandkey_arm64(ctx->key_enc, ctx->key_dec, in_key, key_len,
                            6 + key_len / 4);
        return 0;
}
EXPORT_SYMBOL(ce_aes_expandkey);

static void aes_encrypt_arch(const struct aes_enckey *key,
                             u8 out[AES_BLOCK_SIZE],
                             const u8 in[AES_BLOCK_SIZE])
{
        if (IS_ENABLED(CONFIG_KERNEL_MODE_NEON) &&
            static_branch_likely(&have_aes) && likely(may_use_simd())) {
                scoped_ksimd()
                        __aes_ce_encrypt(key->k.rndkeys, out, in, key->nrounds);
        } else {
                __aes_arm64_encrypt(key->k.rndkeys, out, in, key->nrounds);
        }
}

static void aes_decrypt_arch(const struct aes_key *key,
                             u8 out[AES_BLOCK_SIZE],
                             const u8 in[AES_BLOCK_SIZE])
{
        if (IS_ENABLED(CONFIG_KERNEL_MODE_NEON) &&
            static_branch_likely(&have_aes) && likely(may_use_simd())) {
                scoped_ksimd()
                        __aes_ce_decrypt(key->inv_k.inv_rndkeys, out, in,
                                         key->nrounds);
        } else {
                __aes_arm64_decrypt(key->inv_k.inv_rndkeys, out, in,
                                    key->nrounds);
        }
}

#ifdef CONFIG_KERNEL_MODE_NEON
#define aes_mod_init_arch aes_mod_init_arch
static void aes_mod_init_arch(void)
{
        if (cpu_have_named_feature(AES))
                static_branch_enable(&have_aes);
}
#endif /* CONFIG_KERNEL_MODE_NEON */