From ca7627f2a44b6815b2f23f4da2d0fc996bd20424 Mon Sep 17 00:00:00 2001
From: Hisping Lin <hisping.lin@rock-chips.com>
Date: Fri, 22 Dec 2023 09:52:47 +0800
Subject: [PATCH] tee: optee: interrupt an RPC depend on shutdown flag

Fixes: e6c7ea7d4da7 ("tee: optee: interrupt an RPC when supplicant has been killed")
Change-Id: I33f33af02afb75bda5d68c9731fe636bb97132fe
Signed-off-by: Hisping Lin <hisping.lin@rock-chips.com>
---
 drivers/tee/optee/core.c          |  8 ++++++++
 drivers/tee/optee/optee_private.h |  1 +
 drivers/tee/optee/supp.c          | 13 +++++--------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/drivers/tee/optee/core.c b/drivers/tee/optee/core.c
index 6ea80add7378..6abf2acec14c 100644
--- a/drivers/tee/optee/core.c
+++ b/drivers/tee/optee/core.c
@@ -7,6 +7,7 @@
 
 #include <linux/arm-smccc.h>
 #include <linux/crash_dump.h>
+#include <linux/delay.h>
 #include <linux/errno.h>
 #include <linux/io.h>
 #include <linux/module.h>
@@ -624,6 +625,13 @@ static int optee_remove(struct platform_device *pdev)
  */
 static void optee_shutdown(struct platform_device *pdev)
 {
+	struct optee *optee = platform_get_drvdata(pdev);
+
+	/* Tell requesting thread to interrupt an RPC */
+	smp_store_mb(optee->supp.shutdown, true);
+	/* Wait requesting thread to release resources */
+	mdelay(200);
+
 	optee_disable_shm_cache(platform_get_drvdata(pdev));
 }
 
diff --git a/drivers/tee/optee/optee_private.h b/drivers/tee/optee/optee_private.h
index ea09533e30cd..8af2c97c8c09 100644
--- a/drivers/tee/optee/optee_private.h
+++ b/drivers/tee/optee/optee_private.h
@@ -64,6 +64,7 @@ struct optee_supp {
 	struct list_head reqs;
 	struct idr idr;
 	struct completion reqs_c;
+	bool shutdown;
 };
 
 /**
diff --git a/drivers/tee/optee/supp.c b/drivers/tee/optee/supp.c
index 2d556b79a67e..a8656a72d83a 100644
--- a/drivers/tee/optee/supp.c
+++ b/drivers/tee/optee/supp.c
@@ -82,7 +82,6 @@ u32 optee_supp_thrd_req(struct tee_context *ctx, u32 func, size_t num_params,
 	struct optee_supp_req *req;
 	bool interruptable;
 	u32 ret;
-	unsigned long timeleft;
 	int id;
 	struct optee_supp_req *get_req;
 
@@ -117,14 +116,12 @@ u32 optee_supp_thrd_req(struct tee_context *ctx, u32 func, size_t num_params,
 	 * exclusive access again.
 	 */
 	while (wait_for_completion_interruptible(&req->c)) {
-		pr_err("Warning, Interrupting an RPC to supplicant!\n");
-		timeleft = wait_for_completion_timeout(&req->c, msecs_to_jiffies(2000));
-		if (timeleft) {
-			/* get completion, it means tee-supplicant is alive. */
-			break;
-		} else {
-			/* timeout, it means tee-supplicant is dead, interrupting an RPC. */
+		if (supp->shutdown) {
+			/* Reboot happen, tee-supplicant is dead, interrupt an RPC */
 			interruptable = true;
+		} else {
+			/* Deep sleep, tee-supplicant is freeze, wait tee-supplicant */
+			continue;
 		}
 
 		mutex_lock(&supp->mutex);