Set thread affinity during ioctl calls (#580)

Signed-off-by: Alan Jowett <alanjo@microsoft.com>

Co-authored-by: Dave Thaler <dthaler@microsoft.com>

libs/execution_context/ebpf_core.c

@@ -1384,6 +1384,11 @@ ebpf_core_invoke_protocol_handler(
     if (operation_id >= EBPF_COUNT_OF(_ebpf_protocol_handlers) || operation_id < EBPF_OPERATION_RESOLVE_HELPER) {
         return EBPF_OPERATION_NOT_SUPPORTED;
     }
+    uintptr_t old_affinity_mask = 0;
+
+    retval = ebpf_set_current_thread_affinity((uintptr_t)1 << ebpf_get_current_cpu(), &old_affinity_mask);
+    if (retval != EBPF_SUCCESS)
+        return retval;
 
     retval = ebpf_epoch_enter();
     if (retval != EBPF_SUCCESS)
@@ -1396,5 +1401,7 @@ ebpf_core_invoke_protocol_handler(
             input_buffer, output_buffer, output_buffer_length);
 
     ebpf_epoch_exit();
+
+    ebpf_restore_current_thread_affinity(old_affinity_mask);
     return retval;
 }

libs/platform/ebpf_platform.h

@@ -875,6 +875,12 @@ extern "C"
     uint64_t
     ebpf_query_time_since_boot(bool include_suspended_time);
 
+    ebpf_result_t
+    ebpf_set_current_thread_affinity(uintptr_t new_thread_affinity_mask, _Out_ uintptr_t* old_thread_affinity_mask);
+
+    void
+    ebpf_restore_current_thread_affinity(uintptr_t old_thread_affinity_mask);
+
 #ifdef __cplusplus
 }
 #endif

libs/platform/kernel/ebpf_platform_kernel.c

@@ -255,6 +255,28 @@ ebpf_interlocked_compare_exchange_pointer(
     return InterlockedCompareExchangePointer((void* volatile*)destination, (void*)exchange, (void*)comperand);
 }
 
+ebpf_result_t
+ebpf_set_current_thread_affinity(uintptr_t new_thread_affinity_mask, _Out_ uintptr_t* old_thread_affinity_mask)
+{
+    if (KeGetCurrentIrql() > PASSIVE_LEVEL) {
+        return EBPF_OPERATION_NOT_SUPPORTED;
+    }
+
+    KAFFINITY old_affinity = KeSetSystemAffinityThreadEx(new_thread_affinity_mask);
+    if (old_affinity == 0) {
+        return EBPF_OPERATION_NOT_SUPPORTED;
+    } else {
+        *old_thread_affinity_mask = old_affinity;
+        return EBPF_SUCCESS;
+    }
+}
+
+void
+ebpf_restore_current_thread_affinity(uintptr_t old_thread_affinity_mask)
+{
+    KeRevertToUserAffinityThreadEx(old_thread_affinity_mask);
+}
+
 _Ret_range_(>, 0) uint32_t ebpf_get_cpu_count() { return KeQueryMaximumProcessorCount(); }
 
 bool

libs/platform/user/ebpf_platform_user.cpp

@@ -58,7 +58,11 @@ ebpf_free(_Frees_ptr_opt_ void* memory)
 __drv_allocatesMem(Mem) _Must_inspect_result_ _Ret_maybenull_
     _Post_writable_byte_size_(size) void* ebpf_allocate_cache_aligned(size_t size)
 {
-    return _aligned_malloc(size, EBPF_CACHE_LINE_SIZE);
+    void* memory = _aligned_malloc(size, EBPF_CACHE_LINE_SIZE);
+    if (memory) {
+        memset(memory, 0, size);
+    }
+    return memory;
 }
 
 void
@@ -246,6 +250,24 @@ ebpf_query_time_since_boot(bool include_suspended_time)
     return interrupt_time;
 }
 
+ebpf_result_t
+ebpf_set_current_thread_affinity(uintptr_t new_thread_affinity_mask, _Out_ uintptr_t* old_thread_affinity_mask)
+{
+    uintptr_t old_mask = SetThreadAffinityMask(GetCurrentThread(), new_thread_affinity_mask);
+    if (old_mask == 0) {
+        return EBPF_OPERATION_NOT_SUPPORTED;
+    } else {
+        *old_thread_affinity_mask = old_mask;
+        return EBPF_SUCCESS;
+    }
+}
+
+void
+ebpf_restore_current_thread_affinity(uintptr_t old_thread_affinity_mask)
+{
+    SetThreadAffinityMask(GetCurrentThread(), old_thread_affinity_mask);
+}
+
 _Ret_range_(>, 0) uint32_t ebpf_get_cpu_count()
 {
     SYSTEM_INFO system_info;