From 095e05e927d44b323ba85cccd5ca031ef2dc5fe0 Mon Sep 17 00:00:00 2001
From: Jordan Bancino <jordan@bancino.net>
Date: Sun, 4 Jun 2023 18:44:37 +0000
Subject: [PATCH] Make Memory use a recursive mutex.

This allows some additional operations to be performed inside the memory
hooks, although it's still a bad idea to allocate or free memory while
inside the hook.
---
 Cytoplasm/src/Memory.c | 32 +++++++++++++++++++++++++--
 Cytoplasm/src/RtStub.c | 50 ++++++++++++++++++++++++++++++++----------
 2 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/Cytoplasm/src/Memory.c b/Cytoplasm/src/Memory.c
index 99fcb25..7fb2f64 100644
--- a/Cytoplasm/src/Memory.c
+++ b/Cytoplasm/src/Memory.c
@@ -57,14 +57,42 @@ struct MemoryInfo
 #define MEM_BOUND_UPPER(p, x) *(((MEM_BOUND_TYPE *) (((UInt8 *) p) + x)) + 1)
 #define MEM_SIZE_ACTUAL(x) (((x) * sizeof(UInt8)) + (2 * sizeof(MEM_BOUND_TYPE)))
 
-static pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
-static void (*hook) (MemoryAction, MemoryInfo *, void *) = NULL;
+static pthread_mutex_t lock;
+static void (*hook) (MemoryAction, MemoryInfo *, void *) = MemoryDefaultHook;
 static void *hookArgs = NULL;
 
 static MemoryInfo **allocations = NULL;
 static size_t allocationsSize = 0;
 static size_t allocationsLen = 0;
 
+int
+MemoryRuntimeInit(void)
+{
+    pthread_mutexattr_t attr;
+    int ret = 0;
+
+    if (pthread_mutexattr_init(&attr) != 0)
+    {
+        goto finish;
+    }
+
+    pthread_mutexattr_settype(&attr, PTHREAD_MUTEX_RECURSIVE);
+    ret = pthread_mutex_init(&lock, &attr);
+    pthread_mutexattr_destroy(&attr);
+
+    ret = (ret == 0);
+
+finish:
+    return ret;
+}
+
+int
+MemoryRuntimeDestroy(void)
+{
+    MemoryFreeAll();
+    return pthread_mutex_destroy(&lock) == 0;
+}
+
 static size_t
 MemoryHash(void *p)
 {
diff --git a/Cytoplasm/src/RtStub.c b/Cytoplasm/src/RtStub.c
index a38620a..6eac77f 100644
--- a/Cytoplasm/src/RtStub.c
+++ b/Cytoplasm/src/RtStub.c
@@ -33,12 +33,18 @@
 #include <string.h>
 #include <pthread.h>
 
+#define ERR_BUF_SIZE 128
+
 /* Specified by POSIX to contain environment variables */
 extern char **environ;
 
 /* The linking program is expected to provide Main */
 extern int Main(Array *, HashMap *);
 
+/* Functions in the Memory API that are not exported via header */
+extern int MemoryRuntimeInit(void);
+extern int MemoryRuntimeDestroy(void);
+
 typedef struct MainArgs
 {
     Array *args;
@@ -69,32 +75,46 @@ main(int argc, char **argv)
 
     MainArgs args;
 
-    MemoryHook(MemoryDefaultHook, NULL);
+    char errBuf[ERR_BUF_SIZE];
+    int errLen;
 
     args.args = NULL;
     args.env = NULL;
     args.ret = EXIT_FAILURE;
 
+    if (!MemoryRuntimeInit())
+    {
+	errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to initialize Memory runtime.\n");
+        write(STDERR_FILENO, errBuf, errLen);
+        goto finish;
+    }
+
     args.args = ArrayCreate();
 
     if (!args.args)
     {
-        Log(LOG_ERR, "Bootstrap error: Unable to allocate memory for arguments.");
-        args.ret = EXIT_FAILURE;
+        errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to allocate heap memory for program arguments.\n");
+        write(STDERR_FILENO, errBuf, errLen);
         goto finish;
     }
 
     args.env = HashMapCreate();
     if (!args.env)
     {
-        Log(LOG_ERR, "Bootstrap error: Unable to allocate memory for environment.");
-        args.ret = EXIT_FAILURE;
+        errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to allocate heap memory for program environment.\n");
+        write(STDERR_FILENO, errBuf, errLen);
         goto finish;
     }
 
     for (i = 0; i < (size_t) argc; i++)
     {
-        ArrayAdd(args.args, StrDuplicate(argv[i]));
+        char *arg = StrDuplicate(argv[i]);
+        if (!arg || !ArrayAdd(args.args, arg))
+        {
+            errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to allocate heap memory for program argument.\n");
+            write(STDERR_FILENO, errBuf, errLen);
+            goto finish;
+        }
     }
 
     envp = environ;
@@ -105,6 +125,14 @@ main(int argc, char **argv)
         /* It is unclear whether or not envp strings are writable, so
          * we make our own copy to manipulate it */
         key = StrDuplicate(*envp);
+
+        if (!key)
+        {
+            errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to allocate heap memory for program environment variable.\n");
+            write(STDERR_FILENO, errBuf, errLen);
+            goto finish;
+        }
+
         valInd = strcspn(key, "=");
 
         key[valInd] = '\0';
@@ -116,15 +144,15 @@ main(int argc, char **argv)
 
     if (pthread_create(&mainThread, NULL, MainThread, &args) != 0)
     {
-        Log(LOG_ERR, "Bootstrap error: Unable to create main thread.");
-        args.ret = EXIT_FAILURE;
+        errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to create main thread.\n");
+        write(STDERR_FILENO, errBuf, errLen);
         goto finish;
     }
 
     if (pthread_join(mainThread, NULL) != 0)
     {
-        /* Should never happen */
-        Log(LOG_ERR, "Unable to join main thread.");
+        errLen = snprintf(errBuf, ERR_BUF_SIZE, "Fatal: Unable to join main thread.\n");
+        write(STDERR_FILENO, errBuf, errLen);
         args.ret = EXIT_FAILURE;
         goto finish;
     }
@@ -156,7 +184,7 @@ finish:
 
     GenerateMemoryReport(argc, argv);
 
-    MemoryFreeAll();
+    MemoryRuntimeDestroy();
 
     return args.ret;
 }