Loading...
--- libmalloc/libmalloc-657.80.3/src/early_malloc.c
+++ libmalloc/libmalloc-792.60.6/src/early_malloc.c
@@ -47,7 +47,9 @@
 #if MALLOC_TARGET_EXCLAVES || MALLOC_TARGET_EXCLAVES_INTROSPECTOR
 #define MFM_ARENA_SIZE      (1ul << 20)
 #else
-#define MFM_ARENA_SIZE      (8ul << 20)
+// Originally this was set to 8M: we shrank it down to 4M to accomodate for
+// the guarded range reservation (as the arena was heavily underused anyway).
+#define MFM_ARENA_SIZE      (4ul << 20)
 #endif /* MALLOC_TARGET_EXCLAVES || MALLOC_TARGET_EXCLAVES_INTROSPECTOR */
 #define MFM_QUANTUM         16ul
 #define MFM_SIZE_CLASSES    (__builtin_ctz(MFM_ALLOC_SIZE_MAX / MFM_QUANTUM) + 1)
@@ -95,7 +97,7 @@
 	void *__ptrauth(ptrauth_key_process_dependent_data, true,
 			ptrauth_string_discriminator("mfmb_next"),
 			"authenticates-null-values")
-	                        mfmb_next;
+							mfmb_next;
 #else
 	uint64_t                mfmb_next;
 #endif
@@ -157,6 +159,9 @@
 #define mfmh_freelist       mfm_header.mfm_freelist
 
 static struct mfm_arena    *mfm_arena;
+#if CONFIG_MTE
+static bool mfm_memtag_enabled = false;
+#endif // CONFIG_MTE
 
 #pragma mark validation and helper functions
 
@@ -530,6 +535,23 @@
 #endif
 }
 
+#if CONFIG_MTE
+/*!
+ * @function __mfm_block_fixup_ptr()
+ *
+ * @brief
+ * Loads the tag for a given block pointer, and returns the
+ * tagged pointer.
+ *
+ * @discussion
+ * This must be called with mfm_memtag_enabled set to true.
+ */
+static struct mfm_block *
+__mfm_block_fixup_ptr(struct mfm_block *blk)
+{
+	return (struct mfm_block *)memtag_fixup_ptr((uint8_t *)blk);
+}
+#endif // CONFIG_MTE
 
 /*!
  * @function __mfm_block_insert_head()
@@ -551,6 +573,12 @@
 	offs = __mfm_block_offset(arena, blk);
 	next_blk = &arena->mfm_base[next];
 
+#if CONFIG_MTE
+	if (mfm_memtag_enabled) {
+		blk = __mfm_block_fixup_ptr(blk);
+		next_blk = __mfm_block_fixup_ptr(next_blk);
+	}
+#endif // CONFIG_MTE
 
 	blk->mfmb_prev = head;
 	__mfm_block_set_next(blk, next);
@@ -570,11 +598,22 @@
 	uint64_t next, prev;
 	struct mfm_block *next_blk, *prev_blk;
 
+#if CONFIG_MTE
+	if (mfm_memtag_enabled) {
+		blk = __mfm_block_fixup_ptr(blk);
+	}
+#endif // CONFIG_MTE
 
 	next = __mfm_block_next(blk);
 	prev = blk->mfmb_prev;
 	next_blk = &arena->mfm_base[next];
 	prev_blk = &arena->mfm_base[prev];
+#if CONFIG_MTE
+	if (mfm_memtag_enabled) {
+		next_blk = __mfm_block_fixup_ptr(next_blk);
+		prev_blk = __mfm_block_fixup_ptr(prev_blk);
+	}
+#endif // CONFIG_MTE
 	next_blk->mfmb_prev = prev;
 	__mfm_block_set_next(prev_blk, next);
 	__builtin_bzero(blk, sizeof(struct mfm_block));
@@ -642,33 +681,46 @@
 mfm_initialize(void)
 {
 	struct mfm_arena *arena;
+	int debug_flags;
 #if MALLOC_TARGET_EXCLAVES
 	plat_map_t map = {0};
+	debug_flags = MALLOC_NO_POPULATE;
+#else
+	int alloc_flags = 0;
+	debug_flags = DISABLE_ASLR | MALLOC_ADD_GUARD_PAGE_FLAGS;
 #endif // MALLOC_TARGET_EXCLAVES
 
-
+#if CONFIG_MTE
+	// Tie the enablement of memory tagging support in MFM to
+	// malloc_has_sec_transition: we rely on mfm_initialize
+	// being called *after* this has been setup by malloc proper.
+	mfm_memtag_enabled = malloc_has_sec_transition;
 #if MALLOC_TARGET_EXCLAVES
-	arena = mvm_allocate_pages_plat(MFM_ARENA_SIZE, 0, MALLOC_NO_POPULATE,
+	if (mfm_memtag_enabled) {
+		debug_flags |= MALLOC_MTE_TAGGABLE;
+	}
+#else
+	mfm_memtag_enabled =
+			mfm_memtag_enabled && malloc_sec_transition_early_malloc_support;
+
+	if (mfm_memtag_enabled) {
+		alloc_flags |= VM_FLAGS_MTE;
+	}
+#endif // !MALLOC_TARGET_EXCLAVES
+#endif // CONFIG_MTE
+
+	/* this is called early, which means the address space _does_ have 4M */
+	arena = mvm_allocate_pages_plat(MFM_ARENA_SIZE, 0, debug_flags,
 			VM_MEMORY_MALLOC, mvm_plat_map(map));
-#else
-	/* this is called early, which means the address space _does_ have 8M */
-	arena = mvm_allocate_pages_plat(MFM_ARENA_SIZE, 0,
-			DISABLE_ASLR | MALLOC_ADD_GUARD_PAGE_FLAGS, VM_MEMORY_MALLOC,
-			NULL);
-#endif // MALLOC_TARGET_EXCLAVES
-
 	if (arena == NULL) {
 		MFM_INTERNAL_CRASH(arena, "failed to allocate memory");
 	}
 
 #if MALLOC_TARGET_EXCLAVES
 	/* populate the header up to the block storage */
-	const uintptr_t addr = (uintptr_t)mvm_allocate_plat((uintptr_t)arena,
-			roundup(offsetof(struct mfm_arena, mfm_blocks), PAGE_SIZE), 0,
-			VM_FLAGS_FIXED, 0, 0, mvm_plat_map(map));
-	if (addr != (uintptr_t)arena) {
-		MFM_INTERNAL_CRASH(addr, "populate of header failed");
-	}
+	mvm_madvise_plat(arena,
+			roundup(offsetof(struct mfm_arena, mfm_blocks), PAGE_SIZE),
+			MADV_FAULTABLE, MALLOC_ABORT_ON_ERROR, mvm_plat_map(map));
 
 	arena->mfm_header.mfm_map = map;
 #else
@@ -678,8 +730,7 @@
 	 * originally because the kernel would have placed it in the heap range */
 	mach_vm_address_t vm_addr = (mach_vm_address_t)arena;
 	mach_vm_size_t vm_size = (mach_vm_size_t)MFM_ARENA_SIZE;
-	int alloc_flags = VM_FLAGS_OVERWRITE | VM_MAKE_TAG(VM_MEMORY_MALLOC_TINY);
-
+	alloc_flags |= VM_FLAGS_OVERWRITE | VM_MAKE_TAG(VM_MEMORY_MALLOC_TINY);
 
 	kern_return_t kr = mach_vm_map(mach_task_self(), &vm_addr, vm_size,
 			/* mask */ 0, alloc_flags, MEMORY_OBJECT_NULL, /* offset */ 0,
@@ -719,6 +770,9 @@
 	struct mfm_arena *arena = os_atomic_load(&mfm_arena, dependency);
 	size_t index;
 
+#if CONFIG_MTE
+	ptr = memtag_strip_address((void *)ptr);
+#endif // CONFIG_MTE
 
 	if (!__mfm_address_owned(arena, ptr)) {
 		return 0ul;
@@ -769,7 +823,7 @@
 			if (blk_size > size) {
 				__mfm_block_mark_start(arena, blk_index + size);
 				__mfm_free_block(arena, blk_index + size,
-				    blk_size - size);
+					blk_size - size);
 			}
 
 			__mfm_block_mark_allocated(arena, blk_index, size);
@@ -799,12 +853,9 @@
 			const uintptr_t end = roundup((uintptr_t)ptr + alloc_size, PAGE_SIZE);
 			const size_t bytes = end - begin;
 			if (bytes) {
-				const uintptr_t addr = (uintptr_t)mvm_allocate_plat(begin,
-						bytes, 0, VM_FLAGS_FIXED, 0, 0,
+				mvm_madvise_plat((void*)begin, bytes, MADV_FAULTABLE,
+						MALLOC_ABORT_ON_ERROR,
 						mvm_plat_map(arena->mfm_header.mfm_map));
-				if (addr != begin) {
-					MFM_INTERNAL_CRASH(ptr, "populate of pages failed");
-				}
 			}
 #endif
 
@@ -820,6 +871,14 @@
 #endif
 	__mfm_unlock(arena);
 
+#if CONFIG_MTE
+	// Set the tag for the block we are returning to the caller.
+	if (mfm_memtag_enabled) {
+		ptr = memtag_assign_tag(ptr, size * MFM_QUANTUM);
+		// Blocks from the early arena may have arbitrary 16B-alignment
+		memtag_set_tag_unaligned(ptr, size * MFM_QUANTUM);
+	}
+#endif // CONFIG_MTE
 
 	return ptr;
 }
@@ -835,6 +894,9 @@
 	dprintf(STDERR_FILENO, "{ -1, %p },\n", ptr);
 #endif
 
+#if CONFIG_MTE
+	addr = memtag_strip_address((void *)ptr);
+#endif // CONFIG_MTE
 
 	if (!__mfm_address_owned(arena, addr)) {
 		MFM_INTERNAL_CRASH(ptr, "not MFM owned");
@@ -846,6 +908,14 @@
 	}
 	size = __mfm_block_size(arena, index);
 
+#if CONFIG_MTE
+	// Retag the block we are freeing.
+	if (mfm_memtag_enabled) {
+		ptr = memtag_assign_tag(ptr, MFM_QUANTUM * size);
+		// Blocks from the early arena may have arbitrary 16B-alignment
+		memtag_set_tag_unaligned(ptr, MFM_QUANTUM * size);
+	}
+#endif // CONFIG_MTE
 
 	bzero(ptr, MFM_QUANTUM * size);
 
@@ -872,7 +942,7 @@
 	}
 
 	if (index + size < arena->mfmh_bump &&
-	    !__mfm_block_is_allocated(arena, index + size)) {
+			!__mfm_block_is_allocated(arena, index + size)) {
 		size_t next  = index + size;
 		size_t nsize = __mfm_block_size(arena, next);
 
@@ -897,8 +967,31 @@
 {
 	struct mfm_arena *arena = os_atomic_load(&mfm_arena, dependency);
 
-
+#if CONFIG_MTE
+	// After checking that this pointer belongs to us, verify that the logical
+	// tag of the pointer matches the physical tag stored in memory. If that is
+	// not the case, report this pointer as not claimed.
+	//
+	// The layers above us (xzone and malloc proper) rely on this behaviour.
+	// This is required to properly handle (i.e. fail) the lookup of a pointer
+	// with a mismatching tag: in the free() and realloc() paths, the dispatch
+	// layer (malloc proper) will try to validate the tag of the pointer, and
+	// raise a fatal exception if it doesn't match the tag stored in memory.
+	//
+	// Note: we need to validate the tag after ensuring that we own the
+	// pointer, otherwise we might incur in passing complete garbage to ldg,
+	// which would lead to a crash even in paths that need to possibly operate
+	// on completely invalid pointers (e.g. malloc_size()).
+	bool claimed = __mfm_address_owned(arena, memtag_strip_address(ptr));
+	if (claimed && mfm_memtag_enabled) {
+		if (!memtag_tags_match(ptr, memtag_fixup_ptr(ptr))) {
+			return false;
+		}
+	}
+	return claimed;
+#else
 	return __mfm_address_owned(arena, ptr);
+#endif
 }
 
 void *
@@ -934,7 +1027,7 @@
 			size_t size  = __mfm_block_size(arena, index);
 
 			P("  [%p, %p) size=%zd\n",
-			    blk, blk + size, size * MFM_QUANTUM);
+					blk, blk + size, size * MFM_QUANTUM);
 		}
 	}
 	P("\n");