ntdll: Avoid overflows in heap allocations. Based on a patch by Rob Shearman.

dlls/kernel32/tests/heap.c

@@ -61,6 +61,17 @@ START_TEST(heap)
         HeapFree(GetProcessHeap(), 0, mem);
     }
 
+    /* test some border cases of HeapAlloc and HeapReAlloc */
+    mem = HeapAlloc(GetProcessHeap(), 0, 0);
+    ok(mem != NULL, "memory not allocated for size 0\n");
+    msecond = HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, mem, ~0UL - 7);
+    ok(msecond == NULL, "HeapReAlloc(0xfffffff8) should have failed\n");
+    msecond = HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, mem, ~0UL);
+    ok(msecond == NULL, "HeapReAlloc(0xffffffff) should have failed\n");
+    HeapFree(GetProcessHeap(), 0, mem);
+    mem = HeapAlloc(GetProcessHeap(), 0, ~0UL);
+    ok(mem == NULL, "memory allocated for size ~0UL\n");
+
     /* Global*() functions */
     gbl = GlobalAlloc(GMEM_MOVEABLE, 0);
     ok(gbl != NULL, "global memory not allocated for size 0\n");

dlls/ntdll/heap.c

@@ -735,6 +735,7 @@ static ARENA_FREE *HEAP_FindFreeBlock( HEAP *heap, SIZE_T size,
 {
     SUBHEAP *subheap;
     struct list *ptr;
+    SIZE_T total_size;
     FREE_LIST_ENTRY *pEntry = heap->freeList + get_freelist_index( size + sizeof(ARENA_INUSE) );
 
     /* Find a suitable free list, and in it find a block large enough */
@@ -766,13 +767,15 @@ static ARENA_FREE *HEAP_FindFreeBlock( HEAP *heap, SIZE_T size,
      * So just one heap struct, one first free arena which will eventually
      * get used, and a second free arena that might get assigned all remaining
      * free space in HEAP_ShrinkBlock() */
-    size += ROUND_SIZE(sizeof(SUBHEAP)) + sizeof(ARENA_INUSE) + sizeof(ARENA_FREE);
-    if (!(subheap = HEAP_CreateSubHeap( heap, NULL, heap->flags, size,
-                                        max( HEAP_DEF_SIZE, size ) )))
+    total_size = size + ROUND_SIZE(sizeof(SUBHEAP)) + sizeof(ARENA_INUSE) + sizeof(ARENA_FREE);
+    if (total_size < size) return NULL;  /* overflow */
+
+    if (!(subheap = HEAP_CreateSubHeap( heap, NULL, heap->flags, total_size,
+                                        max( HEAP_DEF_SIZE, total_size ) )))
         return NULL;
 
     TRACE("created new sub-heap %p of %08lx bytes for heap %p\n",
-            subheap, size, heap );
+            subheap, total_size, heap );
 
     *ppSubHeap = subheap;
     return (ARENA_FREE *)(subheap + 1);
@@ -1178,6 +1181,11 @@ PVOID WINAPI RtlAllocateHeap( HANDLE heap, ULONG flags, SIZE_T size )
     flags &= HEAP_GENERATE_EXCEPTIONS | HEAP_NO_SERIALIZE | HEAP_ZERO_MEMORY;
     flags |= heapPtr->flags;
     rounded_size = ROUND_SIZE(size);
+    if (rounded_size < size)  /* overflow */
+    {
+        if (flags & HEAP_GENERATE_EXCEPTIONS) RtlRaiseStatus( STATUS_NO_MEMORY );
+        return NULL;
+    }
     if (rounded_size < HEAP_MIN_DATA_SIZE) rounded_size = HEAP_MIN_DATA_SIZE;
 
     if (!(flags & HEAP_NO_SERIALIZE)) RtlEnterCriticalSection( &heapPtr->critSection );
@@ -1320,6 +1328,12 @@ PVOID WINAPI RtlReAllocateHeap( HANDLE heap, ULONG flags, PVOID ptr, SIZE_T size
              HEAP_REALLOC_IN_PLACE_ONLY;
     flags |= heapPtr->flags;
     rounded_size = ROUND_SIZE(size);
+    if (rounded_size < size)  /* overflow */
+    {
+        if (flags & HEAP_GENERATE_EXCEPTIONS) RtlRaiseStatus( STATUS_NO_MEMORY );
+        RtlSetLastWin32ErrorAndNtStatusFromNtStatus( STATUS_NO_MEMORY );
+        return NULL;
+    }
     if (rounded_size < HEAP_MIN_DATA_SIZE) rounded_size = HEAP_MIN_DATA_SIZE;
 
     if (!(flags & HEAP_NO_SERIALIZE)) RtlEnterCriticalSection( &heapPtr->critSection );