Merge pull request #102 from devnexen/sched_getcpu_as_util

Allowing to enabled CPU_PIN on other oses than Linux (FreeBSD at the …

Author: struct

Makefile

@@ -93,6 +93,8 @@ HUGE_PAGES = -DHUGE_PAGES=1
 ## and has negative performance implications
 CPU_PIN = -DCPU_PIN=0
 
+SCHED_GETCPU =
+
 ## Enable the allocation sanity feature. This works a lot
 ## like GWP-ASAN does. It samples calls to iso_alloc and
 ## randomly swaps them out for raw page allocations that
@@ -169,12 +171,16 @@ endif
 
 ifeq ($(UNAME), Linux)
 STRIP = strip -s $(BUILD_DIR)/$(LIBNAME)
+SCHED_GETCPU = -DSCHED_GETCPU
 endif
 
 ifeq ($(UNAME), FreeBSD)
 STRIP = strip -s $(BUILD_DIR)/$(LIBNAME)
 ## Using spinlocks to avoid recursive locks contentions with calloc
 USE_SPINLOCK = -DUSE_SPINLOCK=1
+## Once FreeBSD 13.1 becomes the minimal non EOL version
+## it can be enabled
+## SCHED_GETCPU = -DSCHED_GETCPU
 endif
 
 HOOKS = $(MALLOC_HOOK)
@@ -187,7 +193,7 @@ else
 BUILD_ERROR_FLAGS := $(BUILD_ERROR_FLAGS) -Wno-attributes -Wno-unused-variable
 endif
 CFLAGS = $(COMMON_CFLAGS) $(SECURITY_FLAGS) $(BUILD_ERROR_FLAGS) $(HOOKS) $(HEAP_PROFILER) -fvisibility=hidden \
-	-std=c11 $(SANITIZER_SUPPORT) $(ALLOC_SANITY) $(UNINIT_READ_SANITY) $(CPU_PIN) $(EXPERIMENTAL) $(UAF_PTR_PAGE) \
+	-std=c11 $(SANITIZER_SUPPORT) $(ALLOC_SANITY) $(UNINIT_READ_SANITY) $(CPU_PIN) $(SCHED_GETCPU) $(EXPERIMENTAL) $(UAF_PTR_PAGE) \
 	$(VERIFY_BIT_SLOT_CACHE) $(NAMED_MAPPINGS) $(ABORT_ON_NULL) $(NO_ZERO_ALLOCATIONS) $(ABORT_NO_ENTROPY) \
 	$(ISO_DTOR_CLEANUP) $(SHUFFLE_BIT_SLOT_CACHE) $(USE_SPINLOCK) $(HUGE_PAGES) $(USE_MLOCK) $(MEMORY_TAGGING)
 CXXFLAGS = $(COMMON_CFLAGS) -DCPP_SUPPORT=1 -std=c++17 $(SANITIZER_SUPPORT) $(HOOKS)

include/iso_alloc_internal.h

@@ -605,3 +605,7 @@ INTERNAL_HIDDEN void _iso_alloc_search_stack(uint8_t *stack_start);
 #if UNIT_TESTING
 EXTERNAL_API iso_alloc_root *_get_root(void);
 #endif
+
+#if CPU_PIN
+INTERNAL_HIDDEN INLINE int _iso_getcpu(void);
+#endif

src/iso_alloc.c

@@ -808,7 +808,7 @@ INTERNAL_HIDDEN iso_alloc_zone_t *_iso_new_zone(size_t size, bool internal, int3
     get_next_free_bit_slot(new_zone);
 
 #if CPU_PIN
-    new_zone->cpu_core = sched_getcpu();
+    new_zone->cpu_core = _iso_getcpu();
 #endif
 
     POISON_ZONE(new_zone);
@@ -899,7 +899,7 @@ INTERNAL_HIDDEN bit_slot_t iso_scan_zone_free_slot_slow(iso_alloc_zone_t *zone)
 
 INTERNAL_HIDDEN iso_alloc_zone_t *is_zone_usable(iso_alloc_zone_t *zone, size_t size) {
 #if CPU_PIN
-    if(zone->cpu_core != sched_getcpu()) {
+    if(zone->cpu_core != _iso_getcpu()) {
         return false;
     }
 #endif
@@ -2102,3 +2102,26 @@ EXTERNAL_API iso_alloc_root *_get_root(void) {
     return _root;
 }
 #endif
+
+#if CPU_PIN
+/* sched_getcpu's performance depends on the
+ * architecture/kernel version, so we lower
+ * the cost of feature's abstraction here.
+ */
+INTERNAL_HIDDEN INLINE int _iso_getcpu(void) {
+#if defined(SCHED_GETCPU)
+    return sched_getcpu();
+#elif defined(__x86_64__)
+    /* rdtscp is not always available and is pretty slow
+     * we instead load from the global descriptor table
+     * then "mov" it to a.
+     */
+    unsigned int a;
+    const unsigned int cpunodesegment = 15 * 8 + 3;
+    __asm__ volatile("lsl %1, %0" : "=r"(a) : "r"(cpunodesegment));
+    return (int)(a & 0xfff);
+#else
+    return -1;
+#endif
+}
+#endif