[PATCH] tee: optee: smc_abi: Fix use-after-free in optee_smc_open

[Yoochan Lee <yoochan1026@gmail.com>]

A race condition may occur if the user physically removes the
smc_abi device while calling open().

This is a race condition between optee_smc_open() function and
the optee_smc_remove() function, which may lead to Use-After-Free.

Therefore, add a refcount check to optee_smc_remove() function
to free the "optee" structure after the device is close()d.

---------------CPU 0--------------------CPU 1-----------------
         optee_smc_open()     |     optee_smc_remove()
--------------------------------------------------------------
struct optee *optee = tee_get_|
drvdata(ctx->teedev); — (1)   |
                              | struct optee *optee = platform_
                              | get_drvdata(pdev);
                              | ...
                              | kfree(optee); — (2)
u32 sec_caps = optee->smc.sec_|
caps; — (3)

Signed-off-by: Yoochan Lee <yoochan1026@gmail.com>
---
 drivers/tee/optee/optee_private.h |  1 +
 drivers/tee/optee/smc_abi.c       | 66 ++++++++++++++++++++++---------
 2 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/drivers/tee/optee/optee_private.h b/drivers/tee/optee/optee_private.h
index 04ae58892608..f52b1cf20eab 100644
--- a/drivers/tee/optee/optee_private.h
+++ b/drivers/tee/optee/optee_private.h
@@ -175,6 +175,7 @@ struct optee {
 	bool   scan_bus_done;
 	struct workqueue_struct *scan_bus_wq;
 	struct work_struct scan_bus_work;
+	struct kref refcnt;
 };
 
 struct optee_session {
diff --git a/drivers/tee/optee/smc_abi.c b/drivers/tee/optee/smc_abi.c
index a1c1fa1a9c28..4fbec2acc255 100644
--- a/drivers/tee/optee/smc_abi.c
+++ b/drivers/tee/optee/smc_abi.c
@@ -1077,18 +1077,61 @@ static void optee_get_version(struct tee_device *teedev,
 	*vers = v;
 }
 
+static void optee_smc_delete(struct kref *kref)
+{
+	struct optee *optee = container_of(kref, struct optee, refcnt);
+
+	/*
+	 * Ask OP-TEE to free all cached shared memory objects to decrease
+	 * reference counters and also avoid wild pointers in secure world
+	 * into the old shared memory range.
+	 */
+	if (!optee->rpc_param_count)
+		optee_disable_shm_cache(optee);
+
+	optee_smc_notif_uninit_irq(optee);
+
+	optee_remove_common(optee);
+
+	if (optee->smc.memremaped_shm)
+		memunmap(optee->smc.memremaped_shm);
+
+	kfree(optee);
+}
+
+static void optee_smc_release_supp(struct tee_context *ctx)
+{
+	struct optee *optee = tee_get_drvdata(ctx->teedev);
+
+	optee_release_helper(ctx, optee_close_session_helper);
+	if (optee->scan_bus_wq) {
+		destroy_workqueue(optee->scan_bus_wq);
+		optee->scan_bus_wq = NULL;
+	}
+	optee_supp_release(&optee->supp);
+	kref_put(&optee->refcnt, optee_smc_delete);
+}
+
+static void optee_smc_release(struct tee_context *ctx)
+{
+	struct optee *optee = tee_get_drvdata(ctx->teedev);
+
+	optee_release_helper(ctx, optee_close_session_helper);
+	kref_put(&optee->refcnt, optee_smc_delete);
+}
+
 static int optee_smc_open(struct tee_context *ctx)
 {
 	struct optee *optee = tee_get_drvdata(ctx->teedev);
 	u32 sec_caps = optee->smc.sec_caps;
-
+	kref_get(&optee->refcnt);
 	return optee_open(ctx, sec_caps & OPTEE_SMC_SEC_CAP_MEMREF_NULL);
 }
 
 static const struct tee_driver_ops optee_clnt_ops = {
 	.get_version = optee_get_version,
 	.open = optee_smc_open,
-	.release = optee_release,
+	.release = optee_smc_release,
 	.open_session = optee_open_session,
 	.close_session = optee_close_session,
 	.invoke_func = optee_invoke_func,
@@ -1106,7 +1149,7 @@ static const struct tee_desc optee_clnt_desc = {
 static const struct tee_driver_ops optee_supp_ops = {
 	.get_version = optee_get_version,
 	.open = optee_smc_open,
-	.release = optee_release_supp,
+	.release = optee_smc_release_supp,
 	.supp_recv = optee_supp_recv,
 	.supp_send = optee_supp_send,
 	.shm_register = optee_shm_register_supp,
@@ -1319,22 +1362,7 @@ static int optee_smc_remove(struct platform_device *pdev)
 {
 	struct optee *optee = platform_get_drvdata(pdev);
 
-	/*
-	 * Ask OP-TEE to free all cached shared memory objects to decrease
-	 * reference counters and also avoid wild pointers in secure world
-	 * into the old shared memory range.
-	 */
-	if (!optee->rpc_param_count)
-		optee_disable_shm_cache(optee);
-
-	optee_smc_notif_uninit_irq(optee);
-
-	optee_remove_common(optee);
-
-	if (optee->smc.memremaped_shm)
-		memunmap(optee->smc.memremaped_shm);
-
-	kfree(optee);
+	kref_put(&optee->refcnt, optee_smc_delete);
 
 	return 0;
 }
-- 
2.39.0