[PATCH v5 1/2] drm/rocket: Check allocations before use

[Triet Hoang <triet.hoang.dev@gmail.com>]

Check the result of kvmalloc_array() in rocket_job_push() and
kmalloc_objs() in rocket_job_open() before using the allocated
buffers.

Fixes: 0810d5ad88a1 ("accel/rocket: Add job submission IOCTL")

Signed-off-by: Triet Hoang <triet.hoang.dev@gmail.com>
---
Changes in v2:
- Free scheds when drm_sched_entity_init() fails.
- Initialize ret to 0.

Changes in v3:
- Move patch version descriptions below the '---' marker.

Changes in v4:
- Remove unnecessary initialization of ret to 0.
- Adjust commit message word wrapping.
- Add Fixes tag.

Changes in v5:
- Add check overflow before kvmalloc_array() in rocket_job_push().
---
 drivers/accel/rocket/rocket_job.c | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/drivers/accel/rocket/rocket_job.c b/drivers/accel/rocket/rocket_job.c
index ac51bff39833..15e93c355c38 100644
--- a/drivers/accel/rocket/rocket_job.c
+++ b/drivers/accel/rocket/rocket_job.c
@@ -189,9 +189,15 @@ static int rocket_job_push(struct rocket_job *job)
 	struct drm_gem_object **bos;
 	struct ww_acquire_ctx acquire_ctx;
 	int ret = 0;
+	size_t bos_count;
+
+	if (check_add_overflow(job->in_bo_count, job->out_bo_count, &bos_count))
+		return -EOVERFLOW;
+
+	bos = kvmalloc_array(bos_count, sizeof(void *), GFP_KERNEL);
+	if (!bos)
+		return -ENOMEM;
 
-	bos = kvmalloc_array(job->in_bo_count + job->out_bo_count, sizeof(void *),
-			     GFP_KERNEL);
 	memcpy(bos, job->in_bos, job->in_bo_count * sizeof(void *));
 	memcpy(&bos[job->in_bo_count], job->out_bos, job->out_bo_count * sizeof(void *));
 
@@ -501,6 +507,9 @@ int rocket_job_open(struct rocket_file_priv *rocket_priv)
 	unsigned int core;
 	int ret;
 
+	if (!scheds)
+		return -ENOMEM;
+
 	for (core = 0; core < rdev->num_cores; core++)
 		scheds[core] = &rdev->cores[core].sched;
 
@@ -509,9 +518,9 @@ int rocket_job_open(struct rocket_file_priv *rocket_priv)
 				    scheds,
 				    rdev->num_cores, NULL);
 	if (WARN_ON(ret))
-		return ret;
+		kfree(scheds);
 
-	return 0;
+	return ret;
 }
 
 void rocket_job_close(struct rocket_file_priv *rocket_priv)
-- 
2.53.0