msm: vidc: Fix use after free memory issue

One thread might free the video buffer structure in
msm_comm_flush() and other thread might still be using it
in msm_comm_qbuf() which raises use after free memory issue.
Fix the issue by using kref functions.

CRs-Fixed: 2075754
Change-Id: I27dc83031ec761b8fd5719a5f6ed3c2acf47e18b
Signed-off-by: Maheshwar Ajja <majja@codeaurora.org>
diff --git a/drivers/media/platform/msm/vidc/msm_vidc.c b/drivers/media/platform/msm/vidc/msm_vidc.c
index 98a59a5..427568c 100644
--- a/drivers/media/platform/msm/vidc/msm_vidc.c
+++ b/drivers/media/platform/msm/vidc/msm_vidc.c
@@ -429,7 +429,7 @@
 		print_vidc_buffer(VIDC_DBG, "release buf", inst, mbuf);
 		msm_comm_unmap_vidc_buffer(inst, mbuf);
 		list_del(&mbuf->list);
-		kfree(mbuf);
+		kref_put_mbuf(mbuf);
 	}
 	mutex_unlock(&inst->registeredbufs.lock);
 
@@ -998,7 +998,7 @@
 			}
 			msm_comm_unmap_vidc_buffer(inst, temp);
 			list_del(&temp->list);
-			kfree(temp);
+			kref_put_mbuf(temp);
 		}
 		mutex_unlock(&inst->registeredbufs.lock);
 	}
@@ -1075,10 +1075,16 @@
 				inst, vb2);
 		return;
 	}
+	if (!kref_get_mbuf(inst, mbuf)) {
+		dprintk(VIDC_ERR, "%s: mbuf not found\n", __func__);
+		return;
+	}
 
 	rc = msm_comm_qbuf(inst, mbuf);
 	if (rc)
 		print_vidc_buffer(VIDC_ERR, "failed qbuf", inst, mbuf);
+
+	kref_put_mbuf(mbuf);
 }
 
 static const struct vb2_ops msm_vidc_vb2q_ops = {
@@ -1620,7 +1626,7 @@
 		print_vidc_buffer(VIDC_ERR, "undequeud buf", inst, temp);
 		msm_comm_unmap_vidc_buffer(inst, temp);
 		list_del(&temp->list);
-		kfree(temp);
+		kref_put_mbuf(temp);
 	}
 	mutex_unlock(&inst->registeredbufs.lock);
 
diff --git a/drivers/media/platform/msm/vidc/msm_vidc_common.c b/drivers/media/platform/msm/vidc/msm_vidc_common.c
index 844a4e1..219f155 100644
--- a/drivers/media/platform/msm/vidc/msm_vidc_common.c
+++ b/drivers/media/platform/msm/vidc/msm_vidc_common.c
@@ -1492,6 +1492,7 @@
 		break;
 	case HAL_EVENT_RELEASE_BUFFER_REFERENCE:
 	{
+		struct msm_vidc_buffer *mbuf;
 		u32 planes[VIDEO_MAX_PLANES] = {0};
 
 		dprintk(VIDC_DBG,
@@ -1501,8 +1502,15 @@
 
 		planes[0] = event_notify->packet_buffer;
 		planes[1] = event_notify->extra_data_buffer;
-		handle_release_buffer_reference(inst, planes);
-
+		mbuf = msm_comm_get_buffer_using_device_planes(inst, planes);
+		if (!mbuf || !kref_get_mbuf(inst, mbuf)) {
+			dprintk(VIDC_ERR,
+				"%s: data_addr %x, extradata_addr %x not found\n",
+				__func__, planes[0], planes[1]);
+		} else {
+			handle_release_buffer_reference(inst, mbuf);
+			kref_put_mbuf(mbuf);
+		}
 		goto err_bad_event;
 	}
 	default:
@@ -2200,7 +2208,7 @@
 	planes[1] = empty_buf_done->extra_data_buffer;
 
 	mbuf = msm_comm_get_buffer_using_device_planes(inst, planes);
-	if (!mbuf) {
+	if (!mbuf || !kref_get_mbuf(inst, mbuf)) {
 		dprintk(VIDC_ERR,
 			"%s: data_addr %x, extradata_addr %x not found\n",
 			__func__, planes[0], planes[1]);
@@ -2242,12 +2250,11 @@
 	/*
 	 * put_buffer should be done before vb2_buffer_done else
 	 * client might queue the same buffer before it is unmapped
-	 * in put_buffer. also don't use mbuf after put_buffer
-	 * as it may be freed in put_buffer.
+	 * in put_buffer.
 	 */
 	msm_comm_put_vidc_buffer(inst, mbuf);
 	msm_comm_vb2_buffer_done(inst, vb);
-
+	kref_put_mbuf(mbuf);
 exit:
 	put_inst(inst);
 }
@@ -2326,7 +2333,7 @@
 	buffer_type = msm_comm_get_hal_output_buffer(inst);
 	if (fill_buf_done->buffer_type == buffer_type) {
 		mbuf = msm_comm_get_buffer_using_device_planes(inst, planes);
-		if (!mbuf) {
+		if (!mbuf || !kref_get_mbuf(inst, mbuf)) {
 			dprintk(VIDC_ERR,
 				"%s: data_addr %x, extradata_addr %x not found\n",
 				__func__, planes[0], planes[1]);
@@ -2429,6 +2436,7 @@
 	 */
 	msm_comm_put_vidc_buffer(inst, mbuf);
 	msm_comm_vb2_buffer_done(inst, vb);
+	kref_put_mbuf(mbuf);
 
 exit:
 	put_inst(inst);
@@ -4749,8 +4757,7 @@
 
 		/* remove from list */
 		list_del(&mbuf->list);
-		kfree(mbuf);
-		mbuf = NULL;
+		kref_put_mbuf(mbuf);
 	}
 	mutex_unlock(&inst->registeredbufs.lock);
 
@@ -5923,6 +5930,7 @@
 			rc = -ENOMEM;
 			goto exit;
 		}
+		kref_init(&mbuf->kref);
 	}
 
 	vbuf = to_vb2_v4l2_buffer(vb2);
@@ -5986,11 +5994,11 @@
 	return mbuf;
 
 exit:
-	mutex_unlock(&inst->registeredbufs.lock);
 	dprintk(VIDC_ERR, "%s: rc %d\n", __func__, rc);
 	msm_comm_unmap_vidc_buffer(inst, mbuf);
 	if (!found)
-		kfree(mbuf);
+		kref_put_mbuf(mbuf);
+	mutex_unlock(&inst->registeredbufs.lock);
 
 	return ERR_PTR(rc);
 }
@@ -6042,24 +6050,26 @@
 	 */
 	if (!mbuf->smem[0].refcount) {
 		list_del(&mbuf->list);
-		kfree(mbuf);
-		mbuf = NULL;
+		kref_put_mbuf(mbuf);
 	}
 unlock:
 	mutex_unlock(&inst->registeredbufs.lock);
 }
 
-void handle_release_buffer_reference(struct msm_vidc_inst *inst, u32 *planes)
+void handle_release_buffer_reference(struct msm_vidc_inst *inst,
+		struct msm_vidc_buffer *mbuf)
 {
 	int rc = 0;
-	struct msm_vidc_buffer *mbuf = NULL;
+	struct msm_vidc_buffer *temp;
 	bool found = false;
 	int i = 0;
 
 	mutex_lock(&inst->registeredbufs.lock);
 	found = false;
-	list_for_each_entry(mbuf, &inst->registeredbufs.list, list) {
-		if (msm_comm_compare_device_planes(mbuf, planes)) {
+	/* check if mbuf was not removed by any chance */
+	list_for_each_entry(temp, &inst->registeredbufs.list, list) {
+		if (msm_comm_compare_vb2_planes(inst, mbuf,
+				&temp->vvb.vb2_buf)) {
 			found = true;
 			break;
 		}
@@ -6077,13 +6087,10 @@
 		/* refcount is not zero if client queued the same buffer */
 		if (!mbuf->smem[0].refcount) {
 			list_del(&mbuf->list);
-			kfree(mbuf);
-			mbuf = NULL;
+			kref_put_mbuf(mbuf);
 		}
 	} else {
-		dprintk(VIDC_ERR,
-			"%s: data_addr %x extradata_addr %x not found\n",
-			__func__, planes[0], planes[1]);
+		print_vidc_buffer(VIDC_ERR, "mbuf not found", inst, mbuf);
 		goto unlock;
 	}
 
@@ -6097,8 +6104,9 @@
 	 *    and if found queue it to video hw (if not flushing).
 	 */
 	found = false;
-	list_for_each_entry(mbuf, &inst->registeredbufs.list, list) {
-		if (msm_comm_compare_device_plane(mbuf, planes, 0)) {
+	list_for_each_entry(temp, &inst->registeredbufs.list, list) {
+		if (msm_comm_compare_vb2_plane(inst, mbuf,
+				&temp->vvb.vb2_buf, 0)) {
 			found = true;
 			break;
 		}
@@ -6113,8 +6121,7 @@
 		msm_comm_unmap_vidc_buffer(inst, mbuf);
 		/* remove from list */
 		list_del(&mbuf->list);
-		kfree(mbuf);
-		mbuf = NULL;
+		kref_put_mbuf(mbuf);
 
 		/* don't queue the buffer */
 		found = false;
@@ -6161,3 +6168,41 @@
 	return rc;
 }
 
+static void kref_free_mbuf(struct kref *kref)
+{
+	struct msm_vidc_buffer *mbuf = container_of(kref,
+			struct msm_vidc_buffer, kref);
+
+	kfree(mbuf);
+}
+
+void kref_put_mbuf(struct msm_vidc_buffer *mbuf)
+{
+	if (!mbuf)
+		return;
+
+	kref_put(&mbuf->kref, kref_free_mbuf);
+}
+
+bool kref_get_mbuf(struct msm_vidc_inst *inst, struct msm_vidc_buffer *mbuf)
+{
+	struct msm_vidc_buffer *temp;
+	bool matches = false;
+	bool ret = false;
+
+	if (!inst || !mbuf)
+		return false;
+
+	mutex_lock(&inst->registeredbufs.lock);
+	list_for_each_entry(temp, &inst->registeredbufs.list, list) {
+		if (temp == mbuf) {
+			matches = true;
+			break;
+		}
+	}
+	ret = (matches && kref_get_unless_zero(&mbuf->kref)) ? true : false;
+	mutex_unlock(&inst->registeredbufs.lock);
+
+	return ret;
+}
+
diff --git a/drivers/media/platform/msm/vidc/msm_vidc_common.h b/drivers/media/platform/msm/vidc/msm_vidc_common.h
index bc881a0..18ba4a5 100644
--- a/drivers/media/platform/msm/vidc/msm_vidc_common.h
+++ b/drivers/media/platform/msm/vidc/msm_vidc_common.h
@@ -116,7 +116,8 @@
 		struct vb2_buffer *vb2);
 void msm_comm_put_vidc_buffer(struct msm_vidc_inst *inst,
 		struct msm_vidc_buffer *mbuf);
-void handle_release_buffer_reference(struct msm_vidc_inst *inst, u32 *planes);
+void handle_release_buffer_reference(struct msm_vidc_inst *inst,
+		struct msm_vidc_buffer *mbuf);
 int msm_comm_vb2_buffer_done(struct msm_vidc_inst *inst,
 		struct vb2_buffer *vb);
 int msm_comm_flush_vidc_buffer(struct msm_vidc_inst *inst,
@@ -145,4 +146,7 @@
 		struct vb2_buffer *vb2);
 void print_v4l2_buffer(u32 tag, const char *str, struct msm_vidc_inst *inst,
 		struct v4l2_buffer *v4l2);
+void kref_put_mbuf(struct msm_vidc_buffer *mbuf);
+bool kref_get_mbuf(struct msm_vidc_inst *inst, struct msm_vidc_buffer *mbuf);
+
 #endif
diff --git a/drivers/media/platform/msm/vidc/msm_vidc_internal.h b/drivers/media/platform/msm/vidc/msm_vidc_internal.h
index 49e6c3ec..e554a46 100644
--- a/drivers/media/platform/msm/vidc/msm_vidc_internal.h
+++ b/drivers/media/platform/msm/vidc/msm_vidc_internal.h
@@ -387,6 +387,7 @@
 
 struct msm_vidc_buffer {
 	struct list_head list;
+	struct kref kref;
 	struct msm_smem smem[VIDEO_MAX_PLANES];
 	struct vb2_v4l2_buffer vvb;
 	bool deferred;