From a8ee531704b607d5f023315d34f78d4fb7ef421f Mon Sep 17 00:00:00 2001
From: Gao Wang <quic_gaowang@quicinc.com>
Date: Mon, 23 Dec 2024 17:14:59 +0800
Subject: [PATCH] msm: npu: Fix use after free issue

There is possibility that network will be used after free.
This change is to fix this issue.

Change-Id: I12205b750450bee36f85dff3f620f8f0689a4e46
Signed-off-by: Gao Wang <quic_gaowang@quicinc.com>
---
 drivers/media/platform/msm/npu/npu_dev.c |  2 +-
 drivers/media/platform/msm/npu/npu_mgr.c | 30 ++++++++++++++++++++++++
 drivers/media/platform/msm/npu/npu_mgr.h |  3 +++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/drivers/media/platform/msm/npu/npu_dev.c b/drivers/media/platform/msm/npu/npu_dev.c
index e5b2d93d9db0..704ccddd09e2 100644
--- a/drivers/media/platform/msm/npu/npu_dev.c
+++ b/drivers/media/platform/msm/npu/npu_dev.c
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-only
 /*
  * Copyright (c) 2017-2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
@@ -1552,7 +1553,6 @@ static int npu_set_fw_state(struct npu_client *client, uint32_t enable)
 
 	if (host_ctx->network_num > 0) {
 		pr_err("Need to unload network first\n");
-		mutex_unlock(&npu_dev->dev_lock);
 		return -EINVAL;
 	}
 
diff --git a/drivers/media/platform/msm/npu/npu_mgr.c b/drivers/media/platform/msm/npu/npu_mgr.c
index 6b017934f985..e21ba85a839d 100644
--- a/drivers/media/platform/msm/npu/npu_mgr.c
+++ b/drivers/media/platform/msm/npu/npu_mgr.c
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-only
 /*
  * Copyright (c) 2018-2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
@@ -1593,6 +1594,13 @@ int32_t npu_host_unload_network(struct npu_client *client,
 		return -EINVAL;
 	}
 
+	if (network->is_unloading) {
+		pr_err("network is unloading\n");
+		network_put(network);
+		mutex_unlock(&host_ctx->lock);
+		return -EINVAL;
+	}
+
 	if (!network->is_active) {
 		pr_err("network is not active\n");
 		network_put(network);
@@ -1600,10 +1608,18 @@ int32_t npu_host_unload_network(struct npu_client *client,
 		return -EINVAL;
 	}
 
+	if (network->is_executing) {
+		pr_err("network is in execution\n");
+		network_put(network);
+		mutex_unlock(&host_ctx->lock);
+		return -EINVAL;
+	}
+
 	if (network->fw_error) {
 		pr_err("fw in error state, skip unload network in fw\n");
 		goto free_network;
 	}
+	network->is_unloading = true;
 
 	pr_debug("Unload network %lld\n", network->id);
 	/* prepare IPC packet for UNLOAD */
@@ -1832,12 +1848,24 @@ int32_t npu_host_exec_network_v2(struct npu_client *client,
 	if (atomic_inc_return(&host_ctx->network_execute_cnt) == 1)
 		npu_notify_cdsprm_cxlimit_activity(npu_dev, true);
 
+	if (network->is_unloading) {
+		pr_err("network is unloading\n");
+		ret = -EINVAL;
+		goto exec_v2_done;
+	}
+
 	if (!network->is_active) {
 		pr_err("network is not active\n");
 		ret = -EINVAL;
 		goto exec_v2_done;
 	}
 
+	if (network->is_executing) {
+		pr_err("network is already in execution\n");
+		ret = -EINVAL;
+		goto exec_v2_done;
+	}
+
 	if (network->fw_error) {
 		pr_err("fw is in error state\n");
 		ret = -EIO;
@@ -1855,6 +1883,7 @@ int32_t npu_host_exec_network_v2(struct npu_client *client,
 		goto exec_v2_done;
 	}
 
+	network->is_executing = true;
 	for (i = 0; i < num_patch_params; i++) {
 		exec_packet->patch_params[i].id = patch_buf_info[i].buf_id;
 		pr_debug("%d: patch_id: %x\n", i,
@@ -1936,6 +1965,7 @@ int32_t npu_host_exec_network_v2(struct npu_client *client,
 
 free_exec_packet:
 	kfree(exec_packet);
+	network->is_executing = false;
 exec_v2_done:
 	network_put(network);
 	mutex_unlock(&host_ctx->lock);
diff --git a/drivers/media/platform/msm/npu/npu_mgr.h b/drivers/media/platform/msm/npu/npu_mgr.h
index f0ab23267787..b6bf46557dbd 100644
--- a/drivers/media/platform/msm/npu/npu_mgr.h
+++ b/drivers/media/platform/msm/npu/npu_mgr.h
@@ -1,6 +1,7 @@
 /* SPDX-License-Identifier: GPL-2.0-only */
 /*
  * Copyright (c) 2018-2021, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #ifndef _NPU_MGR_H
@@ -52,6 +53,8 @@ struct npu_network {
 	atomic_t ref_cnt;
 	bool is_valid;
 	bool is_active;
+	bool is_unloading;
+	bool is_executing;
 	bool fw_error;
 	bool cmd_pending;
 	bool cmd_async;