// 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/gcm.h>
#include <crypto/internal/aead.h>
#include <crypto/internal/skcipher.h>

#include "dthev2-common.h"

#include <linux/bitfield.h>
#include <linux/delay.h>
#include <linux/dmaengine.h>
#include <linux/dma-mapping.h>
#include <linux/io.h>
#include <linux/iopoll.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_TAG_OUT	0x0070

#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_CTR_MASK = BIT(6),
	AES_CTRL_XTS_MASK = BIT(12) | BIT(11),
	AES_CTRL_GCM_MASK = BIT(17) | BIT(16) | BIT(6),
	AES_CTRL_CCM_MASK = BIT(18) | BIT(6),
};

#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_CTR_WIDTH_128B		(BIT(7) | BIT(8))

#define DTHE_AES_CCM_L_FROM_IV_MASK		GENMASK(2, 0)
#define DTHE_AES_CCM_M_BITS			GENMASK(2, 0)
#define DTHE_AES_CTRL_CCM_L_FIELD_MASK		GENMASK(21, 19)
#define DTHE_AES_CTRL_CCM_M_FIELD_MASK		GENMASK(24, 22)

#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
#define DTHE_AES_GCM_AAD_MAXLEN			(BIT_ULL(32) - 1)
#define DTHE_AES_CCM_AAD_MAXLEN			(BIT(16) - BIT(8))
#define DTHE_AES_CCM_CRYPT_MAXLEN		(BIT_ULL(61) - 1)
#define POLL_TIMEOUT_INTERVAL			HZ

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_init_tfm_fallback(struct crypto_skcipher *tfm)
{
	struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);
	struct dthe_data *dev_data = dthe_get_dev(ctx);
	const char *alg_name = crypto_tfm_alg_name(crypto_skcipher_tfm(tfm));

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

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

	return 0;
}

static void dthe_cipher_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_ctr_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
{
	struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);
	int ret = dthe_aes_setkey(tfm, key, keylen);

	if (ret)
		return ret;

	ctx->aes_mode = DTHE_AES_CTR;

	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 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_CTR:
		ctrl_val |= AES_CTRL_CTR_MASK;
		ctrl_val |= DTHE_AES_CTRL_CTR_WIDTH_128B;
		break;
	case DTHE_AES_XTS:
		ctrl_val |= AES_CTRL_XTS_MASK;
		break;
	case DTHE_AES_GCM:
		ctrl_val |= AES_CTRL_GCM_MASK;
		break;
	case DTHE_AES_CCM:
		ctrl_val |= AES_CTRL_CCM_MASK;
		ctrl_val |= FIELD_PREP(DTHE_AES_CTRL_CCM_L_FIELD_MASK,
				       (iv_in[0] & DTHE_AES_CCM_L_FROM_IV_MASK));
		ctrl_val |= FIELD_PREP(DTHE_AES_CTRL_CCM_M_FIELD_MASK,
				       ((ctx->authsize - 2) >> 1) & DTHE_AES_CCM_M_BITS);
		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 int dthe_aes_do_fallback(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);

	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);
}

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 = sg_nents_for_len(dst, len);

	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;
	}

	/*
	 * CTR mode can operate on any input length, but the hardware
	 * requires input length to be a multiple of the block size.
	 * We need to handle the padding in the driver.
	 */
	if (ctx->aes_mode == DTHE_AES_CTR && req->cryptlen % AES_BLOCK_SIZE) {
		unsigned int pad_size = AES_BLOCK_SIZE - (req->cryptlen % AES_BLOCK_SIZE);
		u8 *pad_buf = rctx->padding;
		struct scatterlist *sg;

		len += pad_size;
		src_nents++;
		dst_nents++;

		src = kmalloc_array(src_nents, sizeof(*src), GFP_ATOMIC);
		if (!src) {
			ret = -ENOMEM;
			goto aes_ctr_src_alloc_err;
		}

		sg_init_table(src, src_nents);
		sg = dthe_copy_sg(src, req->src, req->cryptlen);
		memzero_explicit(pad_buf, AES_BLOCK_SIZE);
		sg_set_buf(sg, pad_buf, pad_size);

		if (diff_dst) {
			dst = kmalloc_array(dst_nents, sizeof(*dst), GFP_ATOMIC);
			if (!dst) {
				ret = -ENOMEM;
				goto aes_ctr_dst_alloc_err;
			}

			sg_init_table(dst, dst_nents);
			sg = dthe_copy_sg(dst, req->dst, req->cryptlen);
			sg_set_buf(sg, pad_buf, pad_size);
		} else {
			dst = src;
		}
	}

	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_map_src_err;
	}

	if (!diff_dst) {
		dst_mapped_nents = src_mapped_nents;
	} else {
		dst_mapped_nents = dma_map_sg(rx_dev, dst, dst_nents, dst_dir);
		if (dst_mapped_nents == 0) {
			ret = -EINVAL;
			goto aes_map_dst_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(len), aes_base_reg + DTHE_P_AES_C_LENGTH_0);
	writel_relaxed(upper_32_bits(len), 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:
	if (dst_dir != DMA_BIDIRECTIONAL)
		dma_unmap_sg(rx_dev, dst, dst_nents, dst_dir);
aes_map_dst_err:
	dma_unmap_sg(tx_dev, src, src_nents, src_dir);

aes_map_src_err:
	if (ctx->aes_mode == DTHE_AES_CTR && req->cryptlen % AES_BLOCK_SIZE) {
		memzero_explicit(rctx->padding, AES_BLOCK_SIZE);
		if (diff_dst)
			kfree(dst);
aes_ctr_dst_alloc_err:
		kfree(src);
aes_ctr_src_alloc_err:
		/*
		 * Fallback to software if ENOMEM
		 */
		if (ret == -ENOMEM)
			ret = dthe_aes_do_fallback(req);
	}

	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_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
	 * - do nothing for CTR
	 */
	if (req->cryptlen % AES_BLOCK_SIZE) {
		if (ctx->aes_mode == DTHE_AES_XTS)
			return dthe_aes_do_fallback(req);

		if (ctx->aes_mode != DTHE_AES_CTR)
			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 int dthe_aead_init_tfm(struct crypto_aead *tfm)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);
	struct dthe_data *dev_data = dthe_get_dev(ctx);

	ctx->dev_data = dev_data;

	const char *alg_name = crypto_tfm_alg_name(crypto_aead_tfm(tfm));

	ctx->aead_fb = crypto_alloc_sync_aead(alg_name, 0,
					      CRYPTO_ALG_NEED_FALLBACK);
	if (IS_ERR(ctx->aead_fb)) {
		dev_err(dev_data->dev, "fallback driver %s couldn't be loaded\n",
			alg_name);
		return PTR_ERR(ctx->aead_fb);
	}

	return 0;
}

static void dthe_aead_exit_tfm(struct crypto_aead *tfm)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);

	crypto_free_sync_aead(ctx->aead_fb);
}

/**
 * dthe_aead_prep_aad - Prepare AAD scatterlist from input request
 * @sg: Input scatterlist containing AAD
 * @assoclen: Length of AAD
 * @pad_buf: Buffer to hold AAD padding if needed
 *
 * Description:
 *   Creates a scatterlist containing only the AAD portion with padding
 *   to align to AES_BLOCK_SIZE. This simplifies DMA handling by allowing
 *   AAD to be sent separately via TX-only DMA.
 *
 * Return:
 *   Pointer to the AAD scatterlist, or ERR_PTR(error) on failure.
 *   The calling function needs to free the returned scatterlist when done.
 **/
static struct scatterlist *dthe_aead_prep_aad(struct scatterlist *sg,
					      unsigned int assoclen,
					      u8 *pad_buf)
{
	struct scatterlist *aad_sg;
	struct scatterlist *to_sg;
	int aad_nents;

	if (assoclen == 0)
		return NULL;

	aad_nents = sg_nents_for_len(sg, assoclen);
	if (assoclen % AES_BLOCK_SIZE)
		aad_nents++;

	aad_sg = kmalloc_array(aad_nents, sizeof(struct scatterlist), GFP_ATOMIC);
	if (!aad_sg)
		return ERR_PTR(-ENOMEM);

	sg_init_table(aad_sg, aad_nents);
	to_sg = dthe_copy_sg(aad_sg, sg, assoclen);
	if (assoclen % AES_BLOCK_SIZE) {
		unsigned int pad_len = AES_BLOCK_SIZE - (assoclen % AES_BLOCK_SIZE);

		memset(pad_buf, 0, pad_len);
		sg_set_buf(to_sg, pad_buf, pad_len);
	}

	return aad_sg;
}

/**
 * dthe_aead_prep_crypt - Prepare crypt scatterlist from req->src/req->dst
 * @sg: Input req->src/req->dst scatterlist
 * @assoclen: Length of AAD (to skip)
 * @cryptlen: Length of ciphertext/plaintext (minus the size of TAG in decryption)
 * @pad_buf: Zeroed buffer to hold crypt padding if needed
 *
 * Description:
 *   Creates a scatterlist containing only the ciphertext/plaintext portion
 *   (skipping AAD) with padding to align to AES_BLOCK_SIZE.
 *
 * Return:
 *   Pointer to the ciphertext scatterlist, or ERR_PTR(error) on failure.
 *   The calling function needs to free the returned scatterlist when done.
 **/
static struct scatterlist *dthe_aead_prep_crypt(struct scatterlist *sg,
						unsigned int assoclen,
						unsigned int cryptlen,
						u8 *pad_buf)
{
	struct scatterlist *out_sg[1];
	struct scatterlist *crypt_sg;
	struct scatterlist *to_sg;
	size_t split_sizes[1] = {cryptlen};
	int out_mapped_nents[1];
	int crypt_nents;
	int err;

	if (cryptlen == 0)
		return NULL;

	/* Skip AAD, extract ciphertext portion */
	err = sg_split(sg, 0, assoclen, 1, split_sizes, out_sg, out_mapped_nents, GFP_ATOMIC);
	if (err)
		goto dthe_aead_prep_crypt_split_err;

	crypt_nents = sg_nents_for_len(out_sg[0], cryptlen);
	if (cryptlen % AES_BLOCK_SIZE)
		crypt_nents++;

	crypt_sg = kmalloc_array(crypt_nents, sizeof(struct scatterlist), GFP_ATOMIC);
	if (!crypt_sg) {
		err = -ENOMEM;
		goto dthe_aead_prep_crypt_mem_err;
	}

	sg_init_table(crypt_sg, crypt_nents);
	to_sg = dthe_copy_sg(crypt_sg, out_sg[0], cryptlen);
	if (cryptlen % AES_BLOCK_SIZE) {
		unsigned int pad_len = AES_BLOCK_SIZE - (cryptlen % AES_BLOCK_SIZE);

		sg_set_buf(to_sg, pad_buf, pad_len);
	}

dthe_aead_prep_crypt_mem_err:
	kfree(out_sg[0]);

dthe_aead_prep_crypt_split_err:
	if (err)
		return ERR_PTR(err);
	return crypt_sg;
}

static int dthe_aead_read_tag(struct dthe_tfm_ctx *ctx, u32 *tag)
{
	struct dthe_data *dev_data = dthe_get_dev(ctx);
	void __iomem *aes_base_reg = dev_data->regs + DTHE_P_AES_BASE;
	u32 val;
	int ret;

	ret = readl_relaxed_poll_timeout(aes_base_reg + DTHE_P_AES_CTRL, val,
					 (val & DTHE_AES_CTRL_SAVED_CTX_READY),
					 0, POLL_TIMEOUT_INTERVAL);
	if (ret)
		return ret;

	for (int i = 0; i < AES_BLOCK_WORDS; ++i)
		tag[i] = readl_relaxed(aes_base_reg +
				       DTHE_P_AES_TAG_OUT +
				       DTHE_REG_SIZE * i);
	return 0;
}

static int dthe_aead_enc_get_tag(struct aead_request *req)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
	u32 tag[AES_BLOCK_WORDS];
	int nents;
	int ret;

	ret = dthe_aead_read_tag(ctx, tag);
	if (ret)
		return ret;

	nents = sg_nents_for_len(req->dst, req->cryptlen + req->assoclen + ctx->authsize);

	sg_pcopy_from_buffer(req->dst, nents, tag, ctx->authsize,
			     req->assoclen + req->cryptlen);

	return 0;
}

static int dthe_aead_dec_verify_tag(struct aead_request *req)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
	u32 tag_out[AES_BLOCK_WORDS];
	u32 tag_in[AES_BLOCK_WORDS];
	int nents;
	int ret;

	ret = dthe_aead_read_tag(ctx, tag_out);
	if (ret)
		return ret;

	nents = sg_nents_for_len(req->src, req->assoclen + req->cryptlen);

	sg_pcopy_to_buffer(req->src, nents, tag_in, ctx->authsize,
			   req->assoclen + req->cryptlen - ctx->authsize);

	if (crypto_memneq(tag_in, tag_out, ctx->authsize))
		return -EBADMSG;
	else
		return 0;
}

static int dthe_aead_setkey(struct crypto_aead *tfm, const u8 *key, unsigned int keylen)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);

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

	crypto_sync_aead_clear_flags(ctx->aead_fb, CRYPTO_TFM_REQ_MASK);
	crypto_sync_aead_set_flags(ctx->aead_fb,
				   crypto_aead_get_flags(tfm) &
				   CRYPTO_TFM_REQ_MASK);

	return crypto_sync_aead_setkey(ctx->aead_fb, key, keylen);
}

static int dthe_gcm_aes_setkey(struct crypto_aead *tfm, const u8 *key, unsigned int keylen)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);
	int ret;

	ret = dthe_aead_setkey(tfm, key, keylen);
	if (ret)
		return ret;

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

	return ret;
}

static int dthe_ccm_aes_setkey(struct crypto_aead *tfm, const u8 *key, unsigned int keylen)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);
	int ret;

	ret = dthe_aead_setkey(tfm, key, keylen);
	if (ret)
		return ret;

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

	return ret;
}

static int dthe_aead_setauthsize(struct crypto_aead *tfm, unsigned int authsize)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);

	/* Invalid auth size will be handled by crypto_aead_setauthsize() */
	ctx->authsize = authsize;

	return crypto_sync_aead_setauthsize(ctx->aead_fb, authsize);
}

static int dthe_aead_do_fallback(struct aead_request *req)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);

	SYNC_AEAD_REQUEST_ON_STACK(subreq, ctx->aead_fb);

	aead_request_set_callback(subreq, req->base.flags,
				  req->base.complete, req->base.data);
	aead_request_set_crypt(subreq, req->src, req->dst, req->cryptlen, req->iv);
	aead_request_set_ad(subreq, req->assoclen);

	return rctx->enc ? crypto_aead_encrypt(subreq) :
		crypto_aead_decrypt(subreq);
}

static void dthe_aead_dma_in_callback(void *data)
{
	struct aead_request *req = (struct aead_request *)data;
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);

	complete(&rctx->aes_compl);
}

static int dthe_aead_run(struct crypto_engine *engine, void *areq)
{
	struct aead_request *req = container_of(areq, struct aead_request, base);
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);
	struct dthe_data *dev_data = dthe_get_dev(ctx);

	unsigned int cryptlen = req->cryptlen;
	unsigned int assoclen = req->assoclen;
	unsigned int authsize = ctx->authsize;
	unsigned int unpadded_cryptlen;
	struct scatterlist *src = NULL;
	struct scatterlist *dst = NULL;
	struct scatterlist *aad_sg = NULL;
	u32 iv_in[AES_IV_WORDS];

	int aad_nents = 0;
	int src_nents = 0;
	int dst_nents = 0;
	int aad_mapped_nents = 0;
	int src_mapped_nents = 0;
	int dst_mapped_nents = 0;

	u8 *src_assoc_padbuf = rctx->padding;
	u8 *src_crypt_padbuf = rctx->padding + AES_BLOCK_SIZE;
	u8 *dst_crypt_padbuf = rctx->padding + AES_BLOCK_SIZE;

	bool diff_dst;
	enum dma_data_direction aad_dir, src_dir, dst_dir;

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

	int ret;
	int err;

	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);

	/* In decryption, the last authsize bytes are the TAG */
	if (!rctx->enc)
		cryptlen -= authsize;
	unpadded_cryptlen = cryptlen;

	memset(src_assoc_padbuf, 0, AES_BLOCK_SIZE);
	memset(src_crypt_padbuf, 0, AES_BLOCK_SIZE);
	memset(dst_crypt_padbuf, 0, AES_BLOCK_SIZE);

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

	if (req->src == req->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;
	}
	aad_dir = DMA_TO_DEVICE;

	/* Prep AAD scatterlist (always from req->src) */
	aad_sg = dthe_aead_prep_aad(req->src, req->assoclen, src_assoc_padbuf);
	if (IS_ERR(aad_sg)) {
		ret = PTR_ERR(aad_sg);
		goto aead_prep_aad_err;
	}

	/* Prep ciphertext src scatterlist */
	src = dthe_aead_prep_crypt(req->src, req->assoclen, cryptlen, src_crypt_padbuf);
	if (IS_ERR(src)) {
		ret = PTR_ERR(src);
		goto aead_prep_src_err;
	}

	/* Prep ciphertext dst scatterlist (only if separate dst) */
	if (diff_dst) {
		dst = dthe_aead_prep_crypt(req->dst, req->assoclen, unpadded_cryptlen,
					   dst_crypt_padbuf);
		if (IS_ERR(dst)) {
			ret = PTR_ERR(dst);
			goto aead_prep_dst_err;
		}
	} else {
		dst = src;
	}

	/* Calculate padded lengths for nents calculations */
	if (req->assoclen % AES_BLOCK_SIZE)
		assoclen += AES_BLOCK_SIZE - (req->assoclen % AES_BLOCK_SIZE);
	if (cryptlen % AES_BLOCK_SIZE)
		cryptlen += AES_BLOCK_SIZE - (cryptlen % AES_BLOCK_SIZE);

	if (assoclen != 0) {
		/* Map AAD for TX only */
		aad_nents = sg_nents_for_len(aad_sg, assoclen);
		aad_mapped_nents = dma_map_sg(tx_dev, aad_sg, aad_nents, aad_dir);
		if (aad_mapped_nents == 0) {
			dev_err(dev_data->dev, "Failed to map AAD for TX\n");
			ret = -EINVAL;
			goto aead_dma_map_aad_err;
		}

		/* Prepare DMA descriptors for AAD TX */
		desc_aad_out = dmaengine_prep_slave_sg(dev_data->dma_aes_tx, aad_sg,
						       aad_mapped_nents, DMA_MEM_TO_DEV,
						       DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
		if (!desc_aad_out) {
			dev_err(dev_data->dev, "AAD TX prep_slave_sg() failed\n");
			ret = -EINVAL;
			goto aead_dma_prep_aad_err;
		}
	}

	if (cryptlen != 0) {
		/* Map ciphertext src for TX (BIDIRECTIONAL if in-place) */
		src_nents = sg_nents_for_len(src, cryptlen);
		src_mapped_nents = dma_map_sg(tx_dev, src, src_nents, src_dir);
		if (src_mapped_nents == 0) {
			dev_err(dev_data->dev, "Failed to map ciphertext src for TX\n");
			ret = -EINVAL;
			goto aead_dma_prep_aad_err;
		}

		/* Prepare DMA descriptors for ciphertext TX */
		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, "Ciphertext TX prep_slave_sg() failed\n");
			ret = -EINVAL;
			goto aead_dma_prep_src_err;
		}

		/* Map ciphertext dst for RX (only if separate dst) */
		if (diff_dst) {
			dst_nents = sg_nents_for_len(dst, cryptlen);
			dst_mapped_nents = dma_map_sg(rx_dev, dst, dst_nents, dst_dir);
			if (dst_mapped_nents == 0) {
				dev_err(dev_data->dev, "Failed to map ciphertext dst for RX\n");
				ret = -EINVAL;
				goto aead_dma_prep_src_err;
			}
		} else {
			dst_nents = src_nents;
			dst_mapped_nents = src_mapped_nents;
		}

		/* Prepare DMA descriptor for ciphertext RX */
		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, "Ciphertext RX prep_slave_sg() failed\n");
			ret = -EINVAL;
			goto aead_dma_prep_dst_err;
		}

		desc_in->callback = dthe_aead_dma_in_callback;
		desc_in->callback_param = req;
	} else if (assoclen != 0) {
		/* AAD-only operation */
		desc_aad_out->callback = dthe_aead_dma_in_callback;
		desc_aad_out->callback_param = req;
	}

	init_completion(&rctx->aes_compl);

	/*
	 * HACK: There is an unknown hw issue where if the previous operation had alen = 0 and
	 * plen != 0, the current operation's tag calculation is incorrect in the case where
	 * plen = 0 and alen != 0 currently. This is a workaround for now which somehow works;
	 * by resetting the context by writing a 1 to the C_LENGTH_0 and AUTH_LENGTH registers.
	 */
	if (cryptlen == 0) {
		writel_relaxed(1, aes_base_reg + DTHE_P_AES_C_LENGTH_0);
		writel_relaxed(1, aes_base_reg + DTHE_P_AES_AUTH_LENGTH);
	}

	if (ctx->aes_mode == DTHE_AES_GCM) {
		if (req->iv) {
			memcpy(iv_in, req->iv, GCM_AES_IV_SIZE);
		} else {
			iv_in[0] = 0;
			iv_in[1] = 0;
			iv_in[2] = 0;
		}
		iv_in[3] = 0x01000000;
	} else {
		memcpy(iv_in, req->iv, AES_IV_SIZE);
	}

	/* Clear key2 to reset previous GHASH intermediate data */
	for (int i = 0; i < AES_KEYSIZE_256 / sizeof(u32); ++i)
		writel_relaxed(0, aes_base_reg + DTHE_P_AES_KEY2_6 + DTHE_REG_SIZE * i);

	dthe_aes_set_ctrl_key(ctx, rctx, iv_in);

	writel_relaxed(lower_32_bits(unpadded_cryptlen), aes_base_reg + DTHE_P_AES_C_LENGTH_0);
	writel_relaxed(upper_32_bits(unpadded_cryptlen), aes_base_reg + DTHE_P_AES_C_LENGTH_1);
	writel_relaxed(req->assoclen, aes_base_reg + DTHE_P_AES_AUTH_LENGTH);

	/* Submit DMA descriptors: AAD TX, ciphertext TX, ciphertext RX */
	if (assoclen != 0)
		dmaengine_submit(desc_aad_out);
	if (cryptlen != 0) {
		dmaengine_submit(desc_out);
		dmaengine_submit(desc_in);
	}

	if (cryptlen != 0)
		dma_async_issue_pending(dev_data->dma_aes_rx);
	dma_async_issue_pending(dev_data->dma_aes_tx);

	/* Need to do 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;
		if (cryptlen != 0)
			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;
	}

	if (cryptlen != 0)
		dma_sync_sg_for_cpu(rx_dev, dst, dst_nents, dst_dir);

	if (rctx->enc)
		err = dthe_aead_enc_get_tag(req);
	else
		err = dthe_aead_dec_verify_tag(req);

	ret = (ret) ? ret : err;

aead_dma_prep_dst_err:
	if (diff_dst && cryptlen != 0)
		dma_unmap_sg(rx_dev, dst, dst_nents, dst_dir);
aead_dma_prep_src_err:
	if (cryptlen != 0)
		dma_unmap_sg(tx_dev, src, src_nents, src_dir);
aead_dma_prep_aad_err:
	if (assoclen != 0)
		dma_unmap_sg(tx_dev, aad_sg, aad_nents, aad_dir);

aead_dma_map_aad_err:
	if (diff_dst && cryptlen != 0)
		kfree(dst);
aead_prep_dst_err:
	if (cryptlen != 0)
		kfree(src);
aead_prep_src_err:
	if (assoclen != 0)
		kfree(aad_sg);

aead_prep_aad_err:
	memzero_explicit(rctx->padding, 2 * AES_BLOCK_SIZE);

	if (ret)
		ret = dthe_aead_do_fallback(req);

	local_bh_disable();
	crypto_finalize_aead_request(engine, req, ret);
	local_bh_enable();
	return 0;
}

static int dthe_aead_crypt(struct aead_request *req)
{
	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);
	struct dthe_data *dev_data = dthe_get_dev(ctx);
	struct crypto_engine *engine;
	unsigned int cryptlen = req->cryptlen;
	bool is_zero_ctr = true;

	/* In decryption, last authsize bytes are the TAG */
	if (!rctx->enc)
		cryptlen -= ctx->authsize;

	if (ctx->aes_mode == DTHE_AES_CCM) {
		/*
		 * For CCM Mode, the 128-bit IV contains the following:
		 * | 0 .. 2 | 3 .. 7 | 8 .. (127-8*L) | (128-8*L) .. 127 |
		 * |   L-1  |  Zero  |     Nonce      |      Counter     |
		 * L needs to be between 2-8 (inclusive), i.e. 1 <= (L-1) <= 7
		 * and the next 5 bits need to be zeroes. Else return -EINVAL
		 */
		u8 *iv = req->iv;
		u8 L = iv[0];

		/* variable L stores L-1 here */
		if (L < 1 || L > 7)
			return -EINVAL;
		/*
		 * DTHEv2 HW can only work with zero initial counter in CCM mode.
		 * Check if the initial counter value is zero or not
		 */
		for (int i = 0; i < L + 1; ++i) {
			if (iv[AES_IV_SIZE - 1 - i] != 0) {
				is_zero_ctr = false;
				break;
			}
		}
	}

	/*
	 * Need to fallback to software in the following cases due to HW restrictions:
	 * - Both AAD and plaintext/ciphertext are zero length
	 * - For AES-GCM, AAD length is more than 2^32 - 1 bytes
	 * - For AES-CCM, AAD length is more than 2^16 - 2^8 bytes
	 * - For AES-CCM, plaintext/ciphertext length is more than 2^61 - 1 bytes
	 * - For AES-CCM, AAD length is non-zero but plaintext/ciphertext length is zero
	 * - For AES-CCM, the initial counter (last L+1 bytes of IV) is not all zeroes
	 *
	 * PS: req->cryptlen is currently unsigned int type, which causes the second and fourth
	 * cases above tautologically false. If req->cryptlen is to be changed to a 64-bit
	 * type, the check for these would also need to be added below.
	 */
	if ((req->assoclen == 0 && cryptlen == 0) ||
	    (ctx->aes_mode == DTHE_AES_CCM && req->assoclen > DTHE_AES_CCM_AAD_MAXLEN) ||
	    (ctx->aes_mode == DTHE_AES_CCM && cryptlen == 0) ||
	    (ctx->aes_mode == DTHE_AES_CCM && !is_zero_ctr))
		return dthe_aead_do_fallback(req);

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

static int dthe_aead_encrypt(struct aead_request *req)
{
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);

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

static int dthe_aead_decrypt(struct aead_request *req)
{
	struct dthe_aes_req_ctx *rctx = aead_request_ctx(req);

	rctx->enc = 0;
	return dthe_aead_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_init_tfm_fallback,
		.base.exit			= dthe_cipher_exit_tfm,
		.base.setkey			= dthe_aes_ctr_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.chunksize			= AES_BLOCK_SIZE,
		.base.base = {
			.cra_name		= "ctr(aes)",
			.cra_driver_name	= "ctr-aes-dthev2",
			.cra_priority		= 299,
			.cra_flags		= CRYPTO_ALG_TYPE_SKCIPHER |
						  CRYPTO_ALG_ASYNC |
						  CRYPTO_ALG_KERN_DRIVER_ONLY |
						  CRYPTO_ALG_NEED_FALLBACK,
			.cra_blocksize		= 1,
			.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,
	}, /* CTR AES */
	{
		.base.init			= dthe_cipher_init_tfm_fallback,
		.base.exit			= dthe_cipher_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 */
};

static struct aead_engine_alg aead_algs[] = {
	{
		.base.init			= dthe_aead_init_tfm,
		.base.exit			= dthe_aead_exit_tfm,
		.base.setkey			= dthe_gcm_aes_setkey,
		.base.setauthsize		= dthe_aead_setauthsize,
		.base.maxauthsize		= AES_BLOCK_SIZE,
		.base.encrypt			= dthe_aead_encrypt,
		.base.decrypt			= dthe_aead_decrypt,
		.base.chunksize			= AES_BLOCK_SIZE,
		.base.ivsize			= GCM_AES_IV_SIZE,
		.base.base = {
			.cra_name		= "gcm(aes)",
			.cra_driver_name	= "gcm-aes-dthev2",
			.cra_priority		= 299,
			.cra_flags		= CRYPTO_ALG_TYPE_AEAD |
						  CRYPTO_ALG_KERN_DRIVER_ONLY |
						  CRYPTO_ALG_ASYNC |
						  CRYPTO_ALG_NEED_FALLBACK,
			.cra_blocksize		= 1,
			.cra_ctxsize		= sizeof(struct dthe_tfm_ctx),
			.cra_reqsize		= sizeof(struct dthe_aes_req_ctx),
			.cra_module		= THIS_MODULE,
		},
		.op.do_one_request = dthe_aead_run,
	}, /* GCM AES */
	{
		.base.init			= dthe_aead_init_tfm,
		.base.exit			= dthe_aead_exit_tfm,
		.base.setkey			= dthe_ccm_aes_setkey,
		.base.setauthsize		= dthe_aead_setauthsize,
		.base.maxauthsize		= AES_BLOCK_SIZE,
		.base.encrypt			= dthe_aead_encrypt,
		.base.decrypt			= dthe_aead_decrypt,
		.base.chunksize			= AES_BLOCK_SIZE,
		.base.ivsize			= AES_IV_SIZE,
		.base.base = {
			.cra_name		= "ccm(aes)",
			.cra_driver_name	= "ccm-aes-dthev2",
			.cra_priority		= 299,
			.cra_flags		= CRYPTO_ALG_TYPE_AEAD |
						  CRYPTO_ALG_KERN_DRIVER_ONLY |
						  CRYPTO_ALG_ASYNC |
						  CRYPTO_ALG_NEED_FALLBACK,
			.cra_blocksize		= 1,
			.cra_ctxsize		= sizeof(struct dthe_tfm_ctx),
			.cra_reqsize		= sizeof(struct dthe_aes_req_ctx),
			.cra_module		= THIS_MODULE,
		},
		.op.do_one_request = dthe_aead_run,
	}, /* CCM AES */
};

int dthe_register_aes_algs(void)
{
	int ret = 0;

	ret = crypto_engine_register_skciphers(cipher_algs, ARRAY_SIZE(cipher_algs));
	if (ret)
		return ret;
	ret = crypto_engine_register_aeads(aead_algs, ARRAY_SIZE(aead_algs));
	if (ret)
		crypto_engine_unregister_skciphers(cipher_algs, ARRAY_SIZE(cipher_algs));

	return ret;
}

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