From f436315e02102cf25e37175b3485f5bcc86bc7c8 Mon Sep 17 00:00:00 2001
From: Nicolas Pitre <npitre@baylibre.com>
Date: Tue, 19 Jan 2021 14:38:56 -0500
Subject: [PATCH] lib/os/heap: guard against arithmetic overflows

Let's do it upfront only once for each entry point and dispense
with overflow checks later to keep the code simple.

Signed-off-by: Nicolas Pitre <npitre@baylibre.com>
---
 lib/os/heap.c | 14 ++++++++++----
 lib/os/heap.h | 10 ++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lib/os/heap.c b/lib/os/heap.c
index a0cc7c87369..7bb3635a709 100644
--- a/lib/os/heap.c
+++ b/lib/os/heap.c
@@ -223,11 +223,12 @@ static chunkid_t alloc_chunk(struct z_heap *h, size_t sz)
 
 void *sys_heap_alloc(struct sys_heap *heap, size_t bytes)
 {
-	if (bytes == 0U) {
+	struct z_heap *h = heap->heap;
+
+	if (bytes == 0U || size_too_big(h, bytes)) {
 		return NULL;
 	}
 
-	struct z_heap *h = heap->heap;
 	size_t chunk_sz = bytes_to_chunksz(h, bytes);
 	chunkid_t c = alloc_chunk(h, chunk_sz);
 	if (c == 0U) {
@@ -253,7 +254,7 @@ void *sys_heap_aligned_alloc(struct sys_heap *heap, size_t align, size_t bytes)
 	if (align <= chunk_header_bytes(h)) {
 		return sys_heap_alloc(heap, bytes);
 	}
-	if (bytes == 0) {
+	if (bytes == 0 || size_too_big(h, bytes)) {
 		return NULL;
 	}
 
@@ -297,6 +298,8 @@ void *sys_heap_aligned_alloc(struct sys_heap *heap, size_t align, size_t bytes)
 
 void *sys_heap_realloc(struct sys_heap *heap, void *ptr, size_t bytes)
 {
+	struct z_heap *h = heap->heap;
+
 	/* special realloc semantics */
 	if (ptr == NULL) {
 		return sys_heap_alloc(heap, bytes);
@@ -306,7 +309,10 @@ void *sys_heap_realloc(struct sys_heap *heap, void *ptr, size_t bytes)
 		return NULL;
 	}
 
-	struct z_heap *h = heap->heap;
+	if (size_too_big(h, bytes)) {
+		return NULL;
+	}
+
 	chunkid_t c = mem_to_chunkid(h, ptr);
 	chunkid_t rc = right_chunk(h, c);
 	size_t chunks_need = bytes_to_chunksz(h, bytes);
diff --git a/lib/os/heap.h b/lib/os/heap.h
index 5ee490e88e1..d2b0d6d049c 100644
--- a/lib/os/heap.h
+++ b/lib/os/heap.h
@@ -229,6 +229,16 @@ static inline int bucket_idx(struct z_heap *h, size_t sz)
 	return 31 - __builtin_clz(usable_sz);
 }
 
+static inline bool size_too_big(struct z_heap *h, size_t bytes)
+{
+	/*
+	 * Quick check to bail out early if size is too big.
+	 * Also guards against potential arithmetic overflows elsewhere.
+	 * There is a minimum of one chunk always in use by the heap header.
+	 */
+	return (bytes / CHUNK_UNIT) >= h->len;
+}
+
 /* For debugging */
 void heap_dump(struct z_heap *h);