Make memguard(9) capable to guard uma(9) allocations.

1 files changed, 38 insertions, 2 deletions

diff --git a/sys/vm/uma_core.c b/sys/vm/uma_core.c
index 8da5b8e..9fbea55 100644
--- a/sys/vm/uma_core.c
+++ b/sys/vm/uma_core.c
@@ -59,6 +59,7 @@ __FBSDID("$FreeBSD$");
 
 #include "opt_ddb.h"
 #include "opt_param.h"
+#include "opt_vm.h"
 
 #include <sys/param.h>
 #include <sys/systm.h>
@@ -88,6 +89,10 @@ __FBSDID("$FreeBSD$");
 
 #include <ddb/ddb.h>
 
+#ifdef DEBUG_MEMGUARD
+#include <vm/memguard.h>
+#endif
+
 /*
  * This is the zone and keg from which all zones are spawned.  The idea is that
  * even the zone & keg heads are allocated from the allocator, so we use the
@@ -1978,7 +1983,29 @@ uma_zalloc_arg(uma_zone_t zone, void *udata, int flags)
 		WITNESS_WARN(WARN_GIANTOK | WARN_SLEEPOK, NULL,
 		    "uma_zalloc_arg: zone \"%s\"", zone->uz_name);
 	}
-
+#ifdef DEBUG_MEMGUARD
+	if (memguard_cmp_zone(zone)) {
+		item = memguard_alloc(zone->uz_size, flags);
+		if (item != NULL) {
+			/*
+			 * Avoid conflict with the use-after-free
+			 * protecting infrastructure from INVARIANTS.
+			 */
+			if (zone->uz_init != NULL &&
+			    zone->uz_init != mtrash_init &&
+			    zone->uz_init(item, zone->uz_size, flags) != 0)
+				return (NULL);
+			if (zone->uz_ctor != NULL &&
+			    zone->uz_ctor != mtrash_ctor &&
+			    zone->uz_ctor(item, zone->uz_size, udata, flags) != 0) {
+			    	zone->uz_fini(item, zone->uz_size);
+				return (NULL);
+			}
+			return (item);
+		}
+		/* This is unfortunate but should not be fatal. */
+	}
+#endif
 	/*
 	 * If possible, allocate from the per-CPU cache.  There are two
 	 * requirements for safe access to the per-CPU cache: (1) the thread
@@ -2544,7 +2571,16 @@ uma_zfree_arg(uma_zone_t zone, void *item, void *udata)
         /* uma_zfree(..., NULL) does nothing, to match free(9). */
         if (item == NULL)
                 return;
-
+#ifdef DEBUG_MEMGUARD
+	if (is_memguard_addr(item)) {
+		if (zone->uz_dtor != NULL && zone->uz_dtor != mtrash_dtor)
+			zone->uz_dtor(item, zone->uz_size, udata);
+		if (zone->uz_fini != NULL && zone->uz_fini != mtrash_fini)
+			zone->uz_fini(item, zone->uz_size);
+		memguard_free(item);
+		return;
+	}
+#endif
 	if (zone->uz_dtor)
 		zone->uz_dtor(item, zone->uz_size, udata);