[malloc] Avoid integer overflow for excessively large memory allocations

Signed-off-by: Michael Brown <mcb30@ipxe.org>

1 files changed, 49 insertions, 48 deletions

diff --git a/src/core/malloc.c b/src/core/malloc.c
index b120c032..d7c67823 100644
--- a/src/core/malloc.c
+++ b/src/core/malloc.c
@@ -275,7 +275,7 @@ void * alloc_memblock ( size_t size, size_t align, size_t offset ) {
 	size_t align_mask;
 	size_t actual_size;
 	size_t pre_size;
-	ssize_t post_size;
+	size_t post_size;
 	struct memory_block *pre;
 	struct memory_block *post;
 	void *ptr;
@@ -291,7 +291,9 @@ void * alloc_memblock ( size_t size, size_t align, size_t offset ) {
 	 */
 	actual_size = ( ( size + MIN_MEMBLOCK_SIZE - 1 ) &
 			~( MIN_MEMBLOCK_SIZE - 1 ) );
+	assert ( actual_size >= size );
 	align_mask = ( ( align - 1 ) | ( MIN_MEMBLOCK_SIZE - 1 ) );
+	assert ( ( actual_size + align_mask ) > actual_size );
 
 	DBGC2 ( &heap, "Allocating %#zx (aligned %#zx+%zx)\n",
 		size, align, offset );
@@ -300,55 +302,54 @@ void * alloc_memblock ( size_t size, size_t align, size_t offset ) {
 		list_for_each_entry ( block, &free_blocks, list ) {
 			pre_size = ( ( offset - virt_to_phys ( block ) )
 				     & align_mask );
+			if ( block->size < ( pre_size + actual_size ) )
+				continue;
 			post_size = ( block->size - pre_size - actual_size );
-			if ( post_size >= 0 ) {
-				/* Split block into pre-block, block, and
-				 * post-block.  After this split, the "pre"
-				 * block is the one currently linked into the
-				 * free list.
-				 */
-				pre   = block;
-				block = ( ( ( void * ) pre   ) + pre_size );
-				post  = ( ( ( void * ) block ) + actual_size );
-				DBGC2 ( &heap, "[%p,%p) -> [%p,%p) + [%p,%p)\n",
-					pre, ( ( ( void * ) pre ) + pre->size ),
-					pre, block, post,
-					( ( ( void * ) pre ) + pre->size ) );
-				/* If there is a "post" block, add it in to
-				 * the free list.  Leak it if it is too small
-				 * (which can happen only at the very end of
-				 * the heap).
-				 */
-				if ( (size_t) post_size >= MIN_MEMBLOCK_SIZE ) {
-					VALGRIND_MAKE_MEM_UNDEFINED
-						( post, sizeof ( *post ) );
-					post->size = post_size;
-					list_add ( &post->list, &pre->list );
-				}
-				/* Shrink "pre" block, leaving the main block
-				 * isolated and no longer part of the free
-				 * list.
-				 */
-				pre->size = pre_size;
-				/* If there is no "pre" block, remove it from
-				 * the list.  Also remove it (i.e. leak it) if
-				 * it is too small, which can happen only at
-				 * the very start of the heap.
-				 */
-				if ( pre_size < MIN_MEMBLOCK_SIZE ) {
-					list_del ( &pre->list );
-					VALGRIND_MAKE_MEM_NOACCESS
-						( pre, sizeof ( *pre ) );
-				}
-				/* Update total free memory */
-				freemem -= actual_size;
-				/* Return allocated block */
-				DBGC2 ( &heap, "Allocated [%p,%p)\n", block,
-					( ( ( void * ) block ) + size ) );
-				ptr = block;
-				VALGRIND_MAKE_MEM_UNDEFINED ( ptr, size );
-				goto done;
+			/* Split block into pre-block, block, and
+			 * post-block.  After this split, the "pre"
+			 * block is the one currently linked into the
+			 * free list.
+			 */
+			pre   = block;
+			block = ( ( ( void * ) pre   ) + pre_size );
+			post  = ( ( ( void * ) block ) + actual_size );
+			DBGC2 ( &heap, "[%p,%p) -> [%p,%p) + [%p,%p)\n", pre,
+				( ( ( void * ) pre ) + pre->size ), pre, block,
+				post, ( ( ( void * ) pre ) + pre->size ) );
+			/* If there is a "post" block, add it in to
+			 * the free list.  Leak it if it is too small
+			 * (which can happen only at the very end of
+			 * the heap).
+			 */
+			if ( post_size >= MIN_MEMBLOCK_SIZE ) {
+				VALGRIND_MAKE_MEM_UNDEFINED ( post,
+							      sizeof ( *post ));
+				post->size = post_size;
+				list_add ( &post->list, &pre->list );
 			}
+			/* Shrink "pre" block, leaving the main block
+			 * isolated and no longer part of the free
+			 * list.
+			 */
+			pre->size = pre_size;
+			/* If there is no "pre" block, remove it from
+			 * the list.  Also remove it (i.e. leak it) if
+			 * it is too small, which can happen only at
+			 * the very start of the heap.
+			 */
+			if ( pre_size < MIN_MEMBLOCK_SIZE ) {
+				list_del ( &pre->list );
+				VALGRIND_MAKE_MEM_NOACCESS ( pre,
+							     sizeof ( *pre ) );
+			}
+			/* Update total free memory */
+			freemem -= actual_size;
+			/* Return allocated block */
+			DBGC2 ( &heap, "Allocated [%p,%p)\n", block,
+				( ( ( void * ) block ) + size ) );
+			ptr = block;
+			VALGRIND_MAKE_MEM_UNDEFINED ( ptr, size );
+			goto done;
 		}
 
 		/* Try discarding some cached data to free up memory */