3 files changed, 41 insertions, 9 deletions

diff --git a/drivers/tee/tee_core.c b/drivers/tee/tee_core.c
index 295910f5cdd0..3d49ac2e3c84 100644
--- a/drivers/tee/tee_core.c
+++ b/drivers/tee/tee_core.c
@@ -54,6 +54,7 @@ static int tee_open(struct inode *inode, struct file *filp)
 		goto err;
 	}
 
+	kref_init(&ctx->refcount);
 	ctx->teedev = teedev;
 	INIT_LIST_HEAD(&ctx->list_shm);
 	filp->private_data = ctx;
@@ -68,19 +69,40 @@ err:
 	return rc;
 }
 
-static int tee_release(struct inode *inode, struct file *filp)
+void teedev_ctx_get(struct tee_context *ctx)
 {
-	struct tee_context *ctx = filp->private_data;
-	struct tee_device *teedev = ctx->teedev;
-	struct tee_shm *shm;
+	if (ctx->releasing)
+		return;
+
+	kref_get(&ctx->refcount);
+}
 
+static void teedev_ctx_release(struct kref *ref)
+{
+	struct tee_context *ctx = container_of(ref, struct tee_context,
+					       refcount);
+	ctx->releasing = true;
 	ctx->teedev->desc->ops->release(ctx);
-	mutex_lock(&ctx->teedev->mutex);
-	list_for_each_entry(shm, &ctx->list_shm, link)
-		shm->ctx = NULL;
-	mutex_unlock(&ctx->teedev->mutex);
 	kfree(ctx);
-	tee_device_put(teedev);
+}
+
+void teedev_ctx_put(struct tee_context *ctx)
+{
+	if (ctx->releasing)
+		return;
+
+	kref_put(&ctx->refcount, teedev_ctx_release);
+}
+
+static void teedev_close_context(struct tee_context *ctx)
+{
+	tee_device_put(ctx->teedev);
+	teedev_ctx_put(ctx);
+}
+
+static int tee_release(struct inode *inode, struct file *filp)
+{
+	teedev_close_context(filp->private_data);
 	return 0;
 }
 
diff --git a/drivers/tee/tee_private.h b/drivers/tee/tee_private.h
index 2bc2b5ab1661..85d99d621603 100644
--- a/drivers/tee/tee_private.h
+++ b/drivers/tee/tee_private.h
@@ -73,4 +73,7 @@ int tee_shm_get_fd(struct tee_shm *shm);
 bool tee_device_get(struct tee_device *teedev);
 void tee_device_put(struct tee_device *teedev);
 
+void teedev_ctx_get(struct tee_context *ctx);
+void teedev_ctx_put(struct tee_context *ctx);
+
 #endif /*TEE_PRIVATE_H*/
diff --git a/drivers/tee/tee_shm.c b/drivers/tee/tee_shm.c
index 11d11a46d86e..ba02a15eefcb 100644
--- a/drivers/tee/tee_shm.c
+++ b/drivers/tee/tee_shm.c
@@ -53,6 +53,9 @@ static void tee_shm_release(struct tee_shm *shm)
 		kfree(shm->pages);
 	}
 
+	if (shm->ctx)
+		teedev_ctx_put(shm->ctx);
+
 	kfree(shm);
 
 	tee_device_put(teedev);
@@ -187,6 +190,7 @@ struct tee_shm *__tee_shm_alloc(struct tee_context *ctx,
 	}
 
 	if (ctx) {
+		teedev_ctx_get(ctx);
 		mutex_lock(&teedev->mutex);
 		list_add_tail(&shm->link, &ctx->list_shm);
 		mutex_unlock(&teedev->mutex);
@@ -253,6 +257,8 @@ struct tee_shm *tee_shm_register(struct tee_context *ctx, unsigned long addr,
 		return ERR_PTR(-ENOTSUPP);
 	}
 
+	teedev_ctx_get(ctx);
+
 	shm = kzalloc(sizeof(*shm), GFP_KERNEL);
 	if (!shm) {
 		ret = ERR_PTR(-ENOMEM);
@@ -334,6 +340,7 @@ err:
 		kfree(shm->pages);
 	}
 	kfree(shm);
+	teedev_ctx_put(ctx);
 	tee_device_put(teedev);
 	return ret;
 }