purecap-kernel: Enforce stricter bounds in generic non-large malloc a…

…llocations

Bounds are tightened in malloc and malloc_domainset to the requested size.
Bounds aren't tightened in malloc_large (yet).
Capabilities are rederived in free, zfree, realloc, when required, to restore
the bounds of the underlying allocation, except in large allocations.
Capabilities are rederived from the associated slab data pointer.

sys/kern/kern_malloc.c

@@ -110,7 +110,7 @@ dtrace_malloc_probe_func_t __read_mostly	dtrace_malloc_probe;
 #define	MALLOC_DEBUG	1
 #endif
 
-#if defined(KASAN) || defined(DEBUG_REDZONE)
+#if defined(KASAN) || defined(DEBUG_REDZONE) || defined(__CHERI_PURE_CAPABILITY__)
 #define	DEBUG_REDZONE_ARG_DEF	, unsigned long osize
 #define	DEBUG_REDZONE_ARG	, osize
 #else
@@ -665,7 +665,7 @@ void *
 	int indx;
 	caddr_t va;
 	uma_zone_t zone;
-#if defined(DEBUG_REDZONE) || defined(KASAN)
+#if defined(DEBUG_REDZONE) || defined(KASAN) || defined(__CHERI_PURE_CAPABILITY__)
 	unsigned long osize = size;
 #endif
 
@@ -707,9 +707,10 @@ void *
 		kasan_mark((void *)va, osize, size, KASAN_MALLOC_REDZONE);
 #endif
 #ifdef __CHERI_PURE_CAPABILITY__
-	KASSERT(cheri_getlen(va) <= CHERI_REPRESENTABLE_LENGTH(size),
+	va = cheri_setbounds(va, osize);
+	KASSERT(cheri_getlen(va) <= CHERI_REPRESENTABLE_LENGTH(osize),
 	    ("Invalid bounds: expected %zx found %zx",
-	        (size_t)CHERI_REPRESENTABLE_LENGTH(size),
+	        (size_t)CHERI_REPRESENTABLE_LENGTH(osize),
 	        (size_t)cheri_getlen(va)));
 #endif
 	return ((void *) va);
@@ -750,7 +751,7 @@ malloc_domainset(size_t size, struct malloc_type *mtp, struct domainset *ds,
 	caddr_t va;
 	int domain;
 	int indx;
-#if defined(KASAN) || defined(DEBUG_REDZONE)
+#if defined(KASAN) || defined(DEBUG_REDZONE) || defined(__CHERI_PURE_CAPABILITY__)
 	unsigned long osize = size;
 #endif
 
@@ -788,6 +789,13 @@ malloc_domainset(size_t size, struct malloc_type *mtp, struct domainset *ds,
 		kmsan_mark(va, size, KMSAN_STATE_UNINIT);
 		kmsan_orig(va, size, KMSAN_TYPE_MALLOC, KMSAN_RET_ADDR);
 	}
+#endif
+#ifdef __CHERI_PURE_CAPABILITY__
+	va = cheri_setbounds(va, osize);
+	KASSERT(cheri_getlen(va) <= CHERI_REPRESENTABLE_LENGTH(osize),
+	    ("Invalid bounds: expected %zx found %zx",
+	        (size_t)CHERI_REPRESENTABLE_LENGTH(osize),
+	        (size_t)cheri_getlen(va)));
 #endif
 	return (va);
 }
@@ -806,7 +814,7 @@ void *
 malloc_domainset_exec(size_t size, struct malloc_type *mtp, struct domainset *ds,
     int flags)
 {
-#if defined(DEBUG_REDZONE) || defined(KASAN)
+#if defined(DEBUG_REDZONE) || defined(KASAN) || defined(__CHERI_PURE_CAPABILITY__)
 	unsigned long osize = size;
 #endif
 #ifdef MALLOC_DEBUG
@@ -963,6 +971,10 @@ free(void *addr, struct malloc_type *mtp)
 	uma_zone_t zone;
 	uma_slab_t slab;
 	u_long size;
+#ifdef __CHERI_PURE_CAPABILITY__
+	uma_keg_t keg;
+	int i;
+#endif
 
 #ifdef MALLOC_DEBUG
 	if (free_dbg(&addr, mtp) != 0)
@@ -987,6 +999,10 @@ free(void *addr, struct malloc_type *mtp)
 	if (__predict_true(!malloc_large_slab(slab))) {
 		size = zone->uz_size;
 #ifdef __CHERI_PURE_CAPABILITY__
+		keg = zone->uz_keg;
+		i = slab_item_index(slab, keg, addr);
+		addr = slab_item(slab, keg, i);
+		addr = cheri_setboundsexact(addr, size);
 		if (__predict_false(cheri_getlen(addr) !=
 		    CHERI_REPRESENTABLE_LENGTH(size)))
 			panic("Invalid bounds: expected %zx found %zx",
@@ -1024,6 +1040,10 @@ zfree(void *addr, struct malloc_type *mtp)
 	uma_zone_t zone;
 	uma_slab_t slab;
 	u_long size;
+#ifdef __CHERI_PURE_CAPABILITY__
+	uma_keg_t keg;
+	int i;
+#endif
 
 #ifdef MALLOC_DEBUG
 	if (free_dbg(&addr, mtp) != 0)
@@ -1047,6 +1067,12 @@ zfree(void *addr, struct malloc_type *mtp)
 
 	if (__predict_true(!malloc_large_slab(slab))) {
 		size = zone->uz_size;
+#ifdef __CHERI_PURE_CAPABILITY__
+		keg = zone->uz_keg;
+		i = slab_item_index(slab, keg, addr);
+		addr = slab_item(slab, keg, i);
+		addr = cheri_setboundsexact(addr, size);
+#endif
 #if defined(INVARIANTS) && !defined(KASAN)
 		free_save_type(addr, mtp, size);
 #endif
@@ -1074,6 +1100,10 @@ realloc(void *addr, size_t size, struct malloc_type *mtp, int flags)
 #endif
 	unsigned long alloc;
 	void *newaddr;
+#ifdef __CHERI_PURE_CAPABILITY__
+	uma_keg_t keg;
+	int i;
+#endif
 
 	KASSERT(mtp->ks_version == M_VERSION,
 	    ("realloc: bad malloc type version"));
@@ -1111,7 +1141,17 @@ realloc(void *addr, size_t size, struct malloc_type *mtp, int flags)
 
 	/* Get the size of the original block */
 	if (!malloc_large_slab(slab))
+#ifdef __CHERI_PURE_CAPABILITY__
+	{
 		alloc = zone->uz_size;
+		keg = zone->uz_keg;
+		i = slab_item_index(slab, keg, addr);
+		addr = slab_item(slab, keg, i);
+		addr = cheri_setboundsexact(addr, alloc);
+	}
+#else
+		alloc = zone->uz_size;
+#endif
 	else
 		alloc = malloc_large_size(slab);