better alloc-inl.h

1 files changed, 193 insertions, 0 deletions

diff --git a/include/alloc-inl.h b/include/alloc-inl.h
index a3cc63ca..19ac86e9 100644
--- a/include/alloc-inl.h
+++ b/include/alloc-inl.h
@@ -34,6 +34,197 @@
 #include "types.h"
 #include "debug.h"
 
+#ifndef _WANT_ORIGINAL_AFL_ALLOC
+// afl++ stuff without memory corruption checks - for speed
+
+/* User-facing macro to sprintf() to a dynamically allocated buffer. */
+
+#define alloc_printf(_str...)                        \
+  ({                                                 \
+                                                     \
+    u8* _tmp;                                        \
+    s32 _len = snprintf(NULL, 0, _str);              \
+    if (_len < 0) FATAL("Whoa, snprintf() fails?!"); \
+    _tmp = ck_alloc(_len + 1);                       \
+    snprintf((char*)_tmp, _len + 1, _str);           \
+    _tmp;                                            \
+                                                     \
+  })
+
+/* Macro to enforce allocation limits as a last-resort defense against
+   integer overflows. */
+
+#define ALLOC_CHECK_SIZE(_s)                                          \
+  do {                                                                \
+                                                                      \
+    if ((_s) > MAX_ALLOC) ABORT("Bad alloc request: %u bytes", (_s)); \
+                                                                      \
+  } while (0)
+
+/* Macro to check malloc() failures and the like. */
+
+#define ALLOC_CHECK_RESULT(_r, _s)                                    \
+  do {                                                                \
+                                                                      \
+    if (!(_r)) ABORT("Out of memory: can't allocate %u bytes", (_s)); \
+                                                                      \
+  } while (0)
+
+/* Allocator increments for ck_realloc_block(). */
+
+#define ALLOC_BLK_INC 256
+
+/* Allocate a buffer, explicitly not zeroing it. Returns NULL for zero-sized
+   requests. */
+
+static inline void* DFL_ck_alloc_nozero(u32 size) {
+
+  u8* ret;
+
+  if (!size) return NULL;
+
+  ALLOC_CHECK_SIZE(size);
+  ret = malloc(size);
+  ALLOC_CHECK_RESULT(ret, size);
+
+  return (void*)ret;
+
+}
+
+/* Allocate a buffer, returning zeroed memory. */
+
+static inline void* DFL_ck_alloc(u32 size) {
+
+  void* mem;
+
+  if (!size) return NULL;
+  mem = DFL_ck_alloc_nozero(size);
+
+  return memset(mem, 0, size);
+
+}
+
+/* Free memory, checking for double free and corrupted heap. When DEBUG_BUILD
+   is set, the old memory will be also clobbered with 0xFF. */
+
+static inline void DFL_ck_free(void* mem) {
+
+  if (!mem) return;
+
+  free(mem);
+
+}
+
+/* Re-allocate a buffer, checking for issues and zeroing any newly-added tail.
+   With DEBUG_BUILD, the buffer is always reallocated to a new addresses and the
+   old memory is clobbered with 0xFF. */
+
+static inline void* DFL_ck_realloc(void* orig, u32 size) {
+
+  u8* ret;
+
+  if (!size) {
+
+    DFL_ck_free(orig);
+    return NULL;
+
+  }
+
+  ALLOC_CHECK_SIZE(size);
+
+  /* Catch pointer issues sooner: force relocation and make sure that the
+     original buffer is wiped. */
+
+  ret = realloc(orig, size);
+
+  ALLOC_CHECK_RESULT(ret, size);
+
+  return (void*)ret;
+
+}
+
+/* Re-allocate a buffer with ALLOC_BLK_INC increments (used to speed up
+   repeated small reallocs without complicating the user code). */
+
+static inline void* DFL_ck_realloc_block(void* orig, u32 size) {
+
+  return DFL_ck_realloc(orig, size);
+
+}
+
+/* Create a buffer with a copy of a string. Returns NULL for NULL inputs. */
+
+static inline u8* DFL_ck_strdup(u8* str) {
+
+  u8* ret;
+  u32 size;
+
+  if (!str) return NULL;
+
+  size = strlen((char*)str) + 1;
+
+  ALLOC_CHECK_SIZE(size);
+  ret = malloc(size);
+  ALLOC_CHECK_RESULT(ret, size);
+
+  return memcpy(ret, str, size);
+
+}
+
+/* Create a buffer with a copy of a memory block. Returns NULL for zero-sized
+   or NULL inputs. */
+
+static inline void* DFL_ck_memdup(void* mem, u32 size) {
+
+  u8* ret;
+
+  if (!mem || !size) return NULL;
+
+  ALLOC_CHECK_SIZE(size);
+  ret = malloc(size);
+  ALLOC_CHECK_RESULT(ret, size);
+
+  return memcpy(ret, mem, size);
+
+}
+
+/* Create a buffer with a block of text, appending a NUL terminator at the end.
+   Returns NULL for zero-sized or NULL inputs. */
+
+static inline u8* DFL_ck_memdup_str(u8* mem, u32 size) {
+
+  u8* ret;
+
+  if (!mem || !size) return NULL;
+
+  ALLOC_CHECK_SIZE(size);
+  ret = malloc(size + 1);
+  ALLOC_CHECK_RESULT(ret, size);
+
+  memcpy(ret, mem, size);
+  ret[size] = 0;
+
+  return ret;
+
+}
+
+/* In non-debug mode, we just do straightforward aliasing of the above functions
+   to user-visible names such as ck_alloc(). */
+
+#define ck_alloc DFL_ck_alloc
+#define ck_alloc_nozero DFL_ck_alloc_nozero
+#define ck_realloc DFL_ck_realloc
+#define ck_realloc_block DFL_ck_realloc_block
+#define ck_strdup DFL_ck_strdup
+#define ck_memdup DFL_ck_memdup
+#define ck_memdup_str DFL_ck_memdup_str
+#define ck_free DFL_ck_free
+
+#define alloc_report()
+
+#else
+// This is the original alloc-inl of stock afl
+
 /* User-facing macro to sprintf() to a dynamically allocated buffer. */
 
 #define alloc_printf(_str...)                        \
@@ -568,5 +759,7 @@ static inline void TRK_ck_free(void* ptr, const char* file, const char* func,
 
 #endif                                                     /* ^!DEBUG_BUILD */
 
+#endif                                          /* _WANT_ORIGINAL_AFL_ALLOC */
+
 #endif                                               /* ! _HAVE_ALLOC_INL_H */