From 6cb585b13ad196ca2e4588ce984c269f3fdb4cea Mon Sep 17 00:00:00 2001
From: Alex Lapenkou <lapenkov@fb.com>
Date: Tue, 2 Nov 2021 15:56:36 -0700
Subject: [PATCH] San: Unguard guarded slabs during arena destruction

When opt_retain is on, slab extents remain guarded in all states, even
retained. This works well if arena is never destroyed, because we
anticipate those slabs will be eventually reused. But if the arena is
destroyed, the slabs must be unguarded to prevent leaking guard pages.
---
 include/jemalloc/internal/guard.h |  9 +++++++-
 src/extent.c                      |  8 ++++---
 src/guard.c                       | 35 +++++++++++++++++++++++++------
 3 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/include/jemalloc/internal/guard.h b/include/jemalloc/internal/guard.h
index 31f98c5f..8e578168 100644
--- a/include/jemalloc/internal/guard.h
+++ b/include/jemalloc/internal/guard.h
@@ -14,7 +14,14 @@ extern size_t opt_san_guard_large;
 extern size_t opt_san_guard_small;
 
 void guard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap);
-void unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap);
+void unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata,
+    emap_t *emap);
+/*
+ * Unguard the extent, but don't modify emap boundaries. Must be called on an
+ * extent that has been erased from emap and shouldn't be placed back.
+ */
+void unguard_pages_pre_destroy(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata,
+    emap_t *emap);
 void tsd_san_init(tsd_t *tsd);
 
 static inline bool
diff --git a/src/extent.c b/src/extent.c
index 84ecd6b2..a79e1c72 100644
--- a/src/extent.c
+++ b/src/extent.c
@@ -1057,12 +1057,14 @@ extent_destroy_wrapper(tsdn_t *tsdn, pac_t *pac, ehooks_t *ehooks,
     edata_t *edata) {
 	assert(edata_base_get(edata) != NULL);
 	assert(edata_size_get(edata) != 0);
+	assert(edata_state_get(edata) == extent_state_retained);
+	assert(emap_edata_is_acquired(tsdn, pac->emap, edata));
 	witness_assert_depth_to_rank(tsdn_witness_tsdp_get(tsdn),
 	    WITNESS_RANK_CORE, 0);
 
-	/* Deregister first to avoid a race with other allocating threads. */
-	extent_deregister(tsdn, pac, edata);
-
+	if (edata_guarded_get(edata)) {
+		unguard_pages_pre_destroy(tsdn, ehooks, edata, pac->emap);
+	}
 	edata_addr_set(edata, edata_base_get(edata));
 
 	/* Try to destroy; silently fail otherwise. */
diff --git a/src/guard.c b/src/guard.c
index 07232199..4dadc970 100644
--- a/src/guard.c
+++ b/src/guard.c
@@ -32,10 +32,16 @@ guard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap) {
 	/* The new boundary will be registered on the pa_alloc path. */
 }
 
-void
-unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap) {
+static void
+unguard_pages_impl(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap,
+    bool reg_emap) {
 	/* Remove the inner boundary which no longer exists. */
-	emap_deregister_boundary(tsdn, emap, edata);
+	if (reg_emap) {
+		assert(edata_state_get(edata) == extent_state_active);
+		emap_deregister_boundary(tsdn, emap, edata);
+	} else {
+		assert(edata_state_get(edata) == extent_state_retained);
+	}
 
 	size_t size = edata_size_get(edata);
 	size_t size_with_guards = size + PAGE_GUARDS_SIZE;
@@ -44,7 +50,6 @@ unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap) {
 	uintptr_t guard1 = addr - PAGE;
 	uintptr_t guard2 = addr + size;
 
-	assert(edata_state_get(edata) == extent_state_active);
 	ehooks_unguard(tsdn, ehooks, (void *)guard1, (void *)guard2);
 
 	/* Update the true addr and usable size of the edata. */
@@ -52,8 +57,26 @@ unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap) {
 	edata_addr_set(edata, (void *)guard1);
 	edata_guarded_set(edata, false);
 
-	/* Then re-register the outer boundary including the guards. */
-	emap_register_boundary(tsdn, emap, edata, SC_NSIZES, /* slab */ false);
+	/*
+	 * Then re-register the outer boundary including the guards, if
+	 * requested.
+	 */
+	if (reg_emap) {
+		emap_register_boundary(tsdn, emap, edata, SC_NSIZES,
+		    /* slab */ false);
+	}
+}
+
+void
+unguard_pages(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata, emap_t *emap) {
+	unguard_pages_impl(tsdn, ehooks, edata, emap, /* reg_emap */ true);
+}
+
+void
+unguard_pages_pre_destroy(tsdn_t *tsdn, ehooks_t *ehooks, edata_t *edata,
+    emap_t *emap) {
+	emap_assert_not_mapped(tsdn, emap, edata);
+	unguard_pages_impl(tsdn, ehooks, edata, emap, /* reg_emap */ false);
 }
 
 void