Fix ZwUnloadDriver failure issue (#965)

* fix zwunloaddriver failure issue

* build break

* fix failing unit test

libs/api/ebpf_api.cpp

@@ -2292,6 +2292,7 @@ _ebpf_program_load_native(
     std::wstring service_name;
     std::string file_name_string(file_name);
     SC_HANDLE service_handle = nullptr;
+    SERVICE_STATUS status = {0};
     std::wstring service_path(SERVICE_PATH_PREFIX);
     std::wstring paramaters_path(PARAMETERS_PATH_PREFIX);
     ebpf_protocol_buffer_t request_buffer;
@@ -2436,10 +2437,18 @@ Done:
     free(map_handles);
     free(program_handles);
 
-    // https://github.com/microsoft/ebpf-for-windows/issues/867
-    // TODO: (Isse# 867) On Server 2019, ZwUnloadDriver fails for services which have
-    // been marked for deletion. As a workaround for this, not deleting the service.
-    // Platform::_delete_service(service_handle);
+    // Workaround: Querying service status hydrates service reference count in SCM.
+    // This ensures that when _delete_service() is called, the service is marked
+    // pending for delete, and a later call to ZwUnloadDriver() by ebpfcore does not
+    // fail. One side effect of this approach still is that the stale service entries
+    // in the registry will not be cleaned up till the next reboot.
+    Platform::_query_service_status(service_handle, &status);
+    EBPF_LOG_MESSAGE_WSTRING(
+        EBPF_TRACELOG_LEVEL_INFO,
+        EBPF_TRACELOG_KEYWORD_API,
+        "_ebpf_program_load_native: Deleting service",
+        service_name.c_str());
+    Platform::_delete_service(service_handle);
     EBPF_RETURN_RESULT(result);
 }
 

libs/platform/ebpf_platform.h

@@ -332,7 +332,9 @@ extern "C"
      * @returns - The previous lock_state required for unlock.
      */
     _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPATCH_LEVEL) _IRQL_saves_
-        _IRQL_raises_(DISPATCH_LEVEL) ebpf_lock_state_t ebpf_lock_lock(_In_ ebpf_lock_t* lock);
+        _IRQL_raises_(DISPATCH_LEVEL)
+    ebpf_lock_state_t
+    ebpf_lock_lock(_In_ ebpf_lock_t* lock);
 
     /**
      * @brief Release exclusive access to the lock.
@@ -1285,6 +1287,15 @@ extern "C"
         TraceLoggingString(message, "Message"),                        \
         TraceLoggingString(string, #string));
 
+#define EBPF_LOG_MESSAGE_WSTRING(trace_level, keyword, message, wstring) \
+    TraceLoggingWrite(                                                   \
+        ebpf_tracelog_provider,                                          \
+        EBPF_TRACELOG_EVENT_GENERIC_MESSAGE,                             \
+        TraceLoggingLevel(trace_level),                                  \
+        TraceLoggingKeyword((keyword)),                                  \
+        TraceLoggingString(message, "Message"),                          \
+        TraceLoggingWideString(wstring, #wstring));
+
 #define EBPF_LOG_WIN32_API_FAILURE(keyword, api)          \
     do {                                                  \
         DWORD last_error = GetLastError();                \

libs/thunk/mock/mock.cpp

@@ -161,6 +161,15 @@ _delete_service(SC_HANDLE service_handle)
     return delete_service_handler(service_handle);
 }
 
+bool
+_query_service_status(SC_HANDLE service_handle, _Inout_ SERVICE_STATUS* status)
+{
+    UNREFERENCED_PARAMETER(service_handle);
+    UNREFERENCED_PARAMETER(status);
+
+    return true;
+}
+
 uint32_t
 _stop_service(SC_HANDLE service_handle)
 {

libs/thunk/platform.h

@@ -75,4 +75,7 @@ _delete_service(SC_HANDLE service_handle);
 uint32_t
 _stop_service(SC_HANDLE service_handle);
 
+bool
+_query_service_status(SC_HANDLE service_handle, _Inout_ SERVICE_STATUS* status);
+
 } // namespace Platform

libs/thunk/windows/platform.cpp

@@ -256,6 +256,12 @@ Done:
     return error;
 }
 
+bool
+_query_service_status(SC_HANDLE service_handle, _Inout_ SERVICE_STATUS* status)
+{
+    return QueryServiceStatus(service_handle, status);
+}
+
 uint32_t
 _delete_service(SC_HANDLE service_handle)
 {

tests/api_test/api_test.cpp

@@ -35,6 +35,8 @@
 #define DROP_PACKET_MAP_COUNT 2
 #define BIND_MONITOR_MAP_COUNT 2
 
+#define WAIT_TIME_IN_MS 5000
+
 static service_install_helper
     _ebpf_core_driver_helper(EBPF_CORE_DRIVER_NAME, EBPF_CORE_DRIVER_BINARY_NAME, SERVICE_KERNEL_DRIVER);
 
@@ -238,6 +240,15 @@ TEST_CASE("pinned_map_enum", "[pinned_map_enum]") { ebpf_test_pinned_map_enum();
         _test_program_load(file, program_type, execution_type, expected_result);     \
     }
 
+// Duplicate tests sleep for WAIT_TIME_IN_MS seconds. This ensures the previous driver is
+// unloaded by the time the test is re-run.
+#define DECLARE_DUPLICATE_LOAD_TEST_CASE(file, program_type, execution_type, instance, expected_result) \
+    TEST_CASE("test_ebpf_program_load-" #file "-" #program_type "-" #execution_type "-" #instance)      \
+    {                                                                                                   \
+        Sleep(WAIT_TIME_IN_MS);                                                                         \
+        _test_program_load(file, program_type, execution_type, expected_result);                        \
+    }
+
 #if defined(CONFIG_BPF_JIT_ALWAYS_ON)
 #define INTERPRET_LOAD_RESULT EBPF_PROGRAM_LOAD_FAILED
 #else
@@ -249,6 +260,10 @@ DECLARE_LOAD_TEST_CASE("droppacket.o", nullptr, EBPF_EXECUTION_JIT, EBPF_SUCCESS
 
 DECLARE_LOAD_TEST_CASE("droppacket_km.sys", nullptr, EBPF_EXECUTION_NATIVE, EBPF_SUCCESS);
 
+// Declare a duplicate test case. This will ensure that the earlier driver is actually unloaded,
+// else this test case will fail.
+DECLARE_DUPLICATE_LOAD_TEST_CASE("droppacket_km.sys", nullptr, EBPF_EXECUTION_NATIVE, 2, EBPF_SUCCESS);
+
 // Load droppacket (ANY) without providing expected program type.
 DECLARE_LOAD_TEST_CASE("droppacket.o", nullptr, EBPF_EXECUTION_ANY, EBPF_SUCCESS);
 

tests/end_to_end/end_to_end.cpp

@@ -604,7 +604,7 @@ bindmonitor_tailcall_test(ebpf_execution_type_t execution_type)
     fd_t process_map_fd = bpf_object__find_map_fd_by_name(object, "process_map");
     REQUIRE(process_map_fd > 0);
 
-    // Setup tail calls.
+    // Set up tail calls.
     struct bpf_program* callee0 = bpf_object__find_program_by_name(object, "BindMonitor_Callee0");
     REQUIRE(callee0 != nullptr);
     fd_t callee0_fd = bpf_program__fd(callee0);
@@ -625,14 +625,14 @@ bindmonitor_tailcall_test(ebpf_execution_type_t execution_type)
 
     // Validate various maps.
 
-    // Validate map-in-maps with "inner_id"
+    // Validate map-in-maps with "inner_id".
     struct bpf_map* outer_map = bpf_object__find_map_by_name(object, "dummy_outer_map");
     REQUIRE(outer_map != nullptr);
 
     int outer_map_fd = bpf_map__fd(outer_map);
     REQUIRE(outer_map_fd > 0);
 
-    // Validate map-in-maps with "inner_idx"
+    // Validate map-in-maps with "inner_idx".
     struct bpf_map* outer_idx_map = bpf_object__find_map_by_name(object, "dummy_outer_idx_map");
     REQUIRE(outer_idx_map != nullptr);
 

tests/end_to_end/test_helper.cpp

@@ -49,6 +49,7 @@ typedef struct _service_context
     HMODULE dll;
     bool loaded;
     ebpf_extension_client_t* binding_context;
+    bool delete_pending = false;
 } service_context_t;
 
 static uint64_t _ebpf_file_descriptor_counter = 0;
@@ -226,6 +227,8 @@ _unload_all_native_modules()
             // Deregister client.
             ebpf_extension_unload(context->binding_context);
         }
+        // The service should have been marked for deletion till now.
+        REQUIRE(context->delete_pending);
         if (context->dll != nullptr) {
             FreeLibrary(context->dll);
         }
@@ -236,14 +239,14 @@ _unload_all_native_modules()
 #pragma warning(pop)
 
 static void
-_load_native_module(_Inout_ service_context_t* context)
+_preprocess_load_native_module(_Inout_ service_context_t* context)
 {
     // For user mode tests, most of the sample programs are compiled into a single DLL file.
     // First try loading from there.
     context->dll = LoadLibraryA(NATIVE_DLL_NAME);
     REQUIRE(context->dll != nullptr);
 
-    // Get file name from the file path
+    // Get file name from the file path.
     std::string file_name = std::filesystem::path(context->file_path).filename().stem().string();
 
     auto get_function_common =
@@ -305,7 +308,7 @@ _preprocess_ioctl(_In_ const ebpf_operation_header_t* user_request)
                 context->second->module_id = request->module_id;
 
                 // Load the module.
-                _load_native_module(context->second);
+                _preprocess_load_native_module(context->second);
             }
         } catch (...) {
             // Ignore.
@@ -474,7 +477,14 @@ Glue_delete_service(SC_HANDLE handle)
 {
     for (auto& [path, context] : _service_path_to_context_map) {
         if (context->handle == (intptr_t)handle) {
-            _service_path_to_context_map.erase(path);
+            // Delete the service if it has not been loaded yet. Otherwise
+            // mark it pending for delete.
+            if (!context->loaded) {
+                _service_path_to_context_map.erase(path);
+                ebpf_free(context);
+            } else {
+                context->delete_pending = true;
+            }
             break;
         }
     }