Add helper function to set redirect context (#2823)

* add new helper function to set redirect context

* WIP

* update typo

* build failure fix

* update expected programs, update tcp_udp_listener to initialize to null terminated

* update expected program with proper clang version

* disable expected failing test

* update comments with periods

* address comment regarding potentially leaked memory

* Add checks for invalid parameters

* update comment

* add tracing for errors

* address CR

include/ebpf_nethooks.h

@@ -156,6 +156,7 @@ typedef struct bpf_sock_addr
 typedef enum
 {
     BPF_FUNC_sock_addr_get_current_pid_tgid = SOCK_ADDR_EXT_HELPER_FN_BASE + 1,
+    BPF_FUNC_sock_addr_set_redirect_context = SOCK_ADDR_EXT_HELPER_FN_BASE + 2,
 } ebpf_sock_addr_helper_id_t;
 
 /**
@@ -175,6 +176,23 @@ EBPF_HELPER(uint64_t, bpf_sock_addr_get_current_pid_tgid, (bpf_sock_addr_t * ctx
     ((bpf_sock_addr_get_current_pid_tgid_t)BPF_FUNC_sock_addr_get_current_pid_tgid)
 #endif
 
+/**
+ * @brief Set a context for consumption by a user-mode application (sock_addr specific only).
+ * This function is not supported for the recv_accept hooks.
+ *
+ * @param[in] ctx Pointer to bpf_sock_addr_t context.
+ * @param[in] data Pointer to data to store.
+ * @param[in] data_size The size of the data to store.
+ *
+ * @retval 0 The operation was successful.
+ * @retval <0 A failure occurred.
+ */
+EBPF_HELPER(int, bpf_sock_addr_set_redirect_context, (bpf_sock_addr_t * ctx, void* data, uint32_t data_size));
+#ifndef __doxygen
+#define bpf_sock_addr_set_redirect_context \
+    ((bpf_sock_addr_set_redirect_context_t)BPF_FUNC_sock_addr_set_redirect_context)
+#endif
+
 /**
  * @brief Handle socket operation. Currently supports ingress/egress connection initialization.
  *

netebpfext/net_ebpf_ext_program_info.h

@@ -46,7 +46,11 @@ static const ebpf_helper_function_prototype_t _sock_addr_ebpf_extension_helper_f
     {BPF_FUNC_sock_addr_get_current_pid_tgid,
      "bpf_sock_addr_get_current_pid_tgid",
      EBPF_RETURN_TYPE_INTEGER,
-     {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}}};
+     {EBPF_ARGUMENT_TYPE_PTR_TO_CTX}},
+    {BPF_FUNC_sock_addr_set_redirect_context,
+     "bpf_sock_addr_set_redirect_context",
+     EBPF_RETURN_TYPE_INTEGER,
+     {EBPF_ARGUMENT_TYPE_PTR_TO_CTX, EBPF_ARGUMENT_TYPE_PTR_TO_READABLE_MEM, EBPF_ARGUMENT_TYPE_CONST_SIZE}}};
 
 // CGROUP_SOCK_ADDR global helper function prototypes.
 static const ebpf_helper_function_prototype_t _ebpf_sock_addr_global_helper_function_prototype[] = {

netebpfext/net_ebpf_ext_sock_addr.c

@@ -179,6 +179,9 @@ typedef struct _net_ebpf_bpf_sock_addr
     TOKEN_ACCESS_INFORMATION* access_information;
     uint64_t process_id;
     uint32_t flags;
+    net_ebpf_extension_hook_id_t hook_id;
+    void* redirect_context;
+    uint32_t redirect_context_size;
 } net_ebpf_sock_addr_t;
 
 /**
@@ -258,6 +261,59 @@ _ebpf_sock_addr_get_current_logon_id(_In_ const bpf_sock_addr_t* ctx)
     return logon_id;
 }
 
+static int
+_ebpf_sock_addr_set_redirect_context(_In_ const bpf_sock_addr_t* ctx, _In_ void* data, _In_ uint32_t data_size)
+{
+    int return_value = 0;
+    net_ebpf_sock_addr_t* sock_addr_ctx = CONTAINING_RECORD(ctx, net_ebpf_sock_addr_t, base);
+    void* redirect_context = NULL;
+
+    // Check for invalid parameters.
+    if (data == NULL || data_size == 0) {
+        return_value = -1;
+        goto Exit;
+    }
+
+    // This function is only supported at the connect redirect layer.
+    if (sock_addr_ctx->hook_id != EBPF_HOOK_ALE_CONNECT_REDIRECT_V4 &&
+        sock_addr_ctx->hook_id != EBPF_HOOK_ALE_CONNECT_REDIRECT_V6) {
+        NET_EBPF_EXT_LOG_MESSAGE(
+            NET_EBPF_EXT_TRACELOG_LEVEL_ERROR,
+            NET_EBPF_EXT_TRACELOG_KEYWORD_SOCK_ADDR,
+            "_ebpf_sock_addr_set_redirect_context invoked at incorrect hook.");
+        return_value = -1;
+        goto Exit;
+    }
+
+    // Allocate buffer to store redirect context.
+    redirect_context = ExAllocatePoolUninitialized(NonPagedPoolNx, data_size, NET_EBPF_EXTENSION_POOL_TAG);
+    if (redirect_context == NULL) {
+        NET_EBPF_EXT_LOG_MESSAGE(
+            NET_EBPF_EXT_TRACELOG_LEVEL_ERROR,
+            NET_EBPF_EXT_TRACELOG_KEYWORD_SOCK_ADDR,
+            "_ebpf_sock_addr_set_redirect_context failed to allocate memory for the redirect context.");
+        return_value = -1;
+        goto Exit;
+    }
+    memcpy(redirect_context, data, data_size);
+
+    // If a redirect context already exists, free the existing buffer.
+    if (sock_addr_ctx->redirect_context != NULL) {
+        ExFreePool(sock_addr_ctx->redirect_context);
+    }
+
+    // Set the redirect context.
+    sock_addr_ctx->redirect_context = redirect_context;
+    sock_addr_ctx->redirect_context_size = data_size;
+
+Exit:
+    if (return_value == -1) {
+        NET_EBPF_EXT_LOG_FUNCTION_ERROR(return_value);
+    }
+
+    return return_value;
+}
+
 _IRQL_requires_max_(DISPATCH_LEVEL) static NTSTATUS _perform_access_check(
     _In_ SECURITY_DESCRIPTOR* security_descriptor,
     _In_ TOKEN_ACCESS_INFORMATION* access_information,
@@ -410,7 +466,8 @@ _ebpf_sock_addr_context_destroy(
 // SOCK_ADDR Program Information NPI Provider.
 //
 
-static const void* _ebpf_sock_addr_specific_helper_functions[] = {(void*)_ebpf_sock_addr_get_current_pid_tgid};
+static const void* _ebpf_sock_addr_specific_helper_functions[] = {
+    (void*)_ebpf_sock_addr_get_current_pid_tgid, (void*)_ebpf_sock_addr_set_redirect_context};
 
 static ebpf_helper_function_addresses_t _ebpf_sock_addr_specific_helper_function_address_table = {
     EBPF_COUNT_OF(_ebpf_sock_addr_specific_helper_functions), (uint64_t*)_ebpf_sock_addr_specific_helper_functions};
@@ -1053,6 +1110,8 @@ _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;
+
     // Copy IP address fields.
     if ((hook_id == EBPF_HOOK_ALE_AUTH_CONNECT_V4) || (hook_id == EBPF_HOOK_ALE_AUTH_RECV_ACCEPT_V4) ||
         (hook_id == EBPF_HOOK_ALE_CONNECT_REDIRECT_V4)) {
@@ -1307,6 +1366,7 @@ _net_ebpf_ext_process_redirect_verdict(
     NTSTATUS status = STATUS_SUCCESS;
     FWPS_CONNECT_REQUEST* connect_request = NULL;
     BOOLEAN commit_layer_data = FALSE;
+    net_ebpf_sock_addr_t* sock_addr_ctx = CONTAINING_RECORD(redirected_context, net_ebpf_sock_addr_t, base);
 
     *redirected = FALSE;
 
@@ -1351,6 +1411,12 @@ _net_ebpf_ext_process_redirect_verdict(
 
         connect_request->localRedirectTargetPID = TARGET_PROCESS_ID;
         connect_request->localRedirectHandle = redirect_handle;
+
+        connect_request->localRedirectContext = sock_addr_ctx->redirect_context;
+        connect_request->localRedirectContextSize = sock_addr_ctx->redirect_context_size;
+        // Ownership transferred to WFP.
+        sock_addr_ctx->redirect_context = NULL;
+        sock_addr_ctx->redirect_context_size = 0;
     }
 
 Exit:
@@ -1648,6 +1714,10 @@ net_ebpf_extension_sock_addr_redirect_connection_classify(
         net_ebpf_extension_hook_client_leave_rundown(attached_client);
     }
 
+    if (net_ebpf_sock_addr_ctx.redirect_context != NULL) {
+        ExFreePool(net_ebpf_sock_addr_ctx.redirect_context);
+    }
+
     NET_EBPF_EXT_LOG_EXIT();
 }