Loading...
--- libmalloc/libmalloc-317.140.5/src/malloc.c
+++ libmalloc/libmalloc-317.40.8/src/malloc.c
@@ -50,7 +50,6 @@
 malloc_zone_t **malloc_zones = (malloc_zone_t **)0xdeaddeaddeaddead;
 malloc_logger_t *malloc_logger = NULL;
 
-static uint32_t initial_num_zones;
 static malloc_zone_t *initial_scalable_zone;
 static malloc_zone_t *initial_nano_zone;
 static malloc_zone_t *initial_default_zone = NULL;
@@ -122,7 +121,7 @@
 #define DEFAULT_MALLOC_ZONE_STRING "DefaultMallocZone"
 #define DEFAULT_PUREGEABLE_ZONE_STRING "DefaultPurgeableMallocZone"
 #define MALLOC_HELPER_ZONE_STRING "MallocHelperZone"
-#define MALLOC_PGUARD_ZONE_STRING "ProbGuardMallocZone"
+#define MALLOC_PGUARD_ZONE_STRING "PGuardMallocZone"
 
 MALLOC_NOEXPORT
 unsigned int phys_ncpus;
@@ -201,24 +200,6 @@
 	return (*(uint64_t*)_COMM_PAGE_CPU_CAPABILITIES64) & kIsTranslated;
 }
 #endif /* TARGET_OS_OSX */
-
-
-#define LIBMALLOC_EXPERIMENT_FACTORS_KEY "MallocExperiment="
-#define LIBMALLOC_EXPERIMENT_DISABLE_MEDIUM (1ULL)
-static void
-__malloc_init_experiments(const char *str)
-{
-	uint64_t experiment_factors = 0;
-	str = strchr(str, '=');
-	if (str) {
-		experiment_factors = strtoull_l(str + 1, NULL, 16, NULL);
-	}
-	switch (experiment_factors) {
-	case LIBMALLOC_EXPERIMENT_DISABLE_MEDIUM:
-		magazine_medium_enabled = false;
-		break;
-	}
-}
 
 static void
 __malloc_init_from_bootargs(const char *bootargs)
@@ -341,7 +322,6 @@
 	}
 
 	const char **p;
-	const char *malloc_experiments = NULL;
 	for (p = apple; p && *p; p++) {
 		if (strstr(*p, "malloc_entropy") == *p) {
 			int count = __entropy_from_kernel(*p);
@@ -350,9 +330,7 @@
 			if (sizeof(malloc_entropy) / sizeof(malloc_entropy[0]) == count) {
 				_malloc_entropy_initialized = true;
 			}
-		}
-		if (strstr(*p, LIBMALLOC_EXPERIMENT_FACTORS_KEY) == *p) {
-			malloc_experiments = *p;
+			break;
 		}
 	}
 	if (!_malloc_entropy_initialized) {
@@ -360,9 +338,6 @@
 		_malloc_entropy_initialized = true;
 	}
 
-	if (malloc_experiments) {
-		__malloc_init_experiments(malloc_experiments);
-	}
 	__malloc_init_from_bootargs(bootargs);
 	mvm_aslr_init();
 
@@ -377,19 +352,6 @@
 	force_asan_init_if_present();
 
 	_malloc_initialize(apple, bootargs);
-}
-
-static void register_pgm_zone(bool internal_diagnostics);
-static void stack_logging_early_finished(const struct _malloc_late_init *funcs);
-
-// WARNING: The passed _malloc_late_init is a stack variable in
-// libSystem_initializer().  We must not hold on to it.
-void
-__malloc_late_init(const struct _malloc_late_init *mli)
-{
-	register_pgm_zone(mli->internal_diagnostics);
-	stack_logging_early_finished(mli);
-	initial_num_zones = malloc_num_zones;
 }
 
 MALLOC_NOEXPORT malloc_zone_t* lite_zone = NULL;
@@ -672,37 +634,32 @@
 			return default_zone;
 		}
 	}
-
-	malloc_zone_t *zone;
-	size_t size;
-
-	// We assume that the initial zones will never be unregistered concurrently while this code is running so we can have
-	// a fast path without locking.  Callers who really do unregister these (to install their own default zone) need to
-	// ensure they establish their zone setup during initialization and before entering a multi-threaded environment.
-	for (uint32_t i = 0; i < initial_num_zones; i++) {
-		zone = malloc_zones[i];
-		size = zone->size(zone, ptr);
-
-		if (size) { // Claimed by this zone?
-			if (returned_size) {
-				*returned_size = size;
-			}
-
-			// Asan and others replace the zone at position 0 with their own zone.
-			// In that case just return that zone as they need this information.
-			// Otherwise return the virtual default zone, not the actual zone in position 0.
-			if (i == 0 && has_default_zone0()) {
-				return default_zone;
-			}
-
+	
+	// The default zone is registered in malloc_zones[0]. There's no danger that it will ever be unregistered.
+	// So don't advance the FRZ counter yet.
+	malloc_zone_t *zone = malloc_zones[0];
+	size_t size = zone->size(zone, ptr);
+	if (size) { // Claimed by this zone?
+		if (returned_size) {
+			*returned_size = size;
+		}
+
+		// Asan and others replace the zone at position 0 with their own zone.
+		// In that case just return that zone as they need this information.
+		// Otherwise return the virtual default zone, not the actual zone in position 0.
+		if (!has_default_zone0()) {
 			return zone;
+		} else {
+			return default_zone;
 		}
 	}
 
 	int32_t volatile *pFRZCounter = pFRZCounterLive;   // Capture pointer to the counter of the moment
 	OSAtomicIncrement32Barrier(pFRZCounter); // Advance this counter -- our thread is in FRZ
 
+	unsigned index;
 	int32_t limit = *(int32_t volatile *)&malloc_num_zones;
+	malloc_zone_t **zones = &malloc_zones[1];
 
 	// From this point on, FRZ is accessing the malloc_zones[] array without locking
 	// in order to avoid contention on common operations (such as non-default-zone free()).
@@ -719,8 +676,8 @@
 	//      are still valid). It also ensures that all the pointers in the zones array are
 	//      valid until it returns, so that a stale value in limit is not dangerous.
 
-	for (uint32_t i = initial_num_zones; i < limit; i++) {
-		zone = malloc_zones[i];
+	for (index = 1; index < limit; ++index, ++zones) {
+		zone = *zones;
 		size = zone->size(zone, ptr);
 		if (size) { // Claimed by this zone?
 			goto out;
@@ -907,7 +864,7 @@
 	nano_common_init(envp, apple, bootargs);
 #endif
 
-	const uint32_t k_max_zones = 2;
+	const uint32_t k_max_zones = 3;
 	malloc_zone_t *zone_stack[k_max_zones];
 	const char *name_stack[k_max_zones];
 	uint32_t num_zones = 0;
@@ -938,6 +895,14 @@
 	}
 #endif
 
+	if (pguard_enabled()) {
+		malloc_zone_t *wrapped_zone = zone_stack[num_zones - 1];
+		zone_stack[num_zones] = pguard_create_zone(wrapped_zone, malloc_debug_flags);
+		name_stack[num_zones] = MALLOC_PGUARD_ZONE_STRING;
+		// TODO(yln): what is the external contract for zone names?
+		num_zones++;
+	}
+
 	MALLOC_ASSERT(num_zones <= k_max_zones);
 	initial_default_zone = zone_stack[num_zones - 1];
 
@@ -945,25 +910,9 @@
 	for (int i = num_zones - 1; i >= 0; i--) malloc_zone_register_while_locked(zone_stack[i]);
 	for (int i = num_zones - 1; i >= 0; i--) malloc_set_zone_name(zone_stack[i], name_stack[i]);
 
-	initial_num_zones = malloc_num_zones;
-
 	// malloc_report(ASL_LEVEL_INFO, "%d registered zones\n", malloc_num_zones);
 	// malloc_report(ASL_LEVEL_INFO, "malloc_zones is at %p; malloc_num_zones is at %p\n", (unsigned)&malloc_zones,
 	// (unsigned)&malloc_num_zones);
-}
-
-static void make_last_zone_default_zone(void);
-static void
-register_pgm_zone(bool internal_diagnostics)
-{
-	if (pguard_enabled(internal_diagnostics)) {
-		malloc_zone_t *wrapped_zone = malloc_zones[0];
-		malloc_zone_t *pgm_zone = pguard_create_zone(wrapped_zone);
-		malloc_zone_register_while_locked(pgm_zone);
-		make_last_zone_default_zone();
-		initial_default_zone = pgm_zone;
-		malloc_set_zone_name(pgm_zone, MALLOC_PGUARD_ZONE_STRING);
-	}
 }
 
 static inline malloc_zone_t *
@@ -1395,23 +1344,6 @@
 	return zone;
 }
 
-static void
-make_last_zone_default_zone(void)
-{
-	unsigned protect_size = malloc_num_zones_allocated * sizeof(malloc_zone_t *);
-	mprotect(malloc_zones, protect_size, PROT_READ | PROT_WRITE);
-
-	malloc_zone_t *last_zone = malloc_zones[malloc_num_zones - 1];
-
-	// assert(zone == malloc_zones[malloc_num_zones - 1];
-	for (int i = malloc_num_zones - 1; i > 0; --i) {
-		malloc_zones[i] = malloc_zones[i - 1];
-	}
-	malloc_zones[0] = last_zone;
-
-	mprotect(malloc_zones, protect_size, PROT_READ);
-}
-
 /*
  * For use by CheckFix: establish a new default zone whose behavior is, apart from
  * the use of death-row and per-CPU magazines, that of Leopard.
@@ -1420,6 +1352,7 @@
 malloc_create_legacy_default_zone(void)
 {
 	malloc_zone_t *zone;
+	int i;
 
 	zone = create_legacy_scalable_zone(0, malloc_debug_flags);
 
@@ -1435,7 +1368,16 @@
 	}
 	malloc_set_zone_name(zone, DEFAULT_MALLOC_ZONE_STRING);
 
-	make_last_zone_default_zone();
+	unsigned protect_size = malloc_num_zones_allocated * sizeof(malloc_zone_t *);
+	mprotect(malloc_zones, protect_size, PROT_READ | PROT_WRITE);
+
+	// assert(zone == malloc_zones[malloc_num_zones - 1];
+	for (i = malloc_num_zones - 1; i > 0; --i) {
+		malloc_zones[i] = malloc_zones[i - 1];
+	}
+	malloc_zones[0] = zone;
+
+	mprotect(malloc_zones, protect_size, PROT_READ);
 	MALLOC_UNLOCK();
 }
 
@@ -1830,10 +1772,6 @@
 		--malloc_num_zones;
 
 		mprotect(malloc_zones, protect_size, PROT_READ);
-
-		// MAX(num_zones, 1) enables the fast path in find_registered_zone() for zone 0 even
-		// if it is a custom zone, e.g., ASan and user zones.
-		initial_num_zones = MIN(MAX(malloc_num_zones, 1), initial_num_zones);
 
 		// Exchange the roles of the FRZ counters. The counter that has captured the number of threads presently
 		// executing *inside* find_registered_zone is swapped with the counter drained to zero last time through.
@@ -2055,11 +1993,9 @@
 		return true;
 	}
 
-	// Next, try the initial zones.
-	for (uint32_t i = 0; i < initial_num_zones; i++) {
-		if (malloc_zone_claimed_address(malloc_zones[i], ptr)) {
-			return true;
-		}
+	// Next, try the default zone, which is always present.
+	if (malloc_zone_claimed_address(malloc_zones[0], ptr)) {
+		return true;
 	}
 
 	// Try all the other zones. Increment the FRZ barrier so that we can
@@ -2069,9 +2005,10 @@
 	OSAtomicIncrement32Barrier(pFRZCounter);
 
 	int32_t limit = *(int32_t volatile *)&malloc_num_zones;
+	malloc_zone_t **zones = &malloc_zones[1];
 	boolean_t result = false;
-	for (uint32_t i = initial_num_zones; i < limit; i++) {
-		malloc_zone_t *zone = malloc_zones[i];
+	for (unsigned index = 1; index < limit; ++index, ++zones) {
+		malloc_zone_t *zone = *zones;
 		if (malloc_zone_claimed_address(zone, ptr)) {
 			result = true;
 			break;
@@ -2660,8 +2597,8 @@
 
 
 /* this is called from libsystem during initialization. */
-static void
-stack_logging_early_finished(const struct _malloc_late_init *funcs)
+void
+__stack_logging_early_finished(const struct _malloc_functions *funcs)
 {
 #if !TARGET_OS_DRIVERKIT
 	_dlopen = funcs->dlopen;