Allow sockops BPF programs to query the WFP flow_id

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

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};
 

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,

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<bind_context_header_t*>(ebpf_allocate_with_tag(sizeof(bind_context_header_t), EBPF_POOL_TAG_DEFAULT));
+    bind_context_header_t* bind_context_header = reinterpret_cast<bind_context_header_t*>(
+        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<sock_addr_context_header_t*>(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<sock_addr_context_header_t*>(
+        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,13 +818,27 @@ _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 = {
     EBPF_HELPER_FUNCTION_ADDRESSES_HEADER,
     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<sock_ops_context_header_t*>(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<sock_ops_context_header_t*>(
+        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,

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 <build_output_path>
+// 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";