Introduce xfree_sized and xrealloc_sized

This will allow prism to pass buffer sizes to the Ruby GC.

It also helps avoid buffer overflow as it confirms the size was
correctly tracked all the way until the buffer is freed.

Author: byroot

Doxyfile

@@ -24,6 +24,7 @@ OUTPUT_DIRECTORY       = doc
 JAVADOC_AUTOBRIEF      = YES
 OPTIMIZE_OUTPUT_FOR_C  = YES
 INPUT                  = src src/util include include/prism include/prism/util
+EXCLUDE                = include/prism/debug_allocator.h
 HTML_OUTPUT            = c
 SORT_MEMBER_DOCS       = NO
 GENERATE_LATEX         = NO

ext/prism/extension.c

@@ -413,7 +413,7 @@ dump(int argc, VALUE *argv, VALUE self) {
     if (options.freeze) rb_obj_freeze(value);
 
 #ifdef PRISM_BUILD_DEBUG
-    xfree(dup);
+    xfree_sized(dup, length);
 #endif
 
     pm_string_free(&input);
@@ -929,7 +929,7 @@ parse(int argc, VALUE *argv, VALUE self) {
     VALUE value = parse_input(&input, &options);
 
 #ifdef PRISM_BUILD_DEBUG
-    xfree(dup);
+    xfree_sized(dup, length);
 #endif
 
     pm_string_free(&input);

include/prism/debug_allocator.h

@@ -0,0 +1,99 @@
+/**
+ * @file debug_allocator.h
+ *
+ * Decorate allocation function to ensure sizes are correct.
+ */
+#ifndef PRISM_DEBUG_ALLOCATOR_H
+#define PRISM_DEBUG_ALLOCATOR_H
+
+#include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+static inline void *
+pm_debug_malloc(size_t size)
+{
+    size_t *memory = xmalloc(size + sizeof(size_t));
+    memory[0] = size;
+    return memory + 1;
+}
+
+static inline void *
+pm_debug_calloc(size_t nmemb, size_t size)
+{
+    size_t total_size = nmemb * size;
+    void *ptr = pm_debug_malloc(total_size);
+    memset(ptr, 0, total_size);
+    return ptr;
+}
+
+static inline void *
+pm_debug_realloc(void *ptr, size_t size)
+{
+    if (ptr == NULL) {
+        return pm_debug_malloc(size);
+    }
+
+    size_t *memory = (size_t *)ptr;
+    void *raw_memory = memory - 1;
+    memory = (size_t *)xrealloc(raw_memory, size + sizeof(size_t));
+    memory[0] = size;
+    return memory + 1;
+}
+
+static inline void
+pm_debug_free(void *ptr)
+{
+    if (ptr != NULL) {
+        size_t *memory = (size_t *)ptr;
+        xfree(memory - 1);
+    }
+}
+
+static inline void
+pm_debug_free_sized(void *ptr, size_t old_size)
+{
+    if (ptr != NULL) {
+        size_t *memory = (size_t *)ptr;
+        if (old_size != memory[-1]) {
+            fprintf(stderr, "[BUG] buffer %p was allocated with size %lu but freed with size %lu\n", ptr, memory[-1], old_size);
+            abort();
+        }
+        xfree_sized(memory - 1, old_size + sizeof(size_t));
+    }
+}
+
+static inline void *
+pm_debug_realloc_sized(void *ptr, size_t size, size_t old_size)
+{
+    if (ptr == NULL) {
+        if (old_size != 0) {
+            fprintf(stderr, "[BUG] realloc_sized called with NULL pointer and old size %lu\n", old_size);
+            abort();
+        }
+        return pm_debug_malloc(size);
+    }
+
+    size_t *memory = (size_t *)ptr;
+    if (old_size != memory[-1]) {
+        fprintf(stderr, "[BUG] buffer %p was allocated with size %lu but realloced with size %lu\n", ptr, memory[-1], old_size);
+        abort();
+    }
+    return pm_debug_realloc(ptr, size);
+}
+
+#undef xmalloc
+#undef xrealloc
+#undef xcalloc
+#undef xfree
+#undef xrealloc_sized
+#undef xfree_sized
+
+#define xmalloc          pm_debug_malloc
+#define xrealloc         pm_debug_realloc
+#define xcalloc          pm_debug_calloc
+#define xfree            pm_debug_free
+#define xrealloc_sized   pm_debug_realloc_sized
+#define xfree_sized      pm_debug_free_sized
+
+#endif

include/prism/defines.h

@@ -161,10 +161,12 @@
  * ```
  * #ifndef PRISM_XALLOCATOR_H
  * #define PRISM_XALLOCATOR_H
- * #define xmalloc      my_malloc
- * #define xrealloc     my_realloc
- * #define xcalloc      my_calloc
- * #define xfree        my_free
+ * #define xmalloc          my_malloc
+ * #define xrealloc         my_realloc
+ * #define xcalloc          my_calloc
+ * #define xfree            my_free
+ * #define xrealloc_sized   my_realloc_sized // (optional)
+ * #define xfree_sized      my_free_sized    // (optional)
  * #endif
  * ```
  */
@@ -204,6 +206,28 @@
     #endif
 #endif
 
+#ifndef xfree_sized
+/**
+ * The free_sized function that should be used. This can be overridden with the
+ * PRISM_XALLOCATOR define.
+ * If not defined, defaults to calling xfree.
+ */
+    #define xfree_sized(p, s) xfree(((void)(s), (p)))
+#endif
+
+#ifndef xrealloc_sized
+/**
+ * The xrealloc_sized function that should be used. This can be overridden with the
+ * PRISM_XALLOCATOR define.
+ * If not defined, defaults to calling xrealloc.
+ */
+    #define xrealloc_sized(p, ns, os) xrealloc((p), ((void)(os), (ns)))
+#endif
+
+#ifdef PRISM_BUILD_DEBUG
+    #include "prism/debug_allocator.h"
+#endif
+
 /**
  * If PRISM_BUILD_MINIMAL is defined, then we're going to define every possible
  * switch that will turn off certain features of prism.

prism.gemspec

@@ -48,6 +48,7 @@ Gem::Specification.new do |spec|
     "ext/prism/extension.h",
     "include/prism.h",
     "include/prism/ast.h",
+    "include/prism/debug_allocator.h",
     "include/prism/defines.h",
     "include/prism/diagnostic.h",
     "include/prism/encoding.h",

src/options.c

@@ -226,10 +226,10 @@ pm_options_free(pm_options_t *options) {
             pm_string_free(&scope->locals[local_index]);
         }
 
-        xfree(scope->locals);
+        xfree_sized(scope->locals, scope->locals_count * sizeof(pm_string_t));
     }
 
-    xfree(options->scopes);
+    xfree_sized(options->scopes, options->scopes_count * sizeof(pm_options_scope_t));
 }
 
 /**