[netebpfext] Add per-provider WFP handles to avoid improper use from …

…parallel invocations (#3866)

* add locking

* fix build error

* move to per-provider handle

* move to per provider handle, move to static handle, add cleanup, add error resilience

* remove version

* temporarily move to branch submodule for testing

* update submodule

* reset gitmodules

* PR part 1: Code cleanup, nits

* move engine handle into filter context

* CR - move filter unregister to after filter delete, replace null checks with asserts, remove unneded helper function

* CR

netebpfext/net_ebpf_ext.c

@@ -238,6 +238,7 @@ net_ebpf_extension_wfp_filter_context_create(
     _In_ const net_ebpf_extension_hook_provider_t* provider_context,
     _Outptr_ net_ebpf_extension_wfp_filter_context_t** filter_context)
 {
+    NTSTATUS status = STATUS_SUCCESS;
     ebpf_result_t result = EBPF_SUCCESS;
     net_ebpf_extension_wfp_filter_context_t* local_filter_context = NULL;
     uint32_t client_context_count_max = NET_EBPF_EXT_MAX_CLIENTS_PER_HOOK_SINGLE_ATTACH;
@@ -287,6 +288,14 @@ net_ebpf_extension_wfp_filter_context_create(
     // Set the provider context.
     local_filter_context->provider_context = provider_context;
 
+    // Open the WFP engine handle.
+    status = FwpmEngineOpen(NULL, RPC_C_AUTHN_WINNT, NULL, NULL, &local_filter_context->wfp_engine_handle);
+    if (!NT_SUCCESS(status)) {
+        NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmEngineOpen", status);
+        result = EBPF_FAILED;
+        goto Exit;
+    }
+
     *filter_context = local_filter_context;
     local_filter_context = NULL;
 
@@ -376,13 +385,17 @@ net_ebpf_extension_get_callout_id_for_hook(net_ebpf_extension_hook_id_t hook_id)
 }
 void
 net_ebpf_extension_delete_wfp_filters(
-    uint32_t filter_count, _Frees_ptr_ _In_count_(filter_count) net_ebpf_ext_wfp_filter_id_t* filter_ids)
+    _In_ HANDLE wfp_engine_handle,
+    uint32_t filter_count,
+    _Frees_ptr_ _In_count_(filter_count) net_ebpf_ext_wfp_filter_id_t* filter_ids)
 {
     NET_EBPF_EXT_LOG_ENTRY();
     NTSTATUS status = STATUS_SUCCESS;
 
+    ASSERT(wfp_engine_handle != NULL);
+
     for (uint32_t index = 0; index < filter_count; index++) {
-        status = FwpmFilterDeleteById(_fwp_engine_handle, filter_ids[index].id);
+        status = FwpmFilterDeleteById(wfp_engine_handle, filter_ids[index].id);
         if (!NT_SUCCESS(status)) {
             NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
                 NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmFilterDeleteById", status);
@@ -403,6 +416,7 @@ net_ebpf_extension_delete_wfp_filters(
 
 _Must_inspect_result_ ebpf_result_t
 net_ebpf_extension_add_wfp_filters(
+    _In_ HANDLE wfp_engine_handle,
     uint32_t filter_count,
     _In_count_(filter_count) const net_ebpf_extension_wfp_filter_parameters_t* parameters,
     uint32_t condition_count,
@@ -418,6 +432,8 @@ net_ebpf_extension_add_wfp_filters(
 
     NET_EBPF_EXT_LOG_ENTRY();
 
+    ASSERT(wfp_engine_handle != NULL);
+
     if (filter_count == 0) {
         NET_EBPF_EXT_LOG_MESSAGE(
             NET_EBPF_EXT_TRACELOG_LEVEL_ERROR, NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "Filter count is 0");
@@ -432,7 +448,7 @@ net_ebpf_extension_add_wfp_filters(
 
     memset(local_filter_ids, 0, (sizeof(net_ebpf_ext_wfp_filter_id_t) * filter_count));
 
-    status = FwpmTransactionBegin(_fwp_engine_handle, 0);
+    status = FwpmTransactionBegin(wfp_engine_handle, 0);
     if (!NT_SUCCESS(status)) {
         NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmTransactionBegin", status);
         result = EBPF_INVALID_ARGUMENT;
@@ -463,7 +479,7 @@ net_ebpf_extension_add_wfp_filters(
         REFERENCE_FILTER_CONTEXT(filter_context);
         filter.rawContext = (uint64_t)(uintptr_t)filter_context;
 
-        status = FwpmFilterAdd(_fwp_engine_handle, &filter, NULL, &local_filter_id);
+        status = FwpmFilterAdd(wfp_engine_handle, &filter, NULL, &local_filter_id);
         if (!NT_SUCCESS(status)) {
             NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE_MESSAGE_STRING(
                 NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION,
@@ -486,7 +502,7 @@ net_ebpf_extension_add_wfp_filters(
         }
     }
 
-    status = FwpmTransactionCommit(_fwp_engine_handle);
+    status = FwpmTransactionCommit(wfp_engine_handle);
     if (!NT_SUCCESS(status)) {
         NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmTransactionCommit", status);
         result = EBPF_INVALID_ARGUMENT;
@@ -502,7 +518,7 @@ net_ebpf_extension_add_wfp_filters(
             ExFreePool(local_filter_ids);
         }
         if (is_in_transaction) {
-            status = FwpmTransactionAbort(_fwp_engine_handle);
+            status = FwpmTransactionAbort(wfp_engine_handle);
             if (!NT_SUCCESS(status)) {
                 NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
                     NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmTransactionAbort", status);
@@ -527,9 +543,7 @@ _net_ebpf_ext_register_wfp_callout(_Inout_ net_ebpf_ext_wfp_callout_state_t* cal
 
     FWPS_CALLOUT callout_register_state = {0};
     FWPM_CALLOUT callout_add_state = {0};
-
     FWPM_DISPLAY_DATA display_data = {0};
-
     BOOLEAN was_callout_registered = FALSE;
 
     callout_register_state.calloutKey = *callout_state->callout_guid;
@@ -651,17 +665,12 @@ net_ebpf_extension_initialize_wfp_components(_Inout_ void* device_object)
 
 -- */
 {
+    UNREFERENCED_PARAMETER(device_object);
     NTSTATUS status = STATUS_SUCCESS;
     FWPM_PROVIDER ebpf_wfp_provider = {0};
     FWPM_SUBLAYER ebpf_hook_sub_layer;
-
-    UNREFERENCED_PARAMETER(device_object);
-
     BOOLEAN is_engine_opened = FALSE;
     BOOLEAN is_in_transaction = FALSE;
-
-    FWPM_SESSION session = {0};
-
     size_t index;
 
     NET_EBPF_EXT_LOG_ENTRY();
@@ -671,9 +680,7 @@ net_ebpf_extension_initialize_wfp_components(_Inout_ void* device_object)
         goto Exit;
     }
 
-    session.flags = FWPM_SESSION_FLAG_DYNAMIC;
-
-    status = FwpmEngineOpen(NULL, RPC_C_AUTHN_WINNT, NULL, &session, &_fwp_engine_handle);
+    status = FwpmEngineOpen(NULL, RPC_C_AUTHN_WINNT, NULL, NULL, &_fwp_engine_handle);
     NET_EBPF_EXT_BAIL_ON_API_FAILURE_STATUS(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmEngineOpen", status);
     is_engine_opened = TRUE;
 
@@ -750,19 +757,43 @@ net_ebpf_extension_uninitialize_wfp_components(void)
     NTSTATUS status;
 
     if (_fwp_engine_handle != NULL) {
-        status = FwpmEngineClose(_fwp_engine_handle);
-        if (!NT_SUCCESS(status)) {
-            NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmEngineClose", status);
-        }
-        _fwp_engine_handle = NULL;
-
+        // Clean up the callouts
         for (index = 0; index < EBPF_COUNT_OF(_net_ebpf_ext_wfp_callout_states); index++) {
             status = FwpsCalloutUnregisterById(_net_ebpf_ext_wfp_callout_states[index].assigned_callout_id);
-            if (!NT_SUCCESS(status)) {
+            if (!NT_SUCCESS(status) && !WFP_ERROR(status, CALLOUT_NOT_FOUND)) {
                 NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
                     NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpsCalloutUnregisterById", status);
             }
+
+            status = FwpmCalloutDeleteByKey(_fwp_engine_handle, _net_ebpf_ext_wfp_callout_states[index].callout_guid);
+            if (!NT_SUCCESS(status) && !WFP_ERROR(status, CALLOUT_NOT_FOUND)) {
+                NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
+                    NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmCalloutDeleteByKey", status);
+            }
         }
+
+        // Clean up the sub layers.
+        for (index = 0; index < EBPF_COUNT_OF(_net_ebpf_ext_sublayers); index++) {
+            status = FwpmSubLayerDeleteByKey(_fwp_engine_handle, _net_ebpf_ext_sublayers[index].sublayer_guid);
+            if (!NT_SUCCESS(status) && !WFP_ERROR(status, SUBLAYER_NOT_FOUND)) {
+                NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
+                    NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmSubLayerDeleteByKey", status);
+            }
+        }
+
+        // Clean up the providers.
+        status = FwpmProviderDeleteByKey(_fwp_engine_handle, &EBPF_WFP_PROVIDER);
+        if (!NT_SUCCESS(status) && !WFP_ERROR(status, PROVIDER_NOT_FOUND)) {
+            NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
+                NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmProviderDeleteByKey", status);
+        }
+
+        // Close the engine handle.
+        status = FwpmEngineClose(_fwp_engine_handle);
+        if (!NT_SUCCESS(status)) {
+            NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmEngineClose", status);
+        }
+        _fwp_engine_handle = NULL;
     }
 
     // FwpsInjectionHandleCreate can fail. So, check for NULL.

netebpfext/net_ebpf_ext.h

@@ -27,10 +27,10 @@
 
 // Note: The maximum number of clients that can attach per-hook in multi-attach case has been currently capped to
 // a constant value to keep the implementation simple. Keeping the max limit constant allows allocating the memory
-// required for creating a copy of list of clients on the stack itself. In the future, if there is a need to increase this
-// maximum count, the value can be simply increased as long as the required memory can still be allocated on stack. If
-// the required memory becomes too large, we may need to switch to a different design to handle this. One option is to
-// use epoch based memory management for the list of clients. This eliminates the need to create a copy of programs
+// required for creating a copy of list of clients on the stack itself. In the future, if there is a need to increase
+// this maximum count, the value can be simply increased as long as the required memory can still be allocated on stack.
+// If the required memory becomes too large, we may need to switch to a different design to handle this. One option is
+// to use epoch based memory management for the list of clients. This eliminates the need to create a copy of programs
 // per-invocation. Another option can be to always invoke the programs while holding the socket context lock, but that
 // comes with a side effect of every program invocation now happening at DISPATCH_LEVEL.
 #define NET_EBPF_EXT_MAX_CLIENTS_PER_HOOK_MULTI_ATTACH 16
@@ -140,17 +140,21 @@ typedef struct _net_ebpf_extension_wfp_filter_context
     bool context_deleting : 1; ///< True if all the clients have been detached and the context is being deleted.
     bool wildcard : 1;         ///< True if the filter context is for wildcard filters.
     bool initialized : 1;      ///< True if the filter context has been successfully initialized.
+    HANDLE wfp_engine_handle;  ///< WFP engine handle.
 } net_ebpf_extension_wfp_filter_context_t;
 
-#define CLEAN_UP_FILTER_CONTEXT(filter_context)            \
-    if ((filter_context) != NULL) {                        \
-        if ((filter_context)->filter_ids != NULL) {        \
-            ExFreePool((filter_context)->filter_ids);      \
-        }                                                  \
-        if ((filter_context)->client_contexts != NULL) {   \
-            ExFreePool((filter_context)->client_contexts); \
-        }                                                  \
-        ExFreePool((filter_context));                      \
+#define CLEAN_UP_FILTER_CONTEXT(filter_context)                   \
+    if ((filter_context) != NULL) {                               \
+        if ((filter_context)->filter_ids != NULL) {               \
+            ExFreePool((filter_context)->filter_ids);             \
+        }                                                         \
+        if ((filter_context)->client_contexts != NULL) {          \
+            ExFreePool((filter_context)->client_contexts);        \
+        }                                                         \
+        if ((filter_context)->wfp_engine_handle != NULL) {        \
+            FwpmEngineClose((filter_context)->wfp_engine_handle); \
+        }                                                         \
+        ExFreePool((filter_context));                             \
     }
 
 #define REFERENCE_FILTER_CONTEXT(filter_context)                  \
@@ -247,6 +251,7 @@ net_ebpf_extension_get_callout_id_for_hook(net_ebpf_extension_hook_id_t hook_id)
 /**
  * @brief Add WFP filters with specified conditions at specified layers.
  *
+ * @param[in] wfp_filter_handle The WFP filter handle used to add the filters.
  * @param[in] filter_count Count of filters to be added.
  * @param[in] parameters Filter parameters.
  * @param[in] condition_count Count of filter conditions.
@@ -259,6 +264,7 @@ net_ebpf_extension_get_callout_id_for_hook(net_ebpf_extension_hook_id_t hook_id)
  */
 _Must_inspect_result_ ebpf_result_t
 net_ebpf_extension_add_wfp_filters(
+    _In_ HANDLE wfp_engine_handle,
     uint32_t filter_count,
     _In_count_(filter_count) const net_ebpf_extension_wfp_filter_parameters_t* parameters,
     uint32_t condition_count,
@@ -269,12 +275,15 @@ net_ebpf_extension_add_wfp_filters(
 /**
  * @brief Deletes WFP filters with specified filter IDs.
  *
+ * @param[in] wfp_filter_handle The WFP filter handle used to delete the filters.
  * @param[in]  filter_count Count of filters to be added.
  * @param[in]  filter_ids ID of the filter being deleted.
  */
 void
 net_ebpf_extension_delete_wfp_filters(
-    uint32_t filter_count, _Frees_ptr_ _In_count_(filter_count) net_ebpf_ext_wfp_filter_id_t* filter_ids);
+    _In_ HANDLE wfp_engine_handle,
+    uint32_t filter_count,
+    _Frees_ptr_ _In_count_(filter_count) net_ebpf_ext_wfp_filter_id_t* filter_ids);
 
 // eBPF WFP Provider GUID.
 // ddb851f5-841a-4b77-8a46-bb7063e9f162

netebpfext/net_ebpf_ext_bind.c

@@ -121,6 +121,7 @@ _net_ebpf_ext_bind_create_filter_context(
 
     // Add WFP filters at appropriate layers and set the hook NPI client as the filter's raw context.
     result = net_ebpf_extension_add_wfp_filters(
+        local_filter_context->wfp_engine_handle,
         EBPF_COUNT_OF(_net_ebpf_extension_bind_wfp_filter_parameters),
         _net_ebpf_extension_bind_wfp_filter_parameters,
         0,
@@ -151,7 +152,8 @@ _net_ebpf_ext_bind_delete_filter_context(
     }
 
     // Delete the WFP filters.
-    net_ebpf_extension_delete_wfp_filters(filter_context->filter_ids_count, filter_context->filter_ids);
+    net_ebpf_extension_delete_wfp_filters(
+        filter_context->wfp_engine_handle, filter_context->filter_ids_count, filter_context->filter_ids);
     net_ebpf_extension_wfp_filter_context_cleanup((net_ebpf_extension_wfp_filter_context_t*)filter_context);
 
 Exit:
@@ -214,11 +216,11 @@ net_ebpf_ext_bind_register_providers()
 void
 net_ebpf_ext_bind_unregister_providers()
 {
-    if (_ebpf_bind_hook_provider_context) {
+    if (_ebpf_bind_hook_provider_context != NULL) {
         net_ebpf_extension_hook_provider_unregister(_ebpf_bind_hook_provider_context);
         _ebpf_bind_hook_provider_context = NULL;
     }
-    if (_ebpf_bind_program_info_provider_context) {
+    if (_ebpf_bind_program_info_provider_context != NULL) {
         net_ebpf_extension_program_info_provider_unregister(_ebpf_bind_program_info_provider_context);
         _ebpf_bind_program_info_provider_context = NULL;
     }

netebpfext/net_ebpf_ext_sock_addr.c

@@ -690,6 +690,7 @@ _net_ebpf_extension_sock_addr_create_filter_context(
     // Add a single WFP filter at the WFP layer corresponding to the hook type, and set the hook NPI client as the
     // filter's raw context.
     result = net_ebpf_extension_add_wfp_filters(
+        local_filter_context->base.wfp_engine_handle,
         filter_parameters_array->count, // filter_count
         filter_parameters_array->filter_parameters,
         (compartment_id == UNSPECIFIED_COMPARTMENT_ID) ? 0 : 1,
@@ -720,7 +721,8 @@ _net_ebpf_extension_sock_addr_delete_filter_context(
     }
     sock_addr_filter_context = (net_ebpf_extension_sock_addr_wfp_filter_context_t*)filter_context;
 
-    net_ebpf_extension_delete_wfp_filters(filter_context->filter_ids_count, filter_context->filter_ids);
+    net_ebpf_extension_delete_wfp_filters(
+        filter_context->wfp_engine_handle, filter_context->filter_ids_count, filter_context->filter_ids);
     if (sock_addr_filter_context->redirect_handle != NULL) {
         FwpsRedirectHandleDestroy(sock_addr_filter_context->redirect_handle);
     }
@@ -1286,7 +1288,7 @@ net_ebpf_ext_sock_addr_unregister_providers()
             _ebpf_sock_addr_hook_provider_context[i] = NULL;
         }
     }
-    if (_ebpf_sock_addr_program_info_provider_context) {
+    if (_ebpf_sock_addr_program_info_provider_context != NULL) {
         net_ebpf_extension_program_info_provider_unregister(_ebpf_sock_addr_program_info_provider_context);
         _ebpf_sock_addr_program_info_provider_context = NULL;
     }

netebpfext/net_ebpf_ext_sock_ops.c

@@ -186,6 +186,7 @@ _net_ebpf_extension_sock_ops_create_filter_context(
     // Add WFP filters at appropriate layers and set the hook NPI client as the filter's raw context.
     filter_count = NET_EBPF_SOCK_OPS_FILTER_COUNT;
     result = net_ebpf_extension_add_wfp_filters(
+        local_filter_context->base.wfp_engine_handle,
         filter_count,
         _net_ebpf_extension_sock_ops_wfp_filter_parameters,
         (compartment_id == UNSPECIFIED_COMPARTMENT_ID) ? 0 : 1,
@@ -258,7 +259,9 @@ _net_ebpf_extension_sock_ops_delete_filter_context(
 
     InitializeListHead(&local_list_head);
     net_ebpf_extension_delete_wfp_filters(
-        local_filter_context->base.filter_ids_count, local_filter_context->base.filter_ids);
+        filter_context->wfp_engine_handle,
+        local_filter_context->base.filter_ids_count,
+        local_filter_context->base.filter_ids);
 
     KeAcquireSpinLock(&local_filter_context->lock, &irql);
     if (local_filter_context->flow_context_list.count > 0) {
@@ -353,11 +356,11 @@ net_ebpf_ext_sock_ops_register_providers()
 void
 net_ebpf_ext_sock_ops_unregister_providers()
 {
-    if (_ebpf_sock_ops_hook_provider_context) {
+    if (_ebpf_sock_ops_hook_provider_context != NULL) {
         net_ebpf_extension_hook_provider_unregister(_ebpf_sock_ops_hook_provider_context);
         _ebpf_sock_ops_hook_provider_context = NULL;
     }
-    if (_ebpf_sock_ops_program_info_provider_context) {
+    if (_ebpf_sock_ops_program_info_provider_context != NULL) {
         net_ebpf_extension_program_info_provider_unregister(_ebpf_sock_ops_program_info_provider_context);
         _ebpf_sock_ops_program_info_provider_context = NULL;
     }

netebpfext/net_ebpf_ext_xdp.c

@@ -167,6 +167,7 @@ _net_ebpf_extension_xdp_create_filter_context(
     // Add WFP filters at appropriate layers and set the hook NPI client as the filter's raw context.
     filter_count = NET_EBPF_XDP_FILTER_COUNT;
     result = net_ebpf_extension_add_wfp_filters(
+        xdp_filter_context->base.wfp_engine_handle,
         filter_count,
         _net_ebpf_extension_xdp_wfp_filter_parameters,
         (if_index == 0) ? 0 : 1,
@@ -242,7 +243,9 @@ _net_ebpf_extension_xdp_delete_filter_context(
     xdp_filter_context = (net_ebpf_extension_xdp_wfp_filter_context_t*)filter_context;
 
     net_ebpf_extension_delete_wfp_filters(
-        xdp_filter_context->base.filter_ids_count, xdp_filter_context->base.filter_ids);
+        filter_context->wfp_engine_handle,
+        xdp_filter_context->base.filter_ids_count,
+        xdp_filter_context->base.filter_ids);
     net_ebpf_extension_wfp_filter_context_cleanup((net_ebpf_extension_wfp_filter_context_t*)xdp_filter_context);
 
 Exit:
@@ -303,11 +306,11 @@ net_ebpf_ext_xdp_register_providers()
 void
 net_ebpf_ext_xdp_unregister_providers()
 {
-    if (_ebpf_xdp_test_hook_provider_context) {
+    if (_ebpf_xdp_test_hook_provider_context != NULL) {
         net_ebpf_extension_hook_provider_unregister(_ebpf_xdp_test_hook_provider_context);
         _ebpf_xdp_test_hook_provider_context = NULL;
     }
-    if (_ebpf_xdp_test_program_info_provider_context) {
+    if (_ebpf_xdp_test_program_info_provider_context != NULL) {
         net_ebpf_extension_program_info_provider_unregister(_ebpf_xdp_test_program_info_provider_context);
         _ebpf_xdp_test_program_info_provider_context = NULL;
     }

netebpfext/sys/netebpfext_platform.h

@@ -9,3 +9,5 @@
 #pragma warning(disable : 4201) // unnamed struct/union
 #include <fwpsk.h>
 #pragma warning(pop)
+
+#define WFP_ERROR(status, error) ((status) == (STATUS_FWP_##error))