[PATCH] lkdtm: heap: Add missing NULL checks after memory allocation

[Jiangshan Yi]

lkdtm_VMALLOC_LINEAR_OVERFLOW() and lkdtm_SLAB_FREE_PAGE() do not check
the return value of vzalloc() / __get_free_page(). If the allocation
fails, the test dereferences a NULL pointer and crashes in an
unintended way rather than testing the intended overflow or invalid
free scenario.

Add the missing checks. In lkdtm_VMALLOC_LINEAR_OVERFLOW(), the checks
are placed after OPTIMIZER_HIDE_VAR() so the compiler cannot use the
non-NULL proof to defeat the variable-hiding.

Other functions in the same file (e.g. lkdtm_WRITE_BUDDY_AFTER_FREE,
lkdtm_READ_AFTER_FREE) already check for allocation failure, so this
makes the error handling consistent.

Signed-off-by: Jiangshan Yi <yijiangshan@xxxxxxxxxx>
---
 drivers/misc/lkdtm/heap.c | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/drivers/misc/lkdtm/heap.c b/drivers/misc/lkdtm/heap.c
index c1a05b935894..76edc0d8124b 100644
--- a/drivers/misc/lkdtm/heap.c
+++ b/drivers/misc/lkdtm/heap.c
@@ -32,8 +32,14 @@ static void lkdtm_VMALLOC_LINEAR_OVERFLOW(void)
 	char *one, *two;
 
 	one = vzalloc(PAGE_SIZE);
 	OPTIMIZER_HIDE_VAR(one);
+	if (!one)
+		return;
 	two = vzalloc(PAGE_SIZE);
+	if (!two) {
+		vfree(one);
+		return;
+	}
 
 	pr_info("Attempting vmalloc linear overflow ...\n");
 	memset(one, 0xAA, PAGE_SIZE + __offset);
@@ -350,6 +356,11 @@ static void lkdtm_SLAB_FREE_PAGE(void)
 {
 	unsigned long p = __get_free_page(GFP_KERNEL);
 
+	if (!p) {
+		pr_info("Unable to allocate free page\n");
+		return;
+	}
+
 	pr_info("Attempting non-Slab slab free ...\n");
 	kmem_cache_free(NULL, (void *)p);
 	free_page(p);
-- 
2.25.1