Repository: microsoft/ebpf-for-windows (C)

Issue:
Implement bpf_get_socket_cookie(bpf_sock_addr*)
### Describe the feature you'd like supported

Implement `bpf_get_socket_cookie()` helper function for CGROUP SOCK_ADDR program type.

Per bpf-helper man page:

```
u64 bpf_get_socket_cookie(struct bpf_sock_addr *ctx)

              Description
                     Equivalent to bpf_get_socket_cookie() helper that
                     accepts skb, but gets socket from struct
                     bpf_sock_addr context.

              Return A 8-byte long unique number.
```

Description from sk_buff version:
```
              Description
                     If the struct sk_buff pointed by skb has a known
                     socket, retrieve the cookie (generated by the
                     kernel) of this socket.  If no cookie has been set
                     yet, generate a new cookie. Once generated, the
                     socket cookie remains stable for the life of the
                     socket. This helper can be useful for monitoring
                     per socket networking traffic statistics as it
                     provides a global socket identifier that can be
                     assumed unique.
```

### Proposed solution

On Windows, transportEndpointHandle is a 64 bit unique identifier for a socket. This is equivalent to the Linux socket cookie.
The proposal is for the helper function to return the transport endpoint handle for the socket.

### Additional context

_No response_



Fix (patch):
diff --git a/include/bpf_helper_defs.h b/include/bpf_helper_defs.h
index bedd249556..982cb632eb 100644
--- a/include/bpf_helper_defs.h
+++ b/include/bpf_helper_defs.h
@@ -406,6 +406,19 @@ EBPF_HELPER(
 #define bpf_memmove ((bpf_memmove_t)BPF_FUNC_memmove)
 #endif
 
+/**
+ * @brief Get the socket cookie associated with the socket context.
+ * The context can be *bpf_sock_addr* struct, *bpf_sock_ops* struct, or *bpf_sock* struct.
+ *
+ * @param[in] ctx Context passed to the eBPF program.
+ *
+ * @returns The socket cookie.
+ */
+EBPF_HELPER(int64_t, bpf_get_socket_cookie, (const void* ctx));
+#ifndef __doxygen
+#define bpf_get_socket_cookie ((bpf_get_socket_cookie_t)BPF_FUNC_get_socket_cookie)
+#endif
+
 #if __clang__
 #define memcpy(dest, src, dest_size) bpf_memcpy(dest, dest_size, src, dest_size)
 #define memcmp(mem1, mem2, mem1_size) bpf_memcmp(mem1, mem1_size, mem2, mem1_size)
diff --git a/include/ebpf_structs.h b/include/ebpf_structs.h
index f88ce20b0d..84610d527a 100644
--- a/include/ebpf_structs.h
+++ b/include/ebpf_structs.h
@@ -160,6 +160,7 @@ typedef enum
     BPF_FUNC_memcmp = 23,                    ///< \ref bpf_memcmp
     BPF_FUNC_memset = 24,                    ///< \ref bpf_memset
     BPF_FUNC_memmove = 25,                   ///< \ref bpf_memmove
+    BPF_FUNC_get_socket_cookie = 26,         ///< \ref bpf_get_socket_cookie
 } ebpf_helper_id_t;
 
 // Cross-platform BPF program types.
diff --git a/libs/execution_context/ebpf_core.c b/libs/execution_context/ebpf_core.c
index 1d5f3728c3..648130fdbb 100644
--- a/libs/execution_context/ebpf_core.c
+++ b/libs/execution_context/ebpf_core.c
@@ -148,6 +148,8 @@ static const void* _ebpf_general_helpers[] = {
     (void*)&_ebpf_core_memcmp,
     (void*)&_ebpf_core_memset,
     (void*)&_ebpf_core_memmove,
+    // No default implementation of bpf_get_socket_cookie
+    (void*)NULL, // bpf_get_socket_cookie
 };
 
 static const ebpf_helper_function_addresses_t _ebpf_global_helper_function_dispatch_table = {
diff --git a/libs/execution_context/ebpf_general_helpers.c b/libs/execution_context/ebpf_general_helpers.c
index 5a15fea84a..51be31bc43 100644
--- a/libs/execution_context/ebpf_general_helpers.c
+++ b/libs/execution_context/ebpf_general_helpers.c
@@ -168,6 +168,11 @@ ebpf_helper_function_prototype_t ebpf_core_helper_function_prototype_array[] = {
          EBPF_ARGUMENT_TYPE_PTR_TO_READABLE_MEM,
          EBPF_ARGUMENT_TYPE_CONST_SIZE_OR_ZERO,
      }},
+    {{EBPF_HELPER_FUNCTION_PROTOTYPE_CURRENT_VERSION, EBPF_HELPER_FUNCTION_PROTOTYPE_CURRENT_VERSION_SIZE},
+     BPF_FUNC_get_socket_cookie,
+     "bpf_get_socket_cookie",
+     EBPF_RETURN_TYPE_INTEGER,
+     {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}},
 };
 
 #ifdef __cplusplus
diff --git a/libs/execution_context/ebpf_program.c b/libs/execution_context/ebpf_program.c
index 3c4de59ec1..4c44daec15 100644
--- a/libs/execution_context/ebpf_program.c
+++ b/libs/execution_context/ebpf_program.c
@@ -1514,7 +1514,8 @@ _Requires_lock_held_(program->lock) static ebpf_result_t _ebpf_program_get_helpe
     }
 
     if (helper_function_id < EBPF_MAX_GENERAL_HELPER_FUNCTION) {
-        // Check the general helper function table of the program type.
+        // First check the global helper function table of the program type for any helper functions that are
+        // overwritten.
         if (!found) {
             for (size_t index = 0; index < program_data->program_info->count_of_global_helpers; index++) {
                 if (program_data->program_info->global_helper_prototype[index].helper_id == helper_function_id) {
@@ -1526,7 +1527,7 @@ _Requires_lock_held_(program->lock) static ebpf_result_t _ebpf_program_get_helpe
             }
         }
 
-        // Check the general helper function table of the general program type.
+        // Next check the general helper function table of the general program type.
         if (!found) {
             for (size_t index = 0; index < general_program_data->program_info->count_of_program_type_specific_helpers;
                  index++) {
@@ -1534,7 +1535,17 @@ _Requires_lock_held_(program->lock) static ebpf_result_t _ebpf_program_get_helpe
                     helper_function_id) {
                     function_address = (void*)general_program_data->program_type_specific_helper_function_addresses
                                            ->helper_function_address[index];
-                    found = true;
+                    // In case of helper functions that do not have any default / general implementation, the function
+                    // address will be NULL.
+                    if (function_address != NULL) {
+                        found = true;
+                    } else {
+                        EBPF_LOG_MESSAGE_UINT64(
+                            EBPF_TRACELOG_LEVEL_ERROR,
+                            EBPF_TRACELOG_KEYWORD_PROGRAM,
+                            "No override implementation found for helper ID",
+                            helper_function_id);
+                    }
                     break;
                 }
             }
diff --git a/netebpfext/net_ebpf_ext_program_info.h b/netebpfext/net_ebpf_ext_program_info.h
index fa14e6188e..273f3e4a1c 100644
--- a/netebpfext/net_ebpf_ext_program_info.h
+++ b/netebpfext/net_ebpf_ext_program_info.h
@@ -100,6 +100,11 @@ static const ebpf_helper_function_prototype_t _ebpf_sock_addr_global_helper_func
      BPF_FUNC_is_current_admin,
      "bpf_is_current_admin",
      EBPF_RETURN_TYPE_INTEGER,
+     {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}},
+    {{EBPF_HELPER_FUNCTION_PROTOTYPE_CURRENT_VERSION, EBPF_HELPER_FUNCTION_PROTOTYPE_CURRENT_VERSION_SIZE},
+     BPF_FUNC_get_socket_cookie,
+     "bpf_get_socket_cookie",
+     EBPF_RETURN_TYPE_INTEGER,
      {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}}};
 
 // CGROUP_SOCK_ADDR program information.
diff --git a/netebpfext/net_ebpf_ext_sock_addr.c b/netebpfext/net_ebpf_ext_sock_addr.c
index 9bf07fb29c..eb951968b0 100644
--- a/netebpfext/net_ebpf_ext_sock_addr.c
+++ b/netebpfext/net_ebpf_ext_sock_addr.c
@@ -182,6 +182,7 @@ typedef struct _net_ebpf_bpf_sock_addr
     net_ebpf_extension_hook_id_t hook_id;
     void* redirect_context;
     uint32_t redirect_context_size;
+    uint64_t transport_endpoint_handle;
 } net_ebpf_sock_addr_t;
 
 /**
@@ -261,6 +262,13 @@ _ebpf_sock_addr_get_current_logon_id(_In_ const bpf_sock_addr_t* ctx)
     return logon_id;
 }
 
+static uint64_t
+_ebpf_sock_addr_get_socket_cookie(_In_ const bpf_sock_addr_t* ctx)
+{
+    net_ebpf_sock_addr_t* sock_addr_ctx = CONTAINING_RECORD(ctx, net_ebpf_sock_addr_t, base);
+    return sock_addr_ctx->transport_endpoint_handle;
+}
+
 static int
 _ebpf_sock_addr_set_redirect_context(_In_ const bpf_sock_addr_t* ctx, _In_ void* data, _In_ uint32_t data_size)
 {
@@ -475,7 +483,9 @@ static ebpf_helper_function_addresses_t _ebpf_sock_addr_specific_helper_function
     (uint64_t*)_ebpf_sock_addr_specific_helper_functions};
 
 static const void* _ebpf_sock_addr_global_helper_functions[] = {
-    (void*)_ebpf_sock_addr_get_current_logon_id, (void*)_ebpf_sock_addr_is_current_admin};
+    (void*)_ebpf_sock_addr_get_current_logon_id,
+    (void*)_ebpf_sock_addr_is_current_admin,
+    (void*)_ebpf_sock_addr_get_socket_cookie};
 
 static ebpf_helper_function_addresses_t _ebpf_sock_addr_global_helper_function_address_table = {
     {EBPF_HELPER_FUNCTION_ADDRESSES_CURRENT_VERSION, EBPF_HELPER_FUNCTION_ADDRESSES_CURRENT_VERSION_SIZE},
@@ -1084,6 +1094,7 @@ _net_ebpf_extension_sock_addr_copy_wfp_connection_fields(
     FWPS_INCOMING_VALUE0* incoming_values = incoming_fixed_values->incomingValue;
 
     sock_addr_ctx->hook_id = hook_id;
+    sock_addr_ctx->transport_endpoint_handle = incoming_metadata_values->transportEndpointHandle;
 
     // Copy IP address fields.
     if ((hook_id == EBPF_HOOK_ALE_AUTH_CONNECT_V4) || (hook_id == EBPF_HOOK_ALE_AUTH_RECV_ACCEPT_V4) ||