From 79ad9f23304d27415c9e9d9021d01c2af8a1ed2f Mon Sep 17 00:00:00 2001
From: Aashvij Shenai <a-shenai@ti.com>
Date: Mon, 19 Aug 2024 15:17:40 +0530
Subject: [PATCH 1/6] sa2ul stuff

Signed-off-by: Aashvij Shenai <a-shenai@ti.com>
---
 crypto/testmgr.c       | 1437 +---------------------------------------
 drivers/crypto/sa2ul.c |  759 ++++++++++++++++++++-
 drivers/crypto/sa2ul.h |    7 +
 3 files changed, 753 insertions(+), 1450 deletions(-)
 
diff --git a/drivers/crypto/sa2ul.c b/drivers/crypto/sa2ul.c
index 461eca40e878..e0524f4e57a1 100644
--- a/drivers/crypto/sa2ul.c
+++ b/drivers/crypto/sa2ul.c
@@ -29,7 +29,7 @@
 #include <crypto/scatterwalk.h>
 #include <crypto/sha1.h>
 #include <crypto/sha2.h>
-
+#include <linux/delay.h>
 #include "sa2ul.h"
 
 /* Byte offset for key in encryption security context */
@@ -45,6 +45,7 @@
 
 #define SA_AUTH_SUBKEY_LEN	16
 #define SA_CMDL_PAYLOAD_LENGTH_MASK	0xFFFF
+#define SA_CMDL_LABEL_LENGTH_MASK	0xFF0000
 #define SA_CMDL_SOP_BYPASS_LEN_MASK	0xFF000000
 
 #define MODE_CONTROL_BYTES	27
@@ -90,7 +91,7 @@ struct sa_match_data {
 };
 
 static struct device *sa_k3_dev;
-
+static int count = 0;
 /**
  * struct sa_cmdl_cfg - Command label configuration descriptor
  * @aalg: authentication algorithm ID
@@ -191,12 +192,33 @@ struct sa_rx_data {
 	void *req;
 	struct device *ddev;
 	struct dma_async_tx_descriptor *tx_in;
+        struct dma_async_tx_descriptor *tx_in_extra;
+	struct sa_mapped_sg mapped_sg[2];
+	u8 enc;
+	u8 enc_iv_size;
+	u8 iv_idx;
+  	dma_addr_t dmaAddr;
+        dma_addr_t dmaAddrr;
+	u64 len;
+        struct completion done;
+};
+
+struct sa_tx_data {
+	void *req;
+	struct device *ddev;
+	struct dma_async_tx_descriptor *tx_in;
+        struct dma_async_tx_descriptor *tx_in_extra;
 	struct sa_mapped_sg mapped_sg[2];
 	u8 enc;
 	u8 enc_iv_size;
 	u8 iv_idx;
+  	dma_addr_t dmaAddr;
+        dma_addr_t dmaAddrr;
+	u64 len;
+        struct completion done;
 };
 
+
 /**
  * struct sa_req: SA request definition
  * @dev: device for the request
@@ -219,12 +241,12 @@ struct sa_rx_data {
  */
 struct sa_req {
 	struct device *dev;
-	u16 size;
+	u32 size;
 	u8 enc_offset;
 	u16 enc_size;
 	u8 *enc_iv;
 	u8 auth_offset;
-	u16 auth_size;
+	u32 auth_size;
 	u8 *auth_iv;
 	u32 type;
 	u32 *cmdl;
@@ -234,6 +256,7 @@ struct sa_req {
 	struct scatterlist *src;
 	struct scatterlist *dst;
 	dma_async_tx_callback callback;
+        dma_async_tx_callback callback_extra;
 	u16 mdata_size;
 };
 
@@ -362,6 +385,14 @@ static u8 mci_ecb_3des_dec_array[MODE_CONTROL_BYTES] = {
 	0x00, 0x00, 0x00,
 };
 
+static void sa_sha_dma_in_callback_tx_out(void *data)
+{
+	struct sa_tx_data *txd = data;
+
+        complete(&txd->done);
+        printk("%s: %d\n",__func__,__LINE__);
+}
+
 /*
  * Perform 16 byte or 128 bit swizzling
  * The SA2UL Expects the security context to
@@ -773,7 +804,7 @@ int sa_init_sc(struct sa_ctx_info *ctx, const struct sa_match_data *match_data,
 	sa_swiz_128(sc_buf, SA_CTX_MAX_SZ);
 
 	sa_set_swinfo(first_engine, ctx->sc_id, ctx->sc_phys, 1, 0,
-		      SA_SW_INFO_FLAG_EVICT, ad->hash_size, swinfo);
+		      SA_SW_INFO_FLAG_EVICT |SA_SW_INFO_FLAG_TEAR, ad->hash_size, swinfo);
 
 	sa_dump_sc(sc_buf, ctx->sc_phys);
 
@@ -1019,6 +1050,24 @@ static void sa_sync_from_device(struct sa_rx_data *rxd)
 	dma_sync_sgtable_for_cpu(rxd->ddev, sgt, DMA_FROM_DEVICE);
 }
 
+
+static void sa_free_sa_rx_data_sha512(struct sa_rx_data *rxd)
+{
+	int i;
+
+	for (i = 0; i < ARRAY_SIZE(rxd->mapped_sg); i++) {
+		struct sa_mapped_sg *mapped_sg = &rxd->mapped_sg[i];
+
+		if (mapped_sg->mapped) {
+			dma_unmap_sgtable(rxd->ddev, &mapped_sg->sgt,
+					  mapped_sg->dir, 0);
+			kfree(mapped_sg->split_sg);
+		}
+	}
+
+        complete(&rxd->done);
+}
+
 static void sa_free_sa_rx_data(struct sa_rx_data *rxd)
 {
 	int i;
@@ -1077,6 +1126,222 @@ sa_prepare_tx_desc(u32 *mdptr, u32 pslen, u32 *psdata, u32 epiblen, u32 *epib)
 		*out++ = *in++;
 }
 
+static int sa_run_sha512(struct sa_req *req, int flags, u8 *buf, u8 *bufr, size_t len, u32 skip)
+{
+
+	struct sa_rx_data *rxd;
+        struct sa_tx_data *txd;
+	gfp_t gfp_flags;
+	u32 cmdl[SA_MAX_CMDL_WORDS] = {0};
+	struct sa_crypto_data *pdata = dev_get_drvdata(sa_k3_dev);
+	struct device *ddev;
+	struct dma_chan *dma_rx;
+        struct dma_chan *dma_rx_extra = NULL;
+	int sg_nents, src_nents, dst_nents;
+	struct scatterlist *src, *dst;
+	size_t pl, ml, split_size;
+	struct sa_ctx_info *sa_ctx = req->enc ? &req->ctx->enc : &req->ctx->dec;
+	int ret;
+	struct dma_async_tx_descriptor *tx_out;
+	u32 *mdptr;
+	bool diff_dst;
+	enum dma_data_direction dir_src;
+	struct sa_mapped_sg *mapped_sg;
+
+	gfp_flags = req->base->flags & CRYPTO_TFM_REQ_MAY_SLEEP ?
+		GFP_KERNEL : GFP_ATOMIC;
+
+	rxd = kzalloc(sizeof(*rxd), gfp_flags);
+	if (!rxd)
+		return -ENOMEM;
+
+        
+	txd = kzalloc(sizeof(*txd), gfp_flags);
+	if (!txd)
+		return -ENOMEM;
+
+	if (req->src != req->dst) {
+		diff_dst = true;
+		dir_src = DMA_TO_DEVICE;
+	} else {
+		diff_dst = false;
+		dir_src = DMA_BIDIRECTIONAL;
+	}
+        dir_src = DMA_TO_DEVICE;
+	/*
+	 * SA2UL has an interesting feature where the receive DMA channel
+	 * is selected based on the data passed to the engine. Within the
+	 * transition range, there is also a space where it is impossible
+	 * to determine where the data will end up, and this should be
+	 * avoided. This will be handled by the SW fallback mechanism by
+	 * the individual algorithm implementations.
+	 */
+	if (req->size >= 256) {
+		dma_rx = pdata->dma_rx2;
+        }
+	else if (req->size >= SA_UNSAFE_DATA_SZ_MIN &&
+                 req->size <= SA_UNSAFE_DATA_SZ_MAX) {
+                dma_rx_extra = pdata->dma_rx1;
+                dma_rx = pdata->dma_rx2;
+        }
+	else {
+                dma_rx = pdata->dma_rx1;
+        }
+        
+	ddev = dmaengine_get_dma_device(pdata->dma_tx);
+	rxd->ddev = ddev;
+
+	memcpy(cmdl, sa_ctx->cmdl, sa_ctx->cmdl_size);
+
+	sa_update_cmdl(req, cmdl, &sa_ctx->cmdl_upd_info);
+
+	if (req->type != CRYPTO_ALG_TYPE_AHASH) {
+		if (req->enc)
+			req->type |=
+				(SA_REQ_SUBTYPE_ENC << SA_REQ_SUBTYPE_SHIFT);
+		else
+			req->type |=
+				(SA_REQ_SUBTYPE_DEC << SA_REQ_SUBTYPE_SHIFT);
+	}
+
+	cmdl[0] &= ~SA_CMDL_LABEL_LENGTH_MASK;
+	cmdl[0] |= ((16u << 16) & SA_CMDL_LABEL_LENGTH_MASK);
+
+	cmdl[1] |= (0x1 << 19) | (0x1 << 16);
+	cmdl[2] = 0;
+	cmdl[3] = (req->auth_size<<3);
+
+	/*
+	 * Map the packets, first we check if the data fits into a single
+	 * sg entry and use that if possible. If it does not fit, we check
+	 * if we need to do sg_split to align the scatterlist data on the
+	 * actual data size being processed by the crypto engine.
+	 */
+	
+	buf += skip;
+	rxd->len = len;
+	rxd->dmaAddr = dma_map_single(ddev, buf, len, dir_src);
+	if (rxd->dmaAddr == DMA_MAPPING_ERROR) {
+		printk("AHHHHHHHHHHHHHHHHHHHHH\n");
+		goto err_cleanup;
+	}
+
+        rxd->dmaAddrr = dma_map_single(ddev, bufr, len, DMA_FROM_DEVICE);
+	if (rxd->dmaAddr == DMA_MAPPING_ERROR) {
+		printk("AHHHHHHHHHHHHHHHHHHHHH\n");
+		goto err_cleanup;
+	}
+
+
+	rxd->tx_in = dmaengine_prep_slave_single(dma_rx, rxd->dmaAddrr, len,
+			DMA_DEV_TO_MEM,
+			DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
+	if (!rxd->tx_in) {
+		dev_err(pdata->dev, "IN prep_slave_sg() extra failed\n");
+		ret = -EINVAL;
+		goto err_cleanup;
+	}
+
+        if(dma_rx_extra) {
+                rxd->tx_in_extra = dmaengine_prep_slave_single(dma_rx_extra, rxd->dmaAddrr, len,
+                                                         DMA_DEV_TO_MEM,
+                                                         DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
+                if (!rxd->tx_in_extra) {
+                        dev_err(pdata->dev, "IN prep_slave_sg() xetra failed\n");
+                        ret = -EINVAL;
+                        goto err_cleanup;
+                }
+        }
+
+	rxd->req = (void *)req->base;
+	rxd->enc = req->enc;
+	rxd->iv_idx = req->ctx->iv_idx;
+	rxd->enc_iv_size = sa_ctx->cmdl_upd_info.enc_iv.size;
+	rxd->tx_in->callback = req->callback;
+	rxd->tx_in->callback_param = rxd;
+        if(dma_rx_extra){
+                rxd->tx_in_extra->callback = req->callback_extra;
+                rxd->tx_in_extra->callback_param = rxd;
+        }
+        
+	tx_out = dmaengine_prep_slave_single(pdata->dma_tx, rxd->dmaAddr,
+			len, DMA_MEM_TO_DEV,
+			DMA_PREP_INTERRUPT | DMA_CTRL_ACK);
+
+	if (!tx_out) {
+		dev_err(pdata->dev, "OUT prep_slave_sg() failed\n");
+		ret = -EINVAL;
+		goto err_cleanup;
+	}
+
+        txd->tx_in = tx_out;
+        txd->tx_in->callback = sa_sha_dma_in_callback_tx_out;
+	txd->tx_in->callback_param = txd;
+	/*
+	 * Prepare metadata for DMA engine. This essentially describes the
+	 * crypto algorithm to be used, data sizes, different keys etc.
+	 */
+	mdptr = (u32 *)dmaengine_desc_get_metadata_ptr(tx_out, &pl, &ml);
+
+	if(flags) {
+		sa_ctx->epib[1] |= (1U << 19U) & 0x80000; //FRAG BIT
+		sa_ctx->epib[1] |= (1U << 31U) & 0x80000000; //CMD_LBL_BYPASS
+
+		sa_ctx->epib[1] &= ~((1U << 16U) & 0x10000); //EVICT
+		sa_ctx->epib[1] &= ~((1U << 17U) & 0x20000); //TEARDOWN
+		cmdl[3] = 0;
+                printk("update length = %d\n", req->auth_size);
+	}
+	else {
+		sa_ctx->epib[1] &= ~((1U << 19U) & 0x80000); //FRAG BIT
+		sa_ctx->epib[1] &= ~((1U << 31U) & 0x80000000); //CMD_LBL_BYPASS
+
+		sa_ctx->epib[1] |= (1U << 16U) & 0x10000; //EVICT
+		sa_ctx->epib[1] |= (1U << 17U) & 0x20000; //TEARDOWN
+                printk("final len = %d\n", req->auth_size);
+		cmdl[3] = (req->auth_size<<3);
+	}
+	sa_prepare_tx_desc(mdptr, (sa_ctx->cmdl_size + (SA_PSDATA_CTX_WORDS *
+				   sizeof(u32))), cmdl, sizeof(sa_ctx->epib),
+			   sa_ctx->epib);
+
+	ml = sa_ctx->cmdl_size + (SA_PSDATA_CTX_WORDS * sizeof(u32));
+	dmaengine_desc_set_metadata_len(tx_out, req->mdata_size);
+
+	dmaengine_submit(tx_out);
+	dmaengine_submit(rxd->tx_in);
+        if(dma_rx_extra) {
+                dmaengine_submit(rxd->tx_in_extra);
+        }
+
+	dma_async_issue_pending(dma_rx);
+        if(dma_rx_extra) {
+                dma_async_issue_pending(dma_rx_extra);
+        }
+
+	dma_async_issue_pending(pdata->dma_tx);
+
+        init_completion(&txd->done);
+        init_completion(&rxd->done);
+        
+        wait_for_completion(&txd->done);
+        wait_for_completion(&rxd->done);
+
+        kfree(txd);
+
+        dma_unmap_single(rxd->ddev, rxd->dmaAddr, rxd->len, DMA_MEM_TO_DEV/* DMA_BIDIRECTIONAL */);
+        dma_unmap_single(rxd->ddev, rxd->dmaAddrr, rxd->len, DMA_DEV_TO_MEM/* DMA_BIDIRECTIONAL */);
+                
+	kfree(rxd);
+
+	return -EINPROGRESS;
+
+err_cleanup:
+	sa_free_sa_rx_data_sha512(rxd);
+
+	return ret;
+}
+
 static int sa_run(struct sa_req *req)
 {
 	struct sa_rx_data *rxd;
@@ -1350,6 +1615,164 @@ static int sa_decrypt(struct skcipher_request *req)
 	return sa_cipher_run(req, req->iv, 0);
 }
 
+static void sa_sha_dma_in_callback_sha512_final(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+	mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in, &pl, &ml);
+	result = (u32 *)req->result;
+
+        for (i = 0; i < (authsize / 4); i++)
+                result[i] = be32_to_cpu(mdptr[i + 4]);	
+
+        sa_free_sa_rx_data_sha512(rxd);
+        
+        ahash_request_complete(req, 0);
+}
+
+
+static void sa_sha_dma_in_callback_sha512_final_exception_rx2(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+	mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in, &pl, &ml);
+	result = (u32 *)req->result;
+
+        for (i = 0; i < (authsize / 4); i++)
+                result[i] = be32_to_cpu(mdptr[i + 4]);
+
+        dmaengine_terminate_async(rxd->tx_in_extra->chan);
+
+        sa_free_sa_rx_data_sha512(rxd);
+        
+        ahash_request_complete(req, 0);
+}
+
+static void sa_sha_dma_in_callback_sha512_final_exception_rx1(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+	mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in_extra, &pl, &ml);
+	result = (u32 *)req->result;
+
+        for (i = 0; i < (authsize / 4); i++)
+                result[i] = be32_to_cpu(mdptr[i + 4]);
+
+        dmaengine_terminate_async(rxd->tx_in->chan);
+
+        sa_free_sa_rx_data_sha512(rxd);
+        
+        ahash_request_complete(req, 0);
+}
+
+
+static void sa_sha_dma_in_callback_sha512_update(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	printk("%s\n", __func__);
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+        mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in, &pl, &ml);
+
+        sa_free_sa_rx_data_sha512(rxd);
+}
+
+static void sa_sha_dma_in_callback_sha512_update_exception_rx2(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+        mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in, &pl, &ml);
+
+        dmaengine_terminate_async(rxd->tx_in_extra->chan);
+        
+        sa_free_sa_rx_data_sha512(rxd);
+        
+}
+
+static void sa_sha_dma_in_callback_sha512_update_exception_rx1(void *data)
+{
+	struct sa_rx_data *rxd = data;
+	struct ahash_request *req;
+	struct crypto_ahash *tfm;
+	unsigned int authsize;
+	int i;
+	size_t ml, pl;
+	u32 *result;
+	__be32 *mdptr;
+
+	sa_sync_from_device(rxd);
+	req = container_of(rxd->req, struct ahash_request, base);
+	tfm = crypto_ahash_reqtfm(req);
+	authsize = crypto_ahash_digestsize(tfm);
+
+        mdptr = (__be32 *)dmaengine_desc_get_metadata_ptr(rxd->tx_in_extra, &pl, &ml);
+
+        dmaengine_terminate_async(rxd->tx_in->chan);
+        
+        sa_free_sa_rx_data_sha512(rxd);
+}
+
+
 static void sa_sha_dma_in_callback(void *data)
 {
 	struct sa_rx_data *rxd = data;
@@ -1407,7 +1830,6 @@ static int sa_sha_run(struct ahash_request *req)
 	size_t auth_len;
 
 	auth_len = req->nbytes;
-
 	if (!auth_len)
 		return zero_message_process(req);
 
@@ -1533,9 +1955,257 @@ static int sa_sha_cra_init_alg(struct crypto_tfm *tfm, const char *alg_base)
 
 static int sa_sha_digest(struct ahash_request *req)
 {
+        struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
+
+        struct sa_tfm_ctx *ctx = crypto_ahash_ctx(tfm);
+	struct algo_data ad = { 0 };
+
+	printk("count = %d\n",count++);
+
+        ad.aalg_id = SA_AALG_ID_SHA2_512;
+        ad.hash_size = SHA512_DIGEST_SIZE;
+        ad.auth_ctrl = SA_AUTH_SW_CTRL_SHA512;
+        sa_sha_setup(ctx, &ad);
+	
 	return sa_sha_run(req);
 }
 
+static int globule_count = 0;
+
+static int sa_sha_init_sha512(struct ahash_request *req)
+{
+	struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
+	struct sa_sha_req_ctx *rctx = ahash_request_ctx(req);
+	struct sa_tfm_ctx *ctx = crypto_ahash_ctx(tfm);
+	struct algo_data ad = { 0 };
+
+	printk("count = %d\n",count++);
+	size_t auth_len = req->nbytes;
+
+	dev_dbg(sa_k3_dev, "init: digest size: %u, rctx=%p\n",
+		crypto_ahash_digestsize(tfm), rctx);
+
+	/* if(globule_count++){ */
+        ad.aalg_id = SA_AALG_ID_SHA2_512;
+        ad.hash_size = SHA512_DIGEST_SIZE;
+        ad.auth_ctrl = SA_AUTH_SW_CTRL_SHA512;
+        sa_sha_setup(ctx, &ad);
+	/* } */
+
+	ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+	rctx->fallback_req.base.flags =
+		req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+
+	crypto_ahash_init(&rctx->fallback_req);
+
+	ctx->first_pkt = true;
+	ctx->total_len = 0;
+	ctx->go_to_fallback = 0;
+
+	count = 0;
+
+	return 0;
+}
+
+void inspect_sg_list(struct scatterlist *sg_list, int sg_len)
+{
+    int i;
+    struct scatterlist *sg;
+
+    pr_info("Inspecting scatter-gather list:\n");
+
+    for_each_sg(sg_list, sg, sg_len, i) {
+        void *vaddr;
+        char data[8] = {0}; // Buffer to hold the first 8 bytes
+
+        // Check if the scatterlist entry is valid
+        if (!sg) {
+            pr_warn("Entry %d is NULL\n", i);
+            continue;
+        }
+
+        if (!sg_page(sg)) {
+            pr_warn("Entry %d has a NULL page pointer\n", i);
+            continue;
+        }
+
+        // Map the page
+        vaddr = kmap_atomic(sg_page(sg)) + sg->offset;
+        if (!vaddr) {
+            pr_warn("Entry %d: Unable to map page\n", i);
+            continue;
+        }
+
+        // Copy the first 8 bytes safely
+        memcpy(data, vaddr, min_t(size_t, 8, sg->length));
+        kunmap_atomic(vaddr);
+
+        // Print the details
+        pr_info("Entry %d:\n", i);
+        pr_info("  Page: %p\n", sg_page(sg));
+        pr_info("  Offset: %u\n", sg->offset);
+        pr_info("  Length: %u\n", sg->length);
+        pr_info("  DMA address: %pad\n", &sg->dma_address);
+        pr_info("  First 8 bytes: %*ph\n", 8, data); // Print the first 8 bytes in hex
+    }
+}
+
+static int sa_sha_update_sha512(struct ahash_request *req)
+{
+	struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
+	struct sa_sha_req_ctx *rctx = ahash_request_ctx(req);
+	struct sa_tfm_ctx *ctx = crypto_ahash_ctx(tfm);
+	struct sa_req sa_req = { 0 };
+	size_t auth_len = req->nbytes;
+
+	ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+	rctx->fallback_req.base.flags =
+		req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+	rctx->fallback_req.nbytes = req->nbytes;
+	rctx->fallback_req.src = req->src;
+	count = 0;
+
+	mdelay(1);
+
+	if(ctx->first_pkt) 
+	  {
+	    if (!auth_len || auth_len < 128 || ctx->go_to_fallback ||
+	    (auth_len >= SA_UNSAFE_DATA_SZ_MIN &&
+	     auth_len <= SA_UNSAFE_DATA_SZ_MAX)) {
+			ctx->go_to_fallback = 1;
+
+		ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+		rctx->fallback_req.base.flags =
+			req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+		rctx->fallback_req.nbytes = req->nbytes;
+		rctx->fallback_req.src = req->src;
+
+                return crypto_ahash_update(&rctx->fallback_req);
+	    }
+	  }
+
+	if(!ctx->go_to_fallback) {
+
+		if(ctx->first_pkt) {
+			ctx->first_pkt = false;
+
+			ctx->prev_len = auth_len;
+			ctx->prev_nents = sg_nents_for_len(req->src, auth_len);
+			sg_copy_to_buffer(req->src, ctx->prev_nents, ctx->buf, auth_len);
+
+			return 0;
+		}
+		else {
+			ctx->total_len += ctx->prev_len;
+			sa_req.size = ctx->prev_len;
+			sa_req.auth_size = ctx->prev_len;
+
+			sa_req.enc = true;
+			sa_req.type = CRYPTO_ALG_TYPE_AHASH;
+			sa_req.callback = sa_sha_dma_in_callback_sha512_update;
+                        if(sa_req.size >= SA_UNSAFE_DATA_SZ_MIN &&
+                           sa_req.size <= SA_UNSAFE_DATA_SZ_MAX) {
+                                sa_req.callback = sa_sha_dma_in_callback_sha512_update_exception_rx2;
+                                sa_req.callback_extra = sa_sha_dma_in_callback_sha512_update_exception_rx1;
+                        }
+			sa_req.mdata_size = 28+8;
+			sa_req.ctx = ctx;
+			sa_req.base = &req->base;
+			sa_req.dst = sa_req.src;
+			sa_run_sha512(&sa_req, 1, ctx->buf, ctx->bufr, ctx->prev_len, 0);
+
+			ctx->prev_len = auth_len;
+			ctx->prev_nents = sg_nents_for_len(req->src, auth_len);
+			sg_copy_to_buffer(req->src, ctx->prev_nents, ctx->buf, auth_len);
+
+			return 0;
+		}
+	}
+	return -EINVAL;
+}
+
+static int sa_sha_final_sha512(struct ahash_request *req)
+{
+	struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
+	struct sa_sha_req_ctx *rctx = ahash_request_ctx(req);
+	struct sa_tfm_ctx *ctx = crypto_ahash_ctx(tfm);
+	struct sa_req sa_req = { 0 };
+
+	size_t auth_len = req->nbytes;
+
+	mdelay(1);
+
+	if(ctx->first_pkt) {
+                if (!auth_len || ctx->go_to_fallback ||
+                    (auth_len >= SA_UNSAFE_DATA_SZ_MIN &&
+                     auth_len <= SA_UNSAFE_DATA_SZ_MAX)) {
+                        ctx->go_to_fallback = 0;
+                        
+                        ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+                        rctx->fallback_req.base.flags =
+                                req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+                        rctx->fallback_req.result = req->result;
+                        
+                        return crypto_ahash_final(&rctx->fallback_req);
+                }
+        }
+	ctx->total_len += ctx->prev_len;
+	sa_req.size = ctx->prev_len;
+	sa_req.auth_size = ctx->total_len;
+	sa_req.dst = sa_req.src;
+
+	count = 1;
+
+	sa_req.enc = true;
+	sa_req.type = CRYPTO_ALG_TYPE_AHASH;
+	sa_req.callback = sa_sha_dma_in_callback_sha512_final;
+        if(sa_req.size >= SA_UNSAFE_DATA_SZ_MIN &&
+           sa_req.size <= SA_UNSAFE_DATA_SZ_MAX) {
+                sa_req.callback = sa_sha_dma_in_callback_sha512_final_exception_rx2;
+                sa_req.callback_extra = sa_sha_dma_in_callback_sha512_final_exception_rx1;
+        }
+	sa_req.mdata_size = 28+8;
+	sa_req.ctx = ctx;
+	sa_req.base = &req->base;
+
+	sa_run_sha512(&sa_req, 0, ctx->buf,ctx->bufr, auth_len,0);
+	
+	ctx->go_to_fallback = 0;
+
+	return -EINPROGRESS;
+}
+
+static int sa_sha_finup_sha512(struct ahash_request *req)
+{
+	struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
+	struct sa_sha_req_ctx *rctx = ahash_request_ctx(req);
+	struct sa_tfm_ctx *ctx = crypto_ahash_ctx(tfm);
+
+	ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+	rctx->fallback_req.base.flags =
+		req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+	size_t auth_len = req->nbytes;
+
+        if(ctx->first_pkt) {
+                if (!auth_len ||
+                    (auth_len >= SA_UNSAFE_DATA_SZ_MIN &&
+                     auth_len <= SA_UNSAFE_DATA_SZ_MAX)) {
+
+                        rctx->fallback_req.nbytes = req->nbytes;
+                        rctx->fallback_req.src = req->src;
+                        rctx->fallback_req.result = req->result;
+                        ahash_request_set_tfm(&rctx->fallback_req, ctx->fallback.ahash);
+                        rctx->fallback_req.base.flags =
+                                req->base.flags & CRYPTO_TFM_REQ_MAY_SLEEP;
+                
+                        return crypto_ahash_finup(&rctx->fallback_req);
+
+                }
+        }
+	crypto_ahash_update(req);
+	return crypto_ahash_final(req);
+}
+
 static int sa_sha_init(struct ahash_request *req)
 {
 	struct crypto_ahash *tfm = crypto_ahash_reqtfm(req);
@@ -1672,6 +2342,56 @@ static int sa_sha512_cra_init(struct crypto_tfm *tfm)
 	return 0;
 }
 
+static int sa_sha512_cra_init_sha512(struct crypto_tfm *tfm)
+{
+	struct algo_data ad = { 0 };
+	struct sa_tfm_ctx *ctx = crypto_tfm_ctx(tfm);
+
+	sa_sha_cra_init_alg(tfm, "sha512");
+
+	ad.aalg_id = SA_AALG_ID_SHA2_512;
+	ad.hash_size = SHA512_DIGEST_SIZE;
+	ad.auth_ctrl = SA_AUTH_SW_CTRL_SHA512;
+
+	sa_sha_setup(ctx, &ad);
+
+	ctx->buf = kmalloc(SA_MAX_DATA_SZ, GFP_KERNEL);
+	if(!ctx->buf) {
+		pr_err("Failed to allocate storage buffer\n");
+		return -ENOMEM;
+	}
+
+        ctx->bufr = kmalloc(SA_MAX_DATA_SZ, GFP_KERNEL);
+	if(!ctx->bufr) {
+		pr_err("Failed to allocate storage buffer\n");
+		return -ENOMEM;
+	}
+	globule_count = 0;
+
+	return 0;
+}
+
+static void sa_sha_cra_exit_sha512(struct crypto_tfm *tfm)
+{
+	struct sa_tfm_ctx *ctx = crypto_tfm_ctx(tfm);
+	struct sa_crypto_data *data = dev_get_drvdata(sa_k3_dev);
+
+	dev_dbg(sa_k3_dev, "%s(0x%p) sc-ids(0x%x(0x%pad), 0x%x(0x%pad))\n",
+		__func__, tfm, ctx->enc.sc_id, &ctx->enc.sc_phys,
+		ctx->dec.sc_id, &ctx->dec.sc_phys);
+
+	if (crypto_tfm_alg_type(tfm) == CRYPTO_ALG_TYPE_AHASH)
+		sa_free_ctx_info(&ctx->enc, data);
+
+	if(ctx->buf)
+		kfree(ctx->buf);
+
+        if(ctx->bufr)
+		kfree(ctx->bufr);
+	crypto_free_shash(ctx->shash);
+	crypto_free_ahash(ctx->fallback.ahash);
+}
+
 static void sa_sha_cra_exit(struct crypto_tfm *tfm)
 {
 	struct sa_tfm_ctx *ctx = crypto_tfm_ctx(tfm);
@@ -1807,8 +2527,8 @@ static int sa_aead_setkey(struct crypto_aead *authenc,
 	struct sa_cmdl_cfg cfg;
 	int key_idx;
 
-	if (crypto_authenc_extractkeys(&keys, key, keylen) != 0)
-		return -EINVAL;
+	/* if (crypto_authenc_extractkeys(&keys, key, keylen) != 0) */
+		/* return -EINVAL; */
 
 	/* Convert the key size (16/24/32) to the key size index (0/1/2) */
 	key_idx = (keys.enckeylen >> 3) - 2;
@@ -2131,16 +2851,16 @@ static struct sa_alg_tmpl sa_algs[] = {
 				.cra_blocksize	= SHA512_BLOCK_SIZE,
 				.cra_ctxsize	= sizeof(struct sa_tfm_ctx),
 				.cra_module	= THIS_MODULE,
-				.cra_init	= sa_sha512_cra_init,
-				.cra_exit	= sa_sha_cra_exit,
+				.cra_init	= sa_sha512_cra_init_sha512,
+				.cra_exit	= sa_sha_cra_exit_sha512,
 			},
 			.halg.digestsize	= SHA512_DIGEST_SIZE,
 			.halg.statesize		= sizeof(struct sa_sha_req_ctx) +
 						  sizeof(struct sha512_state),
-			.init			= sa_sha_init,
-			.update			= sa_sha_update,
-			.final			= sa_sha_final,
-			.finup			= sa_sha_finup,
+			.init			= sa_sha_init_sha512,
+			.update			= sa_sha_update_sha512,
+			.final			= sa_sha_final_sha512,
+			.finup			= sa_sha_finup_sha512,
 			.digest			= sa_sha_digest,
 			.export			= sa_sha_export,
 			.import			= sa_sha_import,
@@ -2374,13 +3094,14 @@ static struct sa_match_data am654_match_data = {
 static struct sa_match_data am64_match_data = {
 	.priv = 0,
 	.priv_id = 0,
-	.supported_algos = BIT(SA_ALG_CBC_AES) |
-			   BIT(SA_ALG_EBC_AES) |
-			   BIT(SA_ALG_SHA256) |
-			   BIT(SA_ALG_SHA512) |
-			   BIT(SA_ALG_AUTHENC_SHA256_AES),
+	.supported_algos = /* BIT(SA_ALG_CBC_AES) | */
+			   /* BIT(SA_ALG_EBC_AES) |
+			    * BIT(SA_ALG_SHA256) | */
+        BIT(SA_ALG_SHA512),
+			   /* BIT(SA_ALG_AUTHENC_SHA256_AES), */
 };
 
+
 static const struct of_device_id of_match[] = {
 	{ .compatible = "ti,j721e-sa2ul", .data = &am654_match_data, },
 	{ .compatible = "ti,am654-sa2ul", .data = &am654_match_data, },
diff --git a/drivers/crypto/sa2ul.h b/drivers/crypto/sa2ul.h
index 12c17a68d350..bc674c316d74 100644
--- a/drivers/crypto/sa2ul.h
+++ b/drivers/crypto/sa2ul.h
@@ -319,6 +319,13 @@ struct sa_tfm_ctx {
 		struct crypto_ahash		*ahash;
 		struct crypto_aead		*aead;
 	} fallback;
+        bool first_pkt;
+	int prev_len;
+	int prev_nents;
+	int total_len;
+	u8 *buf;
+        u8 *bufr;
+	int go_to_fallback;
 };
 
 /**
-- 
2.34.1

