From 3e01ac9b04f2bf7da1b1be46cc502b8f0423b08d Mon Sep 17 00:00:00 2001
From: Colin Leroy-Mira <colin@colino.net>
Date: Mon, 15 Jan 2024 20:30:20 +0100
Subject: [PATCH] Fix malloc and realloc overflow

If user requests a size >= 65532, adding the heap admin size
overflows size. Fixes #2358.
---
 libsrc/common/malloc.s        |  2 +-
 libsrc/common/realloc.c       |  8 ++--
 test/val/lib_common_malloc.c  | 34 +++++++++++++++
 test/val/lib_common_realloc.c | 81 +++++++++++++++++++++++++++++++++++
 4 files changed, 121 insertions(+), 4 deletions(-)
 create mode 100644 test/val/lib_common_malloc.c
 create mode 100644 test/val/lib_common_realloc.c

diff --git a/libsrc/common/malloc.s b/libsrc/common/malloc.s
index 6872f1f2e..72c4aedaa 100644
--- a/libsrc/common/malloc.s
+++ b/libsrc/common/malloc.s
@@ -131,6 +131,7 @@ _malloc:
         sta     ptr1
         bcc     @L1
         inc     ptr1+1
+        beq     OutOfHeapSpace          ; if high byte's 0, we overflowed!
 @L1:    ldx     ptr1+1
         bne     @L2
         cmp     #HEAP_MIN_BLOCKSIZE+1
@@ -336,4 +337,3 @@ RetUserPtr:
         bcc     @L9
         inx
 @L9:    rts
-
diff --git a/libsrc/common/realloc.c b/libsrc/common/realloc.c
index eeb1eeea5..b5429b3c2 100644
--- a/libsrc/common/realloc.c
+++ b/libsrc/common/realloc.c
@@ -59,6 +59,11 @@ void* __fastcall__ realloc (void* block, register size_t size)
         return 0;
     }
 
+    /* Don't overflow! */
+    if (size > 0xFFFF - HEAP_ADMIN_SPACE) {
+      return 0;
+    }
+
     /* Make the internal used size from the given size */
     size += HEAP_ADMIN_SPACE;
     if (size < sizeof (struct freeblock)) {
@@ -107,6 +112,3 @@ void* __fastcall__ realloc (void* block, register size_t size)
     }
     return newblock;
 }
-
-
-
diff --git a/test/val/lib_common_malloc.c b/test/val/lib_common_malloc.c
new file mode 100644
index 000000000..5c68dc39a
--- /dev/null
+++ b/test/val/lib_common_malloc.c
@@ -0,0 +1,34 @@
+#include <errno.h>
+#include <stdlib.h>
+#include <string.h>
+#include "unittest.h"
+
+TEST
+{
+    char *buf, *buf2;
+    unsigned int i;
+
+    buf = malloc(0);
+    ASSERT_IsTrue (buf == NULL, "malloc (0) returned something");
+
+    for (i = 1; i < 10; i++) {
+      buf = malloc(i);
+      ASSERT_IsTrue (buf != NULL, "small returned nothing");
+    }
+
+    buf = malloc(4096);
+    ASSERT_IsTrue (buf != NULL, "malloc (4096) returned nothing");
+
+    buf = malloc(61000UL);
+    ASSERT_IsTrue (buf == NULL, "malloc (61000) returned something");
+
+    for (i = 65535UL; i > _heapmaxavail(); i--) {
+      buf = malloc(i);
+      ASSERT_IsTrue (buf == NULL, "malloc returned something but shouldn't have");
+    }
+
+    buf = malloc(i);
+    ASSERT_IsTrue (buf != NULL, "malloc returned nothing but should have");
+    ASSERT_IsTrue(_heapmaxavail() == 0, "heapmaxavail should be 0");
+}
+ENDTEST
diff --git a/test/val/lib_common_realloc.c b/test/val/lib_common_realloc.c
new file mode 100644
index 000000000..d1e4fa3eb
--- /dev/null
+++ b/test/val/lib_common_realloc.c
@@ -0,0 +1,81 @@
+#include <errno.h>
+#include <stdlib.h>
+#include <string.h>
+#include "unittest.h"
+
+TEST
+{
+    char *buf, *buf2;
+    unsigned int i;
+
+    buf = realloc(NULL, 0);
+    ASSERT_IsTrue (buf == NULL, "realloc (NULL, 0) returned something");
+
+    for (i = 1; i < 10; i++) {
+      buf2 = realloc(buf, i);
+      ASSERT_IsTrue (buf2 != NULL, "small realloc returned nothing");
+      if (i > 1) {
+        ASSERT_IsTrue (buf2 == buf, "buf shouldn't have moved");
+      }
+      buf = buf2;
+    }
+
+    buf = realloc(NULL, 15);
+    ASSERT_IsTrue (buf != NULL, "realloc (NULL, 15) returned nothing");
+
+    buf = realloc(buf, 0);
+    ASSERT_IsTrue (buf == NULL, "realloc (buf, 0) returned something");
+
+    buf = realloc(buf, 32);
+    memset(buf, 'a', 32);
+    for (i = 0; i < 32; i++) {
+      ASSERT_IsTrue(buf[i] == 'a', "wrong contents in buf");
+    }
+
+    /* Now realloc larger, while there's nothing else in the heap */
+    buf = realloc(buf, 64);
+    memset(buf+32, 'b', 32);
+    for (i = 0; i < 32; i++) {
+      ASSERT_IsTrue(buf[i] == 'a', "wrong contents in start of buf");
+    }
+    for (i = 32; i < 64; i++) {
+      ASSERT_IsTrue(buf[i] == 'b', "wrong contents in end of buf");
+    }
+
+    /* Now realloc smaller, while there's nothing else in the heap */
+    buf = realloc(buf, 40);
+    for (i = 0; i < 32; i++) {
+      ASSERT_IsTrue(buf[i] == 'a', "wrong contents in start of buf");
+    }
+    for (i = 32; i < 40; i++) {
+      ASSERT_IsTrue(buf[i] == 'b', "wrong contents in end of buf");
+    }
+
+    /* Allocate something else, so next realloc has to change block */
+    malloc(50);
+
+    /* Now realloc larger, with something else in the heap */
+    buf = realloc(buf, 128);
+    for (i = 0; i < 32; i++) {
+      ASSERT_IsTrue(buf[i] == 'a', "wrong contents in start of buf");
+    }
+    for (i = 32; i < 40; i++) {
+      ASSERT_IsTrue(buf[i] == 'b', "wrong contents in end of buf");
+    }
+
+    for (i = 129; i < 8192; i++) {
+      buf = realloc(buf, i);
+      ASSERT_IsTrue(buf != NULL, "realloc failed");
+    }
+
+    malloc(4096);
+
+    buf2 = realloc(buf, 58000UL);
+    ASSERT_IsTrue (buf2 == NULL, "realloc (buf, 58000) returned something");
+
+    for (i = 65535UL; i > 65527UL; i--) {
+      buf2 = realloc(buf, i);
+      ASSERT_IsTrue (buf2 == NULL, "realloc returned something but shouldn't have");
+    }
+}
+ENDTEST