[PATCH] tee: optee: ffa_abi: Fix use-after-free in optee_ffa_open

[Yoochan Lee]

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

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

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

---------------CPU 0--------------------CPU 1-----------------
            optee_open()      |     optee_ffa_remove()
--------------------------------------------------------------
struct optee *optee = tee_get_|
drvdata(teedev); — (1)        |
                              | struct optee *optee = ffa_dev_
                              | get_drvdata(ffa_dev);
                              | ...
                              | kfree(optee); — (2)
if (teedev == optee->supp_    |
teedev) { — (3)               |

Signed-off-by: Yoochan Lee <yoochan1026@xxxxxxxxx>
---
 drivers/tee/optee/ffa_abi.c       | 49 +++++++++++++++++++++++++------
 drivers/tee/optee/optee_private.h |  1 +
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/drivers/tee/optee/ffa_abi.c b/drivers/tee/optee/ffa_abi.c
index 0828240f27e6..ea76d7532419 100644
--- a/drivers/tee/optee/ffa_abi.c
+++ b/drivers/tee/optee/ffa_abi.c
@@ -726,15 +726,52 @@ static void optee_ffa_get_version(struct tee_device *teedev,
 	*vers = v;
 }
 
+static void optee_ffa_delete(struct kref *kref)
+{
+	struct optee *optee = container_of(kref, struct optee, refcnt);
+
+	optee_remove_common(optee);
+
+	mutex_destroy(&optee->ffa.mutex);
+	rhashtable_free_and_destroy(&optee->ffa.global_ids, rh_free_fn, NULL);
+
+	kfree(optee);
+
+}
+
+static void optee_ffa_release(struct tee_context *ctx)
+{
+	struct optee *optee = tee_get_drvdata(teedev);
+
+	optee_release_helper(ctx, optee_close_session_helper);
+	kref_put(&optee->refcnt, optee_ffa_delete);
+}
+
+void optee_ffa_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_ffa_delete);
+}
+
 static int optee_ffa_open(struct tee_context *ctx)
 {
+	struct optee *optee = tee_get_drvdata(teedev);
+	kref_get(&optee->refcnt);
+
 	return optee_open(ctx, true);
 }
 
 static const struct tee_driver_ops optee_ffa_clnt_ops = {
 	.get_version = optee_ffa_get_version,
 	.open = optee_ffa_open,
-	.release = optee_release,
+	.release = optee_ffa_release,
 	.open_session = optee_open_session,
 	.close_session = optee_close_session,
 	.invoke_func = optee_invoke_func,
@@ -752,7 +789,7 @@ static const struct tee_desc optee_ffa_clnt_desc = {
 static const struct tee_driver_ops optee_ffa_supp_ops = {
 	.get_version = optee_ffa_get_version,
 	.open = optee_ffa_open,
-	.release = optee_release_supp,
+	.release = optee_ffa_release_supp,
 	.supp_recv = optee_supp_recv,
 	.supp_send = optee_supp_send,
 	.shm_register = optee_ffa_shm_register, /* same as for clnt ops */
@@ -775,13 +812,7 @@ static const struct optee_ops optee_ffa_ops = {
 static void optee_ffa_remove(struct ffa_device *ffa_dev)
 {
 	struct optee *optee = ffa_dev_get_drvdata(ffa_dev);
-
-	optee_remove_common(optee);
-
-	mutex_destroy(&optee->ffa.mutex);
-	rhashtable_free_and_destroy(&optee->ffa.global_ids, rh_free_fn, NULL);
-
-	kfree(optee);
+	kref_put(&optee->refcnt, optee_ffa_delete);
 }
 
 static int optee_ffa_probe(struct ffa_device *ffa_dev)
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 {
-- 
2.39.0