From a7c66cc1cd721c105d301230c4c7dce8c3aa47cc Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 28 Oct 2025 16:40:10 -0700 Subject: [PATCH 1/2] Allow sockops BPF programs to query the WFP flow_id Signed-off-by: Alan Jowett --- include/ebpf_nethooks.h | 19 ++++ netebpfext/net_ebpf_ext_program_info.h | 17 ++- netebpfext/net_ebpf_ext_sock_ops.c | 34 ++++++ tests/end_to_end/helpers.h | 28 +++-- tests/sample/sockops_flow_id.c | 140 +++++++++++++++++++++++++ tests/socket/socket_tests.cpp | 140 +++++++++++++++++++++++++ 6 files changed, 370 insertions(+), 8 deletions(-) create mode 100644 tests/sample/sockops_flow_id.c diff --git a/include/ebpf_nethooks.h b/include/ebpf_nethooks.h index e2d0c80fa6..d462a68cc2 100644 --- a/include/ebpf_nethooks.h +++ b/include/ebpf_nethooks.h @@ -199,6 +199,25 @@ typedef struct _bpf_sock_ops typedef int sock_ops_hook_t(bpf_sock_ops_t* context); +#define SOCK_OPS_EXT_HELPER_FN_BASE 0xFFFE + +typedef enum +{ + BPF_FUNC_sock_ops_get_flow_id = SOCK_OPS_EXT_HELPER_FN_BASE + 1, +} ebpf_sock_ops_helper_id_t; + +/** + * @brief Get the WFP flow ID associated with the current sock_ops context. + * + * @param[in] ctx Pointer to bpf_sock_ops_t context. + * + * @return The WFP flow ID as a 64-bit unsigned integer. + */ +EBPF_HELPER(uint64_t, bpf_sock_ops_get_flow_id, (bpf_sock_ops_t * ctx)); +#ifndef __doxygen +#define bpf_sock_ops_get_flow_id ((bpf_sock_ops_get_flow_id_t)BPF_FUNC_sock_ops_get_flow_id) +#endif + #ifdef _MSC_VER #pragma warning(pop) #endif diff --git a/netebpfext/net_ebpf_ext_program_info.h b/netebpfext/net_ebpf_ext_program_info.h index c4147447c7..4e814c40bb 100644 --- a/netebpfext/net_ebpf_ext_program_info.h +++ b/netebpfext/net_ebpf_ext_program_info.h @@ -144,6 +144,11 @@ enum _sock_ops_global_helper_functions SOCK_OPS_GLOBAL_HELPER_GET_CURRENT_PID_TGID, }; +enum _sock_ops_program_specific_helper_functions +{ + SOCK_OPS_PROGRAM_SPECIFIC_HELPER_GET_FLOW_ID, +}; + // SOCK_OPS global helper function prototypes. static const ebpf_helper_function_prototype_t _ebpf_sock_ops_global_helper_function_prototype[] = { {.header = EBPF_HELPER_FUNCTION_PROTOTYPE_HEADER, @@ -152,11 +157,19 @@ static const ebpf_helper_function_prototype_t _ebpf_sock_ops_global_helper_funct .return_type = EBPF_RETURN_TYPE_INTEGER, .arguments = {}, .implicit_context = true}}; + +// SOCK_OPS program type specific helper function prototypes. +static const ebpf_helper_function_prototype_t _ebpf_sock_ops_program_type_specific_helper_function_prototype[] = { + {EBPF_HELPER_FUNCTION_PROTOTYPE_HEADER, + BPF_FUNC_sock_ops_get_flow_id, + "bpf_sock_ops_get_flow_id", + EBPF_RETURN_TYPE_INTEGER, + {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}}}; static const ebpf_program_info_t _ebpf_sock_ops_program_info = { EBPF_PROGRAM_INFORMATION_HEADER, &_ebpf_sock_ops_program_type_descriptor, - 0, - NULL, + EBPF_COUNT_OF(_ebpf_sock_ops_program_type_specific_helper_function_prototype), + _ebpf_sock_ops_program_type_specific_helper_function_prototype, EBPF_COUNT_OF(_ebpf_sock_ops_global_helper_function_prototype), _ebpf_sock_ops_global_helper_function_prototype}; diff --git a/netebpfext/net_ebpf_ext_sock_ops.c b/netebpfext/net_ebpf_ext_sock_ops.c index ddc382349e..d6789a8836 100644 --- a/netebpfext/net_ebpf_ext_sock_ops.c +++ b/netebpfext/net_ebpf_ext_sock_ops.c @@ -22,6 +22,7 @@ typedef struct _net_ebpf_bpf_sock_ops EBPF_CONTEXT_HEADER; bpf_sock_ops_t context; uint64_t process_id; + uint64_t flow_id; ///< WFP flow ID associated with this connection. } net_ebpf_sock_ops_t; /** @@ -86,6 +87,27 @@ _ebpf_sock_ops_get_current_pid_tgid( return (sock_ops_ctx->process_id << 32 | (uint32_t)(uintptr_t)PsGetCurrentThreadId()); } +// +// SOCK_OPS Program-type specific helper function implementation. +// +static uint64_t +_ebpf_sock_ops_get_flow_id( + uint64_t dummy_param1, + uint64_t dummy_param2, + uint64_t dummy_param3, + uint64_t dummy_param4, + uint64_t dummy_param5, + _In_ const bpf_sock_ops_t* ctx) +{ + UNREFERENCED_PARAMETER(dummy_param1); + UNREFERENCED_PARAMETER(dummy_param2); + UNREFERENCED_PARAMETER(dummy_param3); + UNREFERENCED_PARAMETER(dummy_param4); + UNREFERENCED_PARAMETER(dummy_param5); + net_ebpf_sock_ops_t* sock_ops_ctx = CONTAINING_RECORD(ctx, net_ebpf_sock_ops_t, context); + return sock_ops_ctx->flow_id; +} + // // SOCK_OPS Program Information NPI Provider. // @@ -97,6 +119,13 @@ static ebpf_helper_function_addresses_t _ebpf_sock_ops_global_helper_function_ad EBPF_COUNT_OF(_ebpf_sock_ops_global_helper_functions), (uint64_t*)_ebpf_sock_ops_global_helper_functions}; +static const void* _ebpf_sock_ops_program_type_specific_helper_functions[] = {(void*)_ebpf_sock_ops_get_flow_id}; + +static ebpf_helper_function_addresses_t _ebpf_sock_ops_program_type_specific_helper_function_address_table = { + EBPF_HELPER_FUNCTION_ADDRESSES_HEADER, + EBPF_COUNT_OF(_ebpf_sock_ops_program_type_specific_helper_functions), + (uint64_t*)_ebpf_sock_ops_program_type_specific_helper_functions}; + static ebpf_result_t _ebpf_sock_ops_context_create( _In_reads_bytes_opt_(data_size_in) const uint8_t* data_in, @@ -116,6 +145,8 @@ _ebpf_sock_ops_context_destroy( static ebpf_program_data_t _ebpf_sock_ops_program_data = { .header = EBPF_PROGRAM_DATA_HEADER, .program_info = &_ebpf_sock_ops_program_info, + .program_type_specific_helper_function_addresses = + &_ebpf_sock_ops_program_type_specific_helper_function_address_table, .global_helper_function_addresses = &_ebpf_sock_ops_global_helper_function_address_table, .context_create = &_ebpf_sock_ops_context_create, .context_destroy = &_ebpf_sock_ops_context_destroy, @@ -515,6 +546,9 @@ net_ebpf_extension_sock_ops_flow_established_classify( local_flow_context->parameters.layer_id = incoming_fixed_values->layerId; local_flow_context->parameters.callout_id = net_ebpf_extension_get_callout_id_for_hook(hook_id); + // Store the flow_id in the sock_ops context for the helper function. + local_flow_context->context.flow_id = incoming_metadata_values->flowHandle; + status = FwpsFlowAssociateContext( local_flow_context->parameters.flow_id, local_flow_context->parameters.layer_id, diff --git a/tests/end_to_end/helpers.h b/tests/end_to_end/helpers.h index 7ddbd593bf..0fdf3989c2 100644 --- a/tests/end_to_end/helpers.h +++ b/tests/end_to_end/helpers.h @@ -579,8 +579,8 @@ _ebpf_bind_context_create( ebpf_result_t retval; *context = nullptr; bind_md_t* bind_context = nullptr; - bind_context_header_t* bind_context_header = - reinterpret_cast(ebpf_allocate_with_tag(sizeof(bind_context_header_t), EBPF_POOL_TAG_DEFAULT)); + bind_context_header_t* bind_context_header = reinterpret_cast( + ebpf_allocate_with_tag(sizeof(bind_context_header_t), EBPF_POOL_TAG_DEFAULT)); if (bind_context_header == nullptr) { retval = EBPF_NO_MEMORY; goto Done; @@ -714,8 +714,8 @@ _ebpf_sock_addr_context_create( *context = nullptr; bpf_sock_addr_t* sock_addr_context = nullptr; - sock_addr_context_header_t* sock_addr_context_header = - reinterpret_cast(ebpf_allocate_with_tag(sizeof(sock_addr_context_header_t), EBPF_POOL_TAG_DEFAULT)); + sock_addr_context_header_t* sock_addr_context_header = reinterpret_cast( + ebpf_allocate_with_tag(sizeof(sock_addr_context_header_t), EBPF_POOL_TAG_DEFAULT)); if (sock_addr_context_header == nullptr) { retval = EBPF_NO_MEMORY; goto Done; @@ -818,6 +818,13 @@ _ebpf_sock_ops_get_current_pid_tgid( return 0; } +static uint64_t +_ebpf_sock_ops_get_flow_id(_In_ const bpf_sock_ops_t* ctx) +{ + UNREFERENCED_PARAMETER(ctx); + return 12345; // Mock flow ID for testing. +} + static const void* _ebpf_sock_ops_global_helper_functions[] = {(void*)_ebpf_sock_ops_get_current_pid_tgid}; static ebpf_helper_function_addresses_t _ebpf_sock_ops_global_helper_function_address_table = { @@ -825,6 +832,13 @@ static ebpf_helper_function_addresses_t _ebpf_sock_ops_global_helper_function_ad EBPF_COUNT_OF(_ebpf_sock_ops_global_helper_functions), (uint64_t*)_ebpf_sock_ops_global_helper_functions}; +static const void* _ebpf_sock_ops_program_type_specific_helper_functions[] = {(void*)_ebpf_sock_ops_get_flow_id}; + +static ebpf_helper_function_addresses_t _ebpf_sock_ops_program_type_specific_helper_function_address_table = { + EBPF_HELPER_FUNCTION_ADDRESSES_HEADER, + EBPF_COUNT_OF(_ebpf_sock_ops_program_type_specific_helper_functions), + (uint64_t*)_ebpf_sock_ops_program_type_specific_helper_functions}; + static ebpf_result_t _ebpf_sock_ops_context_create( _In_reads_bytes_opt_(data_size_in) const uint8_t* data_in, @@ -839,8 +853,8 @@ _ebpf_sock_ops_context_create( *context = nullptr; bpf_sock_ops_t* sock_ops_context = nullptr; - sock_ops_context_header_t* sock_ops_context_header = - reinterpret_cast(ebpf_allocate_with_tag(sizeof(sock_ops_context_header_t), EBPF_POOL_TAG_DEFAULT)); + sock_ops_context_header_t* sock_ops_context_header = reinterpret_cast( + ebpf_allocate_with_tag(sizeof(sock_ops_context_header_t), EBPF_POOL_TAG_DEFAULT)); if (sock_ops_context_header == nullptr) { retval = EBPF_NO_MEMORY; goto Done; @@ -898,6 +912,8 @@ _ebpf_sock_ops_context_destroy( static ebpf_program_data_t _ebpf_sock_ops_program_data = { .header = EBPF_PROGRAM_DATA_HEADER, .program_info = &_ebpf_sock_ops_program_info, + .program_type_specific_helper_function_addresses = + &_ebpf_sock_ops_program_type_specific_helper_function_address_table, .global_helper_function_addresses = &_ebpf_sock_ops_global_helper_function_address_table, .context_create = &_ebpf_sock_ops_context_create, .context_destroy = &_ebpf_sock_ops_context_destroy, diff --git a/tests/sample/sockops_flow_id.c b/tests/sample/sockops_flow_id.c new file mode 100644 index 0000000000..302226b698 --- /dev/null +++ b/tests/sample/sockops_flow_id.c @@ -0,0 +1,140 @@ +// Copyright (c) eBPF for Windows contributors +// SPDX-License-Identifier: MIT + +// Whenever this sample program changes, bpf2c_tests will fail unless the +// expected files in tests\bpf2c_tests\expected are updated. The following +// script can be used to regenerate the expected files: +// generate_expected_bpf2c_output.ps1 +// +// Usage: +// .\scripts\generate_expected_bpf2c_output.ps1 +// Example: +// .\scripts\generate_expected_bpf2c_output.ps1 .\x64\Debug\ + +#include "bpf_helpers.h" +#include "ebpf_nethooks.h" +#include "net/ip.h" +#include "socket_tests_common.h" + +struct +{ + __uint(type, BPF_MAP_TYPE_HASH); + __type(key, connection_tuple_t); + __type(value, uint64_t); + __uint(max_entries, 10); +} flow_id_map SEC(".maps"); + +struct +{ + __uint(type, BPF_MAP_TYPE_RINGBUF); + __uint(max_entries, 256 * 1024); +} flow_id_audit_map SEC(".maps"); + +typedef struct _flow_id_audit_entry +{ + connection_tuple_t tuple; + uint64_t flow_id; + uint64_t process_id; + uint32_t operation; + bool outbound; + bool connected; +} flow_id_audit_entry_t; + +inline int +update_flow_id_audit_map(flow_id_audit_entry_t* audit_entry) +{ + return bpf_ringbuf_output(&flow_id_audit_map, audit_entry, sizeof(*audit_entry), 0); +} + +inline int +handle_v4_flow_id(bpf_sock_ops_t* ctx, bool outbound, bool connected) +{ + int result = 0; + flow_id_audit_entry_t audit_entry = {0}; + + // Get the WFP flow ID using the new helper function. + uint64_t flow_id = bpf_sock_ops_get_flow_id(ctx); + + audit_entry.tuple.local_ip.ipv4 = ctx->local_ip4; + audit_entry.tuple.local_port = ctx->local_port; + audit_entry.tuple.remote_ip.ipv4 = ctx->remote_ip4; + audit_entry.tuple.remote_port = ctx->remote_port; + audit_entry.tuple.protocol = ctx->protocol; + audit_entry.tuple.interface_luid = ctx->interface_luid; + audit_entry.process_id = bpf_get_current_pid_tgid(); + // Ignore the thread Id. + audit_entry.process_id >>= 32; + audit_entry.outbound = outbound; + audit_entry.connected = connected; + audit_entry.operation = ctx->op; + audit_entry.flow_id = flow_id; + + // Store the flow ID in our map for later verification. + bpf_map_update_elem(&flow_id_map, &audit_entry.tuple, &flow_id, BPF_ANY); + + return update_flow_id_audit_map(&audit_entry); +} + +inline int +handle_v6_flow_id(bpf_sock_ops_t* ctx, bool outbound, bool connected) +{ + int result = 0; + flow_id_audit_entry_t audit_entry = {0}; + + // Get the WFP flow ID using the new helper function. + uint64_t flow_id = bpf_sock_ops_get_flow_id(ctx); + + // Copy IPv6 addresses. + __builtin_memcpy(&audit_entry.tuple.local_ip.ipv6, &ctx->local_ip6, sizeof(audit_entry.tuple.local_ip.ipv6)); + __builtin_memcpy(&audit_entry.tuple.remote_ip.ipv6, &ctx->remote_ip6, sizeof(audit_entry.tuple.remote_ip.ipv6)); + + audit_entry.tuple.local_port = ctx->local_port; + audit_entry.tuple.remote_port = ctx->remote_port; + audit_entry.tuple.protocol = ctx->protocol; + audit_entry.tuple.interface_luid = ctx->interface_luid; + audit_entry.process_id = bpf_get_current_pid_tgid(); + // Ignore the thread Id. + audit_entry.process_id >>= 32; + audit_entry.outbound = outbound; + audit_entry.connected = connected; + audit_entry.operation = ctx->op; + audit_entry.flow_id = flow_id; + + // Store the flow ID in our map for later verification + bpf_map_update_elem(&flow_id_map, &audit_entry.tuple, &flow_id, BPF_ANY); + + return update_flow_id_audit_map(&audit_entry); +} + +SEC("sockops") +int +flow_id_monitor(bpf_sock_ops_t* ctx) +{ + switch (ctx->op) { + case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB: + if (ctx->family == AF_INET) { + return handle_v4_flow_id(ctx, false, true); + } else if (ctx->family == AF_INET6) { + return handle_v6_flow_id(ctx, false, true); + } + break; + case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB: + if (ctx->family == AF_INET) { + return handle_v4_flow_id(ctx, true, true); + } else if (ctx->family == AF_INET6) { + return handle_v6_flow_id(ctx, true, true); + } + break; + case BPF_SOCK_OPS_CONNECTION_DELETED_CB: + if (ctx->family == AF_INET) { + return handle_v4_flow_id(ctx, false, false); + } else if (ctx->family == AF_INET6) { + return handle_v6_flow_id(ctx, false, false); + } + break; + } + + return 0; +} + +char _license[] SEC("license") = "MIT"; \ No newline at end of file diff --git a/tests/socket/socket_tests.cpp b/tests/socket/socket_tests.cpp index 66406cee48..20b154a87d 100644 --- a/tests/socket/socket_tests.cpp +++ b/tests/socket/socket_tests.cpp @@ -501,6 +501,146 @@ TEST_CASE("attach_sockops_programs", "[sock_ops_tests]") SAFE_REQUIRE(result == 0); } +// Custom event handler for flow ID validation +int +flow_id_ring_buffer_event_handler(_Inout_ void* ctx, _In_opt_ const void* data, size_t size) +{ + ring_buffer_test_event_context_t* event_context = reinterpret_cast(ctx); + + if ((data == nullptr) || (size == 0)) { + return 0; + } + + if (event_context->canceled) { + // Ignore the callback as the subscription is canceled. + // Return error so that no further callback is made. + return -1; + } + + if (event_context->matched_entry_count >= event_context->test_event_count) { + // Required number of event notifications already received. + return 0; + } + + // Parse the flow_id_audit_entry_t from the ring buffer + struct flow_id_audit_entry_t + { + connection_tuple_t tuple; + uint64_t flow_id; + uint64_t process_id; + uint32_t operation; + bool outbound; + bool connected; + }; + + if (size != sizeof(flow_id_audit_entry_t)) { + // Unexpected entry size + return 0; + } + + const flow_id_audit_entry_t* entry = reinterpret_cast(data); + + // Validate that the flow ID is non-zero + if (entry->flow_id != 0) { + event_context->matched_entry_count++; + + if (event_context->matched_entry_count == event_context->test_event_count) { + // If all the expected entries were received, fulfill the promise. + event_context->ring_buffer_event_promise.set_value(); + } + } + + return 0; +} + +TEST_CASE("sock_ops_flow_id_helper_test", "[sock_ops_tests]") +{ + native_module_helper_t helper; + helper.initialize("sockops_flow_id", _is_main_thread); + struct bpf_object* object = bpf_object__open(helper.get_file_name().c_str()); + bpf_object_ptr object_ptr(object); + + SAFE_REQUIRE(object != nullptr); + // Load the programs. + SAFE_REQUIRE(bpf_object__load(object) == 0); + + // Ring buffer event callback context. + std::unique_ptr context = std::make_unique(); + context->test_event_count = 2; // Expect 2 events for TCP connection (outbound and inbound) + + bpf_program* _program = bpf_object__find_program_by_name(object, "flow_id_monitor"); + SAFE_REQUIRE(_program != nullptr); + + bpf_map* flow_id_map = bpf_object__find_map_by_name(object, "flow_id_map"); + SAFE_REQUIRE(flow_id_map != nullptr); + + // Attach the program. + int result = bpf_prog_attach(bpf_program__fd(const_cast(_program)), 0, BPF_CGROUP_SOCK_OPS, 0); + SAFE_REQUIRE(result == 0); + + // Get the std::future from the promise field in ring buffer event context, which should be in ready state + // once notifications for all events are received. + auto ring_buffer_event_callback = context->ring_buffer_event_promise.get_future(); + + // Create a new ring buffer manager and subscribe to ring buffer events (using async mode for automatic callbacks). + bpf_map* ring_buffer_map = bpf_object__find_map_by_name(object, "flow_id_audit_map"); + SAFE_REQUIRE(ring_buffer_map != nullptr); + ebpf_ring_buffer_opts ring_opts = {}; + ring_opts.sz = sizeof(ring_opts); + ring_opts.flags = EBPF_RINGBUF_FLAG_AUTO_CALLBACK; + context->ring_buffer = ebpf_ring_buffer__new( + bpf_map__fd(ring_buffer_map), + (ring_buffer_sample_fn)flow_id_ring_buffer_event_handler, + context.get(), + &ring_opts); + SAFE_REQUIRE(context->ring_buffer != nullptr); + + // Create a basic TCP connection to trigger the helper function. + stream_client_socket_t stream_client_socket(SOCK_STREAM, IPPROTO_TCP, 0); + stream_server_socket_t stream_server_socket(SOCK_STREAM, IPPROTO_TCP, SOCKET_TEST_PORT); + + PSOCKADDR local_address = nullptr; + int local_address_length = 0; + stream_client_socket.get_local_address(local_address, local_address_length); + + connection_tuple_t tuple{}; + tuple.local_ip.ipv4 = htonl(INADDR_LOOPBACK); + tuple.remote_ip.ipv4 = htonl(INADDR_LOOPBACK); + tuple.local_port = INETADDR_PORT(local_address); + tuple.remote_port = htons(SOCKET_TEST_PORT); + tuple.protocol = IPPROTO_TCP; + + // Post an asynchronous receive on the receiver socket. + stream_server_socket.post_async_receive(); + + // Send loopback message to test port. + const char* message = CLIENT_MESSAGE; + sockaddr_storage destination_address{}; + IN6ADDR_SETV4MAPPED((PSOCKADDR_IN6)&destination_address, &in4addr_loopback, scopeid_unspecified, 0); + + stream_client_socket.send_message_to_remote_host(message, destination_address, SOCKET_TEST_PORT); + // Receive the packet on test port. + stream_server_socket.complete_async_receive(); + + // Wait for event handler getting notifications for all flow ID audit events. + SAFE_REQUIRE(ring_buffer_event_callback.wait_for(1s) == std::future_status::ready); + + // Mark the event context as canceled, such that the event callback stops processing events. + context->canceled = true; + + // Unsubscribe. + context->unsubscribe(); + + // Verify that we got a flow ID stored in the map (should be non-zero). + uint64_t stored_flow_id = 0; + result = bpf_map_lookup_elem(bpf_map__fd(flow_id_map), &tuple, &stored_flow_id); + + // Verify we get a non-zero flow ID. + if (result == 0) { + REQUIRE(stored_flow_id != 0); + } +} + // This function populates map policies for multi-attach tests. // It assumes that the destination and proxy are loopback addresses. static void From f3ba484f8df9c53e632ccc4797f5978e8730de82 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 28 Oct 2025 17:11:15 -0700 Subject: [PATCH 2/2] PR feedback Signed-off-by: Alan Jowett --- include/ebpf_nethooks.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/ebpf_nethooks.h b/include/ebpf_nethooks.h index d462a68cc2..306df4c3ac 100644 --- a/include/ebpf_nethooks.h +++ b/include/ebpf_nethooks.h @@ -199,7 +199,7 @@ typedef struct _bpf_sock_ops typedef int sock_ops_hook_t(bpf_sock_ops_t* context); -#define SOCK_OPS_EXT_HELPER_FN_BASE 0xFFFE +#define SOCK_OPS_EXT_HELPER_FN_BASE 0xFFFF typedef enum {