From f087e771cb27c19b21de65fd0c7d7876cd2d8d54 Mon Sep 17 00:00:00 2001
From: komark06 <87216234+komark06@users.noreply.github.com>
Date: Sun, 17 Mar 2024 19:57:36 +0800
Subject: [PATCH] Hook calloc for test harness (#172)

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 | 39 +++++++++++++++++++++++++++++----------
 harness.h |  1 +
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/harness.c b/harness.c
index 475a89455..4e02891f5 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,
+    TEST_CALLOC,
+} alloc_t;
+
 /* Internal functions */
 
 /* Should this allocation fail? */
@@ -115,17 +122,23 @@ 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_t alloc_type, size_t size)
 {
     if (noallocate_mode) {
-        report_event(MSG_FATAL, "Calls to malloc disallowed");
+        char *msg_alloc_forbidden[] = {
+            "Calls to malloc are disallowed",
+            "Calls to calloc are disallowed",
+        };
+        report_event(MSG_FATAL, "%s", msg_alloc_forbidden[alloc_type]);
         return NULL;
     }
 
     if (fail_allocation()) {
-        report_event(MSG_WARN, "Malloc returning NULL");
+        char *msg_alloc_failure[] = {
+            "Malloc returning NULL",
+            "Calloc returning NULL",
+        };
+        report_event(MSG_WARN, "%s", msg_alloc_failure[alloc_type]);
         return NULL;
     }
 
@@ -142,7 +155,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, !alloc_type * FILLCHAR, size);
     // cppcheck-suppress nullPointerRedundantCheck
     new_block->next = allocated;
     // cppcheck-suppress nullPointerRedundantCheck
@@ -156,16 +169,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 */