From 974222c626b351256f071d18994c70b79d10a627 Mon Sep 17 00:00:00 2001
From: Qi Wang <interwq@gwu.edu>
Date: Thu, 30 Jan 2020 14:35:54 -0800
Subject: [PATCH] Add safety check on sdallocx slow / sampled path.

---
 include/jemalloc/internal/safety_check.h |  1 +
 src/jemalloc.c                           |  6 +++++-
 src/safety_check.c                       | 12 ++++++++++++
 src/tcache.c                             |  5 +----
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/include/jemalloc/internal/safety_check.h b/include/jemalloc/internal/safety_check.h
index ec4b3369..a7a44338 100644
--- a/include/jemalloc/internal/safety_check.h
+++ b/include/jemalloc/internal/safety_check.h
@@ -1,6 +1,7 @@
 #ifndef JEMALLOC_INTERNAL_SAFETY_CHECK_H
 #define JEMALLOC_INTERNAL_SAFETY_CHECK_H
 
+void safety_check_fail_sized_dealloc(bool current_dealloc);
 void safety_check_fail(const char *format, ...);
 /* Can set to NULL for a default. */
 void safety_check_set_abort(void (*abort_fn)(const char *));
diff --git a/src/jemalloc.c b/src/jemalloc.c
index bac050a4..5f11fc38 100644
--- a/src/jemalloc.c
+++ b/src/jemalloc.c
@@ -2632,7 +2632,11 @@ isfree(tsd_t *tsd, void *ptr, size_t usize, tcache_t *tcache, bool slow_path) {
 			rtree_szind_slab_read(tsd_tsdn(tsd), &extents_rtree,
 			    rtree_ctx, (uintptr_t)ptr, true, &ctx->szind,
 			    &ctx->slab);
-			assert(ctx->szind == sz_size2index(usize));
+			/* Small alloc may have !slab (sampled). */
+			bool sz_correct = (ctx->szind == sz_size2index(usize));
+			if (config_opt_safety_checks && !sz_correct) {
+				safety_check_fail_sized_dealloc(true);
+			}
 		} else {
 			ctx = NULL;
 		}
diff --git a/src/safety_check.c b/src/safety_check.c
index 804155dc..a83dca76 100644
--- a/src/safety_check.c
+++ b/src/safety_check.c
@@ -3,6 +3,18 @@
 
 static void (*safety_check_abort)(const char *message);
 
+void safety_check_fail_sized_dealloc(bool current_dealloc) {
+	assert(config_opt_safety_checks);
+	char *src = current_dealloc ? "the current pointer being freed" :
+	    "in thread cache, possibly from previous deallocations";
+
+	safety_check_fail("<jemalloc>: size mismatch detected, likely caused by"
+	   " application sized deallocation bugs (source: %s). Suggest building"
+	    "with --enable-debug or address sanitizer for debugging. Abort.\n",
+	    src);
+	abort();
+}
+
 void safety_check_set_abort(void (*abort_fn)(const char *)) {
 	safety_check_abort = abort_fn;
 }
diff --git a/src/tcache.c b/src/tcache.c
index 0a511e2d..2f4ca5a4 100644
--- a/src/tcache.c
+++ b/src/tcache.c
@@ -135,10 +135,7 @@ tbin_edatas_lookup_size_check(tsdn_t *tsdn, cache_bin_t *tbin, szind_t binind,
 		sz_sum -= szind;
 	}
 	if (sz_sum != 0) {
-		safety_check_fail("<jemalloc>: size mismatch in thread cache "
-		    "detected, likely caused by sized deallocation bugs by "
-		    "application. Abort.\n");
-		abort();
+		safety_check_fail_sized_dealloc(false);
 	}
 }