Fix closing native program / maps handles in system worker thread. (#…

…2500) * fix * fix analyze build * fix memory leak * add tests * code reorder * mark new tests for DEBUG only * fix Release build * CR comments * cr comments
microsoft · May 26, 2023 · c471719 · c471719
1 parent 904608d
commit c471719
Show file tree

Hide file tree

Showing 7 changed files with 213 additions and 4 deletions.
diff --git a/libs/execution_context/ebpf_native.c b/libs/execution_context/ebpf_native.c
@@ -58,6 +58,8 @@ typedef enum _ebpf_native_module_state
 
 typedef struct _ebpf_native_handle_cleanup_information
 {
+    intptr_t process_handle;
+    ebpf_process_state_t* process_state;
     size_t count_of_program_handles;
     ebpf_handle_t* program_handles;
     size_t count_of_map_handles;
@@ -1272,6 +1274,10 @@ _ebpf_native_close_handles_work_item(_In_opt_ const void* context)
     }
 
     ebpf_native_handle_cleanup_info_t* handle_info = (ebpf_native_handle_cleanup_info_t*)context;
+
+    // Attach process to this worker thread.
+    ebpf_platform_attach_process(handle_info->process_handle, handle_info->process_state);
+
     for (uint32_t i = 0; i < handle_info->count_of_program_handles; i++) {
         if (handle_info->program_handles[i] != ebpf_handle_invalid) {
             (void)ebpf_handle_close(handle_info->program_handles[i]);
@@ -1285,6 +1291,13 @@ _ebpf_native_close_handles_work_item(_In_opt_ const void* context)
         }
     }
 
+    // Detach process from this worker thread.
+    ebpf_platform_detach_process(handle_info->process_state);
+
+    // Release the reference on the process object.
+    ebpf_platform_dereference_process(handle_info->process_handle);
+
+    ebpf_free(handle_info->process_state);
     ebpf_free(handle_info->program_handles);
     ebpf_free(handle_info->map_handles);
 }
@@ -1299,12 +1312,16 @@ _ebpf_native_clean_up_handle_cleanup_context(_Inout_ ebpf_native_handle_cleanup_
     if (cleanup_context->handle_information != NULL) {
         ebpf_free(cleanup_context->handle_information->map_handles);
         ebpf_free(cleanup_context->handle_information->program_handles);
+        ebpf_free(cleanup_context->handle_information->process_state);
+
+        if (cleanup_context->handle_information->process_handle != 0) {
+            ebpf_platform_dereference_process(cleanup_context->handle_information->process_handle);
+        }
     }
 
     if (cleanup_context->handle_cleanup_work_item != NULL) {
         // ebpf_free_preemptible_work_item() will free the handle_information.
         ebpf_free_preemptible_work_item(cleanup_context->handle_cleanup_work_item);
-
     } else {
         ebpf_free(cleanup_context->handle_information);
     }
@@ -1350,10 +1367,21 @@ _ebpf_native_initialize_handle_cleanup_context(
         cleanup_context->handle_information->program_handles[i] = ebpf_handle_invalid;
     }
 
+    cleanup_context->handle_information->process_state = ebpf_allocate_process_state();
+    if (cleanup_context->handle_information->process_state == NULL) {
+        result = EBPF_NO_MEMORY;
+        goto Done;
+    }
+
     result = ebpf_allocate_preemptible_work_item(
         &cleanup_context->handle_cleanup_work_item,
         _ebpf_native_close_handles_work_item,
         cleanup_context->handle_information);
+    if (result != EBPF_SUCCESS) {
+        goto Done;
+    }
+
+    cleanup_context->handle_information->process_handle = ebpf_platform_reference_process();
 
 Done:
     if (result != EBPF_SUCCESS) {

diff --git a/libs/execution_context/ebpf_program.c b/libs/execution_context/ebpf_program.c
@@ -1795,6 +1795,7 @@ _Must_inspect_result_ ebpf_result_t
 ebpf_program_create_and_initialize(
     _In_ const ebpf_program_parameters_t* parameters, _Out_ ebpf_handle_t* program_handle)
 {
+    EBPF_LOG_ENTRY();
     ebpf_result_t retval;
     ebpf_program_t* program = NULL;
 
@@ -1810,7 +1811,7 @@ ebpf_program_create_and_initialize(
 
 Done:
     EBPF_OBJECT_RELEASE_REFERENCE((ebpf_core_object_t*)program);
-    return retval;
+    EBPF_RETURN_RESULT(retval);
 }
 
 typedef struct _ebpf_helper_id_to_index

diff --git a/libs/platform/ebpf_platform.h b/libs/platform/ebpf_platform.h
@@ -95,6 +95,8 @@ extern "C"
     typedef uintptr_t ebpf_lock_t;
     typedef uint8_t ebpf_lock_state_t;
 
+    typedef struct _ebpf_process_state ebpf_process_state_t;
+
     // A self-relative security descriptor.
     typedef struct _SECURITY_DESCRIPTOR ebpf_security_descriptor_t;
     typedef struct _GENERIC_MAPPING ebpf_security_generic_mapping_t;
@@ -1053,6 +1055,48 @@ extern "C"
     uint32_t
     ebpf_platform_thread_id();
 
+    /**
+     * @brief Allocate memory for process state. Caller needs to call ebpf_free()
+     *  to free the memory.
+     *
+     * @return Pointer to the process state.
+     */
+    _Ret_maybenull_ ebpf_process_state_t*
+    ebpf_allocate_process_state();
+
+    /**
+     * @brief Get a handle to the current process.
+     *
+     * @return Handle to the current process.
+     */
+    intptr_t
+    ebpf_platform_reference_process();
+
+    /**
+     * @brief Dereference a handle to a process.
+     *
+     * @param[in] process_handle to the process.
+     */
+    void
+    ebpf_platform_dereference_process(intptr_t process_handle);
+
+    /**
+     * @brief Attach to the specified process.
+     *
+     * @param[in] handle to the process.
+     * @param[in,out] state Pointer to the process state.
+     */
+    void
+    ebpf_platform_attach_process(intptr_t process_handle, _Inout_ ebpf_process_state_t* state);
+
+    /**
+     * @brief Detach from the current process.
+     *
+     * @param[in] state Pointer to the process state.
+     */
+    void
+    ebpf_platform_detach_process(_In_ ebpf_process_state_t* state);
+
     TRACELOGGING_DECLARE_PROVIDER(ebpf_tracelog_provider);
 
     _Must_inspect_result_ ebpf_result_t
@@ -1286,7 +1330,7 @@ extern "C"
         TraceLoggingLevel(WINEVENT_LEVEL_VERBOSE),                      \
         TraceLoggingKeyword(EBPF_TRACELOG_KEYWORD_FUNCTION_ENTRY_EXIT), \
         TraceLoggingOpcode(WINEVENT_OPCODE_START),                      \
-        TraceLoggingString(__FUNCTION__, "<=="));
+        TraceLoggingString(__FUNCTION__, "==>"));
 
 #define EBPF_LOG_EXIT()                                                 \
     TraceLoggingWrite(                                                  \
@@ -1295,7 +1339,7 @@ extern "C"
         TraceLoggingLevel(WINEVENT_LEVEL_VERBOSE),                      \
         TraceLoggingKeyword(EBPF_TRACELOG_KEYWORD_FUNCTION_ENTRY_EXIT), \
         TraceLoggingOpcode(WINEVENT_OPCODE_STOP),                       \
-        TraceLoggingString(__FUNCTION__, "==>"));
+        TraceLoggingString(__FUNCTION__, "<=="));
 
 #define EBPF_RETURN_ERROR(error)                   \
     do {                                           \

diff --git a/libs/platform/kernel/ebpf_handle_kernel.c b/libs/platform/kernel/ebpf_handle_kernel.c
@@ -67,6 +67,10 @@ ebpf_handle_create(_Out_ ebpf_handle_t* handle, _Inout_ ebpf_base_object_t* obje
     *handle = (ebpf_handle_t)file_handle;
     file_handle = 0;
     return_value = EBPF_SUCCESS;
+
+    EBPF_LOG_MESSAGE_UINT64(
+        EBPF_TRACELOG_LEVEL_VERBOSE, EBPF_TRACELOG_KEYWORD_CORE, "ebpf_handle_create: returning handle", *handle);
+
 Done:
     if (file_object) {
         ObDereferenceObject(file_object);
@@ -83,6 +87,10 @@ _Must_inspect_result_ ebpf_result_t
 ebpf_handle_close(ebpf_handle_t handle)
 {
     EBPF_LOG_ENTRY();
+
+    EBPF_LOG_MESSAGE_UINT64(
+        EBPF_TRACELOG_LEVEL_VERBOSE, EBPF_TRACELOG_KEYWORD_CORE, "ebpf_handle_close: closing handle", (uint64_t)handle);
+
     NTSTATUS status = ObCloseHandle((HANDLE)handle, UserMode);
     if (!NT_SUCCESS(status)) {
         EBPF_LOG_NTSTATUS_API_FAILURE(EBPF_TRACELOG_KEYWORD_BASE, ObCloseHandle, status);

diff --git a/libs/platform/kernel/ebpf_platform_kernel.c b/libs/platform/kernel/ebpf_platform_kernel.c
@@ -15,6 +15,11 @@ static uint32_t _ebpf_platform_maximum_processor_count = 0;
 extern DEVICE_OBJECT*
 ebpf_driver_get_device_object();
 
+typedef struct _ebpf_process_state
+{
+    KAPC_STATE state;
+} ebpf_process_state_t;
+
 typedef struct _ebpf_memory_descriptor
 {
     MDL memory_descriptor_list;
@@ -920,3 +925,37 @@ ebpf_utf8_string_to_unicode(_In_ const ebpf_utf8_string_t* input, _Outptr_ wchar
     ebpf_free(unicode_string);
     return retval;
 }
+
+intptr_t
+ebpf_platform_reference_process()
+{
+    PEPROCESS process = PsGetCurrentProcess();
+    ObReferenceObject(process);
+    return (intptr_t)process;
+}
+
+_Ret_maybenull_ ebpf_process_state_t*
+ebpf_allocate_process_state()
+{
+    // Skipping fault injection as call to ebpf_allocate() covers it.
+    ebpf_process_state_t* state = ebpf_allocate(sizeof(ebpf_process_state_t));
+    return state;
+}
+
+void
+ebpf_platform_dereference_process(intptr_t process_handle)
+{
+    ObDereferenceObject((PEPROCESS)process_handle);
+}
+
+void
+ebpf_platform_attach_process(intptr_t process_handle, _Inout_ ebpf_process_state_t* state)
+{
+    KeStackAttachProcess((PEPROCESS)process_handle, &state->state);
+}
+
+void
+ebpf_platform_detach_process(_In_ ebpf_process_state_t* state)
+{
+    KeUnstackDetachProcess(&state->state);
+}
diff --git a/libs/platform/user/ebpf_platform_user.cpp b/libs/platform/user/ebpf_platform_user.cpp
@@ -37,6 +37,11 @@ static EX_RUNDOWN_REF _ebpf_platform_preemptible_work_items_rundown;
 
 ebpf_leak_detector_ptr _ebpf_leak_detector_ptr;
 
+typedef struct _ebpf_process_state
+{
+    uint8_t unused;
+} ebpf_process_state_t;
+
 /**
  * @brief Environment variable to enable fault injection testing.
  *
@@ -1434,3 +1439,41 @@ ebpf_utf8_string_to_unicode(_In_ const ebpf_utf8_string_t* input, _Outptr_ wchar
     ebpf_free(unicode_string);
     return retval;
 }
+
+_Ret_maybenull_ ebpf_process_state_t*
+ebpf_allocate_process_state()
+{
+    // Skipping fault injection as call to ebpf_allocate() covers it.
+    ebpf_process_state_t* state = (ebpf_process_state_t*)ebpf_allocate(sizeof(ebpf_process_state_t));
+    return state;
+}
+
+intptr_t
+ebpf_platform_reference_process()
+{
+
+    HANDLE process = GetCurrentProcess();
+    return (intptr_t)process;
+}
+
+void
+ebpf_platform_dereference_process(intptr_t process_handle)
+{
+    // This is a no-op for the user mode implementation.
+    UNREFERENCED_PARAMETER(process_handle);
+}
+
+void
+ebpf_platform_attach_process(intptr_t process_handle, _Inout_ ebpf_process_state_t* state)
+{
+    // This is a no-op for the user mode implementation.
+    UNREFERENCED_PARAMETER(process_handle);
+    UNREFERENCED_PARAMETER(state);
+}
+
+void
+ebpf_platform_detach_process(_In_ ebpf_process_state_t* state)
+{
+    // This is a no-op for the user mode implementation.
+    UNREFERENCED_PARAMETER(state);
+}
diff --git a/tests/api_test/api_test.cpp b/tests/api_test/api_test.cpp
@@ -983,3 +983,49 @@ test_sock_addr_native_program_load_attach(const char* file_name)
     }
 
 DECLARE_REGRESSION_TEST_CASE("0.9.0")
+
+// The below tests try to load native drivers for invalid programs (that will fail verification).
+// Since verification can be skipped in bpf2c for only Debug builds, these tests are applicable
+// only for Debug build.
+#ifdef _DEBUG
+
+static void
+_load_invalid_program(_In_z_ const char* file_name, ebpf_execution_type_t execution_type, int expected_result)
+{
+    int result;
+    bpf_object* object = nullptr;
+    fd_t program_fd;
+    uint32_t next_id;
+
+    result = _program_load_helper(file_name, BPF_PROG_TYPE_UNSPEC, execution_type, &object, &program_fd);
+    REQUIRE(result == expected_result);
+
+    if (result != 0) {
+        // If load failed, no programs or maps should be loaded.
+        REQUIRE(bpf_map_get_next_id(0, &next_id) == -ENOENT);
+        REQUIRE(bpf_prog_get_next_id(0, &next_id) == -ENOENT);
+    }
+}
+
+TEST_CASE("load_native_program_invalid1", "[native][negative]")
+{
+    _load_invalid_program("invalid_maps1.sys", EBPF_EXECUTION_NATIVE, -EINVAL);
+}
+TEST_CASE("load_native_program_invalid2", "[native][negative]")
+{
+    _load_invalid_program("invalid_maps2.sys", EBPF_EXECUTION_NATIVE, -EINVAL);
+}
+TEST_CASE("load_native_program_invalid3", "[native][negative]")
+{
+    _load_invalid_program("invalid_helpers.sys", EBPF_EXECUTION_NATIVE, -EINVAL);
+}
+TEST_CASE("load_native_program_invalid4", "[native][negative]")
+{
+    _load_invalid_program("empty.sys", EBPF_EXECUTION_NATIVE, -EINVAL);
+}
+TEST_CASE("load_native_program_invalid5", "[native][negative]")
+{
+    _load_invalid_program("invalid_maps3.sys", EBPF_EXECUTION_NATIVE, -EINVAL);
+}
+
+#endif