From 0a0fcd3e6a0816f0a56fa852416d0ece861c0abb Mon Sep 17 00:00:00 2001
From: David Goldblatt <davidgoldblatt@fb.com>
Date: Tue, 28 Mar 2017 17:30:54 -0700
Subject: [PATCH] Add hooking functionality

This allows us to hook chosen functions and do interesting things there (in
particular: reentrancy checking).
---
 Makefile.in                                   |  2 +
 include/jemalloc/internal/hooks.h             | 19 +++++
 .../jemalloc/internal/jemalloc_internal.h.in  |  7 ++
 include/jemalloc/internal/malloc_io.h         |  2 +-
 include/jemalloc/internal/private_symbols.txt |  1 +
 include/jemalloc/internal/stats_externs.h     |  2 +-
 src/hooks.c                                   | 12 +++
 src/prof.c                                    |  7 ++
 src/tsd.c                                     |  9 +++
 test/include/test/jemalloc_test.h.in          |  2 +
 test/include/test/test.h                      |  6 ++
 test/src/test.c                               | 76 ++++++++++++++++++-
 test/unit/hooks.c                             | 38 ++++++++++
 test/unit/prof_accum.c                        |  2 +-
 test/unit/prof_active.c                       |  2 +-
 test/unit/prof_gdump.c                        |  2 +-
 test/unit/prof_reset.c                        |  2 +-
 test/unit/prof_tctx.c                         |  2 +-
 test/unit/tsd.c                               |  1 +
 19 files changed, 183 insertions(+), 11 deletions(-)
 create mode 100644 include/jemalloc/internal/hooks.h
 create mode 100644 src/hooks.c
 create mode 100644 test/unit/hooks.c

diff --git a/Makefile.in b/Makefile.in
index 4fb852da..26c811c8 100644
--- a/Makefile.in
+++ b/Makefile.in
@@ -98,6 +98,7 @@ C_SRCS := $(srcroot)src/jemalloc.c \
 	$(srcroot)src/extent_dss.c \
 	$(srcroot)src/extent_mmap.c \
 	$(srcroot)src/hash.c \
+	$(srcroot)src/hooks.c \
 	$(srcroot)src/large.c \
 	$(srcroot)src/malloc_io.c \
 	$(srcroot)src/mutex.c \
@@ -161,6 +162,7 @@ TESTS_UNIT := \
 	$(srcroot)test/unit/extent_quantize.c \
 	$(srcroot)test/unit/fork.c \
 	$(srcroot)test/unit/hash.c \
+	$(srcroot)test/unit/hooks.c \
 	$(srcroot)test/unit/junk.c \
 	$(srcroot)test/unit/junk_alloc.c \
 	$(srcroot)test/unit/junk_free.c \
diff --git a/include/jemalloc/internal/hooks.h b/include/jemalloc/internal/hooks.h
new file mode 100644
index 00000000..608b268f
--- /dev/null
+++ b/include/jemalloc/internal/hooks.h
@@ -0,0 +1,19 @@
+#ifndef JEMALLOC_INTERNAL_HOOKS_H
+#define JEMALLOC_INTERNAL_HOOKS_H
+
+extern void (*hooks_arena_new_hook)();
+extern void (*hooks_libc_hook)();
+
+#define JEMALLOC_HOOK(fn, hook) ((void)(hook != NULL && (hook(), 0)), fn)
+
+#define open JEMALLOC_HOOK(open, hooks_libc_hook)
+#define read JEMALLOC_HOOK(read, hooks_libc_hook)
+#define write JEMALLOC_HOOK(write, hooks_libc_hook)
+#define readlink JEMALLOC_HOOK(readlink, hooks_libc_hook)
+#define close JEMALLOC_HOOK(close, hooks_libc_hook)
+#define creat JEMALLOC_HOOK(creat, hooks_libc_hook)
+#define secure_getenv JEMALLOC_HOOK(secure_getenv, hooks_libc_hook)
+/* Note that this is undef'd and re-define'd in src/prof.c. */
+#define _Unwind_Backtrace JEMALLOC_HOOK(_Unwind_Backtrace, hooks_libc_hook)
+
+#endif /* JEMALLOC_INTERNAL_HOOKS_H */
diff --git a/include/jemalloc/internal/jemalloc_internal.h.in b/include/jemalloc/internal/jemalloc_internal.h.in
index c00912bf..1c0bf43a 100644
--- a/include/jemalloc/internal/jemalloc_internal.h.in
+++ b/include/jemalloc/internal/jemalloc_internal.h.in
@@ -23,7 +23,14 @@ extern "C" {
 #  define JEMALLOC_N(n) @private_namespace@##n
 #  include "../jemalloc@install_suffix@.h"
 #endif
+
+/*
+ * Note that the ordering matters here; the hook itself is name-mangled.  We
+ * want the inclusion of hooks to happen early, so that we hook as much as
+ * possible.
+ */
 #include "jemalloc/internal/private_namespace.h"
+#include "jemalloc/internal/hooks.h"
 
 static const bool config_debug =
 #ifdef JEMALLOC_DEBUG
diff --git a/include/jemalloc/internal/malloc_io.h b/include/jemalloc/internal/malloc_io.h
index 7ff3d5b1..8b2fb96f 100644
--- a/include/jemalloc/internal/malloc_io.h
+++ b/include/jemalloc/internal/malloc_io.h
@@ -56,7 +56,7 @@ size_t malloc_snprintf(char *str, size_t size, const char *format, ...)
     JEMALLOC_FORMAT_PRINTF(3, 4);
 void malloc_vcprintf(void (*write_cb)(void *, const char *), void *cbopaque,
     const char *format, va_list ap);
-void malloc_cprintf(void (*write)(void *, const char *), void *cbopaque,
+void malloc_cprintf(void (*write_cb)(void *, const char *), void *cbopaque,
     const char *format, ...) JEMALLOC_FORMAT_PRINTF(3, 4);
 void malloc_printf(const char *format, ...) JEMALLOC_FORMAT_PRINTF(1, 2);
 
diff --git a/include/jemalloc/internal/private_symbols.txt b/include/jemalloc/internal/private_symbols.txt
index e2bb0592..deae8243 100644
--- a/include/jemalloc/internal/private_symbols.txt
+++ b/include/jemalloc/internal/private_symbols.txt
@@ -232,6 +232,7 @@ hash_rotl_64
 hash_x64_128
 hash_x86_128
 hash_x86_32
+hooks_libc_hook
 iaalloc
 ialloc
 iallocztm
diff --git a/include/jemalloc/internal/stats_externs.h b/include/jemalloc/internal/stats_externs.h
index a0a1ab6c..519441c9 100644
--- a/include/jemalloc/internal/stats_externs.h
+++ b/include/jemalloc/internal/stats_externs.h
@@ -3,7 +3,7 @@
 
 extern bool	opt_stats_print;
 
-void	stats_print(void (*write)(void *, const char *), void *cbopaque,
+void	stats_print(void (*write_cb)(void *, const char *), void *cbopaque,
     const char *opts);
 
 #endif /* JEMALLOC_INTERNAL_STATS_EXTERNS_H */
diff --git a/src/hooks.c b/src/hooks.c
new file mode 100644
index 00000000..c32471e9
--- /dev/null
+++ b/src/hooks.c
@@ -0,0 +1,12 @@
+#include "jemalloc/internal/jemalloc_internal.h"
+
+/*
+ * The hooks are a little bit screwy -- they're not genuinely exported in the
+ * sense that we want them available to end-users, but we do want them visible
+ * from outside the generated library, so that we can use them in test code.
+ */
+JEMALLOC_EXPORT
+void (*hooks_arena_new_hook)() = NULL;
+
+JEMALLOC_EXPORT
+void (*hooks_libc_hook)() = NULL;
diff --git a/src/prof.c b/src/prof.c
index a0290b8f..db1ef035 100644
--- a/src/prof.c
+++ b/src/prof.c
@@ -8,7 +8,14 @@
 #endif
 
 #ifdef JEMALLOC_PROF_LIBGCC
+/*
+ * We have a circular dependency -- jemalloc_internal.h tells us if we should
+ * use libgcc's unwinding functionality, but after we've included that, we've
+ * already hooked _Unwind_Backtrace.  We'll temporarily disable hooking.
+ */
+#undef _Unwind_Backtrace
 #include <unwind.h>
+#define _Unwind_Backtrace JEMALLOC_HOOK(_Unwind_Backtrace, hooks_libc_hook)
 #endif
 
 /******************************************************************************/
diff --git a/src/tsd.c b/src/tsd.c
index 8b54770e..0d5de8ea 100644
--- a/src/tsd.c
+++ b/src/tsd.c
@@ -149,6 +149,15 @@ _tls_callback(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
 	return true;
 }
 
+/*
+ * We need to be able to say "read" here (in the "pragma section"), but have
+ * hooked "read". We won't read for the rest of the file, so we can get away
+ * with unhooking.
+ */
+#ifdef read
+#  undef read
+#endif
+
 #ifdef _MSC_VER
 #  ifdef _M_IX86
 #    pragma comment(linker, "/INCLUDE:__tls_used")
diff --git a/test/include/test/jemalloc_test.h.in b/test/include/test/jemalloc_test.h.in
index 0770d020..e3882b29 100644
--- a/test/include/test/jemalloc_test.h.in
+++ b/test/include/test/jemalloc_test.h.in
@@ -45,6 +45,7 @@ extern "C" {
 #  define JEMALLOC_MANGLE
 #  include "jemalloc/internal/jemalloc_internal.h"
 
+
 /******************************************************************************/
 /*
  * For integration tests, expose the public jemalloc interfaces, but only
@@ -68,6 +69,7 @@ static const bool config_debug =
 
 #  define JEMALLOC_N(n) @private_namespace@##n
 #  include "jemalloc/internal/private_namespace.h"
+#  include "jemalloc/internal/hooks.h"
 
 /* Hermetic headers. */
 #  include "jemalloc/internal/assert.h"
diff --git a/test/include/test/test.h b/test/include/test/test.h
index d7f05fad..fd0e5265 100644
--- a/test/include/test/test.h
+++ b/test/include/test/test.h
@@ -310,6 +310,9 @@ label_test_end:								\
 #define test(...)							\
 	p_test(__VA_ARGS__, NULL)
 
+#define test_no_reentrancy(...)							\
+	p_test_no_reentrancy(__VA_ARGS__, NULL)
+
 #define test_no_malloc_init(...)					\
 	p_test_no_malloc_init(__VA_ARGS__, NULL)
 
@@ -321,11 +324,14 @@ label_test_end:								\
 	}								\
 } while (0)
 
+bool test_is_reentrant();
+
 void	test_skip(const char *format, ...) JEMALLOC_FORMAT_PRINTF(1, 2);
 void	test_fail(const char *format, ...) JEMALLOC_FORMAT_PRINTF(1, 2);
 
 /* For private use by macros. */
 test_status_t	p_test(test_t *t, ...);
+test_status_t	p_test_no_reentrancy(test_t *t, ...);
 test_status_t	p_test_no_malloc_init(test_t *t, ...);
 void	p_test_init(const char *name);
 void	p_test_fini(void);
diff --git a/test/src/test.c b/test/src/test.c
index c5101d4e..fe6dc60e 100644
--- a/test/src/test.c
+++ b/test/src/test.c
@@ -1,10 +1,42 @@
 #include "test/jemalloc_test.h"
 
+/* Test status state. */
+
 static unsigned		test_count = 0;
 static test_status_t	test_counts[test_status_count] = {0, 0, 0};
 static test_status_t	test_status = test_status_pass;
 static const char *	test_name = "";
 
+/* Reentrancy testing helpers. */
+
+#define NUM_REENTRANT_ALLOCS 20
+static bool reentrant = false;
+static bool hook_ran = false;
+static void *to_free[NUM_REENTRANT_ALLOCS];
+
+static void
+reentrancy_hook() {
+	hook_ran = true;
+	hooks_libc_hook = NULL;
+
+	void *to_free_local[NUM_REENTRANT_ALLOCS];
+	size_t alloc_size = 1;
+	for (int i = 0; i < NUM_REENTRANT_ALLOCS; i++) {
+		to_free[i] = malloc(alloc_size);
+		to_free_local[i] = malloc(alloc_size);
+		alloc_size *= 2;
+	}
+	for (int i = 0; i < NUM_REENTRANT_ALLOCS; i++) {
+		free(to_free_local[i]);
+	}
+}
+
+/* Actual test infrastructure. */
+bool
+test_is_reentrant() {
+	return reentrant;
+}
+
 JEMALLOC_FORMAT_PRINTF(1, 2)
 void
 test_skip(const char *format, ...) {
@@ -49,11 +81,13 @@ p_test_init(const char *name) {
 void
 p_test_fini(void) {
 	test_counts[test_status]++;
-	malloc_printf("%s: %s\n", test_name, test_status_string(test_status));
+	malloc_printf("%s: %s (%s)\n", test_name,
+	    test_status_string(test_status),
+	    reentrant ? "reentrant" : "non-reentrant");
 }
 
 static test_status_t
-p_test_impl(bool do_malloc_init, test_t *t, va_list ap) {
+p_test_impl(bool do_malloc_init, bool do_reentrant, test_t *t, va_list ap) {
 	test_status_t ret;
 
 	if (do_malloc_init) {
@@ -71,10 +105,27 @@ p_test_impl(bool do_malloc_init, test_t *t, va_list ap) {
 
 	ret = test_status_pass;
 	for (; t != NULL; t = va_arg(ap, test_t *)) {
+		/* Non-reentrant run. */
+		reentrant = false;
 		t();
 		if (test_status > ret) {
 			ret = test_status;
 		}
+		/* Reentrant run. */
+		if (do_reentrant) {
+			reentrant = true;
+			hooks_libc_hook = &reentrancy_hook;
+			t();
+			if (test_status > ret) {
+				ret = test_status;
+			}
+			if (hook_ran) {
+				hook_ran = false;
+				for (int i = 0; i < NUM_REENTRANT_ALLOCS; i++) {
+					free(to_free[i]);
+				}
+			}
+		}
 	}
 
 	malloc_printf("--- %s: %u/%u, %s: %u/%u, %s: %u/%u ---\n",
@@ -95,7 +146,20 @@ p_test(test_t *t, ...) {
 
 	ret = test_status_pass;
 	va_start(ap, t);
-	ret = p_test_impl(true, t, ap);
+	ret = p_test_impl(true, true, t, ap);
+	va_end(ap);
+
+	return ret;
+}
+
+test_status_t
+p_test_no_reentrancy(test_t *t, ...) {
+	test_status_t ret;
+	va_list ap;
+
+	ret = test_status_pass;
+	va_start(ap, t);
+	ret = p_test_impl(true, false, t, ap);
 	va_end(ap);
 
 	return ret;
@@ -108,7 +172,11 @@ p_test_no_malloc_init(test_t *t, ...) {
 
 	ret = test_status_pass;
 	va_start(ap, t);
-	ret = p_test_impl(false, t, ap);
+	/*
+	 * We also omit reentrancy from bootstrapping tests, since we don't
+	 * (yet) care about general reentrancy during bootstrapping.
+	 */
+	ret = p_test_impl(false, false, t, ap);
 	va_end(ap);
 
 	return ret;
diff --git a/test/unit/hooks.c b/test/unit/hooks.c
new file mode 100644
index 00000000..b70172e1
--- /dev/null
+++ b/test/unit/hooks.c
@@ -0,0 +1,38 @@
+#include "test/jemalloc_test.h"
+
+static bool hook_called = false;
+
+static void
+hook() {
+	hook_called = true;
+}
+
+static int
+func_to_hook(int arg1, int arg2) {
+	return arg1 + arg2;
+}
+
+#define func_to_hook JEMALLOC_HOOK(func_to_hook, hooks_libc_hook)
+
+TEST_BEGIN(unhooked_call) {
+	hooks_libc_hook = NULL;
+	hook_called = false;
+	assert_d_eq(3, func_to_hook(1, 2), "Hooking changed return value.");
+	assert_false(hook_called, "Nulling out hook didn't take.");
+}
+TEST_END
+
+TEST_BEGIN(hooked_call) {
+	hooks_libc_hook = &hook;
+	hook_called = false;
+	assert_d_eq(3, func_to_hook(1, 2), "Hooking changed return value.");
+	assert_true(hook_called, "Hook should have executed.");
+}
+TEST_END
+
+int
+main(void) {
+	return test(
+	    unhooked_call,
+	    hooked_call);
+}
diff --git a/test/unit/prof_accum.c b/test/unit/prof_accum.c
index 6ccab82b..25220063 100644
--- a/test/unit/prof_accum.c
+++ b/test/unit/prof_accum.c
@@ -76,6 +76,6 @@ TEST_END
 
 int
 main(void) {
-	return test(
+	return test_no_reentrancy(
 	    test_idump);
 }
diff --git a/test/unit/prof_active.c b/test/unit/prof_active.c
index 275aac89..850a24a7 100644
--- a/test/unit/prof_active.c
+++ b/test/unit/prof_active.c
@@ -112,6 +112,6 @@ TEST_END
 
 int
 main(void) {
-	return test(
+	return test_no_reentrancy(
 	    test_prof_active);
 }
diff --git a/test/unit/prof_gdump.c b/test/unit/prof_gdump.c
index 97ade68c..fcb434cb 100644
--- a/test/unit/prof_gdump.c
+++ b/test/unit/prof_gdump.c
@@ -69,6 +69,6 @@ TEST_END
 
 int
 main(void) {
-	return test(
+	return test_no_reentrancy(
 	    test_gdump);
 }
diff --git a/test/unit/prof_reset.c b/test/unit/prof_reset.c
index 6120714e..7cce42d2 100644
--- a/test/unit/prof_reset.c
+++ b/test/unit/prof_reset.c
@@ -278,7 +278,7 @@ main(void) {
 	/* Intercept dumping prior to running any tests. */
 	prof_dump_open = prof_dump_open_intercept;
 
-	return test(
+	return test_no_reentrancy(
 	    test_prof_reset_basic,
 	    test_prof_reset_cleanup,
 	    test_prof_reset,
diff --git a/test/unit/prof_tctx.c b/test/unit/prof_tctx.c
index 183f7ce0..30c6b178 100644
--- a/test/unit/prof_tctx.c
+++ b/test/unit/prof_tctx.c
@@ -41,6 +41,6 @@ TEST_END
 
 int
 main(void) {
-	return test(
+	return test_no_reentrancy(
 	    test_prof_realloc);
 }
diff --git a/test/unit/tsd.c b/test/unit/tsd.c
index 5bfcdf49..4a0f3185 100644
--- a/test/unit/tsd.c
+++ b/test/unit/tsd.c
@@ -79,6 +79,7 @@ thd_start(void *arg) {
 }
 
 TEST_BEGIN(test_tsd_main_thread) {
+	test_skip_if(test_is_reentrant());
 	thd_start((void *)(uintptr_t)0xa5f3e329);
 }
 TEST_END