// SPDX-License-Identifier: GPL-2.0-only
/*
 * K3 DTHE V2 crypto accelerator driver
 *
 * Copyright (C) Texas Instruments 2025 - https://www.ti.com
 * Author: T Pratham <t-pratham@ti.com>
 */

#include <crypto/aead.h>
#include <crypto/aes.h>
#include <crypto/algapi.h>
#include <crypto/engine.h>
#include <crypto/internal/aead.h>
#include <crypto/internal/skcipher.h>

#include "dthev2-common.h"

#include <linux/delay.h>
#include <linux/dmaengine.h>
#include <linux/dma-mapping.h>
#include <linux/io.h>
#include <linux/scatterlist.h>

/* Registers */

// AES Engine
#define DTHE_P_AES_BASE         0x7000

#define DTHE_P_AES_KEY1_0       0x0038
#define DTHE_P_AES_KEY1_1       0x003C
#define DTHE_P_AES_KEY1_2       0x0030
#define DTHE_P_AES_KEY1_3       0x0034
#define DTHE_P_AES_KEY1_4       0x0028
#define DTHE_P_AES_KEY1_5       0x002C
#define DTHE_P_AES_KEY1_6       0x0020
#define DTHE_P_AES_KEY1_7       0x0024

#define DTHE_P_AES_KEY2_0       0x0018
#define DTHE_P_AES_KEY2_1       0x001C
#define DTHE_P_AES_KEY2_2       0x0010
#define DTHE_P_AES_KEY2_3       0x0014
#define DTHE_P_AES_KEY2_4       0x0008
#define DTHE_P_AES_KEY2_5       0x000C
#define DTHE_P_AES_KEY2_6       0x0000
#define DTHE_P_AES_KEY2_7       0x0004

#define DTHE_P_AES_IV_IN_0      0x0040
#define DTHE_P_AES_IV_IN_1      0x0044
#define DTHE_P_AES_IV_IN_2      0x0048
#define DTHE_P_AES_IV_IN_3      0x004C
#define DTHE_P_AES_CTRL         0x0050
#define DTHE_P_AES_C_LENGTH_0   0x0054
#define DTHE_P_AES_C_LENGTH_1   0x0058
#define DTHE_P_AES_AUTH_LENGTH  0x005C
#define DTHE_P_AES_DATA_IN_OUT  0x0060

#define DTHE_P_AES_SYSCONFIG    0x0084
#define DTHE_P_AES_IRQSTATUS    0x008C
#define DTHE_P_AES_IRQENABLE    0x0090

/* Register write values and macros */

enum aes_ctrl_mode_masks {
        AES_CTRL_ECB_MASK = 0x00,
        AES_CTRL_CBC_MASK = BIT(5),
        AES_CTRL_XTS_MASK = BIT(12) | BIT(11),
};

#define DTHE_AES_CTRL_MODE_CLEAR_MASK           ~GENMASK(28, 5)

#define DTHE_AES_CTRL_DIR_ENC                   BIT(2)

#define DTHE_AES_CTRL_KEYSIZE_16B               BIT(3)
#define DTHE_AES_CTRL_KEYSIZE_24B               BIT(4)
#define DTHE_AES_CTRL_KEYSIZE_32B               (BIT(3) | BIT(4))

#define DTHE_AES_CTRL_SAVE_CTX_SET              BIT(29)

#define DTHE_AES_CTRL_OUTPUT_READY              BIT_MASK(0)
#define DTHE_AES_CTRL_INPUT_READY               BIT_MASK(1)
#define DTHE_AES_CTRL_SAVED_CTX_READY           BIT_MASK(30)
#define DTHE_AES_CTRL_CTX_READY                 BIT_MASK(31)

#define DTHE_AES_SYSCONFIG_DMA_DATA_IN_OUT_EN   GENMASK(6, 5)
#define DTHE_AES_IRQENABLE_EN_ALL               GENMASK(3, 0)

/* Misc */
#define AES_IV_SIZE                             AES_BLOCK_SIZE
#define AES_BLOCK_WORDS                         (AES_BLOCK_SIZE / sizeof(u32))
#define AES_IV_WORDS                            AES_BLOCK_WORDS

static int dthe_cipher_init_tfm(struct crypto_skcipher *tfm)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);
        struct dthe_data *dev_data = dthe_get_dev(ctx);

        ctx->dev_data = dev_data;
        ctx->keylen = 0;

        return 0;
}

static int dthe_cipher_xts_init_tfm(struct crypto_skcipher *tfm)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);
        struct dthe_data *dev_data = dthe_get_dev(ctx);

        ctx->dev_data = dev_data;
        ctx->keylen = 0;

        ctx->skcipher_fb = crypto_alloc_sync_skcipher("xts(aes)", 0,
                                                      CRYPTO_ALG_NEED_FALLBACK);
        if (IS_ERR(ctx->skcipher_fb)) {
                dev_err(dev_data->dev, "fallback driver xts(aes) couldn't be loaded\n");
                return PTR_ERR(ctx->skcipher_fb);
        }

        return 0;
}

static void dthe_cipher_xts_exit_tfm(struct crypto_skcipher *tfm)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);

        crypto_free_sync_skcipher(ctx->skcipher_fb);
}

static int dthe_aes_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);

        if (keylen != AES_KEYSIZE_128 && keylen != AES_KEYSIZE_192 && keylen != AES_KEYSIZE_256)
                return -EINVAL;

        ctx->keylen = keylen;
        memcpy(ctx->key, key, keylen);

        return 0;
}

static int dthe_aes_ecb_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);

        ctx->aes_mode = DTHE_AES_ECB;

        return dthe_aes_setkey(tfm, key, keylen);
}

static int dthe_aes_cbc_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);

        ctx->aes_mode = DTHE_AES_CBC;

        return dthe_aes_setkey(tfm, key, keylen);
}

static int dthe_aes_xts_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);

        if (keylen != 2 * AES_KEYSIZE_128 &&
            keylen != 2 * AES_KEYSIZE_192 &&
            keylen != 2 * AES_KEYSIZE_256)
                return -EINVAL;

        ctx->aes_mode = DTHE_AES_XTS;
        ctx->keylen = keylen / 2;
        memcpy(ctx->key, key, keylen);

        crypto_sync_skcipher_clear_flags(ctx->skcipher_fb, CRYPTO_TFM_REQ_MASK);
        crypto_sync_skcipher_set_flags(ctx->skcipher_fb,
                                  crypto_skcipher_get_flags(tfm) &
                                  CRYPTO_TFM_REQ_MASK);

        return crypto_sync_skcipher_setkey(ctx->skcipher_fb, key, keylen);
}

static void dthe_aes_set_ctrl_key(struct dthe_tfm_ctx *ctx,
                                  struct dthe_aes_req_ctx *rctx,
                                  u32 *iv_in)
{
        struct dthe_data *dev_data = dthe_get_dev(ctx);
        void __iomem *aes_base_reg = dev_data->regs + DTHE_P_AES_BASE;
        u32 ctrl_val = 0;

        writel_relaxed(ctx->key[0], aes_base_reg + DTHE_P_AES_KEY1_0);
        writel_relaxed(ctx->key[1], aes_base_reg + DTHE_P_AES_KEY1_1);
        writel_relaxed(ctx->key[2], aes_base_reg + DTHE_P_AES_KEY1_2);
        writel_relaxed(ctx->key[3], aes_base_reg + DTHE_P_AES_KEY1_3);

        if (ctx->keylen > AES_KEYSIZE_128) {
                writel_relaxed(ctx->key[4], aes_base_reg + DTHE_P_AES_KEY1_4);
                writel_relaxed(ctx->key[5], aes_base_reg + DTHE_P_AES_KEY1_5);
        }
        if (ctx->keylen == AES_KEYSIZE_256) {
                writel_relaxed(ctx->key[6], aes_base_reg + DTHE_P_AES_KEY1_6);
                writel_relaxed(ctx->key[7], aes_base_reg + DTHE_P_AES_KEY1_7);
        }

        if (ctx->aes_mode == DTHE_AES_XTS) {
                size_t key2_offset = ctx->keylen / sizeof(u32);

                writel_relaxed(ctx->key[key2_offset + 0], aes_base_reg + DTHE_P_AES_KEY2_0);
                writel_relaxed(ctx->key[key2_offset + 1], aes_base_reg + DTHE_P_AES_KEY2_1);
                writel_relaxed(ctx->key[key2_offset + 2], aes_base_reg + DTHE_P_AES_KEY2_2);
                writel_relaxed(ctx->key[key2_offset + 3], aes_base_reg + DTHE_P_AES_KEY2_3);

                if (ctx->keylen > AES_KEYSIZE_128) {
                        writel_relaxed(ctx->key[key2_offset + 4], aes_base_reg + DTHE_P_AES_KEY2_4);
                        writel_relaxed(ctx->key[key2_offset + 5], aes_base_reg + DTHE_P_AES_KEY2_5);
                }
                if (ctx->keylen == AES_KEYSIZE_256) {
                        writel_relaxed(ctx->key[key2_offset + 6], aes_base_reg + DTHE_P_AES_KEY2_6);
                        writel_relaxed(ctx->key[key2_offset + 7], aes_base_reg + DTHE_P_AES_KEY2_7);
                }
        }

        if (rctx->enc)
                ctrl_val |= DTHE_AES_CTRL_DIR_ENC;

        if (ctx->keylen == AES_KEYSIZE_128)
                ctrl_val |= DTHE_AES_CTRL_KEYSIZE_16B;
        else if (ctx->keylen == AES_KEYSIZE_192)
                ctrl_val |= DTHE_AES_CTRL_KEYSIZE_24B;
        else
                ctrl_val |= DTHE_AES_CTRL_KEYSIZE_32B;

        // Write AES mode
        ctrl_val &= DTHE_AES_CTRL_MODE_CLEAR_MASK;
        switch (ctx->aes_mode) {
        case DTHE_AES_ECB:
                ctrl_val |= AES_CTRL_ECB_MASK;
                break;
        case DTHE_AES_CBC:
                ctrl_val |= AES_CTRL_CBC_MASK;
                break;
        case DTHE_AES_XTS:
                ctrl_val |= AES_CTRL_XTS_MASK;
                break;
        }

        if (iv_in) {
                ctrl_val |= DTHE_AES_CTRL_SAVE_CTX_SET;
                for (int i = 0; i < AES_IV_WORDS; ++i)
                        writel_relaxed(iv_in[i],
                                       aes_base_reg + DTHE_P_AES_IV_IN_0 + (DTHE_REG_SIZE * i));
        }

        writel_relaxed(ctrl_val, aes_base_reg + DTHE_P_AES_CTRL);
}

static void dthe_aes_dma_in_callback(void *data)
{
        struct skcipher_request *req = (struct skcipher_request *)data;
        struct dthe_aes_req_ctx *rctx = skcipher_request_ctx(req);

        complete(&rctx->aes_compl);
}

static int dthe_aes_run(struct crypto_engine *engine, void *areq)
{
        struct skcipher_request *req = container_of(areq, struct skcipher_request, base);
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(crypto_skcipher_reqtfm(req));
        struct dthe_data *dev_data = dthe_get_dev(ctx);
        struct dthe_aes_req_ctx *rctx = skcipher_request_ctx(req);

        unsigned int len = req->cryptlen;
        struct scatterlist *src = req->src;
        struct scatterlist *dst = req->dst;

        int src_nents = sg_nents_for_len(src, len);
        int dst_nents;

        int src_mapped_nents;
        int dst_mapped_nents;

        bool diff_dst;
        enum dma_data_direction src_dir, dst_dir;

        struct device *tx_dev, *rx_dev;
        struct dma_async_tx_descriptor *desc_in, *desc_out;

        int ret;

        void __iomem *aes_base_reg = dev_data->regs + DTHE_P_AES_BASE;

        u32 aes_irqenable_val = readl_relaxed(aes_base_reg + DTHE_P_AES_IRQENABLE);
        u32 aes_sysconfig_val = readl_relaxed(aes_base_reg + DTHE_P_AES_SYSCONFIG);

        aes_sysconfig_val |= DTHE_AES_SYSCONFIG_DMA_DATA_IN_OUT_EN;
        writel_relaxed(aes_sysconfig_val, aes_base_reg + DTHE_P_AES_SYSCONFIG);

        aes_irqenable_val |= DTHE_AES_IRQENABLE_EN_ALL;
        writel_relaxed(aes_irqenable_val, aes_base_reg + DTHE_P_AES_IRQENABLE);

        if (src == dst) {
                diff_dst = false;
                src_dir = DMA_BIDIRECTIONAL;
                dst_dir = DMA_BIDIRECTIONAL;
        } else {
                diff_dst = true;
                src_dir = DMA_TO_DEVICE;
                dst_dir  = DMA_FROM_DEVICE;
        }

        tx_dev = dmaengine_get_dma_device(dev_data->dma_aes_tx);
        rx_dev = dmaengine_get_dma_device(dev_data->dma_aes_rx);

        src_mapped_nents = dma_map_sg(tx_dev, src, src_nents, src_dir);
        if (src_mapped_nents == 0) {
                ret = -EINVAL;
                goto aes_err;
        }

        if (!diff_dst) {
                dst_nents = src_nents;
                dst_mapped_nents = src_mapped_nents;
        } else {
                dst_nents = sg_nents_for_len(dst, len);
                dst_mapped_nents = dma_map_sg(rx_dev, dst, dst_nents, dst_dir);
                if (dst_mapped_nents == 0) {
                        dma_unmap_sg(tx_dev, src, src_nents, src_dir);
                        ret = -EINVAL;
                        goto aes_err;
                }
        }

        desc_in = dmaengine_prep_slave_sg(dev_data->dma_aes_rx, dst, dst_mapped_nents,
                                          DMA_DEV_TO_MEM, DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
        if (!desc_in) {
                dev_err(dev_data->dev, "IN prep_slave_sg() failed\n");
                ret = -EINVAL;
                goto aes_prep_err;
        }

        desc_out = dmaengine_prep_slave_sg(dev_data->dma_aes_tx, src, src_mapped_nents,
                                           DMA_MEM_TO_DEV, DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
        if (!desc_out) {
                dev_err(dev_data->dev, "OUT prep_slave_sg() failed\n");
                ret = -EINVAL;
                goto aes_prep_err;
        }

        desc_in->callback = dthe_aes_dma_in_callback;
        desc_in->callback_param = req;

        init_completion(&rctx->aes_compl);

        if (ctx->aes_mode == DTHE_AES_ECB)
                dthe_aes_set_ctrl_key(ctx, rctx, NULL);
        else
                dthe_aes_set_ctrl_key(ctx, rctx, (u32 *)req->iv);

        writel_relaxed(lower_32_bits(req->cryptlen), aes_base_reg + DTHE_P_AES_C_LENGTH_0);
        writel_relaxed(upper_32_bits(req->cryptlen), aes_base_reg + DTHE_P_AES_C_LENGTH_1);

        dmaengine_submit(desc_in);
        dmaengine_submit(desc_out);

        dma_async_issue_pending(dev_data->dma_aes_rx);
        dma_async_issue_pending(dev_data->dma_aes_tx);

        // Need to do a timeout to ensure finalise gets called if DMA callback fails for any reason
        ret = wait_for_completion_timeout(&rctx->aes_compl, msecs_to_jiffies(DTHE_DMA_TIMEOUT_MS));
        if (!ret) {
                ret = -ETIMEDOUT;
                dmaengine_terminate_sync(dev_data->dma_aes_rx);
                dmaengine_terminate_sync(dev_data->dma_aes_tx);

                for (int i = 0; i < AES_BLOCK_WORDS; ++i)
                        readl_relaxed(aes_base_reg + DTHE_P_AES_DATA_IN_OUT + (DTHE_REG_SIZE * i));
        } else {
                ret = 0;
        }

        // For modes other than ECB, read IV_OUT
        if (ctx->aes_mode != DTHE_AES_ECB) {
                u32 *iv_out = (u32 *)req->iv;

                for (int i = 0; i < AES_IV_WORDS; ++i)
                        iv_out[i] = readl_relaxed(aes_base_reg +
                                                  DTHE_P_AES_IV_IN_0 +
                                                  (DTHE_REG_SIZE * i));
        }

aes_prep_err:
        dma_unmap_sg(tx_dev, src, src_nents, src_dir);
        if (dst_dir != DMA_BIDIRECTIONAL)
                dma_unmap_sg(rx_dev, dst, dst_nents, dst_dir);

aes_err:
        local_bh_disable();
        crypto_finalize_skcipher_request(dev_data->engine, req, ret);
        local_bh_enable();
        return 0;
}

static int dthe_aes_crypt(struct skcipher_request *req)
{
        struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(crypto_skcipher_reqtfm(req));
        struct dthe_aes_req_ctx *rctx = skcipher_request_ctx(req);
        struct dthe_data *dev_data = dthe_get_dev(ctx);
        struct crypto_engine *engine;

        /*
         * If data is not a multiple of AES_BLOCK_SIZE:
         * - need to return -EINVAL for ECB, CBC as they are block ciphers
         * - need to fallback to software as H/W doesn't support Ciphertext Stealing for XTS
         */
        if (req->cryptlen % AES_BLOCK_SIZE) {
                if (ctx->aes_mode == DTHE_AES_XTS) {
                        SYNC_SKCIPHER_REQUEST_ON_STACK(subreq, ctx->skcipher_fb);

                        skcipher_request_set_callback(subreq, skcipher_request_flags(req),
                                                      req->base.complete, req->base.data);
                        skcipher_request_set_crypt(subreq, req->src, req->dst,
                                                   req->cryptlen, req->iv);

                        return rctx->enc ? crypto_skcipher_encrypt(subreq) :
                                crypto_skcipher_decrypt(subreq);
                }
                return -EINVAL;
        }

        /*
         * If data length input is zero, no need to do any operation.
         * Except for XTS mode, where data length should be non-zero.
         */
        if (req->cryptlen == 0) {
                if (ctx->aes_mode == DTHE_AES_XTS)
                        return -EINVAL;
                return 0;
        }

        engine = dev_data->engine;
        return crypto_transfer_skcipher_request_to_engine(engine, req);
}

static int dthe_aes_encrypt(struct skcipher_request *req)
{
        struct dthe_aes_req_ctx *rctx = skcipher_request_ctx(req);

        rctx->enc = 1;
        return dthe_aes_crypt(req);
}

static int dthe_aes_decrypt(struct skcipher_request *req)
{
        struct dthe_aes_req_ctx *rctx = skcipher_request_ctx(req);

        rctx->enc = 0;
        return dthe_aes_crypt(req);
}

static struct skcipher_engine_alg cipher_algs[] = {
        {
                .base.init                      = dthe_cipher_init_tfm,
                .base.setkey                    = dthe_aes_ecb_setkey,
                .base.encrypt                   = dthe_aes_encrypt,
                .base.decrypt                   = dthe_aes_decrypt,
                .base.min_keysize               = AES_MIN_KEY_SIZE,
                .base.max_keysize               = AES_MAX_KEY_SIZE,
                .base.base = {
                        .cra_name               = "ecb(aes)",
                        .cra_driver_name        = "ecb-aes-dthev2",
                        .cra_priority           = 299,
                        .cra_flags              = CRYPTO_ALG_TYPE_SKCIPHER |
                                                  CRYPTO_ALG_ASYNC |
                                                  CRYPTO_ALG_KERN_DRIVER_ONLY,
                        .cra_alignmask          = AES_BLOCK_SIZE - 1,
                        .cra_blocksize          = AES_BLOCK_SIZE,
                        .cra_ctxsize            = sizeof(struct dthe_tfm_ctx),
                        .cra_reqsize            = sizeof(struct dthe_aes_req_ctx),
                        .cra_module             = THIS_MODULE,
                },
                .op.do_one_request = dthe_aes_run,
        }, /* ECB AES */
        {
                .base.init                      = dthe_cipher_init_tfm,
                .base.setkey                    = dthe_aes_cbc_setkey,
                .base.encrypt                   = dthe_aes_encrypt,
                .base.decrypt                   = dthe_aes_decrypt,
                .base.min_keysize               = AES_MIN_KEY_SIZE,
                .base.max_keysize               = AES_MAX_KEY_SIZE,
                .base.ivsize                    = AES_IV_SIZE,
                .base.base = {
                        .cra_name               = "cbc(aes)",
                        .cra_driver_name        = "cbc-aes-dthev2",
                        .cra_priority           = 299,
                        .cra_flags              = CRYPTO_ALG_TYPE_SKCIPHER |
                                                  CRYPTO_ALG_ASYNC |
                                                  CRYPTO_ALG_KERN_DRIVER_ONLY,
                        .cra_alignmask          = AES_BLOCK_SIZE - 1,
                        .cra_blocksize          = AES_BLOCK_SIZE,
                        .cra_ctxsize            = sizeof(struct dthe_tfm_ctx),
                        .cra_reqsize            = sizeof(struct dthe_aes_req_ctx),
                        .cra_module             = THIS_MODULE,
                },
                .op.do_one_request = dthe_aes_run,
        }, /* CBC AES */
        {
                .base.init                      = dthe_cipher_xts_init_tfm,
                .base.exit                      = dthe_cipher_xts_exit_tfm,
                .base.setkey                    = dthe_aes_xts_setkey,
                .base.encrypt                   = dthe_aes_encrypt,
                .base.decrypt                   = dthe_aes_decrypt,
                .base.min_keysize               = AES_MIN_KEY_SIZE * 2,
                .base.max_keysize               = AES_MAX_KEY_SIZE * 2,
                .base.ivsize                    = AES_IV_SIZE,
                .base.base = {
                        .cra_name               = "xts(aes)",
                        .cra_driver_name        = "xts-aes-dthev2",
                        .cra_priority           = 299,
                        .cra_flags              = CRYPTO_ALG_TYPE_SKCIPHER |
                                                  CRYPTO_ALG_ASYNC |
                                                  CRYPTO_ALG_KERN_DRIVER_ONLY |
                                                  CRYPTO_ALG_NEED_FALLBACK,
                        .cra_alignmask          = AES_BLOCK_SIZE - 1,
                        .cra_blocksize          = AES_BLOCK_SIZE,
                        .cra_ctxsize            = sizeof(struct dthe_tfm_ctx),
                        .cra_reqsize            = sizeof(struct dthe_aes_req_ctx),
                        .cra_module             = THIS_MODULE,
                },
                .op.do_one_request = dthe_aes_run,
        }, /* XTS AES */
};

int dthe_register_aes_algs(void)
{
        return crypto_engine_register_skciphers(cipher_algs, ARRAY_SIZE(cipher_algs));
}

void dthe_unregister_aes_algs(void)
{
        crypto_engine_unregister_skciphers(cipher_algs, ARRAY_SIZE(cipher_algs));
}