Skip to content

Commit

Permalink
Multi-attach support for SOCK_ADDR programs (#3751)
Browse files Browse the repository at this point in the history 
* initial commit

* second commit

* third commit

* update lock logic

* fix analysis errors

* backup

* Add tests

* fix

* fix

* fix

* code refactor

* update test

* Revert "update test"

This reverts commit 6770269.

* Revert "code refactor"

This reverts commit fae36a7.

* Revert "Revert "code refactor""

This reverts commit 3b8e762.

* Revert "Revert "update test""

This reverts commit 11610f0.

* fix capability in xdp

* add tests, remove filter weight

* invoke wildcard programs, update tests

* fix wildcard invocation logic

* fix wildcard invocation logic, remove sleep from concurrency tests

* move invocation out of the lock, fix tests

* code cleanup, update test

* code cleanup

* remove trace

* cleanup

* update wildcard invocation logic

* code cleanup

* code cleanup

* code cleanup, CR comments

* cr comments, code cleanup

* Update netebpfext/net_ebpf_ext_hook_provider.h

Co-authored-by: Dave Thaler <[email protected]>

* cr comments

* add tests

* cr comments

* code cleanup

* tracing changes

* fix sal

* code cleanup

* code cleanup

* Update netebpfext/net_ebpf_ext.h

Co-authored-by: Dave Thaler <[email protected]>

---------

Co-authored-by: Dave Thaler <[email protected]>
  • Loading branch information
@saxena-anurag @dthaler
saxena-anurag and dthaler authored Sep 5, 2024
1 parent 06533c8 commit c73a086
Show file tree
Hide file tree
Showing 13 changed files with 2,506 additions and 744 deletions.
122 changes: 109 additions & 13 deletions netebpfext/net_ebpf_ext.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,39 +235,63 @@ _Must_inspect_result_ ebpf_result_t
net_ebpf_extension_wfp_filter_context_create(
size_t filter_context_size,
_In_ const net_ebpf_extension_hook_client_t* client_context,
_In_ const net_ebpf_extension_hook_provider_t* provider_context,
_Outptr_ net_ebpf_extension_wfp_filter_context_t** filter_context)
{
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;

NET_EBPF_EXT_LOG_ENTRY();

*filter_context = NULL;

if (net_ebpf_extension_hook_provider_get_attach_capability(provider_context) ==
ATTACH_CAPABILITY_MULTI_ATTACH_WITH_WILDCARD) {
client_context_count_max = NET_EBPF_EXT_MAX_CLIENTS_PER_HOOK_MULTI_ATTACH;
}

// Allocate buffer for WFP filter context.
local_filter_context = (net_ebpf_extension_wfp_filter_context_t*)ExAllocatePoolUninitialized(
NonPagedPoolNx, filter_context_size, NET_EBPF_EXTENSION_POOL_TAG);
NET_EBPF_EXT_BAIL_ON_ALLOC_FAILURE_RESULT(
NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, local_filter_context, "local_filter_context", result);

memset(local_filter_context, 0, filter_context_size);

local_filter_context->client_contexts = (net_ebpf_extension_hook_client_t**)ExAllocatePoolUninitialized(
NonPagedPoolNx,
client_context_count_max * sizeof(net_ebpf_extension_hook_client_t*),
NET_EBPF_EXTENSION_POOL_TAG);
NET_EBPF_EXT_BAIL_ON_ALLOC_FAILURE_RESULT(
NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION,
local_filter_context->client_contexts,
"local_filter_context - client_contexts",
result);

memset(
local_filter_context->client_contexts, 0, client_context_count_max * sizeof(net_ebpf_extension_hook_client_t*));
local_filter_context->client_context_count_max = client_context_count_max;
local_filter_context->context_deleting = FALSE;
InitializeListHead(&local_filter_context->link);
local_filter_context->reference_count = 1; // Initial reference.
local_filter_context->client_context = client_context;

if (!net_ebpf_extension_hook_client_enter_rundown(
(net_ebpf_extension_hook_client_t*)local_filter_context->client_context)) {
// Set the first client context.
local_filter_context->client_contexts[0] = (net_ebpf_extension_hook_client_t*)client_context;
local_filter_context->client_context_count = 1;

// We're setting up the filter context here and as this is the very first (and exclusive) attempt to acquire
// rundown, it cannot fail. If it does, this is indicative of a fatal system level error.
__fastfail(FAST_FAIL_INVALID_ARG);
}
// Set filter context as provider data in the hook client.
net_ebpf_extension_hook_client_set_provider_data(
(net_ebpf_extension_hook_client_t*)client_context, local_filter_context);

// Set the provider context.
local_filter_context->provider_context = provider_context;

*filter_context = local_filter_context;
local_filter_context = NULL;

Exit:
if (local_filter_context != NULL) {
ExFreePool(local_filter_context);
}
CLEAN_UP_FILTER_CONTEXT(local_filter_context);

NET_EBPF_EXT_RETURN_RESULT(result);
}
Expand All @@ -276,10 +300,10 @@ void
net_ebpf_extension_wfp_filter_context_cleanup(_Frees_ptr_ net_ebpf_extension_wfp_filter_context_t* filter_context)
{
// Since the hook client is detaching, the eBPF program should not be invoked any further.
// The client_detached field in filter_context is set to false for this reason. This way any
// The context_deleting field in filter_context is set to false for this reason. This way any
// lingering WFP classify callbacks will exit as it would not find any hook client associated
// with the filter context. This is best effort & no locks are held.
filter_context->client_detached = TRUE;
filter_context->context_deleting = TRUE;
DEREFERENCE_FILTER_CONTEXT(filter_context);
}

Expand Down Expand Up @@ -362,13 +386,16 @@ net_ebpf_extension_delete_wfp_filters(
if (!NT_SUCCESS(status)) {
NET_EBPF_EXT_LOG_NTSTATUS_API_FAILURE(
NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION, "FwpmFilterDeleteById", status);
filter_ids[index].state = NET_EBPF_EXT_WFP_FILTER_DELETE_FAILED;
filter_ids[index].error_code = status;
} else {
NET_EBPF_EXT_LOG_MESSAGE_UINT64_UINT64(
NET_EBPF_EXT_TRACELOG_LEVEL_VERBOSE,
NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION,
"Marked WFP filter for deletion: ",
index,
filter_ids[index].id);
filter_ids[index].state = NET_EBPF_EXT_WFP_FILTER_DELETING;
}
}
NET_EBPF_EXT_LOG_EXIT();
Expand Down Expand Up @@ -422,7 +449,8 @@ net_ebpf_extension_add_wfp_filters(
filter.displayData.name = (wchar_t*)filter_parameter->name;
filter.displayData.description = (wchar_t*)filter_parameter->description;
filter.providerKey = (GUID*)&EBPF_WFP_PROVIDER;
filter.action.type = FWP_ACTION_CALLOUT_TERMINATING;
filter.action.type =
filter_parameter->action_type ? filter_parameter->action_type : FWP_ACTION_CALLOUT_TERMINATING;
filter.action.calloutKey = *filter_parameter->callout_guid;
filter.filterCondition = (FWPM_FILTER_CONDITION*)conditions;
filter.numFilterConditions = condition_count;
Expand Down Expand Up @@ -870,3 +898,71 @@ net_ebpf_ext_unregister_providers()
_net_ebpf_sock_ops_providers_registered = false;
}
}

ebpf_result_t
net_ebpf_ext_add_client_context(
_Inout_ net_ebpf_extension_wfp_filter_context_t* filter_context,
_In_ const struct _net_ebpf_extension_hook_client* hook_client)
{
ebpf_result_t result = EBPF_SUCCESS;
KIRQL old_irql;

NET_EBPF_EXT_LOG_ENTRY();

old_irql = ExAcquireSpinLockExclusive(&filter_context->lock);

// Check if we have reached max capacity.
if (filter_context->client_context_count == filter_context->client_context_count_max) {
NET_EBPF_EXT_LOG_MESSAGE(
NET_EBPF_EXT_TRACELOG_LEVEL_ERROR,
NET_EBPF_EXT_TRACELOG_KEYWORD_EXTENSION,
"net_ebpf_ext_add_client_context: Exceeded max client count");
result = EBPF_NO_MEMORY;
goto Exit;
}

// Append the client context to the end.
filter_context->client_contexts[filter_context->client_context_count] =
(struct _net_ebpf_extension_hook_client*)hook_client;
filter_context->client_context_count++;

// Add filter_context as provider data for the client.
net_ebpf_extension_hook_client_set_provider_data(
(struct _net_ebpf_extension_hook_client*)hook_client, (void*)filter_context);

Exit:
ExReleaseSpinLockExclusive(&filter_context->lock, old_irql);
NET_EBPF_EXT_RETURN_RESULT(result);
}

void
net_ebpf_ext_remove_client_context(
_Inout_ net_ebpf_extension_wfp_filter_context_t* filter_context,
_In_ const struct _net_ebpf_extension_hook_client* hook_client)
{
KIRQL old_irql;
uint32_t index;
bool found = FALSE;

old_irql = ExAcquireSpinLockExclusive(&filter_context->lock);

for (index = 0; index < filter_context->client_context_count; index++) {
if (filter_context->client_contexts[index] == hook_client) {
filter_context->client_contexts[index] = NULL;
filter_context->client_context_count--;
found = TRUE;
break;
}
}
ASSERT(found == TRUE);
if (index != filter_context->client_context_count) {
memcpy(
&filter_context->client_contexts[index],
&filter_context->client_contexts[index + 1],
(filter_context->client_context_count - index) * sizeof(net_ebpf_extension_hook_client_t*));

filter_context->client_contexts[filter_context->client_context_count] = NULL;
}

ExReleaseSpinLockExclusive(&filter_context->lock, old_irql);
}
126 changes: 101 additions & 25 deletions netebpfext/net_ebpf_ext.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,37 @@
#define NET_EBPF_EXTENSION_POOL_TAG 'Nfbe'
#define NET_EBPF_EXTENSION_NPI_PROVIDER_VERSION 0

// 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
// 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
#define NET_EBPF_EXT_MAX_CLIENTS_PER_HOOK_SINGLE_ATTACH 1

CONST IN6_ADDR DECLSPEC_SELECTANY in6addr_v4mappedprefix = IN6ADDR_V4MAPPEDPREFIX_INIT;

#define _ACQUIRE_PUSH_LOCK(lock, mode) \
{ \
KeEnterCriticalRegion(); \
ExAcquirePushLock##mode(lock); \
}

#define _RELEASE_PUSH_LOCK(lock, mode) \
{ \
ExReleasePushLock##mode(lock); \
KeLeaveCriticalRegion(); \
}

#define ACQUIRE_PUSH_LOCK_EXCLUSIVE(lock) _ACQUIRE_PUSH_LOCK(lock, Exclusive)
#define ACQUIRE_PUSH_LOCK_SHARED(lock) _ACQUIRE_PUSH_LOCK(lock, Shared)

#define RELEASE_PUSH_LOCK_EXCLUSIVE(lock) _RELEASE_PUSH_LOCK(lock, Exclusive)
#define RELEASE_PUSH_LOCK_SHARED(lock) _RELEASE_PUSH_LOCK(lock, Shared)

#define htonl(x) _byteswap_ulong(x)
#define htons(x) _byteswap_ushort(x)
#define ntohl(x) _byteswap_ulong(x)
Expand All @@ -49,11 +78,12 @@ typedef struct _wfp_ale_layer_fields

typedef struct _net_ebpf_extension_wfp_filter_parameters
{
const GUID* layer_guid; ///< GUID of WFP layer to which this filter is associated.
const GUID* sublayer_guid; ///< GUID of the WFP sublayer to which this filter is associated.
const GUID* callout_guid; ///< GUID of WFP callout to which this filter is associated.
const wchar_t* name; ///< Display name of filter.
const wchar_t* description; ///< Description of filter.
const GUID* layer_guid; ///< GUID of WFP layer to which this filter is associated.
const GUID* sublayer_guid; ///< GUID of the WFP sublayer to which this filter is associated.
const GUID* callout_guid; ///< GUID of WFP callout to which this filter is associated.
const wchar_t* name; ///< Display name of filter.
const wchar_t* description; ///< Description of filter.
const FWP_ACTION_TYPE action_type; ///< Action type for the filter.
} net_ebpf_extension_wfp_filter_parameters_t;

typedef struct _net_ebpf_ext_sublayer_info
Expand All @@ -79,7 +109,9 @@ typedef struct _net_ebpf_extension_wfp_filter_parameters_array
typedef enum _net_ebpf_ext_wfp_filter_state
{
NET_EBPF_EXT_WFP_FILTER_ADDED = 1,
NET_EBPF_EXT_WFP_FILTER_DELETED = 2,
NET_EBPF_EXT_WFP_FILTER_DELETING = 2,
NET_EBPF_EXT_WFP_FILTER_DELETED = 3,
NET_EBPF_EXT_WFP_FILTER_DELETE_FAILED = 4,

} net_ebpf_ext_wfp_filter_state_t;

Expand All @@ -88,42 +120,60 @@ typedef struct _net_ebpf_ext_wfp_filter_id
wchar_t* name;
uint64_t id;
net_ebpf_ext_wfp_filter_state_t state;
NTSTATUS error_code;
} net_ebpf_ext_wfp_filter_id_t;

typedef struct _net_ebpf_extension_wfp_filter_context
{
volatile long reference_count; ///< Reference count.
const struct _net_ebpf_extension_hook_client* client_context; ///< Pointer to hook NPI client.
LIST_ENTRY link; ///< Entry in the list of filter contexts.
volatile long reference_count; ///< Reference count.
EX_SPIN_LOCK lock; ///< Lock to protect the client context array.
uint32_t client_context_count_max; ///< Maximum number of hook NPI clients.
_Guarded_by_(
lock) struct _net_ebpf_extension_hook_client** client_contexts; ///< Array of pointers to hook NPI clients.
_Guarded_by_(lock) uint32_t client_context_count; ///< Current number of hook NPI clients.
const struct _net_ebpf_extension_hook_provider* provider_context; ///< Pointer to provider binding context.

net_ebpf_ext_wfp_filter_id_t* filter_ids; ///< Array of WFP filter Ids.
uint32_t filter_ids_count; ///< Number of WFP filter Ids.

bool client_detached : 1; ///< True if client has detached.
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.
} 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 REFERENCE_FILTER_CONTEXT(filter_context) \
if ((filter_context) != NULL) { \
InterlockedIncrement(&(filter_context)->reference_count); \
}

#define DEREFERENCE_FILTER_CONTEXT(filter_context) \
if ((filter_context) != NULL) { \
if (InterlockedDecrement(&(filter_context)->reference_count) == 0) { \
net_ebpf_extension_hook_client_leave_rundown( \
(net_ebpf_extension_hook_client_t*)(filter_context)->client_context); \
if ((filter_context)->filter_ids != NULL) { \
ExFreePool((filter_context)->filter_ids); \
} \
ExFreePool((filter_context)); \
} \
#define DEREFERENCE_FILTER_CONTEXT(filter_context) \
if ((filter_context) != NULL) { \
if (InterlockedDecrement(&(filter_context)->reference_count) == 0) { \
net_ebpf_extension_hook_provider_leave_rundown( \
(net_ebpf_extension_hook_provider_t*)(filter_context)->provider_context); \
CLEAN_UP_FILTER_CONTEXT((filter_context)); \
} \
}

/**
* @brief This function allocates and initializes a net ebpf extension WFP filter context. This should be invoked when
* the hook client is being attached.
*
* @param[in] filter_context_size Size in bytes of the filter context.
* @param[in] client_context Pointer to hook client being attached. This would be associated with the filter context.
* @param[in] client_context Pointer to hook client being attached.
* @param[in] provider_context Pointer to hook provider context.
* @param[out] filter_context Pointer to created filter_context.
*
* @retval EBPF_SUCCESS The filter context was created successfully.
Expand All @@ -133,6 +183,7 @@ _Must_inspect_result_ ebpf_result_t
net_ebpf_extension_wfp_filter_context_create(
size_t filter_context_size,
_In_ const struct _net_ebpf_extension_hook_client* client_context,
_In_ const struct _net_ebpf_extension_hook_provider* provider_context,
_Outptr_ net_ebpf_extension_wfp_filter_context_t** filter_context);

/**
Expand Down Expand Up @@ -196,11 +247,11 @@ 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] filter_count Count of filters to be added.
* @param[in] parameters Filter parameters.
* @param[in] condition_count Count of filter conditions.
* @param[in] conditions Common filter conditions to be applied to each filter.
* @param[in, out] filter_context Caller supplied context to be associated with the WFP filter.
* @param[in] filter_count Count of filters to be added.
* @param[in] parameters Filter parameters.
* @param[in] condition_count Count of filter conditions.
* @param[in] conditions Common filter conditions to be applied to each filter.
* @param[in, out] filter_context Caller supplied context to be associated with the WFP filter.
* @param[out] filter_ids Output buffer where the added filter IDs are stored.
*
* @retval EBPF_SUCCESS The operation completed successfully.
Expand Down Expand Up @@ -295,3 +346,28 @@ net_ebpf_ext_unregister_providers();
NTSTATUS
net_ebpf_ext_filter_change_notify(
FWPS_CALLOUT_NOTIFY_TYPE callout_notification_type, _In_ const GUID* filter_key, _Inout_ FWPS_FILTER* filter);

/**
* @brief Remove the client context from the filter context.
*
* @param filter_context Filter context to remove the client from.
* @param hook_client Hook client to remove.
*/
void
net_ebpf_ext_remove_client_context(
_Inout_ net_ebpf_extension_wfp_filter_context_t* filter_context,
_In_ const struct _net_ebpf_extension_hook_client* hook_client);

/**
* @brief Add a client context to the filter context.
*
* @param filter_context Filter context to add the client to.
* @param hook_client Hook client to add.
*
* @retval EBPF_SUCCESS The client context was added successfully.
* @retval EBPF_NO_MEMORY No more client contexts can be added.
*/
ebpf_result_t
net_ebpf_ext_add_client_context(
_Inout_ net_ebpf_extension_wfp_filter_context_t* filter_context,
_In_ const struct _net_ebpf_extension_hook_client* hook_client);
Loading

0 comments on commit c73a086

Please sign in to comment.