crypto: keembay - Add support for Keem Bay OCS AES/SM4

Add support for the AES/SM4 crypto engine included in the Offload and
Crypto Subsystem (OCS) of the Intel Keem Bay SoC, thus enabling
hardware-acceleration for the following transformations:

- ecb(aes), cbc(aes), ctr(aes), cts(cbc(aes)), gcm(aes) and cbc(aes);
  supported for 128-bit and 256-bit keys.

- ecb(sm4), cbc(sm4), ctr(sm4), cts(cbc(sm4)), gcm(sm4) and cbc(sm4);
  supported for 128-bit keys.

The driver passes crypto manager self-tests, including the extra tests
(CRYPTO_MANAGER_EXTRA_TESTS=y).

Signed-off-by: Mike Healy <mikex.healy@intel.com>
Co-developed-by: Daniele Alessandrelli <daniele.alessandrelli@intel.com>
Signed-off-by: Daniele Alessandrelli <daniele.alessandrelli@intel.com>
Acked-by: Mark Gross <mgross@linux.intel.com>
Signed-off-by: Herbert Xu <herbert@gondor.apana.org.au>

1 files changed, 1713 insertions, 0 deletions

diff --git a/drivers/crypto/keembay/keembay-ocs-aes-core.c b/drivers/crypto/keembay/keembay-ocs-aes-core.c
new file mode 100644
index 000000000000..b6b25d994af3
--- /dev/null
+++ b/drivers/crypto/keembay/keembay-ocs-aes-core.c
@@ -0,0 +1,1713 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Intel Keem Bay OCS AES Crypto Driver.
+ *
+ * Copyright (C) 2018-2020 Intel Corporation
+ */
+
+#include <linux/clk.h>
+#include <linux/completion.h>
+#include <linux/crypto.h>
+#include <linux/dma-mapping.h>
+#include <linux/interrupt.h>
+#include <linux/io.h>
+#include <linux/module.h>
+#include <linux/of.h>
+#include <linux/platform_device.h>
+#include <linux/types.h>
+
+#include <crypto/aes.h>
+#include <crypto/engine.h>
+#include <crypto/gcm.h>
+#include <crypto/scatterwalk.h>
+
+#include <crypto/internal/aead.h>
+#include <crypto/internal/skcipher.h>
+
+#include "ocs-aes.h"
+
+#define KMB_OCS_PRIORITY	350
+#define DRV_NAME		"keembay-ocs-aes"
+
+#define OCS_AES_MIN_KEY_SIZE	16
+#define OCS_AES_MAX_KEY_SIZE	32
+#define OCS_AES_KEYSIZE_128	16
+#define OCS_AES_KEYSIZE_192	24
+#define OCS_AES_KEYSIZE_256	32
+#define OCS_SM4_KEY_SIZE	16
+
+/**
+ * struct ocs_aes_tctx - OCS AES Transform context
+ * @engine_ctx:		Engine context.
+ * @aes_dev:		The OCS AES device.
+ * @key:		AES/SM4 key.
+ * @key_len:		The length (in bytes) of @key.
+ * @cipher:		OCS cipher to use (either AES or SM4).
+ * @sw_cipher:		The cipher to use as fallback.
+ * @use_fallback:	Whether or not fallback cipher should be used.
+ */
+struct ocs_aes_tctx {
+	struct crypto_engine_ctx engine_ctx;
+	struct ocs_aes_dev *aes_dev;
+	u8 key[OCS_AES_KEYSIZE_256];
+	unsigned int key_len;
+	enum ocs_cipher cipher;
+	union {
+		struct crypto_sync_skcipher *sk;
+		struct crypto_aead *aead;
+	} sw_cipher;
+	bool use_fallback;
+};
+
+/**
+ * struct ocs_aes_rctx - OCS AES Request context.
+ * @instruction:	Instruction to be executed (encrypt / decrypt).
+ * @mode:		Mode to use (ECB, CBC, CTR, CCm, GCM, CTS)
+ * @src_nents:		Number of source SG entries.
+ * @dst_nents:		Number of destination SG entries.
+ * @src_dma_count:	The number of DMA-mapped entries of the source SG.
+ * @dst_dma_count:	The number of DMA-mapped entries of the destination SG.
+ * @in_place:		Whether or not this is an in place request, i.e.,
+ *			src_sg == dst_sg.
+ * @src_dll:		OCS DMA linked list for input data.
+ * @dst_dll:		OCS DMA linked list for output data.
+ * @last_ct_blk:	Buffer to hold last cipher text block (only used in CBC
+ *			mode).
+ * @cts_swap:		Whether or not CTS swap must be performed.
+ * @aad_src_dll:	OCS DMA linked list for input AAD data.
+ * @aad_dst_dll:	OCS DMA linked list for output AAD data.
+ * @in_tag:		Buffer to hold input encrypted tag (only used for
+ *			CCM/GCM decrypt).
+ * @out_tag:		Buffer to hold output encrypted / decrypted tag (only
+ *			used for GCM encrypt / decrypt).
+ */
+struct ocs_aes_rctx {
+	/* Fields common across all modes. */
+	enum ocs_instruction	instruction;
+	enum ocs_mode		mode;
+	int			src_nents;
+	int			dst_nents;
+	int			src_dma_count;
+	int			dst_dma_count;
+	bool			in_place;
+	struct ocs_dll_desc	src_dll;
+	struct ocs_dll_desc	dst_dll;
+
+	/* CBC specific */
+	u8			last_ct_blk[AES_BLOCK_SIZE];
+
+	/* CTS specific */
+	int			cts_swap;
+
+	/* CCM/GCM specific */
+	struct ocs_dll_desc	aad_src_dll;
+	struct ocs_dll_desc	aad_dst_dll;
+	u8			in_tag[AES_BLOCK_SIZE];
+
+	/* GCM specific */
+	u8			out_tag[AES_BLOCK_SIZE];
+};
+
+/* Driver data. */
+struct ocs_aes_drv {
+	struct list_head dev_list;
+	spinlock_t lock;	/* Protects dev_list. */
+};
+
+static struct ocs_aes_drv ocs_aes = {
+	.dev_list = LIST_HEAD_INIT(ocs_aes.dev_list),
+	.lock = __SPIN_LOCK_UNLOCKED(ocs_aes.lock),
+};
+
+static struct ocs_aes_dev *kmb_ocs_aes_find_dev(struct ocs_aes_tctx *tctx)
+{
+	struct ocs_aes_dev *aes_dev;
+
+	spin_lock(&ocs_aes.lock);
+
+	if (tctx->aes_dev) {
+		aes_dev = tctx->aes_dev;
+		goto exit;
+	}
+
+	/* Only a single OCS device available */
+	aes_dev = list_first_entry(&ocs_aes.dev_list, struct ocs_aes_dev, list);
+	tctx->aes_dev = aes_dev;
+
+exit:
+	spin_unlock(&ocs_aes.lock);
+
+	return aes_dev;
+}
+
+/*
+ * Ensure key is 128-bit or 256-bit for AES or 128-bit for SM4 and an actual
+ * key is being passed in.
+ *
+ * Return: 0 if key is valid, -EINVAL otherwise.
+ */
+static int check_key(const u8 *in_key, size_t key_len, enum ocs_cipher cipher)
+{
+	if (!in_key)
+		return -EINVAL;
+
+	/* For AES, only 128-byte or 256-byte keys are supported. */
+	if (cipher == OCS_AES && (key_len == OCS_AES_KEYSIZE_128 ||
+				  key_len == OCS_AES_KEYSIZE_256))
+		return 0;
+
+	/* For SM4, only 128-byte keys are supported. */
+	if (cipher == OCS_SM4 && key_len == OCS_AES_KEYSIZE_128)
+		return 0;
+
+	/* Everything else is unsupported. */
+	return -EINVAL;
+}
+
+/* Save key into transformation context. */
+static int save_key(struct ocs_aes_tctx *tctx, const u8 *in_key, size_t key_len,
+		    enum ocs_cipher cipher)
+{
+	int ret;
+
+	ret = check_key(in_key, key_len, cipher);
+	if (ret)
+		return ret;
+
+	memcpy(tctx->key, in_key, key_len);
+	tctx->key_len = key_len;
+	tctx->cipher = cipher;
+
+	return 0;
+}
+
+/* Set key for symmetric cypher. */
+static int kmb_ocs_sk_set_key(struct crypto_skcipher *tfm, const u8 *in_key,
+			      size_t key_len, enum ocs_cipher cipher)
+{
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+
+	/* Fallback is used for AES with 192-bit key. */
+	tctx->use_fallback = (cipher == OCS_AES &&
+			      key_len == OCS_AES_KEYSIZE_192);
+
+	if (!tctx->use_fallback)
+		return save_key(tctx, in_key, key_len, cipher);
+
+	crypto_sync_skcipher_clear_flags(tctx->sw_cipher.sk,
+					 CRYPTO_TFM_REQ_MASK);
+	crypto_sync_skcipher_set_flags(tctx->sw_cipher.sk,
+				       tfm->base.crt_flags &
+				       CRYPTO_TFM_REQ_MASK);
+
+	return crypto_sync_skcipher_setkey(tctx->sw_cipher.sk, in_key, key_len);
+}
+
+/* Set key for AEAD cipher. */
+static int kmb_ocs_aead_set_key(struct crypto_aead *tfm, const u8 *in_key,
+				size_t key_len, enum ocs_cipher cipher)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(tfm);
+
+	/* Fallback is used for AES with 192-bit key. */
+	tctx->use_fallback = (cipher == OCS_AES &&
+			      key_len == OCS_AES_KEYSIZE_192);
+
+	if (!tctx->use_fallback)
+		return save_key(tctx, in_key, key_len, cipher);
+
+	crypto_aead_clear_flags(tctx->sw_cipher.aead, CRYPTO_TFM_REQ_MASK);
+	crypto_aead_set_flags(tctx->sw_cipher.aead,
+			      crypto_aead_get_flags(tfm) & CRYPTO_TFM_REQ_MASK);
+
+	return crypto_aead_setkey(tctx->sw_cipher.aead, in_key, key_len);
+}
+
+/* Swap two AES blocks in SG lists. */
+static void sg_swap_blocks(struct scatterlist *sgl, unsigned int nents,
+			   off_t blk1_offset, off_t blk2_offset)
+{
+	u8 tmp_buf1[AES_BLOCK_SIZE], tmp_buf2[AES_BLOCK_SIZE];
+
+	/*
+	 * No easy way to copy within sg list, so copy both blocks to temporary
+	 * buffers first.
+	 */
+	sg_pcopy_to_buffer(sgl, nents, tmp_buf1, AES_BLOCK_SIZE, blk1_offset);
+	sg_pcopy_to_buffer(sgl, nents, tmp_buf2, AES_BLOCK_SIZE, blk2_offset);
+	sg_pcopy_from_buffer(sgl, nents, tmp_buf1, AES_BLOCK_SIZE, blk2_offset);
+	sg_pcopy_from_buffer(sgl, nents, tmp_buf2, AES_BLOCK_SIZE, blk1_offset);
+}
+
+/* Initialize request context to default values. */
+static void ocs_aes_init_rctx(struct ocs_aes_rctx *rctx)
+{
+	/* Zero everything. */
+	memset(rctx, 0, sizeof(*rctx));
+
+	/* Set initial value for DMA addresses. */
+	rctx->src_dll.dma_addr = DMA_MAPPING_ERROR;
+	rctx->dst_dll.dma_addr = DMA_MAPPING_ERROR;
+	rctx->aad_src_dll.dma_addr = DMA_MAPPING_ERROR;
+	rctx->aad_dst_dll.dma_addr = DMA_MAPPING_ERROR;
+}
+
+static int kmb_ocs_sk_validate_input(struct skcipher_request *req,
+				     enum ocs_mode mode)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	int iv_size = crypto_skcipher_ivsize(tfm);
+
+	switch (mode) {
+	case OCS_MODE_ECB:
+		/* Ensure input length is multiple of block size */
+		if (req->cryptlen % AES_BLOCK_SIZE != 0)
+			return -EINVAL;
+
+		return 0;
+
+	case OCS_MODE_CBC:
+		/* Ensure input length is multiple of block size */
+		if (req->cryptlen % AES_BLOCK_SIZE != 0)
+			return -EINVAL;
+
+		/* Ensure IV is present and block size in length */
+		if (!req->iv || iv_size != AES_BLOCK_SIZE)
+			return -EINVAL;
+		/*
+		 * NOTE: Since req->cryptlen == 0 case was already handled in
+		 * kmb_ocs_sk_common(), the above two conditions also guarantee
+		 * that: cryptlen >= iv_size
+		 */
+		return 0;
+
+	case OCS_MODE_CTR:
+		/* Ensure IV is present and block size in length */
+		if (!req->iv || iv_size != AES_BLOCK_SIZE)
+			return -EINVAL;
+		return 0;
+
+	case OCS_MODE_CTS:
+		/* Ensure input length >= block size */
+		if (req->cryptlen < AES_BLOCK_SIZE)
+			return -EINVAL;
+
+		/* Ensure IV is present and block size in length */
+		if (!req->iv || iv_size != AES_BLOCK_SIZE)
+			return -EINVAL;
+
+		return 0;
+	default:
+		return -EINVAL;
+	}
+}
+
+/*
+ * Called by encrypt() / decrypt() skcipher functions.
+ *
+ * Use fallback if needed, otherwise initialize context and enqueue request
+ * into engine.
+ */
+static int kmb_ocs_sk_common(struct skcipher_request *req,
+			     enum ocs_cipher cipher,
+			     enum ocs_instruction instruction,
+			     enum ocs_mode mode)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_rctx *rctx = skcipher_request_ctx(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	struct ocs_aes_dev *aes_dev;
+	int rc;
+
+	if (tctx->use_fallback) {
+		SYNC_SKCIPHER_REQUEST_ON_STACK(subreq, tctx->sw_cipher.sk);
+
+		skcipher_request_set_sync_tfm(subreq, tctx->sw_cipher.sk);
+		skcipher_request_set_callback(subreq, req->base.flags, NULL,
+					      NULL);
+		skcipher_request_set_crypt(subreq, req->src, req->dst,
+					   req->cryptlen, req->iv);
+
+		if (instruction == OCS_ENCRYPT)
+			rc = crypto_skcipher_encrypt(subreq);
+		else
+			rc = crypto_skcipher_decrypt(subreq);
+
+		skcipher_request_zero(subreq);
+
+		return rc;
+	}
+
+	/*
+	 * If cryptlen == 0, no processing needed for ECB, CBC and CTR.
+	 *
+	 * For CTS continue: kmb_ocs_sk_validate_input() will return -EINVAL.
+	 */
+	if (!req->cryptlen && mode != OCS_MODE_CTS)
+		return 0;
+
+	rc = kmb_ocs_sk_validate_input(req, mode);
+	if (rc)
+		return rc;
+
+	aes_dev = kmb_ocs_aes_find_dev(tctx);
+	if (!aes_dev)
+		return -ENODEV;
+
+	if (cipher != tctx->cipher)
+		return -EINVAL;
+
+	ocs_aes_init_rctx(rctx);
+	rctx->instruction = instruction;
+	rctx->mode = mode;
+
+	return crypto_transfer_skcipher_request_to_engine(aes_dev->engine, req);
+}
+
+static void cleanup_ocs_dma_linked_list(struct device *dev,
+					struct ocs_dll_desc *dll)
+{
+	if (dll->vaddr)
+		dma_free_coherent(dev, dll->size, dll->vaddr, dll->dma_addr);
+	dll->vaddr = NULL;
+	dll->size = 0;
+	dll->dma_addr = DMA_MAPPING_ERROR;
+}
+
+static void kmb_ocs_sk_dma_cleanup(struct skcipher_request *req)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_rctx *rctx = skcipher_request_ctx(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	struct device *dev = tctx->aes_dev->dev;
+
+	if (rctx->src_dma_count) {
+		dma_unmap_sg(dev, req->src, rctx->src_nents, DMA_TO_DEVICE);
+		rctx->src_dma_count = 0;
+	}
+
+	if (rctx->dst_dma_count) {
+		dma_unmap_sg(dev, req->dst, rctx->dst_nents, rctx->in_place ?
+							     DMA_BIDIRECTIONAL :
+							     DMA_FROM_DEVICE);
+		rctx->dst_dma_count = 0;
+	}
+
+	/* Clean up OCS DMA linked lists */
+	cleanup_ocs_dma_linked_list(dev, &rctx->src_dll);
+	cleanup_ocs_dma_linked_list(dev, &rctx->dst_dll);
+}
+
+static int kmb_ocs_sk_prepare_inplace(struct skcipher_request *req)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_rctx *rctx = skcipher_request_ctx(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	int iv_size = crypto_skcipher_ivsize(tfm);
+	int rc;
+
+	/*
+	 * For CBC decrypt, save last block (iv) to last_ct_blk buffer.
+	 *
+	 * Note: if we are here, we already checked that cryptlen >= iv_size
+	 * and iv_size == AES_BLOCK_SIZE (i.e., the size of last_ct_blk); see
+	 * kmb_ocs_sk_validate_input().
+	 */
+	if (rctx->mode == OCS_MODE_CBC && rctx->instruction == OCS_DECRYPT)
+		scatterwalk_map_and_copy(rctx->last_ct_blk, req->src,
+					 req->cryptlen - iv_size, iv_size, 0);
+
+	/* For CTS decrypt, swap last two blocks, if needed. */
+	if (rctx->cts_swap && rctx->instruction == OCS_DECRYPT)
+		sg_swap_blocks(req->dst, rctx->dst_nents,
+			       req->cryptlen - AES_BLOCK_SIZE,
+			       req->cryptlen - (2 * AES_BLOCK_SIZE));
+
+	/* src and dst buffers are the same, use bidirectional DMA mapping. */
+	rctx->dst_dma_count = dma_map_sg(tctx->aes_dev->dev, req->dst,
+					 rctx->dst_nents, DMA_BIDIRECTIONAL);
+	if (rctx->dst_dma_count == 0) {
+		dev_err(tctx->aes_dev->dev, "Failed to map destination sg\n");
+		return -ENOMEM;
+	}
+
+	/* Create DST linked list */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+					    rctx->dst_dma_count, &rctx->dst_dll,
+					    req->cryptlen, 0);
+	if (rc)
+		return rc;
+	/*
+	 * If descriptor creation was successful, set the src_dll.dma_addr to
+	 * the value of dst_dll.dma_addr, as we do in-place AES operation on
+	 * the src.
+	 */
+	rctx->src_dll.dma_addr = rctx->dst_dll.dma_addr;
+
+	return 0;
+}
+
+static int kmb_ocs_sk_prepare_notinplace(struct skcipher_request *req)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_rctx *rctx = skcipher_request_ctx(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	int rc;
+
+	rctx->src_nents =  sg_nents_for_len(req->src, req->cryptlen);
+	if (rctx->src_nents < 0)
+		return -EBADMSG;
+
+	/* Map SRC SG. */
+	rctx->src_dma_count = dma_map_sg(tctx->aes_dev->dev, req->src,
+					 rctx->src_nents, DMA_TO_DEVICE);
+	if (rctx->src_dma_count == 0) {
+		dev_err(tctx->aes_dev->dev, "Failed to map source sg\n");
+		return -ENOMEM;
+	}
+
+	/* Create SRC linked list */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->src,
+					    rctx->src_dma_count, &rctx->src_dll,
+					    req->cryptlen, 0);
+	if (rc)
+		return rc;
+
+	/* Map DST SG. */
+	rctx->dst_dma_count = dma_map_sg(tctx->aes_dev->dev, req->dst,
+					 rctx->dst_nents, DMA_FROM_DEVICE);
+	if (rctx->dst_dma_count == 0) {
+		dev_err(tctx->aes_dev->dev, "Failed to map destination sg\n");
+		return -ENOMEM;
+	}
+
+	/* Create DST linked list */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+					    rctx->dst_dma_count, &rctx->dst_dll,
+					    req->cryptlen, 0);
+	if (rc)
+		return rc;
+
+	/* If this is not a CTS decrypt operation with swapping, we are done. */
+	if (!(rctx->cts_swap && rctx->instruction == OCS_DECRYPT))
+		return 0;
+
+	/*
+	 * Otherwise, we have to copy src to dst (as we cannot modify src).
+	 * Use OCS AES bypass mode to copy src to dst via DMA.
+	 *
+	 * NOTE: for anything other than small data sizes this is rather
+	 * inefficient.
+	 */
+	rc = ocs_aes_bypass_op(tctx->aes_dev, rctx->dst_dll.dma_addr,
+			       rctx->src_dll.dma_addr, req->cryptlen);
+	if (rc)
+		return rc;
+
+	/*
+	 * Now dst == src, so clean up what we did so far and use in_place
+	 * logic.
+	 */
+	kmb_ocs_sk_dma_cleanup(req);
+	rctx->in_place = true;
+
+	return kmb_ocs_sk_prepare_inplace(req);
+}
+
+static int kmb_ocs_sk_run(struct skcipher_request *req)
+{
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_rctx *rctx = skcipher_request_ctx(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	struct ocs_aes_dev *aes_dev = tctx->aes_dev;
+	int iv_size = crypto_skcipher_ivsize(tfm);
+	int rc;
+
+	rctx->dst_nents = sg_nents_for_len(req->dst, req->cryptlen);
+	if (rctx->dst_nents < 0)
+		return -EBADMSG;
+
+	/*
+	 * If 2 blocks or greater, and multiple of block size swap last two
+	 * blocks to be compatible with other crypto API CTS implementations:
+	 * OCS mode uses CBC-CS2, whereas other crypto API implementations use
+	 * CBC-CS3.
+	 * CBC-CS2 and CBC-CS3 defined by:
+	 * https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38a-add.pdf
+	 */
+	rctx->cts_swap = (rctx->mode == OCS_MODE_CTS &&
+			  req->cryptlen > AES_BLOCK_SIZE &&
+			  req->cryptlen % AES_BLOCK_SIZE == 0);
+
+	rctx->in_place = (req->src == req->dst);
+
+	if (rctx->in_place)
+		rc = kmb_ocs_sk_prepare_inplace(req);
+	else
+		rc = kmb_ocs_sk_prepare_notinplace(req);
+
+	if (rc)
+		goto error;
+
+	rc = ocs_aes_op(aes_dev, rctx->mode, tctx->cipher, rctx->instruction,
+			rctx->dst_dll.dma_addr, rctx->src_dll.dma_addr,
+			req->cryptlen, req->iv, iv_size);
+	if (rc)
+		goto error;
+
+	/* Clean-up DMA before further processing output. */
+	kmb_ocs_sk_dma_cleanup(req);
+
+	/* For CTS Encrypt, swap last 2 blocks, if needed. */
+	if (rctx->cts_swap && rctx->instruction == OCS_ENCRYPT) {
+		sg_swap_blocks(req->dst, rctx->dst_nents,
+			       req->cryptlen - AES_BLOCK_SIZE,
+			       req->cryptlen - (2 * AES_BLOCK_SIZE));
+		return 0;
+	}
+
+	/* For CBC copy IV to req->IV. */
+	if (rctx->mode == OCS_MODE_CBC) {
+		/* CBC encrypt case. */
+		if (rctx->instruction == OCS_ENCRYPT) {
+			scatterwalk_map_and_copy(req->iv, req->dst,
+						 req->cryptlen - iv_size,
+						 iv_size, 0);
+			return 0;
+		}
+		/* CBC decrypt case. */
+		if (rctx->in_place)
+			memcpy(req->iv, rctx->last_ct_blk, iv_size);
+		else
+			scatterwalk_map_and_copy(req->iv, req->src,
+						 req->cryptlen - iv_size,
+						 iv_size, 0);
+		return 0;
+	}
+	/* For all other modes there's nothing to do. */
+
+	return 0;
+
+error:
+	kmb_ocs_sk_dma_cleanup(req);
+
+	return rc;
+}
+
+static int kmb_ocs_aead_validate_input(struct aead_request *req,
+				       enum ocs_instruction instruction,
+				       enum ocs_mode mode)
+{
+	struct crypto_aead *tfm = crypto_aead_reqtfm(req);
+	int tag_size = crypto_aead_authsize(tfm);
+	int iv_size = crypto_aead_ivsize(tfm);
+
+	/* For decrypt crytplen == len(PT) + len(tag). */
+	if (instruction == OCS_DECRYPT && req->cryptlen < tag_size)
+		return -EINVAL;
+
+	/* IV is mandatory. */
+	if (!req->iv)
+		return -EINVAL;
+
+	switch (mode) {
+	case OCS_MODE_GCM:
+		if (iv_size != GCM_AES_IV_SIZE)
+			return -EINVAL;
+
+		return 0;
+
+	case OCS_MODE_CCM:
+		/* Ensure IV is present and block size in length */
+		if (iv_size != AES_BLOCK_SIZE)
+			return -EINVAL;
+
+		return 0;
+
+	default:
+		return -EINVAL;
+	}
+}
+
+/*
+ * Called by encrypt() / decrypt() aead functions.
+ *
+ * Use fallback if needed, otherwise initialize context and enqueue request
+ * into engine.
+ */
+static int kmb_ocs_aead_common(struct aead_request *req,
+			       enum ocs_cipher cipher,
+			       enum ocs_instruction instruction,
+			       enum ocs_mode mode)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
+	struct ocs_aes_rctx *rctx = aead_request_ctx(req);
+	struct ocs_aes_dev *dd;
+	int rc;
+
+	if (tctx->use_fallback) {
+		struct aead_request *subreq = aead_request_ctx(req);
+
+		aead_request_set_tfm(subreq, tctx->sw_cipher.aead);
+		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);
+		rc = crypto_aead_setauthsize(tctx->sw_cipher.aead,
+					     crypto_aead_authsize(crypto_aead_reqtfm(req)));
+		if (rc)
+			return rc;
+
+		return (instruction == OCS_ENCRYPT) ?
+		       crypto_aead_encrypt(subreq) :
+		       crypto_aead_decrypt(subreq);
+	}
+
+	rc = kmb_ocs_aead_validate_input(req, instruction, mode);
+	if (rc)
+		return rc;
+
+	dd = kmb_ocs_aes_find_dev(tctx);
+	if (!dd)
+		return -ENODEV;
+
+	if (cipher != tctx->cipher)
+		return -EINVAL;
+
+	ocs_aes_init_rctx(rctx);
+	rctx->instruction = instruction;
+	rctx->mode = mode;
+
+	return crypto_transfer_aead_request_to_engine(dd->engine, req);
+}
+
+static void kmb_ocs_aead_dma_cleanup(struct aead_request *req)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
+	struct ocs_aes_rctx *rctx = aead_request_ctx(req);
+	struct device *dev = tctx->aes_dev->dev;
+
+	if (rctx->src_dma_count) {
+		dma_unmap_sg(dev, req->src, rctx->src_nents, DMA_TO_DEVICE);
+		rctx->src_dma_count = 0;
+	}
+
+	if (rctx->dst_dma_count) {
+		dma_unmap_sg(dev, req->dst, rctx->dst_nents, rctx->in_place ?
+							     DMA_BIDIRECTIONAL :
+							     DMA_FROM_DEVICE);
+		rctx->dst_dma_count = 0;
+	}
+	/* Clean up OCS DMA linked lists */
+	cleanup_ocs_dma_linked_list(dev, &rctx->src_dll);
+	cleanup_ocs_dma_linked_list(dev, &rctx->dst_dll);
+	cleanup_ocs_dma_linked_list(dev, &rctx->aad_src_dll);
+	cleanup_ocs_dma_linked_list(dev, &rctx->aad_dst_dll);
+}
+
+/**
+ * kmb_ocs_aead_dma_prepare() - Do DMA mapping for AEAD processing.
+ * @req:		The AEAD request being processed.
+ * @src_dll_size:	Where to store the length of the data mapped into the
+ *			src_dll OCS DMA list.
+ *
+ * Do the following:
+ * - DMA map req->src and req->dst
+ * - Initialize the following OCS DMA linked lists: rctx->src_dll,
+ *   rctx->dst_dll, rctx->aad_src_dll and rxtc->aad_dst_dll.
+ *
+ * Return: 0 on success, negative error code otherwise.
+ */
+static int kmb_ocs_aead_dma_prepare(struct aead_request *req, u32 *src_dll_size)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
+	const int tag_size = crypto_aead_authsize(crypto_aead_reqtfm(req));
+	struct ocs_aes_rctx *rctx = aead_request_ctx(req);
+	u32 in_size;	/* The length of the data to be mapped by src_dll. */
+	u32 out_size;	/* The length of the data to be mapped by dst_dll. */
+	u32 dst_size;	/* The length of the data in dst_sg. */
+	int rc;
+
+	/* Get number of entries in input data SG list. */
+	rctx->src_nents = sg_nents_for_len(req->src,
+					   req->assoclen + req->cryptlen);
+	if (rctx->src_nents < 0)
+		return -EBADMSG;
+
+	if (rctx->instruction == OCS_DECRYPT) {
+		/*
+		 * For decrypt:
+		 * - src sg list is:		AAD|CT|tag
+		 * - dst sg list expects:	AAD|PT
+		 *
+		 * in_size == len(CT); out_size == len(PT)
+		 */
+
+		/* req->cryptlen includes both CT and tag. */
+		in_size = req->cryptlen - tag_size;
+
+		/* out_size = PT size == CT size */
+		out_size = in_size;
+
+		/* len(dst_sg) == len(AAD) + len(PT) */
+		dst_size = req->assoclen + out_size;
+
+		/*
+		 * Copy tag from source SG list to 'in_tag' buffer.
+		 *
+		 * Note: this needs to be done here, before DMA mapping src_sg.
+		 */
+		sg_pcopy_to_buffer(req->src, rctx->src_nents, rctx->in_tag,
+				   tag_size, req->assoclen + in_size);
+
+	} else { /* OCS_ENCRYPT */
+		/*
+		 * For encrypt:
+		 *	src sg list is:		AAD|PT
+		 *	dst sg list expects:	AAD|CT|tag
+		 */
+		/* in_size == len(PT) */
+		in_size = req->cryptlen;
+
+		/*
+		 * In CCM mode the OCS engine appends the tag to the ciphertext,
+		 * but in GCM mode the tag must be read from the tag registers
+		 * and appended manually below
+		 */
+		out_size = (rctx->mode == OCS_MODE_CCM) ? in_size + tag_size :
+							  in_size;
+		/* len(dst_sg) == len(AAD) + len(CT) + len(tag) */
+		dst_size = req->assoclen + in_size + tag_size;
+	}
+	*src_dll_size = in_size;
+
+	/* Get number of entries in output data SG list. */
+	rctx->dst_nents = sg_nents_for_len(req->dst, dst_size);
+	if (rctx->dst_nents < 0)
+		return -EBADMSG;
+
+	rctx->in_place = (req->src == req->dst) ? 1 : 0;
+
+	/* Map destination; use bidirectional mapping for in-place case. */
+	rctx->dst_dma_count = dma_map_sg(tctx->aes_dev->dev, req->dst,
+					 rctx->dst_nents,
+					 rctx->in_place ? DMA_BIDIRECTIONAL :
+							  DMA_FROM_DEVICE);
+	if (rctx->dst_dma_count == 0 && rctx->dst_nents != 0) {
+		dev_err(tctx->aes_dev->dev, "Failed to map destination sg\n");
+		return -ENOMEM;
+	}
+
+	/* Create AAD DST list: maps dst[0:AAD_SIZE-1]. */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+					    rctx->dst_dma_count,
+					    &rctx->aad_dst_dll, req->assoclen,
+					    0);
+	if (rc)
+		return rc;
+
+	/* Create DST list: maps dst[AAD_SIZE:out_size] */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+					    rctx->dst_dma_count, &rctx->dst_dll,
+					    out_size, req->assoclen);
+	if (rc)
+		return rc;
+
+	if (rctx->in_place) {
+		/* If this is not CCM encrypt, we are done. */
+		if (!(rctx->mode == OCS_MODE_CCM &&
+		      rctx->instruction == OCS_ENCRYPT)) {
+			/*
+			 * SRC and DST are the same, so re-use the same DMA
+			 * addresses (to avoid allocating new DMA lists
+			 * identical to the dst ones).
+			 */
+			rctx->src_dll.dma_addr = rctx->dst_dll.dma_addr;
+			rctx->aad_src_dll.dma_addr = rctx->aad_dst_dll.dma_addr;
+
+			return 0;
+		}
+		/*
+		 * For CCM encrypt the input and output linked lists contain
+		 * different amounts of data, so, we need to create different
+		 * SRC and AAD SRC lists, even for the in-place case.
+		 */
+		rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+						    rctx->dst_dma_count,
+						    &rctx->aad_src_dll,
+						    req->assoclen, 0);
+		if (rc)
+			return rc;
+		rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->dst,
+						    rctx->dst_dma_count,
+						    &rctx->src_dll, in_size,
+						    req->assoclen);
+		if (rc)
+			return rc;
+
+		return 0;
+	}
+	/* Not in-place case. */
+
+	/* Map source SG. */
+	rctx->src_dma_count = dma_map_sg(tctx->aes_dev->dev, req->src,
+					 rctx->src_nents, DMA_TO_DEVICE);
+	if (rctx->src_dma_count == 0 && rctx->src_nents != 0) {
+		dev_err(tctx->aes_dev->dev, "Failed to map source sg\n");
+		return -ENOMEM;
+	}
+
+	/* Create AAD SRC list. */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->src,
+					    rctx->src_dma_count,
+					    &rctx->aad_src_dll,
+					    req->assoclen, 0);
+	if (rc)
+		return rc;
+
+	/* Create SRC list. */
+	rc = ocs_create_linked_list_from_sg(tctx->aes_dev, req->src,
+					    rctx->src_dma_count,
+					    &rctx->src_dll, in_size,
+					    req->assoclen);
+	if (rc)
+		return rc;
+
+	if (req->assoclen == 0)
+		return 0;
+
+	/* Copy AAD from src sg to dst sg using OCS DMA. */
+	rc = ocs_aes_bypass_op(tctx->aes_dev, rctx->aad_dst_dll.dma_addr,
+			       rctx->aad_src_dll.dma_addr, req->cryptlen);
+	if (rc)
+		dev_err(tctx->aes_dev->dev,
+			"Failed to copy source AAD to destination AAD\n");
+
+	return rc;
+}
+
+static int kmb_ocs_aead_run(struct aead_request *req)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
+	const int tag_size = crypto_aead_authsize(crypto_aead_reqtfm(req));
+	struct ocs_aes_rctx *rctx = aead_request_ctx(req);
+	u32 in_size;	/* The length of the data mapped by src_dll. */
+	int rc;
+
+	rc = kmb_ocs_aead_dma_prepare(req, &in_size);
+	if (rc)
+		goto exit;
+
+	/* For CCM, we just call the OCS processing and we are done. */
+	if (rctx->mode == OCS_MODE_CCM) {
+		rc = ocs_aes_ccm_op(tctx->aes_dev, tctx->cipher,
+				    rctx->instruction, rctx->dst_dll.dma_addr,
+				    rctx->src_dll.dma_addr, in_size,
+				    req->iv,
+				    rctx->aad_src_dll.dma_addr, req->assoclen,
+				    rctx->in_tag, tag_size);
+		goto exit;
+	}
+	/* GCM case; invoke OCS processing. */
+	rc = ocs_aes_gcm_op(tctx->aes_dev, tctx->cipher,
+			    rctx->instruction,
+			    rctx->dst_dll.dma_addr,
+			    rctx->src_dll.dma_addr, in_size,
+			    req->iv,
+			    rctx->aad_src_dll.dma_addr, req->assoclen,
+			    rctx->out_tag, tag_size);
+	if (rc)
+		goto exit;
+
+	/* For GCM decrypt, we have to compare in_tag with out_tag. */
+	if (rctx->instruction == OCS_DECRYPT) {
+		rc = memcmp(rctx->in_tag, rctx->out_tag, tag_size) ?
+		     -EBADMSG : 0;
+		goto exit;
+	}
+
+	/* For GCM encrypt, we must manually copy out_tag to DST sg. */
+
+	/* Clean-up must be called before the sg_pcopy_from_buffer() below. */
+	kmb_ocs_aead_dma_cleanup(req);
+
+	/* Copy tag to destination sg after AAD and CT. */
+	sg_pcopy_from_buffer(req->dst, rctx->dst_nents, rctx->out_tag,
+			     tag_size, req->assoclen + req->cryptlen);
+
+	/* Return directly as DMA cleanup already done. */
+	return 0;
+
+exit:
+	kmb_ocs_aead_dma_cleanup(req);
+
+	return rc;
+}
+
+static int kmb_ocs_aes_sk_do_one_request(struct crypto_engine *engine,
+					 void *areq)
+{
+	struct skcipher_request *req =
+			container_of(areq, struct skcipher_request, base);
+	struct crypto_skcipher *tfm = crypto_skcipher_reqtfm(req);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	int err;
+
+	if (!tctx->aes_dev) {
+		err = -ENODEV;
+		goto exit;
+	}
+
+	err = ocs_aes_set_key(tctx->aes_dev, tctx->key_len, tctx->key,
+			      tctx->cipher);
+	if (err)
+		goto exit;
+
+	err = kmb_ocs_sk_run(req);
+
+exit:
+	crypto_finalize_skcipher_request(engine, req, err);
+
+	return 0;
+}
+
+static int kmb_ocs_aes_aead_do_one_request(struct crypto_engine *engine,
+					   void *areq)
+{
+	struct aead_request *req = container_of(areq,
+						struct aead_request, base);
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(crypto_aead_reqtfm(req));
+	int err;
+
+	if (!tctx->aes_dev)
+		return -ENODEV;
+
+	err = ocs_aes_set_key(tctx->aes_dev, tctx->key_len, tctx->key,
+			      tctx->cipher);
+	if (err)
+		goto exit;
+
+	err = kmb_ocs_aead_run(req);
+
+exit:
+	crypto_finalize_aead_request(tctx->aes_dev->engine, req, err);
+
+	return 0;
+}
+
+static int kmb_ocs_aes_set_key(struct crypto_skcipher *tfm, const u8 *in_key,
+			       unsigned int key_len)
+{
+	return kmb_ocs_sk_set_key(tfm, in_key, key_len, OCS_AES);
+}
+
+static int kmb_ocs_aes_aead_set_key(struct crypto_aead *tfm, const u8 *in_key,
+				    unsigned int key_len)
+{
+	return kmb_ocs_aead_set_key(tfm, in_key, key_len, OCS_AES);
+}
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB
+static int kmb_ocs_aes_ecb_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_ECB);
+}
+
+static int kmb_ocs_aes_ecb_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_ECB);
+}
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB */
+
+static int kmb_ocs_aes_cbc_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_CBC);
+}
+
+static int kmb_ocs_aes_cbc_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_CBC);
+}
+
+static int kmb_ocs_aes_ctr_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_CTR);
+}
+
+static int kmb_ocs_aes_ctr_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_CTR);
+}
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS
+static int kmb_ocs_aes_cts_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_CTS);
+}
+
+static int kmb_ocs_aes_cts_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_CTS);
+}
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS */
+
+static int kmb_ocs_aes_gcm_encrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_GCM);
+}
+
+static int kmb_ocs_aes_gcm_decrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_GCM);
+}
+
+static int kmb_ocs_aes_ccm_encrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_AES, OCS_ENCRYPT, OCS_MODE_CCM);
+}
+
+static int kmb_ocs_aes_ccm_decrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_AES, OCS_DECRYPT, OCS_MODE_CCM);
+}
+
+static int kmb_ocs_sm4_set_key(struct crypto_skcipher *tfm, const u8 *in_key,
+			       unsigned int key_len)
+{
+	return kmb_ocs_sk_set_key(tfm, in_key, key_len, OCS_SM4);
+}
+
+static int kmb_ocs_sm4_aead_set_key(struct crypto_aead *tfm, const u8 *in_key,
+				    unsigned int key_len)
+{
+	return kmb_ocs_aead_set_key(tfm, in_key, key_len, OCS_SM4);
+}
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB
+static int kmb_ocs_sm4_ecb_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_ECB);
+}
+
+static int kmb_ocs_sm4_ecb_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_ECB);
+}
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB */
+
+static int kmb_ocs_sm4_cbc_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_CBC);
+}
+
+static int kmb_ocs_sm4_cbc_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_CBC);
+}
+
+static int kmb_ocs_sm4_ctr_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_CTR);
+}
+
+static int kmb_ocs_sm4_ctr_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_CTR);
+}
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS
+static int kmb_ocs_sm4_cts_encrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_CTS);
+}
+
+static int kmb_ocs_sm4_cts_decrypt(struct skcipher_request *req)
+{
+	return kmb_ocs_sk_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_CTS);
+}
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS */
+
+static int kmb_ocs_sm4_gcm_encrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_GCM);
+}
+
+static int kmb_ocs_sm4_gcm_decrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_GCM);
+}
+
+static int kmb_ocs_sm4_ccm_encrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_SM4, OCS_ENCRYPT, OCS_MODE_CCM);
+}
+
+static int kmb_ocs_sm4_ccm_decrypt(struct aead_request *req)
+{
+	return kmb_ocs_aead_common(req, OCS_SM4, OCS_DECRYPT, OCS_MODE_CCM);
+}
+
+static inline int ocs_common_init(struct ocs_aes_tctx *tctx)
+{
+	tctx->engine_ctx.op.prepare_request = NULL;
+	tctx->engine_ctx.op.do_one_request = kmb_ocs_aes_sk_do_one_request;
+	tctx->engine_ctx.op.unprepare_request = NULL;
+
+	return 0;
+}
+
+static int ocs_aes_init_tfm(struct crypto_skcipher *tfm)
+{
+	const char *alg_name = crypto_tfm_alg_name(&tfm->base);
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+	struct crypto_sync_skcipher *blk;
+
+	/* set fallback cipher in case it will be needed */
+	blk = crypto_alloc_sync_skcipher(alg_name, 0, CRYPTO_ALG_NEED_FALLBACK);
+	if (IS_ERR(blk))
+		return PTR_ERR(blk);
+
+	tctx->sw_cipher.sk = blk;
+
+	crypto_skcipher_set_reqsize(tfm, sizeof(struct ocs_aes_rctx));
+
+	return ocs_common_init(tctx);
+}
+
+static int ocs_sm4_init_tfm(struct crypto_skcipher *tfm)
+{
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+
+	crypto_skcipher_set_reqsize(tfm, sizeof(struct ocs_aes_rctx));
+
+	return ocs_common_init(tctx);
+}
+
+static inline void clear_key(struct ocs_aes_tctx *tctx)
+{
+	memzero_explicit(tctx->key, OCS_AES_KEYSIZE_256);
+
+	/* Zero key registers if set */
+	if (tctx->aes_dev)
+		ocs_aes_set_key(tctx->aes_dev, OCS_AES_KEYSIZE_256,
+				tctx->key, OCS_AES);
+}
+
+static void ocs_exit_tfm(struct crypto_skcipher *tfm)
+{
+	struct ocs_aes_tctx *tctx = crypto_skcipher_ctx(tfm);
+
+	clear_key(tctx);
+
+	if (tctx->sw_cipher.sk) {
+		crypto_free_sync_skcipher(tctx->sw_cipher.sk);
+		tctx->sw_cipher.sk = NULL;
+	}
+}
+
+static inline int ocs_common_aead_init(struct ocs_aes_tctx *tctx)
+{
+	tctx->engine_ctx.op.prepare_request = NULL;
+	tctx->engine_ctx.op.do_one_request = kmb_ocs_aes_aead_do_one_request;
+	tctx->engine_ctx.op.unprepare_request = NULL;
+
+	return 0;
+}
+
+static int ocs_aes_aead_cra_init(struct crypto_aead *tfm)
+{
+	const char *alg_name = crypto_tfm_alg_name(&tfm->base);
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(tfm);
+	struct crypto_aead *blk;
+
+	/* Set fallback cipher in case it will be needed */
+	blk = crypto_alloc_aead(alg_name, 0, CRYPTO_ALG_NEED_FALLBACK);
+	if (IS_ERR(blk))
+		return PTR_ERR(blk);
+
+	tctx->sw_cipher.aead = blk;
+
+	crypto_aead_set_reqsize(tfm,
+				max(sizeof(struct ocs_aes_rctx),
+				    (sizeof(struct aead_request) +
+				     crypto_aead_reqsize(tctx->sw_cipher.aead))));
+
+	return ocs_common_aead_init(tctx);
+}
+
+static int kmb_ocs_aead_ccm_setauthsize(struct crypto_aead *tfm,
+					unsigned int authsize)
+{
+	switch (authsize) {
+	case 4:
+	case 6:
+	case 8:
+	case 10:
+	case 12:
+	case 14:
+	case 16:
+		return 0;
+	default:
+		return -EINVAL;
+	}
+}
+
+static int kmb_ocs_aead_gcm_setauthsize(struct crypto_aead *tfm,
+					unsigned int authsize)
+{
+	return crypto_gcm_check_authsize(authsize);
+}
+
+static int ocs_sm4_aead_cra_init(struct crypto_aead *tfm)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(tfm);
+
+	crypto_aead_set_reqsize(tfm, sizeof(struct ocs_aes_rctx));
+
+	return ocs_common_aead_init(tctx);
+}
+
+static void ocs_aead_cra_exit(struct crypto_aead *tfm)
+{
+	struct ocs_aes_tctx *tctx = crypto_aead_ctx(tfm);
+
+	clear_key(tctx);
+
+	if (tctx->sw_cipher.aead) {
+		crypto_free_aead(tctx->sw_cipher.aead);
+		tctx->sw_cipher.aead = NULL;
+	}
+}
+
+static struct skcipher_alg algs[] = {
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB
+	{
+		.base.cra_name = "ecb(aes)",
+		.base.cra_driver_name = "ecb-aes-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY |
+				  CRYPTO_ALG_NEED_FALLBACK,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_AES_MIN_KEY_SIZE,
+		.max_keysize = OCS_AES_MAX_KEY_SIZE,
+		.setkey = kmb_ocs_aes_set_key,
+		.encrypt = kmb_ocs_aes_ecb_encrypt,
+		.decrypt = kmb_ocs_aes_ecb_decrypt,
+		.init = ocs_aes_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB */
+	{
+		.base.cra_name = "cbc(aes)",
+		.base.cra_driver_name = "cbc-aes-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY |
+				  CRYPTO_ALG_NEED_FALLBACK,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_AES_MIN_KEY_SIZE,
+		.max_keysize = OCS_AES_MAX_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_aes_set_key,
+		.encrypt = kmb_ocs_aes_cbc_encrypt,
+		.decrypt = kmb_ocs_aes_cbc_decrypt,
+		.init = ocs_aes_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+	{
+		.base.cra_name = "ctr(aes)",
+		.base.cra_driver_name = "ctr-aes-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY |
+				  CRYPTO_ALG_NEED_FALLBACK,
+		.base.cra_blocksize = 1,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_AES_MIN_KEY_SIZE,
+		.max_keysize = OCS_AES_MAX_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_aes_set_key,
+		.encrypt = kmb_ocs_aes_ctr_encrypt,
+		.decrypt = kmb_ocs_aes_ctr_decrypt,
+		.init = ocs_aes_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS
+	{
+		.base.cra_name = "cts(cbc(aes))",
+		.base.cra_driver_name = "cts-aes-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY |
+				  CRYPTO_ALG_NEED_FALLBACK,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_AES_MIN_KEY_SIZE,
+		.max_keysize = OCS_AES_MAX_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_aes_set_key,
+		.encrypt = kmb_ocs_aes_cts_encrypt,
+		.decrypt = kmb_ocs_aes_cts_decrypt,
+		.init = ocs_aes_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS */
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB
+	{
+		.base.cra_name = "ecb(sm4)",
+		.base.cra_driver_name = "ecb-sm4-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_SM4_KEY_SIZE,
+		.max_keysize = OCS_SM4_KEY_SIZE,
+		.setkey = kmb_ocs_sm4_set_key,
+		.encrypt = kmb_ocs_sm4_ecb_encrypt,
+		.decrypt = kmb_ocs_sm4_ecb_decrypt,
+		.init = ocs_sm4_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB */
+	{
+		.base.cra_name = "cbc(sm4)",
+		.base.cra_driver_name = "cbc-sm4-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_SM4_KEY_SIZE,
+		.max_keysize = OCS_SM4_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_sm4_set_key,
+		.encrypt = kmb_ocs_sm4_cbc_encrypt,
+		.decrypt = kmb_ocs_sm4_cbc_decrypt,
+		.init = ocs_sm4_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+	{
+		.base.cra_name = "ctr(sm4)",
+		.base.cra_driver_name = "ctr-sm4-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY,
+		.base.cra_blocksize = 1,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_SM4_KEY_SIZE,
+		.max_keysize = OCS_SM4_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_sm4_set_key,
+		.encrypt = kmb_ocs_sm4_ctr_encrypt,
+		.decrypt = kmb_ocs_sm4_ctr_decrypt,
+		.init = ocs_sm4_init_tfm,
+		.exit = ocs_exit_tfm,
+	},
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS
+	{
+		.base.cra_name = "cts(cbc(sm4))",
+		.base.cra_driver_name = "cts-sm4-keembay-ocs",
+		.base.cra_priority = KMB_OCS_PRIORITY,
+		.base.cra_flags = CRYPTO_ALG_ASYNC |
+				  CRYPTO_ALG_KERN_DRIVER_ONLY,
+		.base.cra_blocksize = AES_BLOCK_SIZE,
+		.base.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+		.base.cra_module = THIS_MODULE,
+		.base.cra_alignmask = 0,
+
+		.min_keysize = OCS_SM4_KEY_SIZE,
+		.max_keysize = OCS_SM4_KEY_SIZE,
+		.ivsize = AES_BLOCK_SIZE,
+		.setkey = kmb_ocs_sm4_set_key,
+		.encrypt = kmb_ocs_sm4_cts_encrypt,
+		.decrypt = kmb_ocs_sm4_cts_decrypt,
+		.init = ocs_sm4_init_tfm,
+		.exit = ocs_exit_tfm,
+	}
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS */
+};
+
+static struct aead_alg algs_aead[] = {
+	{
+		.base = {
+			.cra_name = "gcm(aes)",
+			.cra_driver_name = "gcm-aes-keembay-ocs",
+			.cra_priority = KMB_OCS_PRIORITY,
+			.cra_flags = CRYPTO_ALG_ASYNC |
+				     CRYPTO_ALG_KERN_DRIVER_ONLY |
+				     CRYPTO_ALG_NEED_FALLBACK,
+			.cra_blocksize = 1,
+			.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+			.cra_alignmask = 0,
+			.cra_module = THIS_MODULE,
+		},
+		.init = ocs_aes_aead_cra_init,
+		.exit = ocs_aead_cra_exit,
+		.ivsize = GCM_AES_IV_SIZE,
+		.maxauthsize = AES_BLOCK_SIZE,
+		.setauthsize = kmb_ocs_aead_gcm_setauthsize,
+		.setkey = kmb_ocs_aes_aead_set_key,
+		.encrypt = kmb_ocs_aes_gcm_encrypt,
+		.decrypt = kmb_ocs_aes_gcm_decrypt,
+	},
+	{
+		.base = {
+			.cra_name = "ccm(aes)",
+			.cra_driver_name = "ccm-aes-keembay-ocs",
+			.cra_priority = KMB_OCS_PRIORITY,
+			.cra_flags = CRYPTO_ALG_ASYNC |
+				     CRYPTO_ALG_KERN_DRIVER_ONLY |
+				     CRYPTO_ALG_NEED_FALLBACK,
+			.cra_blocksize = 1,
+			.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+			.cra_alignmask = 0,
+			.cra_module = THIS_MODULE,
+		},
+		.init = ocs_aes_aead_cra_init,
+		.exit = ocs_aead_cra_exit,
+		.ivsize = AES_BLOCK_SIZE,
+		.maxauthsize = AES_BLOCK_SIZE,
+		.setauthsize = kmb_ocs_aead_ccm_setauthsize,
+		.setkey = kmb_ocs_aes_aead_set_key,
+		.encrypt = kmb_ocs_aes_ccm_encrypt,
+		.decrypt = kmb_ocs_aes_ccm_decrypt,
+	},
+	{
+		.base = {
+			.cra_name = "gcm(sm4)",
+			.cra_driver_name = "gcm-sm4-keembay-ocs",
+			.cra_priority = KMB_OCS_PRIORITY,
+			.cra_flags = CRYPTO_ALG_ASYNC |
+				     CRYPTO_ALG_KERN_DRIVER_ONLY,
+			.cra_blocksize = 1,
+			.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+			.cra_alignmask = 0,
+			.cra_module = THIS_MODULE,
+		},
+		.init = ocs_sm4_aead_cra_init,
+		.exit = ocs_aead_cra_exit,
+		.ivsize = GCM_AES_IV_SIZE,
+		.maxauthsize = AES_BLOCK_SIZE,
+		.setauthsize = kmb_ocs_aead_gcm_setauthsize,
+		.setkey = kmb_ocs_sm4_aead_set_key,
+		.encrypt = kmb_ocs_sm4_gcm_encrypt,
+		.decrypt = kmb_ocs_sm4_gcm_decrypt,
+	},
+	{
+		.base = {
+			.cra_name = "ccm(sm4)",
+			.cra_driver_name = "ccm-sm4-keembay-ocs",
+			.cra_priority = KMB_OCS_PRIORITY,
+			.cra_flags = CRYPTO_ALG_ASYNC |
+				     CRYPTO_ALG_KERN_DRIVER_ONLY,
+			.cra_blocksize = 1,
+			.cra_ctxsize = sizeof(struct ocs_aes_tctx),
+			.cra_alignmask = 0,
+			.cra_module = THIS_MODULE,
+		},
+		.init = ocs_sm4_aead_cra_init,
+		.exit = ocs_aead_cra_exit,
+		.ivsize = AES_BLOCK_SIZE,
+		.maxauthsize = AES_BLOCK_SIZE,
+		.setauthsize = kmb_ocs_aead_ccm_setauthsize,
+		.setkey = kmb_ocs_sm4_aead_set_key,
+		.encrypt = kmb_ocs_sm4_ccm_encrypt,
+		.decrypt = kmb_ocs_sm4_ccm_decrypt,
+	}
+};
+
+static void unregister_aes_algs(struct ocs_aes_dev *aes_dev)
+{
+	crypto_unregister_aeads(algs_aead, ARRAY_SIZE(algs_aead));
+	crypto_unregister_skciphers(algs, ARRAY_SIZE(algs));
+}
+
+static int register_aes_algs(struct ocs_aes_dev *aes_dev)
+{
+	int ret;
+
+	/*
+	 * If any algorithm fails to register, all preceding algorithms that
+	 * were successfully registered will be automatically unregistered.
+	 */
+	ret = crypto_register_aeads(algs_aead, ARRAY_SIZE(algs_aead));
+	if (ret)
+		return ret;
+
+	ret = crypto_register_skciphers(algs, ARRAY_SIZE(algs));
+	if (ret)
+		crypto_unregister_aeads(algs_aead, ARRAY_SIZE(algs));
+
+	return ret;
+}
+
+/* Device tree driver match. */
+static const struct of_device_id kmb_ocs_aes_of_match[] = {
+	{
+		.compatible = "intel,keembay-ocs-aes",
+	},
+	{}
+};
+
+static int kmb_ocs_aes_remove(struct platform_device *pdev)
+{
+	struct ocs_aes_dev *aes_dev;
+
+	aes_dev = platform_get_drvdata(pdev);
+	if (!aes_dev)
+		return -ENODEV;
+
+	unregister_aes_algs(aes_dev);
+
+	spin_lock(&ocs_aes.lock);
+	list_del(&aes_dev->list);
+	spin_unlock(&ocs_aes.lock);
+
+	crypto_engine_exit(aes_dev->engine);
+
+	return 0;
+}
+
+static int kmb_ocs_aes_probe(struct platform_device *pdev)
+{
+	struct device *dev = &pdev->dev;
+	struct ocs_aes_dev *aes_dev;
+	struct resource *aes_mem;
+	int rc;
+
+	aes_dev = devm_kzalloc(dev, sizeof(*aes_dev), GFP_KERNEL);
+	if (!aes_dev)
+		return -ENOMEM;
+
+	aes_dev->dev = dev;
+
+	platform_set_drvdata(pdev, aes_dev);
+
+	rc = dma_set_mask_and_coherent(dev, DMA_BIT_MASK(32));
+	if (rc) {
+		dev_err(dev, "Failed to set 32 bit dma mask %d\n", rc);
+		return rc;
+	}
+
+	/* Get base register address. */
+	aes_mem = platform_get_resource(pdev, IORESOURCE_MEM, 0);
+	if (!aes_mem) {
+		dev_err(dev, "Could not retrieve io mem resource\n");
+		return -ENODEV;
+	}
+
+	aes_dev->base_reg = devm_ioremap_resource(&pdev->dev, aes_mem);
+	if (IS_ERR(aes_dev->base_reg)) {
+		dev_err(dev, "Failed to get base address\n");
+		return PTR_ERR(aes_dev->base_reg);
+	}
+
+	/* Get and request IRQ */
+	aes_dev->irq = platform_get_irq(pdev, 0);
+	if (aes_dev->irq < 0)
+		return aes_dev->irq;
+
+	rc = devm_request_threaded_irq(dev, aes_dev->irq, ocs_aes_irq_handler,
+				       NULL, 0, "keembay-ocs-aes", aes_dev);
+	if (rc < 0) {
+		dev_err(dev, "Could not request IRQ\n");
+		return rc;
+	}
+
+	INIT_LIST_HEAD(&aes_dev->list);
+	spin_lock(&ocs_aes.lock);
+	list_add_tail(&aes_dev->list, &ocs_aes.dev_list);
+	spin_unlock(&ocs_aes.lock);
+
+	init_completion(&aes_dev->irq_completion);
+
+	/* Initialize crypto engine */
+	aes_dev->engine = crypto_engine_alloc_init(dev, true);
+	if (!aes_dev->engine)
+		goto list_del;
+
+	rc = crypto_engine_start(aes_dev->engine);
+	if (rc) {
+		dev_err(dev, "Could not start crypto engine\n");
+		goto cleanup;
+	}
+
+	rc = register_aes_algs(aes_dev);
+	if (rc) {
+		dev_err(dev,
+			"Could not register OCS algorithms with Crypto API\n");
+		goto cleanup;
+	}
+
+	return 0;
+
+cleanup:
+	crypto_engine_exit(aes_dev->engine);
+list_del:
+	spin_lock(&ocs_aes.lock);
+	list_del(&aes_dev->list);
+	spin_unlock(&ocs_aes.lock);
+
+	return rc;
+}
+
+/* The OCS driver is a platform device. */
+static struct platform_driver kmb_ocs_aes_driver = {
+	.probe = kmb_ocs_aes_probe,
+	.remove = kmb_ocs_aes_remove,
+	.driver = {
+			.name = DRV_NAME,
+			.of_match_table = kmb_ocs_aes_of_match,
+		},
+};
+
+module_platform_driver(kmb_ocs_aes_driver);
+
+MODULE_DESCRIPTION("Intel Keem Bay Offload and Crypto Subsystem (OCS) AES/SM4 Driver");
+MODULE_LICENSE("GPL");
+
+MODULE_ALIAS_CRYPTO("cbc-aes-keembay-ocs");
+MODULE_ALIAS_CRYPTO("ctr-aes-keembay-ocs");
+MODULE_ALIAS_CRYPTO("gcm-aes-keembay-ocs");
+MODULE_ALIAS_CRYPTO("ccm-aes-keembay-ocs");
+
+MODULE_ALIAS_CRYPTO("cbc-sm4-keembay-ocs");
+MODULE_ALIAS_CRYPTO("ctr-sm4-keembay-ocs");
+MODULE_ALIAS_CRYPTO("gcm-sm4-keembay-ocs");
+MODULE_ALIAS_CRYPTO("ccm-sm4-keembay-ocs");
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB
+MODULE_ALIAS_CRYPTO("ecb-aes-keembay-ocs");
+MODULE_ALIAS_CRYPTO("ecb-sm4-keembay-ocs");
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_ECB */
+
+#ifdef CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS
+MODULE_ALIAS_CRYPTO("cts-aes-keembay-ocs");
+MODULE_ALIAS_CRYPTO("cts-sm4-keembay-ocs");
+#endif /* CONFIG_CRYPTO_DEV_KEEMBAY_OCS_AES_SM4_CTS */