libwd/0123-uadk-move-send-and-recv-code-into-wd_util.c.patch

From d72b1bd6a9fae5dadfbaacf64e51353e78c785c5 Mon Sep 17 00:00:00 2001
From: Zhiqi Song <songzhiqi1@huawei.com>
Date: Tue, 7 Jun 2022 14:26:49 +0800
Subject: [PATCH 135/183] uadk: move send and recv code into wd_util.c

In the synchronization scenario, the code for sending and
receiving BD code is same in wd_<alg>.c, so move them
into wd_util.c from wd_<alg>.c.

Signed-off-by: Weili Qian <qianweili@huawei.com>
Signed-off-by: Zhiqi Song <songzhiqi1@huawei.com>
---
 drv/hisi_comp.c             |  6 ++--
 drv/hisi_hpre.c             | 18 +++++++----
 drv/hisi_sec.c              | 44 +++++++++++++++++----------
 include/drv/wd_aead_drv.h   |  4 +--
 include/drv/wd_cipher_drv.h |  4 +--
 include/drv/wd_comp_drv.h   |  4 +--
 include/drv/wd_dh_drv.h     |  4 +--
 include/drv/wd_digest_drv.h |  4 +--
 include/drv/wd_ecc_drv.h    |  4 +--
 include/drv/wd_rsa_drv.h    |  4 +--
 include/wd_util.h           | 18 +++++++++++
 wd_aead.c                   | 35 ++++------------------
 wd_cipher.c                 | 35 ++++------------------
 wd_comp.c                   | 37 ++++-------------------
 wd_dh.c                     | 58 ++++++++---------------------------
 wd_digest.c                 | 45 ++++++++--------------------
 wd_ecc.c                    | 60 ++++++++-----------------------------
 wd_rsa.c                    | 60 ++++++++-----------------------------
 wd_util.c                   | 47 +++++++++++++++++++++++++++++
 19 files changed, 192 insertions(+), 299 deletions(-)

diff --git a/drv/hisi_comp.c b/drv/hisi_comp.c
index 2a2367b..2c0fc41 100644
--- a/drv/hisi_comp.c
+++ b/drv/hisi_comp.c
@@ -865,9 +865,10 @@ static int fill_zip_comp_sqe(struct hisi_qp *qp, struct wd_comp_msg *msg,
 	return 0;
 }
 
-static int hisi_zip_comp_send(handle_t ctx, struct wd_comp_msg *msg)
+static int hisi_zip_comp_send(handle_t ctx, void *comp_msg)
 {
 	struct hisi_qp *qp = wd_ctx_get_priv(ctx);
+	struct wd_comp_msg *msg = comp_msg;
 	handle_t h_qp = (handle_t)qp;
 	struct hisi_zip_sqe sqe = {0};
 	__u16 count = 0;
@@ -1020,9 +1021,10 @@ static int parse_zip_sqe(struct hisi_qp *qp, struct hisi_zip_sqe *sqe,
 	return 0;
 }
 
-static int hisi_zip_comp_recv(handle_t ctx, struct wd_comp_msg *recv_msg)
+static int hisi_zip_comp_recv(handle_t ctx, void *comp_msg)
 {
 	struct hisi_qp *qp = wd_ctx_get_priv(ctx);
+	struct wd_comp_msg *recv_msg = comp_msg;
 	handle_t h_qp = (handle_t)qp;
 	struct hisi_zip_sqe sqe = {0};
 	__u16 count = 0;
diff --git a/drv/hisi_hpre.c b/drv/hisi_hpre.c
index 4e62c0f..7fd4189 100644
--- a/drv/hisi_hpre.c
+++ b/drv/hisi_hpre.c
@@ -499,9 +499,10 @@ static void hpre_exit(void *priv)
 	}
 }
 
-static int rsa_send(handle_t ctx, struct wd_rsa_msg *msg)
+static int rsa_send(handle_t ctx, void *rsa_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_rsa_msg *msg = rsa_msg;
 	struct hisi_hpre_sqe hw_msg;
 	__u16 send_cnt = 0;
 	int ret;
@@ -535,10 +536,11 @@ static int rsa_send(handle_t ctx, struct wd_rsa_msg *msg)
 	return hisi_qm_send(h_qp, &hw_msg, 1, &send_cnt);
 }
 
-static int rsa_recv(handle_t ctx, struct wd_rsa_msg *msg)
+static int rsa_recv(handle_t ctx, void *rsa_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
 	struct hisi_hpre_sqe hw_msg = {0};
+	struct wd_rsa_msg *msg = rsa_msg;
 	__u16 recv_cnt = 0;
 	int ret;
 
@@ -638,9 +640,10 @@ static int dh_out_transfer(struct wd_dh_msg *msg,
 	return WD_SUCCESS;
 }
 
-static int dh_send(handle_t ctx, struct wd_dh_msg *msg)
+static int dh_send(handle_t ctx, void *dh_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_dh_msg *msg = dh_msg;
 	struct wd_dh_req *req = &msg->req;
 	struct hisi_hpre_sqe hw_msg;
 	__u16 send_cnt = 0;
@@ -682,9 +685,10 @@ static int dh_send(handle_t ctx, struct wd_dh_msg *msg)
 	return hisi_qm_send(h_qp, &hw_msg, 1, &send_cnt);
 }
 
-static int dh_recv(handle_t ctx, struct wd_dh_msg *msg)
+static int dh_recv(handle_t ctx, void *dh_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_dh_msg *msg = dh_msg;
 	struct hisi_hpre_sqe hw_msg = {0};
 	__u16 recv_cnt = 0;
 	int ret;
@@ -1774,9 +1778,10 @@ free_dst:
 	return ret;
 }
 
-static int ecc_send(handle_t ctx, struct wd_ecc_msg *msg)
+static int ecc_send(handle_t ctx, void *ecc_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_ecc_msg *msg = ecc_msg;
 
 	hisi_set_msg_id(h_qp, &msg->tag);
 	if (msg->req.op_type == WD_SM2_ENCRYPT)
@@ -2336,9 +2341,10 @@ fail:
 	return ret;
 }
 
-static int ecc_recv(handle_t ctx, struct wd_ecc_msg *msg)
+static int ecc_recv(handle_t ctx, void *ecc_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_ecc_msg *msg = ecc_msg;
 	struct hisi_hpre_sqe hw_msg;
 	__u16 recv_cnt = 0;
 	int ret;
diff --git a/drv/hisi_sec.c b/drv/hisi_sec.c
index b17ce0f..7f12549 100644
--- a/drv/hisi_sec.c
+++ b/drv/hisi_sec.c
@@ -906,9 +906,10 @@ static int fill_cipher_bd2(struct wd_cipher_msg *msg, struct hisi_sec_sqe *sqe)
 	return 0;
 }
 
-int hisi_sec_cipher_send(handle_t ctx, struct wd_cipher_msg *msg)
+int hisi_sec_cipher_send(handle_t ctx, void *cipher_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_cipher_msg *msg = cipher_msg;
 	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	int ret;
@@ -950,10 +951,11 @@ int hisi_sec_cipher_send(handle_t ctx, struct wd_cipher_msg *msg)
 	return 0;
 }
 
-int hisi_sec_cipher_recv(handle_t ctx, struct wd_cipher_msg *recv_msg)
+int hisi_sec_cipher_recv(handle_t ctx, void *cipher_msg)
 {
-	struct hisi_sec_sqe sqe;
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_cipher_msg *recv_msg = cipher_msg;
+	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	int ret;
 
@@ -1112,9 +1114,10 @@ static int fill_cipher_bd3(struct wd_cipher_msg *msg, struct hisi_sec_sqe3 *sqe)
 	return 0;
 }
 
-int hisi_sec_cipher_send_v3(handle_t ctx, struct wd_cipher_msg *msg)
+int hisi_sec_cipher_send_v3(handle_t ctx, void *cipher_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_cipher_msg *msg = cipher_msg;
 	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	int ret;
@@ -1183,10 +1186,11 @@ static void parse_cipher_bd3(struct hisi_sec_sqe3 *sqe,
 	recv_msg->out = rmsg->out;
 }
 
-int hisi_sec_cipher_recv_v3(handle_t ctx, struct wd_cipher_msg *recv_msg)
+int hisi_sec_cipher_recv_v3(handle_t ctx, void *cipher_msg)
 {
-	struct hisi_sec_sqe3 sqe;
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_cipher_msg *recv_msg = cipher_msg;
+	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	int ret;
 
@@ -1332,9 +1336,10 @@ static int digest_len_check(struct wd_digest_msg *msg,  enum sec_bd_type type)
 	return 0;
 }
 
-int hisi_sec_digest_send(handle_t ctx, struct wd_digest_msg *msg)
+int hisi_sec_digest_send(handle_t ctx, void *digest_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_digest_msg *msg = digest_msg;
 	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	__u8 scene;
@@ -1396,9 +1401,10 @@ put_sgl:
 	return ret;
 }
 
-int hisi_sec_digest_recv(handle_t ctx, struct wd_digest_msg *recv_msg)
+int hisi_sec_digest_recv(handle_t ctx, void *digest_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_digest_msg *recv_msg = digest_msg;
 	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	int ret;
@@ -1486,9 +1492,10 @@ static void qm_fill_digest_long_bd3(struct wd_digest_msg *msg,
 	}
 }
 
-int hisi_sec_digest_send_v3(handle_t ctx, struct wd_digest_msg *msg)
+int hisi_sec_digest_send_v3(handle_t ctx, void *digest_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_digest_msg *msg = digest_msg;
 	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	__u16 scene;
@@ -1572,9 +1579,10 @@ static void parse_digest_bd3(struct hisi_sec_sqe3 *sqe,
 	recv_msg->alg_type = WD_DIGEST;
 }
 
-int hisi_sec_digest_recv_v3(handle_t ctx, struct wd_digest_msg *recv_msg)
+int hisi_sec_digest_recv_v3(handle_t ctx, void *digest_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_digest_msg *recv_msg = digest_msg;
 	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	int ret;
@@ -1838,9 +1846,10 @@ static int fill_aead_bd2(struct wd_aead_msg *msg, struct hisi_sec_sqe *sqe)
 	return 0;
 }
 
-int hisi_sec_aead_send(handle_t ctx, struct wd_aead_msg *msg)
+int hisi_sec_aead_send(handle_t ctx, void *aead_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_aead_msg *msg = aead_msg;
 	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	int ret;
@@ -1922,10 +1931,11 @@ static void parse_aead_bd2(struct hisi_sec_sqe *sqe,
 			      sqe->type2.cipher_src_offset;
 }
 
-int hisi_sec_aead_recv(handle_t ctx, struct wd_aead_msg *recv_msg)
+int hisi_sec_aead_recv(handle_t ctx, void *aead_msg)
 {
-	struct hisi_sec_sqe sqe;
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_aead_msg *recv_msg = aead_msg;
+	struct hisi_sec_sqe sqe;
 	__u16 count = 0;
 	int ret;
 
@@ -2103,9 +2113,10 @@ static int fill_aead_bd3(struct wd_aead_msg *msg, struct hisi_sec_sqe3 *sqe)
 	return 0;
 }
 
-int hisi_sec_aead_send_v3(handle_t ctx, struct wd_aead_msg *msg)
+int hisi_sec_aead_send_v3(handle_t ctx, void *aead_msg)
 {
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_aead_msg *msg = aead_msg;
 	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	int ret;
@@ -2181,10 +2192,11 @@ static void parse_aead_bd3(struct hisi_sec_sqe3 *sqe,
 			      sqe->cipher_src_offset;
 }
 
-int hisi_sec_aead_recv_v3(handle_t ctx, struct wd_aead_msg *recv_msg)
+int hisi_sec_aead_recv_v3(handle_t ctx, void *aead_msg)
 {
-	struct hisi_sec_sqe3 sqe;
 	handle_t h_qp = (handle_t)wd_ctx_get_priv(ctx);
+	struct wd_aead_msg *recv_msg = aead_msg;
+	struct hisi_sec_sqe3 sqe;
 	__u16 count = 0;
 	int ret;
 
diff --git a/include/drv/wd_aead_drv.h b/include/drv/wd_aead_drv.h
index cc28ed5..cca28ba 100644
--- a/include/drv/wd_aead_drv.h
+++ b/include/drv/wd_aead_drv.h
@@ -70,8 +70,8 @@ struct wd_aead_driver {
 	__u32	drv_ctx_size;
 	int	(*init)(struct wd_ctx_config_internal *config, void *priv);
 	void	(*exit)(void *priv);
-	int	(*aead_send)(handle_t ctx, struct wd_aead_msg *msg);
-	int	(*aead_recv)(handle_t ctx, struct wd_aead_msg *msg);
+	int	(*aead_send)(handle_t ctx, void *aead_msg);
+	int	(*aead_recv)(handle_t ctx, void *aead_msg);
 };
 
 void wd_aead_set_driver(struct wd_aead_driver *drv);
diff --git a/include/drv/wd_cipher_drv.h b/include/drv/wd_cipher_drv.h
index 0ec758d..84c9844 100644
--- a/include/drv/wd_cipher_drv.h
+++ b/include/drv/wd_cipher_drv.h
@@ -56,8 +56,8 @@ struct wd_cipher_driver {
 	__u32	drv_ctx_size;
 	int	(*init)(struct wd_ctx_config_internal *config, void *priv);
 	void	(*exit)(void *priv);
-	int	(*cipher_send)(handle_t ctx, struct wd_cipher_msg *msg);
-	int	(*cipher_recv)(handle_t ctx, struct wd_cipher_msg *msg);
+	int	(*cipher_send)(handle_t ctx, void *cipher_msg);
+	int	(*cipher_recv)(handle_t ctx, void *cipher_msg);
 };
 
 void wd_cipher_set_driver(struct wd_cipher_driver *drv);
diff --git a/include/drv/wd_comp_drv.h b/include/drv/wd_comp_drv.h
index e3d2269..a84f895 100644
--- a/include/drv/wd_comp_drv.h
+++ b/include/drv/wd_comp_drv.h
@@ -61,8 +61,8 @@ struct wd_comp_driver {
 	__u32 drv_ctx_size;
 	int (*init)(struct wd_ctx_config_internal *config, void *priv);
 	void (*exit)(void *priv);
-	int (*comp_send)(handle_t ctx, struct wd_comp_msg *msg);
-	int (*comp_recv)(handle_t ctx, struct wd_comp_msg *msg);
+	int (*comp_send)(handle_t ctx, void *comp_msg);
+	int (*comp_recv)(handle_t ctx, void *comp_msg);
 };
 
 void wd_comp_set_driver(struct wd_comp_driver *drv);
diff --git a/include/drv/wd_dh_drv.h b/include/drv/wd_dh_drv.h
index 192e7d8..f8e3065 100644
--- a/include/drv/wd_dh_drv.h
+++ b/include/drv/wd_dh_drv.h
@@ -30,8 +30,8 @@ struct wd_dh_driver {
 	int (*init)(struct wd_ctx_config_internal *config, void *priv,
 		    const char *alg_name);
 	void (*exit)(void *priv);
-	int (*send)(handle_t sess, struct wd_dh_msg *msg);
-	int (*recv)(handle_t sess, struct wd_dh_msg *msg);
+	int (*send)(handle_t sess, void *dh_msg);
+	int (*recv)(handle_t sess, void *dh_msg);
 };
 
 void wd_dh_set_driver(struct wd_dh_driver *drv);
diff --git a/include/drv/wd_digest_drv.h b/include/drv/wd_digest_drv.h
index b426ab2..10ce6e1 100644
--- a/include/drv/wd_digest_drv.h
+++ b/include/drv/wd_digest_drv.h
@@ -58,8 +58,8 @@ struct wd_digest_driver {
 	__u32	drv_ctx_size;
 	int	(*init)(struct wd_ctx_config_internal *config, void *priv);
 	void	(*exit)(void *priv);
-	int	(*digest_send)(handle_t ctx, struct wd_digest_msg *msg);
-	int	(*digest_recv)(handle_t ctx, struct wd_digest_msg *msg);
+	int	(*digest_send)(handle_t ctx, void *digest_msg);
+	int	(*digest_recv)(handle_t ctx, void *digest_msg);
 };
 
 void wd_digest_set_driver(struct wd_digest_driver *drv);
diff --git a/include/drv/wd_ecc_drv.h b/include/drv/wd_ecc_drv.h
index ca21759..ef98606 100644
--- a/include/drv/wd_ecc_drv.h
+++ b/include/drv/wd_ecc_drv.h
@@ -184,8 +184,8 @@ struct wd_ecc_driver {
 	int (*init)(struct wd_ctx_config_internal *config, void *priv,
 		    const char *alg_name);
 	void (*exit)(void *priv);
-	int (*send)(handle_t sess, struct wd_ecc_msg *msg);
-	int (*recv)(handle_t sess, struct wd_ecc_msg *msg);
+	int (*send)(handle_t sess, void *ecc_msg);
+	int (*recv)(handle_t sess, void *ecc_msg);
 };
 
 void wd_ecc_set_driver(struct wd_ecc_driver *drv);
diff --git a/include/drv/wd_rsa_drv.h b/include/drv/wd_rsa_drv.h
index bde6bbd..28c3337 100644
--- a/include/drv/wd_rsa_drv.h
+++ b/include/drv/wd_rsa_drv.h
@@ -53,8 +53,8 @@ struct wd_rsa_driver {
 	int (*init)(struct wd_ctx_config_internal *config, void *priv,
 		    const char *alg_name);
 	void (*exit)(void *priv);
-	int (*send)(handle_t sess, struct wd_rsa_msg *msg);
-	int (*recv)(handle_t sess, struct wd_rsa_msg *msg);
+	int (*send)(handle_t sess, void *rsa_msg);
+	int (*recv)(handle_t sess, void *rsa_msg);
 };
 
 void wd_rsa_set_driver(struct wd_rsa_driver *drv);
diff --git a/include/wd_util.h b/include/wd_util.h
index 89dfe87..eb250f0 100644
--- a/include/wd_util.h
+++ b/include/wd_util.h
@@ -99,6 +99,11 @@ struct wd_ctx_attr {
 	__u8 mode;
 };
 
+struct wd_msg_handle {
+	int (*send)(handle_t sess, void *msg);
+	int (*recv)(handle_t sess, void *msg);
+};
+
 /*
  * wd_init_ctx_config() - Init internal ctx configuration.
  * @in:	ctx configuration in global setting.
@@ -327,6 +332,19 @@ int wd_check_ctx(struct wd_ctx_config_internal *config, __u8 mode, __u32 idx);
  */
 int wd_set_epoll_en(const char *var_name, bool *epoll_en);
 
+/**
+ * wd_handle_msg_sync() - recv msg from hardware
+ * @msg_handle: callback of msg handle ops.
+ * @ctx: the handle of context.
+ * @msg: the msg of task.
+ * @balance: estimated number of receiving msg.
+ * @epoll_en: whether to enable epoll.
+ *
+ * Return 0 if successful or less than 0 otherwise.
+ */
+int wd_handle_msg_sync(struct wd_msg_handle *msg_handle, handle_t ctx,
+		void *msg, __u64 *balance, bool epoll_en);
+
 /**
  * wd_init_check() - Check input parameters for wd_<alg>_init.
  * @config: Ctx configuration input by user.
diff --git a/wd_aead.c b/wd_aead.c
index 3f47f8b..7c07271 100644
--- a/wd_aead.c
+++ b/wd_aead.c
@@ -19,8 +19,6 @@
 #define WD_AEAD_CCM_GCM_MIN	4U
 #define WD_AEAD_CCM_GCM_MAX	16
 #define WD_POOL_MAX_ENTRIES	1024
-#define MAX_RETRY_COUNTS	200000000
-
 
 static int g_aead_mac_len[WD_DIGEST_TYPE_MAX] = {
 	WD_DIGEST_SM3_LEN, WD_DIGEST_MD5_LEN, WD_DIGEST_SHA1_LEN,
@@ -496,36 +494,15 @@ static void fill_request_msg(struct wd_aead_msg *msg, struct wd_aead_req *req,
 static int send_recv_sync(struct wd_ctx_internal *ctx,
 			  struct wd_aead_msg *msg)
 {
-	__u64 recv_cnt = 0;
+	struct wd_msg_handle msg_handle;
 	int ret;
 
-	pthread_spin_lock(&ctx->lock);
-	ret = wd_aead_setting.driver->aead_send(ctx->ctx, msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("failed to send aead bd!\n");
-		goto out;
-	}
-
-	do {
-		if (wd_aead_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx->ctx, POLL_TIME);
-			if (unlikely(ret < 0))
-				WD_ERR("wd aead ctx wait timeout(%d)!\n", ret);
-		}
-		ret = wd_aead_setting.driver->aead_recv(ctx->ctx, msg);
-		if (ret == -WD_HW_EACCESS) {
-			WD_ERR("wd aead recv err!\n");
-			goto out;
-		} else if (ret == -WD_EAGAIN) {
-			if (++recv_cnt > MAX_RETRY_COUNTS) {
-				WD_ERR("wd aead recv timeout fail!\n");
-				ret = -WD_ETIMEDOUT;
-				goto out;
-			}
-		}
-	} while (ret < 0);
+	msg_handle.send = wd_aead_setting.driver->aead_send;
+	msg_handle.recv = wd_aead_setting.driver->aead_recv;
 
-out:
+	pthread_spin_lock(&ctx->lock);
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, msg, NULL,
+			  wd_aead_setting.config.epoll_en);
 	pthread_spin_unlock(&ctx->lock);
 	return ret;
 }
diff --git a/wd_cipher.c b/wd_cipher.c
index be8996f..a7e393b 100644
--- a/wd_cipher.c
+++ b/wd_cipher.c
@@ -19,8 +19,6 @@
 
 #define WD_POOL_MAX_ENTRIES	1024
 #define DES_WEAK_KEY_NUM	16
-#define MAX_RETRY_COUNTS	200000000
-
 
 static const unsigned char des_weak_keys[DES_WEAK_KEY_NUM][DES_KEY_SIZE] = {
 	/* weak keys */
@@ -405,36 +403,15 @@ static int wd_cipher_check_params(handle_t h_sess,
 static int send_recv_sync(struct wd_ctx_internal *ctx,
 			  struct wd_cipher_msg *msg)
 {
-	__u64 recv_cnt = 0;
+	struct wd_msg_handle msg_handle;
 	int ret;
 
-	pthread_spin_lock(&ctx->lock);
-	ret = wd_cipher_setting.driver->cipher_send(ctx->ctx, msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("wd cipher send err!\n");
-		goto out;
-	}
-
-	do {
-		if (wd_cipher_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx->ctx, POLL_TIME);
-			if (unlikely(ret < 0))
-				WD_ERR("wd cipher ctx wait timeout(%d)!\n", ret);
-		}
-		ret = wd_cipher_setting.driver->cipher_recv(ctx->ctx, msg);
-		if (ret == -WD_HW_EACCESS) {
-			WD_ERR("wd cipher recv err!\n");
-			goto out;
-		} else if (ret == -WD_EAGAIN) {
-			if (++recv_cnt > MAX_RETRY_COUNTS) {
-				WD_ERR("wd cipher recv timeout fail!\n");
-				ret = -WD_ETIMEDOUT;
-				goto out;
-			}
-		}
-	} while (ret < 0);
+	msg_handle.send = wd_cipher_setting.driver->cipher_send;
+	msg_handle.recv = wd_cipher_setting.driver->cipher_recv;
 
-out:
+	pthread_spin_lock(&ctx->lock);
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, msg, NULL,
+			  wd_cipher_setting.config.epoll_en);
 	pthread_spin_unlock(&ctx->lock);
 	return ret;
 }
diff --git a/wd_comp.c b/wd_comp.c
index 6bf5f82..9107d60 100644
--- a/wd_comp.c
+++ b/wd_comp.c
@@ -18,11 +18,9 @@
 #include "wd_comp.h"
 
 #define WD_POOL_MAX_ENTRIES		1024
-#define MAX_RETRY_COUNTS		200000000
 #define HW_CTX_SIZE			(64 * 1024)
 #define STREAM_CHUNK			(128 * 1024)
 
-
 #define swap_byte(x) \
 	((((x) & 0x000000ff) << 24) | \
 	(((x) & 0x0000ff00) <<  8) | \
@@ -396,8 +394,8 @@ static int wd_comp_sync_job(struct wd_comp_sess *sess,
 {
 	struct wd_ctx_config_internal *config = &wd_comp_setting.config;
 	handle_t h_sched_ctx = wd_comp_setting.sched.h_sched_ctx;
+	struct wd_msg_handle msg_handle;
 	struct wd_ctx_internal *ctx;
-	__u64 recv_count = 0;
 	__u32 idx;
 	int ret;
 
@@ -410,35 +408,12 @@ static int wd_comp_sync_job(struct wd_comp_sess *sess,
 
 	ctx = config->ctxs + idx;
 
-	pthread_spin_lock(&ctx->lock);
-
-	ret = wd_comp_setting.driver->comp_send(ctx->ctx, msg);
-	if (unlikely(ret < 0)) {
-		pthread_spin_unlock(&ctx->lock);
-		WD_ERR("wd comp send error, ret = %d!\n", ret);
-		return ret;
-	}
-
-	do {
-		if (config->epoll_en) {
-			ret = wd_ctx_wait(ctx->ctx, POLL_TIME);
-			if (unlikely(ret < 0))
-				WD_ERR("wd ctx wait timeout, ret = %d!\n", ret);
-		}
-		ret = wd_comp_setting.driver->comp_recv(ctx->ctx, msg);
-		if (unlikely(ret == -WD_HW_EACCESS)) {
-			pthread_spin_unlock(&ctx->lock);
-			WD_ERR("wd comp recv hw error!\n");
-			return ret;
-		} else if (ret == -WD_EAGAIN) {
-			if (++recv_count > MAX_RETRY_COUNTS) {
-				pthread_spin_unlock(&ctx->lock);
-				WD_ERR("wd comp recv timeout!\n");
-				return -WD_ETIMEDOUT;
-			}
-		}
-	} while (ret == -WD_EAGAIN);
+	msg_handle.send = wd_comp_setting.driver->comp_send;
+	msg_handle.recv = wd_comp_setting.driver->comp_recv;
 
+	pthread_spin_lock(&ctx->lock);
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, msg,
+				 NULL, config->epoll_en);
 	pthread_spin_unlock(&ctx->lock);
 
 	return ret;
diff --git a/wd_dh.c b/wd_dh.c
index 5f9dff1..9571ab6 100644
--- a/wd_dh.c
+++ b/wd_dh.c
@@ -18,12 +18,10 @@
 #include "wd_dh.h"
 
 #define WD_POOL_MAX_ENTRIES		1024
-#define DH_BALANCE_THRHD		1280
 #define DH_MAX_KEY_SIZE			512
-#define DH_RECV_MAX_CNT			60000000 // 1 min
 #define WD_DH_G2			2
 
-static __thread int balance;
+static __thread __u64 balance;
 
 struct wd_dh_sess {
 	__u32 alg_type;
@@ -191,45 +189,12 @@ static int fill_dh_msg(struct wd_dh_msg *msg, struct wd_dh_req *req,
 	return 0;
 }
 
-static int dh_recv_sync(handle_t ctx, struct wd_dh_msg *msg)
-{
-	struct wd_dh_req *req = &msg->req;
-	__u32 rx_cnt = 0;
-	int ret;
-
-	do {
-		if (wd_dh_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx, POLL_TIME);
-			if (ret < 0)
-				WD_ERR("wd ctx wait timeout(%d)!\n", ret);
-		}
-
-		ret = wd_dh_setting.driver->recv(ctx, msg);
-		if (ret == -WD_EAGAIN) {
-			if (rx_cnt++ >= DH_RECV_MAX_CNT) {
-				WD_ERR("failed to recv: timeout!\n");
-				return -WD_ETIMEDOUT;
-			}
-
-			if (balance > DH_BALANCE_THRHD)
-				usleep(1);
-		} else if (ret < 0) {
-			WD_ERR("failed to recv: error = %d!\n", ret);
-			return ret;
-		}
-	} while (ret < 0);
-
-	balance = rx_cnt;
-	req->status = msg->result;
-
-	return GET_NEGATIVE(req->status);
-}
-
 int wd_do_dh_sync(handle_t sess, struct wd_dh_req *req)
 {
 	struct wd_ctx_config_internal *config = &wd_dh_setting.config;
 	handle_t h_sched_ctx = wd_dh_setting.sched.h_sched_ctx;
 	struct wd_dh_sess *sess_t = (struct wd_dh_sess *)sess;
+	struct wd_msg_handle msg_handle;
 	struct wd_ctx_internal *ctx;
 	struct wd_dh_msg msg;
 	__u32 idx;
@@ -254,19 +219,20 @@ int wd_do_dh_sync(handle_t sess, struct wd_dh_req *req)
 	if (unlikely(ret))
 		return ret;
 
+	msg_handle.send = wd_dh_setting.driver->send;
+	msg_handle.recv = wd_dh_setting.driver->recv;
+
 	pthread_spin_lock(&ctx->lock);
-	ret = wd_dh_setting.driver->send(ctx->ctx, &msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("failed to send dh BD, ret = %d!\n", ret);
-		goto fail;
-	}
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, &msg, &balance,
+				 wd_dh_setting.config.epoll_en);
+	pthread_spin_unlock(&ctx->lock);
+	if (unlikely(ret))
+		return ret;
 
-	ret = dh_recv_sync(ctx->ctx, &msg);
 	req->pri_bytes = msg.req.pri_bytes;
-fail:
-	pthread_spin_unlock(&ctx->lock);
+	req->status = msg.result;
 
-	return ret;
+	return GET_NEGATIVE(msg.result);
 }
 
 int wd_do_dh_async(handle_t sess, struct wd_dh_req *req)
diff --git a/wd_digest.c b/wd_digest.c
index 46546cb..0fff7c8 100644
--- a/wd_digest.c
+++ b/wd_digest.c
@@ -17,8 +17,6 @@
 
 #define WD_POOL_MAX_ENTRIES	1024
 #define DES_WEAK_KEY_NUM	4
-#define MAX_RETRY_COUNTS	200000000
-
 
 static int g_digest_mac_len[WD_DIGEST_TYPE_MAX] = {
 	WD_DIGEST_SM3_LEN, WD_DIGEST_MD5_LEN, WD_DIGEST_SHA1_LEN,
@@ -294,41 +292,24 @@ static void fill_request_msg(struct wd_digest_msg *msg,
 static int send_recv_sync(struct wd_ctx_internal *ctx, struct wd_digest_sess *dsess,
 			  struct wd_digest_msg *msg)
 {
-	__u64 recv_cnt = 0;
+	struct wd_msg_handle msg_handle;
 	int ret;
 
+	msg_handle.send = wd_digest_setting.driver->digest_send;
+	msg_handle.recv = wd_digest_setting.driver->digest_recv;
+
 	pthread_spin_lock(&ctx->lock);
-	ret = wd_digest_setting.driver->digest_send(ctx->ctx, msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("failed to send bd!\n");
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, msg,
+				 NULL, wd_digest_setting.config.epoll_en);
+	if (unlikely(ret))
 		goto out;
-	}
 
-	do {
-		if (wd_digest_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx->ctx, POLL_TIME);
-			if (unlikely(ret < 0))
-				WD_ERR("wd digest ctx wait timeout(%d)!\n", ret);
-		}
-		ret = wd_digest_setting.driver->digest_recv(ctx->ctx, msg);
-		if (ret == -WD_HW_EACCESS) {
-			WD_ERR("wd digest recv err!\n");
-			goto out;
-		} else if (ret == -WD_EAGAIN) {
-			if (++recv_cnt > MAX_RETRY_COUNTS) {
-				WD_ERR("wd digest recv timeout fail!\n");
-				ret = -WD_ETIMEDOUT;
-				goto out;
-			}
-		}
-
-		/*
-		 * 'out_bytes' can be expressed BD state, non-zero is final BD or
-		 * middle BD as stream mode.
-		 */
-		if (msg->has_next)
-			dsess->bd_state = msg->out_bytes;
-	} while (ret < 0);
+	/*
+	 * 'out_bytes' can be expressed BD state, non-zero is final BD or
+	 * middle BD as stream mode.
+	 */
+	if (msg->has_next)
+		dsess->bd_state = msg->out_bytes;
 
 out:
 	pthread_spin_unlock(&ctx->lock);
diff --git a/wd_ecc.c b/wd_ecc.c
index f22b72d..76c9d5f 100644
--- a/wd_ecc.c
+++ b/wd_ecc.c
@@ -20,8 +20,6 @@
 #define WD_POOL_MAX_ENTRIES		1024
 #define WD_ECC_CTX_MSG_NUM		64
 #define WD_ECC_MAX_CTX			256
-#define ECC_BALANCE_THRHD		1280
-#define ECC_RECV_MAX_CNT		60000000
 #define ECC_MAX_HW_BITS			521
 #define ECC_MAX_KEY_SIZE		BITS_TO_BYTES(ECC_MAX_HW_BITS)
 #define ECC_MAX_IN_NUM			4
@@ -34,7 +32,7 @@
 #define GET_NEGATIVE(val)		(0 - (val))
 #define ZA_PARAM_NUM  			6
 
-static __thread int balance;
+static __thread __u64 balance;
 
 struct curve_param_desc {
 	__u32 type;
@@ -1386,46 +1384,12 @@ static void msg_pack(char *dst, __u64 *out_len,
 	*out_len += src_len;
 }
 
-static int ecc_recv_sync(handle_t ctx, struct wd_ecc_msg *msg)
-{
-	struct wd_ecc_req *req = &msg->req;
-	__u32 rx_cnt = 0;
-	int ret;
-
-	do {
-		if (wd_ecc_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx, POLL_TIME);
-			if (ret < 0)
-				WD_ERR("wd ctx wait timeout(%d)!\n", ret);
-		}
-
-		ret = wd_ecc_setting.driver->recv(ctx, msg);
-		if (ret == -WD_EAGAIN) {
-			if (rx_cnt++ >= ECC_RECV_MAX_CNT) {
-				WD_ERR("failed to recv: timeout!\n");
-				return -WD_ETIMEDOUT;
-			}
-
-			if (balance > ECC_BALANCE_THRHD)
-				usleep(1);
-		} else if (ret < 0) {
-			WD_ERR("failed to recv: error = %d!\n", ret);
-			return ret;
-		}
-	} while (ret < 0);
-
-	balance = rx_cnt;
-	req->status = msg->result;
-	req->dst_bytes = msg->req.dst_bytes;
-
-	return GET_NEGATIVE(req->status);
-}
-
 int wd_do_ecc_sync(handle_t h_sess, struct wd_ecc_req *req)
 {
 	struct wd_ctx_config_internal *config = &wd_ecc_setting.config;
 	handle_t h_sched_ctx = wd_ecc_setting.sched.h_sched_ctx;
 	struct wd_ecc_sess *sess = (struct wd_ecc_sess *)h_sess;
+	struct wd_msg_handle msg_handle;
 	struct wd_ctx_internal *ctx;
 	struct wd_ecc_msg msg;
 	__u32 idx;
@@ -1450,18 +1414,20 @@ int wd_do_ecc_sync(handle_t h_sess, struct wd_ecc_req *req)
 	if (unlikely(ret))
 		return ret;
 
-	pthread_spin_lock(&ctx->lock);
-	ret = wd_ecc_setting.driver->send(ctx->ctx, &msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("failed to send ecc BD, hw is err!\n");
-		goto fail;
-	}
+	msg_handle.send = wd_ecc_setting.driver->send;
+	msg_handle.recv = wd_ecc_setting.driver->recv;
 
-	ret = ecc_recv_sync(ctx->ctx, &msg);
-fail:
+	pthread_spin_lock(&ctx->lock);
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, &msg, &balance,
+				 wd_ecc_setting.config.epoll_en);
 	pthread_spin_unlock(&ctx->lock);
+	if (unlikely(ret))
+		return ret;
 
-	return ret;
+	req->dst_bytes = msg.req.dst_bytes;
+	req->status = msg.result;
+
+	return GET_NEGATIVE(msg.result);
 }
 
 static void get_sign_out_params(struct wd_ecc_out *out,
diff --git a/wd_rsa.c b/wd_rsa.c
index 024f0e0..868e78b 100644
--- a/wd_rsa.c
+++ b/wd_rsa.c
@@ -19,11 +19,9 @@
 #define WD_POOL_MAX_ENTRIES		1024
 #define WD_HW_EACCESS			62
 
-#define RSA_BALANCE_THRHD		1280
 #define RSA_MAX_KEY_SIZE		512
-#define RSA_RECV_MAX_CNT		60000000 // 1 min
 
-static __thread int balance;
+static __thread __u64 balance;
 
 struct wd_rsa_pubkey {
 	struct wd_dtb n;
@@ -251,46 +249,12 @@ static int fill_rsa_msg(struct wd_rsa_msg *msg, struct wd_rsa_req *req,
 	return 0;
 }
 
-static int rsa_recv_sync(handle_t ctx, struct wd_rsa_msg *msg)
-{
-	struct wd_rsa_req *req = &msg->req;
-	__u32 rx_cnt = 0;
-	int ret;
-
-	do {
-		if (wd_rsa_setting.config.epoll_en) {
-			ret = wd_ctx_wait(ctx, POLL_TIME);
-			if (ret < 0)
-				WD_ERR("wd ctx wait timeout(%d)!\n", ret);
-		}
-
-		ret = wd_rsa_setting.driver->recv(ctx, msg);
-		if (ret == -WD_EAGAIN) {
-			if (rx_cnt++ >= RSA_RECV_MAX_CNT) {
-				WD_ERR("failed to recv: timeout!\n");
-				return -WD_ETIMEDOUT;
-			}
-
-			if (balance > RSA_BALANCE_THRHD)
-				usleep(1);
-		} else if (ret < 0) {
-			WD_ERR("failed to recv: error = %d!\n", ret);
-			return ret;
-		}
-	} while (ret < 0);
-
-	balance = rx_cnt;
-	req->status = msg->result;
-	req->dst_bytes = msg->req.dst_bytes;
-
-	return GET_NEGATIVE(req->status);
-}
-
 int wd_do_rsa_sync(handle_t h_sess, struct wd_rsa_req *req)
 {
 	struct wd_ctx_config_internal *config = &wd_rsa_setting.config;
 	handle_t h_sched_ctx = wd_rsa_setting.sched.h_sched_ctx;
 	struct wd_rsa_sess *sess = (struct wd_rsa_sess *)h_sess;
+	struct wd_msg_handle msg_handle;
 	struct wd_ctx_internal *ctx;
 	struct wd_rsa_msg msg;
 	__u32 idx;
@@ -315,18 +279,20 @@ int wd_do_rsa_sync(handle_t h_sess, struct wd_rsa_req *req)
 	if (unlikely(ret))
 		return ret;
 
-	pthread_spin_lock(&ctx->lock);
-	ret = wd_rsa_setting.driver->send(ctx->ctx, &msg);
-	if (unlikely(ret < 0)) {
-		WD_ERR("failed to send rsa BD, ret = %d!\n", ret);
-		goto fail;
-	}
+	msg_handle.send = wd_rsa_setting.driver->send;
+	msg_handle.recv = wd_rsa_setting.driver->recv;
 
-	ret = rsa_recv_sync(ctx->ctx, &msg);
-fail:
+	pthread_spin_lock(&ctx->lock);
+	ret = wd_handle_msg_sync(&msg_handle, ctx->ctx, &msg, &balance,
+				 wd_rsa_setting.config.epoll_en);
 	pthread_spin_unlock(&ctx->lock);
+	if (unlikely(ret))
+		return ret;
 
-	return ret;
+	req->dst_bytes = msg.req.dst_bytes;
+	req->status = msg.result;
+
+	return GET_NEGATIVE(msg.result);
 }
 
 int wd_do_rsa_async(handle_t sess, struct wd_rsa_req *req)
diff --git a/wd_util.c b/wd_util.c
index 4d5faf9..de77d1e 100644
--- a/wd_util.c
+++ b/wd_util.c
@@ -17,6 +17,9 @@
 
 #define WD_ASYNC_DEF_POLL_NUM		1
 #define WD_ASYNC_DEF_QUEUE_DEPTH	1024
+#define WD_BALANCE_THRHD		1280
+#define WD_RECV_MAX_CNT_SLEEP		60000000
+#define WD_RECV_MAX_CNT_NOSLEEP		200000000
 
 struct msg_pool {
 	/* message array allocated dynamically */
@@ -1629,6 +1632,50 @@ int wd_set_epoll_en(const char *var_name, bool *epoll_en)
 	return 0;
 }
 
+int wd_handle_msg_sync(struct wd_msg_handle *msg_handle, handle_t ctx,
+		       void *msg, __u64 *balance, bool epoll_en)
+{
+	__u64 timeout = WD_RECV_MAX_CNT_NOSLEEP;
+	__u64 rx_cnt = 0;
+	int ret;
+
+	if (balance)
+		timeout = WD_RECV_MAX_CNT_SLEEP;
+
+	ret = msg_handle->send(ctx, msg);
+	if (unlikely(ret < 0)) {
+		WD_ERR("failed to send msg to hw, ret = %d!\n", ret);
+		return ret;
+	}
+
+	do {
+		if (epoll_en) {
+			ret = wd_ctx_wait(ctx, POLL_TIME);
+			if (ret < 0)
+				WD_ERR("wd ctx wait timeout(%d)!\n", ret);
+		}
+
+		ret = msg_handle->recv(ctx, msg);
+		if (ret == -WD_EAGAIN) {
+			if (unlikely(rx_cnt++ >= timeout)) {
+				WD_ERR("failed to recv msg: timeout!\n");
+				return -WD_ETIMEDOUT;
+			}
+
+			if (balance && *balance > WD_BALANCE_THRHD)
+				usleep(1);
+		} else if (unlikely(ret < 0)) {
+			WD_ERR("failed to recv msg: error = %d!\n", ret);
+			return ret;
+		}
+	} while (ret < 0);
+
+	if (balance)
+		*balance = rx_cnt;
+
+	return ret;
+}
+
 int wd_init_param_check(struct wd_ctx_config *config, struct wd_sched *sched)
 {
 	if (!config || !config->ctxs || !config->ctxs[0].ctx) {
-- 
2.27.0