habanalabs: check correct vmalloc return code

[ Upstream commit 0839152f8c1efc1cc2d515d1ff1e253ca9402ad3 ]

vmalloc can return different return code than NULL and a valid
pointer. We must validate it in order to dereference a non valid
pointer.

Signed-off-by: Ofir Bitton <obitton@habana.ai>
Signed-off-by: Oded Gabbay <oded.gabbay@gmail.com>
Signed-off-by: Sasha Levin <sashal@kernel.org>

diff --git a/drivers/misc/habanalabs/memory.c b/drivers/misc/habanalabs/memory.c
index 22566b75ca50c20e6ab2509b58b7130a2b01fafe..acfccf32be6b954a7e02f1b75d29ebf86f310ec8 100644 (file)
--- a/drivers/misc/habanalabs/memory.c
+++ b/drivers/misc/habanalabs/memory.c
@@ -67,6 +67,11 @@ static int alloc_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args,
        num_pgs = (args->alloc.mem_size + (page_size - 1)) >> page_shift;
        total_size = num_pgs << page_shift;
 
+       if (!total_size) {
+               dev_err(hdev->dev, "Cannot allocate 0 bytes\n");
+               return -EINVAL;
+       }
+
        contiguous = args->flags & HL_MEM_CONTIGUOUS;
 
        if (contiguous) {
@@ -94,7 +99,7 @@ static int alloc_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args,
        phys_pg_pack->contiguous = contiguous;
 
        phys_pg_pack->pages = kvmalloc_array(num_pgs, sizeof(u64), GFP_KERNEL);
-       if (!phys_pg_pack->pages) {
+       if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
                rc = -ENOMEM;
                goto pages_arr_err;
        }
@@ -689,7 +694,7 @@ static int init_phys_pg_pack_from_userptr(struct hl_ctx *ctx,
 
        phys_pg_pack->pages = kvmalloc_array(total_npages, sizeof(u64),
                                                GFP_KERNEL);
-       if (!phys_pg_pack->pages) {
+       if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
                rc = -ENOMEM;
                goto page_pack_arr_mem_err;
        }

diff --git a/drivers/misc/habanalabs/mmu.c b/drivers/misc/habanalabs/mmu.c
index 176c315836f128d00acab1a6deb42e97e00aeb93..d66e16de4cda32e9f7598ccd5de0cf0a778b1315 100644 (file)
--- a/drivers/misc/habanalabs/mmu.c
+++ b/drivers/misc/habanalabs/mmu.c
@@ -422,7 +422,7 @@ int hl_mmu_init(struct hl_device *hdev)
        hdev->mmu_shadow_hop0 = kvmalloc_array(prop->max_asid,
                                        prop->mmu_hop_table_size,
                                        GFP_KERNEL | __GFP_ZERO);
-       if (!hdev->mmu_shadow_hop0) {
+       if (ZERO_OR_NULL_PTR(hdev->mmu_shadow_hop0)) {
                rc = -ENOMEM;
                goto err_pool_add;
        }