engine: add magic numbers to verify struct integrity

The main focus of this commit is 'bal_safety.h'. This header file
adds magic numbers to prevent Undefined Behavior. Right now this
has only been added to 'bal_assembler_t'.

This introduces 3 magic numbers:

#define BAL_MAGIC_UNINITIALIZED 0x00000000U
#define BAL_ASSEMBLER_MAGIC_ALIVE 0xBA11A550U
#define BAL_ASSEMBLER_MAGIC_DEAD  0xDEADBA11U

'ALIVE' magic numbers will be set only if the struct has been
initialized properly. Every function should check a struct's magic
number to make sure it's not operating on a memory corrupted struct.

'DEAD' magic numbers should be set on a struct's 'destory()' function.
This is done to prevent Double Free and Use-After-Free scenarios.

Convenience macro 'BAL_CHECK_MAGIC' and 'BAL_CHECK_MAGIC_VOID' have
been added to easily validate struct integrity.

Signed-off-by: Ronald Caesar <github43132@proton.me>

include/bal_assembler.h

@@ -55,7 +55,8 @@ extern "C"
         /// be ignored until the assembler is reset.
         bal_error_t status;
 
-        uint32_t pad;
+        /// Integrity check.
+        uint32_t magic;
     } bal_assembler_t;
 
     /// Initializes the assembler with a specific memory buffer and the size of the buffer in
@@ -64,8 +65,10 @@ extern "C"
     /// # Returns
     ///
     /// * [`BAL_SUCCESS`] on success.
-    /// * [`BAL_ERROR_INVALID_ARGUMENT`] if `assembler` or `buffer` is NULL.
+    /// * [`BAL_ERROR_INVALID_ARGUMENT`] if `assembler` or `buffer` is NULL, or if `size` is `0`.
     /// * [`BAL_ERROR_MEMORY_ALIGNMENT`] if `buffer` is not 4-byte aligned.
+    /// * [`BAL_ERROR_CAPACITY_TOO_BIG`] if `size` is large enough to cause a `size_t` integer
+    ///   overflow.
     bal_error_t bal_assembler_init(bal_assembler_t *assembler,
                                    void            *buffer,
                                    size_t           size,

include/bal_safety.h

@@ -0,0 +1,92 @@
+#ifndef BALLISTIC_SAFETY_H
+#define BALLISTIC_SAFETY_H
+
+#ifdef __cplusplus
+extern "C"
+{
+#endif // __cplusplus
+
+#define BAL_MAGIC_UNINITIALIZED 0x00000000U
+
+#define BAL_ASSEMBLER_MAGIC_ALIVE 0xBA11A550U // BALLISTO
+#define BAL_ASSEMBLER_MAGIC_DEAD  0xDEADBA11U // DEADBALL
+
+    static const char *bal_decode_magic(uint32_t magic)
+    {
+        switch (magic)
+        {
+            case BAL_MAGIC_UNINITIALIZED:
+                return "Uninitialized Memory";
+            case BAL_ASSEMBLER_MAGIC_ALIVE:
+                return "BALLISTO";
+            case BAL_ASSEMBLER_MAGIC_DEAD:
+                return "DEADBALL";
+            default:
+                return "Unknown (Likely Buffer Overflow)";
+        }
+    }
+
+    static inline const char *bal_diagnose_magic_failure(uint32_t actual_magic, uint32_t dead_magic)
+    {
+        if (actual_magic == BAL_MAGIC_UNINITIALIZED)
+        {
+            return "Struct was never initialized.";
+        }
+        if (actual_magic == dead_magic)
+        {
+            return "Struct was explicitly destroyed (Double Free?).";
+        }
+        return "Memory corruption or wrong struct type passed.";
+    }
+
+#define BAL_CHECK_MAGIC(ptr, alive_magic, dead_magic, struct_name_str, logger)                   \
+    do                                                                                           \
+    {                                                                                            \
+        if (BAL_UNLIKELY((ptr)->magic != (alive_magic)))                                         \
+        {                                                                                        \
+            BAL_LOG_ERROR(&logger,                                                               \
+                          "\n================================================================\n" \
+                          "%s INTEGRITY CHECK FAILED!\n"                                         \
+                          "  Expected : 0x%08X (%s)\n"                                           \
+                          "  Actual   : 0x%08X (%s)\n"                                           \
+                          "  Reason: %s\n"                                                       \
+                          "================================================================\n",  \
+                          (struct_name_str),                                                     \
+                          (alive_magic),                                                         \
+                          bal_decode_magic(alive_magic),                                         \
+                          (ptr)->magic,                                                          \
+                          bal_decode_magic((ptr)->magic),                                        \
+                          bal_diagnose_magic_failure((ptr)->magic, (dead_magic)));               \
+            (ptr)->status = BAL_ERROR_STRUCT_CORRUPTED;                                          \
+            return (BAL_ERROR_STRUCT_CORRUPTED);                                                 \
+        }                                                                                        \
+    } while (0)
+
+#define BAL_CHECK_MAGIC_VOID(ptr, alive_magic, dead_magic, struct_name_str, logger)              \
+    do                                                                                           \
+    {                                                                                            \
+        if (BAL_UNLIKELY((ptr)->magic != (alive_magic)))                                         \
+        {                                                                                        \
+            BAL_LOG_ERROR(&logger,                                                               \
+                          "\n================================================================\n" \
+                          "%s INTEGRITY CHECK FAILED!\n"                                         \
+                          "  Expected : 0x%08X (%s)\n"                                           \
+                          "  Actual   : 0x%08X (%s)\n"                                           \
+                          "  Reason: %s\n"                                                       \
+                          "================================================================\n",  \
+                          (struct_name_str),                                                     \
+                          (alive_magic),                                                         \
+                          bal_decode_magic(alive_magic),                                         \
+                          (ptr)->magic,                                                          \
+                          bal_decode_magic((ptr)->magic),                                        \
+                          bal_diagnose_magic_failure((ptr)->magic, (dead_magic)));               \
+            (ptr)->status = BAL_ERROR_STRUCT_CORRUPTED;                                          \
+            return;                                                                              \
+        }                                                                                        \
+    } while (0)
+
+#ifdef __cplusplus
+}
+#endif // __cplusplus
+
+#endif // BALLISTIC_SAFETY_H

src/bal_assembler.c

@@ -1,4 +1,5 @@
 #include "bal_assembler.h"
+#include "bal_safety.h"
 #include <stdbool.h>
 #include <string.h>
 
@@ -17,16 +18,40 @@ bal_assembler_init(bal_assembler_t   *assembler,
         return BAL_ERROR_INVALID_ARGUMENT;
     }
 
+    assembler->buffer   = NULL;
+    assembler->capacity = 0;
+    assembler->offset   = 0;
+    assembler->logger   = logger;
+    assembler->status   = BAL_ERROR_INVALID_ARGUMENT;
+    assembler->magic    = 0;
+
     if (NULL == buffer)
     {
-        BAL_LOG_ERROR(&logger, "Buffer is NULL.");
+        BAL_LOG_ERROR(&logger, "Aborting function: Buffer is NULL.");
-        return BAL_ERROR_INVALID_ARGUMENT;
+        return assembler->status;
+    }
+
+    if (0 == size)
+    {
+        BAL_LOG_ERROR(&logger, "Aborting function: Buffer capacity is 0.");
+        return assembler->status;
     }
 
     if ((uintptr_t)buffer % 4 != 0)
     {
-        BAL_LOG_ERROR(&logger, "Buffer %p is not 4-byte aligned.", buffer);
+        BAL_LOG_ERROR(&logger, "Aborting function: Buffer %p is not 4-byte aligned.", buffer);
-        return BAL_ERROR_MEMORY_ALIGNMENT;
+        assembler->status = BAL_ERROR_MEMORY_ALIGNMENT;
+        return assembler->status;
+    }
+
+    if (size > (SIZE_MAX / sizeof(uint32_t)))
+    {
+        BAL_LOG_ERROR(&logger,
+                      "Aborting function: Buffer capacity %zu is too large and would "
+                      "cause integer overflow.",
+                      size);
+        assembler->status = BAL_ERROR_CAPACITY_TOO_BIG;
+        return assembler->status;
     }
 
     assembler->buffer   = (uint32_t *)buffer;
@@ -34,6 +59,7 @@ bal_assembler_init(bal_assembler_t   *assembler,
     assembler->offset   = 0;
     assembler->logger   = logger;
     assembler->status   = BAL_SUCCESS;
+    assembler->magic    = BAL_ASSEMBLER_MAGIC_ALIVE;
 
     BAL_LOG_INFO(
         &logger, "Assembler initialized. Buffer: %p, Capacity: %zu instructions.", buffer, size);
@@ -48,6 +74,26 @@ bal_assembler_reset(bal_assembler_t *assembler)
         return;
     }
 
+    BAL_CHECK_MAGIC_VOID(assembler,
+                         BAL_ASSEMBLER_MAGIC_ALIVE,
+                         BAL_ASSEMBLER_MAGIC_DEAD,
+                         "bal_assembler_t",
+                         assembler->logger);
+
+    if (BAL_UNLIKELY(NULL == assembler->buffer))
+    {
+        BAL_LOG_ERROR(&assembler->logger, "Aborting function: assembler->buffer == NULL");
+        assembler->status = BAL_ERROR_INVALID_ARGUMENT;
+        return;
+    }
+
+    if (BAL_UNLIKELY(0 == assembler->capacity))
+    {
+        BAL_LOG_ERROR(&assembler->logger, "Aborting function: assembler->capacity == 0");
+        assembler->status = BAL_ERROR_INVALID_ARGUMENT;
+        return;
+    }
+
     assembler->offset = 0;
     assembler->status = BAL_SUCCESS;
     (void)memset(assembler->buffer, 0, assembler->capacity * sizeof(uint32_t));