From 2d455fa7f6f429c7fab7dd9448839e56a9465207 Mon Sep 17 00:00:00 2001
From: komark06 <komark06@gmail.com>
Date: Sat, 16 Mar 2024 12:19:15 +0800
Subject: [PATCH] Introduce test_calloc and alloc

This commit introduces a generic alloc function to handle both
test_malloc and test_calloc operations based on the provided
alloc_func_t. Additionally, it implements the test_calloc function to
allocate memory and initialize it to zero.
---
 harness.c | 40 ++++++++++++++++++++++++++++++----------
 harness.h |  1 +
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/harness.c b/harness.c
index 475a89455..c2b32ec9c 100644
--- a/harness.c
+++ b/harness.c
@@ -2,6 +2,7 @@
 
 #include <setjmp.h>
 #include <signal.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -58,6 +59,12 @@ static jmp_buf env;
 static volatile sig_atomic_t jmp_ready = false;
 static bool time_limited = false;
 
+/* For test_malloc and test_calloc */
+typedef enum {
+    TEST_MALLOC = FILLCHAR,
+    TEST_CALLOC = 0,
+} alloc_func_t;
+
 /* Internal functions */
 
 /* Should this allocation fail? */
@@ -115,17 +122,24 @@ static size_t *find_footer(block_element_t *b)
     return p;
 }
 
-/* Implementation of application functions */
-
-void *test_malloc(size_t size)
+static void *alloc(alloc_func_t func_name, size_t size)
 {
+    char *no_alloc_msg, *fail_alloc_msg;
+    if (func_name == TEST_MALLOC) {
+        no_alloc_msg = "Calls to malloc disallowed";
+        fail_alloc_msg = "Malloc returning NULL";
+    } else {
+        no_alloc_msg = "Calls to calloc disallowed";
+        fail_alloc_msg = "Calloc returning NULL";
+    }
+
     if (noallocate_mode) {
-        report_event(MSG_FATAL, "Calls to malloc disallowed");
+        report_event(MSG_FATAL, "%s", no_alloc_msg);
         return NULL;
     }
 
     if (fail_allocation()) {
-        report_event(MSG_WARN, "Malloc returning NULL");
+        report_event(MSG_WARN, "%s", fail_alloc_msg);
         return NULL;
     }
 
@@ -142,7 +156,7 @@ void *test_malloc(size_t size)
     new_block->payload_size = size;
     *find_footer(new_block) = MAGICFOOTER;
     void *p = (void *) &new_block->payload;
-    memset(p, FILLCHAR, size);
+    memset(p, func_name, size);
     // cppcheck-suppress nullPointerRedundantCheck
     new_block->next = allocated;
     // cppcheck-suppress nullPointerRedundantCheck
@@ -156,16 +170,22 @@ void *test_malloc(size_t size)
     return p;
 }
 
+/* Implementation of application functions */
+
+void *test_malloc(size_t size)
+{
+    return alloc(TEST_MALLOC, size);
+}
+
 // cppcheck-suppress unusedFunction
 void *test_calloc(size_t nelem, size_t elsize)
 {
     /* Reference: Malloc tutorial
      * https://danluu.com/malloc-tutorial/
      */
-    size_t size = nelem * elsize;  // TODO: check for overflow
-    void *ptr = test_malloc(size);
-    memset(ptr, 0, size);
-    return ptr;
+    if (!nelem || !elsize || nelem > SIZE_MAX / elsize)
+        return NULL;
+    return alloc(TEST_CALLOC, nelem * elsize);
 }
 
 void test_free(void *p)
diff --git a/harness.h b/harness.h
index 8696b630e..3e2ccec40 100644
--- a/harness.h
+++ b/harness.h
@@ -55,6 +55,7 @@ void trigger_exception(char *msg);
 
 /* Tested program use our versions of malloc and free */
 #define malloc test_malloc
+#define calloc test_calloc
 #define free test_free
 
 /* Use undef to avoid strdup redefined error */