Add support for stack unwind on success. (#367)

Signed-off-by: Alan Jowett alanjo@microsoft.com

external/ubpf

@@ -1 +1 @@
-Subproject commit 843394752c752e38f79d6d49fc704a73ee6490bd
+Subproject commit 9eb26b4bfdec6cafbf629a056155363f12cec972

include/ebpf_helpers.h

@@ -6,6 +6,8 @@
 
 #include "ebpf_structs.h"
 
+#define MAX_TAIL_CALL_CNT 32
+
 // In an execution context, struct bpf_map means struct _ebpf_map_definition,
 // as opposed to for user mode apps, so define the alias here where the execution
 // context and eBPF programs will get it.

include/ebpf_result.h

@@ -94,6 +94,9 @@ extern "C"
 
         // The requested key is already present.
         EBPF_KEY_ALREADY_EXISTS,
+        
+        /// Caller has reached tail call limit.
+        EBPF_NO_MORE_TAIL_CALLS,
     } ebpf_result_t;
 
 #ifdef __cplusplus

libs/execution_context/ebpf_core.c

@@ -10,6 +10,7 @@
 #include "ebpf_program.h"
 #include "ebpf_program_types.h"
 #include "ebpf_serialize.h"
+#include "ebpf_state.h"
 
 GUID ebpf_general_helper_function_interface_id = {/* 8d2a1d3f-9ce6-473d-b48e-17aa5c5581fe */
                                                   0x8d2a1d3f,
@@ -52,11 +53,13 @@ static ebpf_helper_function_prototype_t _ebpf_map_helper_function_prototype[] =
      {EBPF_ARGUMENT_TYPE_PTR_TO_CTX, EBPF_ARGUMENT_TYPE_PTR_TO_MAP_OF_PROGRAMS, EBPF_ARGUMENT_TYPE_ANYTHING}},
 };
 
-static ebpf_program_info_t _ebpf_global_helper_program_info = {{"global_helper", NULL, {0}},
+static ebpf_program_info_t _ebpf_global_helper_program_info = {
+    {"global_helper", NULL, {0}},
     EBPF_COUNT_OF(_ebpf_map_helper_function_prototype),
     _ebpf_map_helper_function_prototype};
 
-static const void* _ebpf_general_helpers[] = {NULL,
+static const void* _ebpf_general_helpers[] = {
+    NULL,
     (void*)&_ebpf_core_map_find_element,
     (void*)&_ebpf_core_map_update_element,
     (void*)&_ebpf_core_map_delete_element,
@@ -65,8 +68,8 @@ static const void* _ebpf_general_helpers[] = {NULL,
 static ebpf_extension_provider_t* _ebpf_global_helper_function_provider_context = NULL;
 static ebpf_helper_function_addresses_t _ebpf_global_helper_function_dispatch_table = {
     EBPF_COUNT_OF(_ebpf_general_helpers), (uint64_t*)_ebpf_general_helpers};
-static ebpf_program_data_t _ebpf_global_helper_function_program_data = {&_ebpf_global_helper_program_info,
-                                                                        &_ebpf_global_helper_function_dispatch_table};
+static ebpf_program_data_t _ebpf_global_helper_function_program_data = {
+    &_ebpf_global_helper_program_info, &_ebpf_global_helper_function_dispatch_table};
 
 static ebpf_extension_data_t _ebpf_global_helper_function_extension_data = {
     EBPF_CORE_GLOBAL_HELPER_EXTENSION_VERSION,
@@ -86,6 +89,10 @@ ebpf_core_initiate()
     if (return_value != EBPF_SUCCESS)
         goto Done;
 
+    return_value = ebpf_state_initiate();
+    if (return_value != EBPF_SUCCESS)
+        goto Done;
+
     ebpf_object_tracking_initiate();
 
     return_value = ebpf_pinning_table_allocate(&_ebpf_core_map_pinning_table);
@@ -96,6 +103,10 @@ ebpf_core_initiate()
     if (return_value != EBPF_SUCCESS)
         goto Done;
 
+    return_value = ebpf_program_initiate();
+    if (return_value != EBPF_SUCCESS)
+        goto Done;
+
     return_value = ebpf_provider_load(
         &_ebpf_global_helper_function_provider_context,
         &ebpf_general_helper_function_interface_id,
@@ -125,10 +136,14 @@ ebpf_core_terminate()
     ebpf_provider_unload(_ebpf_global_helper_function_provider_context);
     _ebpf_global_helper_function_provider_context = NULL;
 
+    ebpf_program_terminate();
+
     ebpf_handle_table_terminate();
 
     ebpf_pinning_table_free(_ebpf_core_map_pinning_table);
 
+    ebpf_state_terminate();
+
     // Shut down the epoch tracker and free any remaining memory or work items.
     // Note: Some objects may only be released on epoch termination.
     ebpf_epoch_flush();
@@ -656,9 +671,9 @@ static ebpf_result_t
 _ebpf_core_protocol_update_pinning(_In_ const struct _ebpf_operation_update_map_pinning_request* request)
 {
     ebpf_result_t retval;
-    const ebpf_utf8_string_t name = {(uint8_t*)request->name,
-                                     request->header.length -
-                                         EBPF_OFFSET_OF(ebpf_operation_update_pinning_request_t, name)};
+    const ebpf_utf8_string_t name = {
+        (uint8_t*)request->name,
+        request->header.length - EBPF_OFFSET_OF(ebpf_operation_update_pinning_request_t, name)};
     ebpf_object_t* object = NULL;
 
     if (name.length == 0) {
@@ -982,19 +997,14 @@ _ebpf_core_map_delete_element(ebpf_map_t* map, const uint8_t* key)
 static int64_t
 _ebpf_core_tail_call(void* context, ebpf_map_t* map, uint32_t index)
 {
+    UNREFERENCED_PARAMETER(context);
+
     // Get program from map[index].
     ebpf_program_t* callee = ebpf_map_get_program_from_entry(map, sizeof(index), (uint8_t*)&index);
     if (callee == NULL) {
         return -EBPF_INVALID_ARGUMENT;
     }
-
-    // TODO(issue #344): Jump to (not call) the callee.
-    uint32_t result;
-    ebpf_program_invoke(callee, context, &result);
-
-    ebpf_object_release_reference((ebpf_object_t*)callee);
-
-    return result;
+    return -ebpf_program_set_tail_call(callee);
 }
 
 typedef struct _ebpf_protocol_handler

libs/execution_context/ebpf_program.c

@@ -3,14 +3,18 @@
 
 #include "ebpf_core.h"
 #include "ebpf_epoch.h"
+#include "ebpf_helpers.h"
 #include "ebpf_link.h"
 #include "ebpf_object.h"
 #include "ebpf_program.h"
 #include "ebpf_program_types.h"
+#include "ebpf_state.h"
 
 typedef struct _FILE FILE;
 #include "ubpf.h"
 
+static size_t _ebpf_program_state_index = MAXUINT64;
+
 typedef struct _ebpf_program
 {
     ebpf_object_t object;
@@ -60,6 +64,16 @@ typedef struct _ebpf_program
 static ebpf_result_t
 _ebpf_program_register_helpers(ebpf_program_t* program);
 
+ebpf_result_t
+ebpf_program_initiate()
+{
+    return ebpf_state_allocate_index(&_ebpf_program_state_index);
+}
+
+void
+ebpf_program_terminate()
+{}
+
 static void
 _ebpf_program_detach_links(_Inout_ ebpf_program_t* program)
 {
@@ -548,27 +562,78 @@ ebpf_program_load_code(
     return result;
 }
 
+typedef struct _ebpf_program_tail_call_state
+{
+    const ebpf_program_t* next_program;
+    uint32_t count;
+} ebpf_program_tail_call_state_t;
+
+ebpf_result_t
+ebpf_program_set_tail_call(_In_ const ebpf_program_t* next_program)
+{
+    ebpf_result_t result;
+    ebpf_program_tail_call_state_t* state = NULL;
+    result = ebpf_state_load(_ebpf_program_state_index, (uintptr_t*)&state);
+    if (result != EBPF_SUCCESS)
+        return result;
+
+    if (state == NULL)
+        return EBPF_INVALID_ARGUMENT;
+
+    if (state->count == MAX_TAIL_CALL_CNT) {
+        return EBPF_NO_MORE_TAIL_CALLS;
+    }
+
+    state->next_program = next_program;
+
+    return EBPF_SUCCESS;
+}
+
 void
 ebpf_program_invoke(_In_ const ebpf_program_t* program, _In_ void* context, _Out_ uint32_t* result)
 {
+    ebpf_program_tail_call_state_t state = {0};
+    const ebpf_program_t* current_program = program;
+
     if (!program || program->program_invalidated) {
         *result = 0;
         return;
     }
 
-    if (program->parameters.code_type == EBPF_CODE_NATIVE) {
+    if (!ebpf_state_store(_ebpf_program_state_index, (uintptr_t)&state) == EBPF_SUCCESS) {
+        *result = 0;
+        return;
+    }
+
+    for (state.count = 0; state.count < MAX_TAIL_CALL_CNT; state.count++) {
+        if (current_program->parameters.code_type == EBPF_CODE_NATIVE) {
             ebpf_program_entry_point_t function_pointer;
-        function_pointer = (ebpf_program_entry_point_t)(program->code_or_vm.code.code_pointer);
+            function_pointer = (ebpf_program_entry_point_t)(current_program->code_or_vm.code.code_pointer);
             *result = (function_pointer)(context);
         } else {
             uint64_t out_value;
-        int ret = (uint32_t)(ubpf_exec(program->code_or_vm.vm, context, 1024, &out_value));
+            int ret = (uint32_t)(ubpf_exec(current_program->code_or_vm.vm, context, 1024, &out_value));
             if (ret < 0) {
                 *result = ret;
             } else {
                 *result = (uint32_t)(out_value);
             }
         }
+
+        if (state.count != 0) {
+            ebpf_object_release_reference((ebpf_object_t*)current_program);
+            current_program = NULL;
+        }
+
+        if (state.next_program == NULL) {
+            break;
+        } else {
+            current_program = state.next_program;
+            state.next_program = NULL;
+        }
+    }
+
+    ebpf_state_store(_ebpf_program_state_index, 0);
 }
 
 static ebpf_result_t

libs/execution_context/ebpf_program.h

@@ -34,6 +34,23 @@ extern "C"
 
     typedef ebpf_result_t (*ebpf_program_entry_point_t)(void* context);
 
+    /**
+     * @brief Initialize global state for the ebpf program module.
+     *
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_NO_MEMORY Unable to allocate resources for this
+     *  operation.
+     */
+    ebpf_result_t
+    ebpf_program_initiate();
+
+    /**
+     * @brief Uninitialize the eBPF state tracking module.
+     *
+     */
+    void
+    ebpf_program_terminate();
+
     /**
      * @brief Create a new program instance.
      *
@@ -187,6 +204,17 @@ extern "C"
     void
     ebpf_program_detach_link(_Inout_ ebpf_program_t* program, _Inout_ ebpf_link_t* link);
 
+    /**
+     * @brief Store the pointer to the program to execute on tail call.
+     *
+     * @param[in] next_program Next program to execute.
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_INVALID_ARGUMENT Internal error.
+     * @retval EBPF_NO_MORE_TAIL_CALLS Program has executed to many tail calls.
+     */
+    ebpf_result_t
+    ebpf_program_set_tail_call(_In_ const ebpf_program_t* next_program);
+
 #ifdef __cplusplus
 }
 #endif

libs/platform/ebpf_state.c

@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+
+#include "ebpf_state.h"
+#include "ebpf_epoch.h"
+
+#define EBPF_MAX_STATE_ENTRIES 64
+
+// Table to track what state for each thread.
+static ebpf_hash_table_t* _ebpf_state_thread_table = NULL;
+
+static int64_t _ebpf_state_next_index = 0;
+
+// Table to track what state for each CPU.
+typedef struct _ebpf_state_entry
+{
+    uintptr_t state[EBPF_MAX_STATE_ENTRIES];
+} ebpf_state_entry_t;
+
+static _Writable_elements_(_ebpf_state_cpu_table_size) ebpf_state_entry_t* _ebpf_state_cpu_table = NULL;
+static uint32_t _ebpf_state_cpu_table_size = 0;
+
+ebpf_result_t
+ebpf_state_initiate()
+{
+    ebpf_result_t return_value = EBPF_SUCCESS;
+
+    if (ebpf_is_non_preemptible_work_item_supported()) {
+        ebpf_get_cpu_count(&_ebpf_state_cpu_table_size);
+        _Analysis_assume_(_ebpf_state_cpu_table_size >= 1);
+
+        _ebpf_state_cpu_table = ebpf_allocate(_ebpf_state_cpu_table_size * sizeof(ebpf_state_entry_t));
+        if (!_ebpf_state_cpu_table) {
+            return_value = EBPF_NO_MEMORY;
+            goto Error;
+        }
+    }
+
+    return_value = ebpf_hash_table_create(
+        &_ebpf_state_thread_table,
+        ebpf_epoch_allocate,
+        ebpf_epoch_free,
+        sizeof(uint64_t),
+        sizeof(ebpf_state_entry_t),
+        ebpf_get_cpu_count(),
+        NULL);
+    if (return_value != EBPF_SUCCESS) {
+        goto Error;
+    }
+
+    return return_value;
+
+Error:
+    ebpf_state_terminate();
+    return return_value;
+}
+
+/**
+ * @brief Uninitialize the eBPF state tracking module.
+ *
+ */
+void
+ebpf_state_terminate()
+{
+    ebpf_hash_table_destroy(_ebpf_state_thread_table);
+    ebpf_free(_ebpf_state_cpu_table);
+}
+
+ebpf_result_t
+ebpf_state_allocate_index(_Out_ size_t* new_index)
+{
+    if (_ebpf_state_next_index >= EBPF_MAX_STATE_ENTRIES) {
+        return EBPF_NO_MEMORY;
+    }
+
+    *new_index = ebpf_interlocked_increment_int64(&_ebpf_state_next_index) - 1;
+    return EBPF_SUCCESS;
+}
+
+ebpf_result_t
+_ebpf_state_get_entry(_Out_ ebpf_state_entry_t** entry)
+{
+    ebpf_state_entry_t* local_entry = NULL;
+
+    if (!ebpf_is_non_preemptible_work_item_supported() || ebpf_is_preemptible()) {
+        ebpf_result_t return_value;
+        uint64_t current_thread_id = ebpf_get_current_thread_id();
+
+        return_value =
+            ebpf_hash_table_find(_ebpf_state_thread_table, (const uint8_t*)&current_thread_id, (uint8_t**)&local_entry);
+
+        if (return_value == EBPF_KEY_NOT_FOUND) {
+            ebpf_state_entry_t new_entry = {0};
+
+            return_value = ebpf_hash_table_update(
+                _ebpf_state_thread_table,
+                (const uint8_t*)&current_thread_id,
+                (const uint8_t*)&new_entry,
+                EBPF_HASH_TABLE_OPERATION_INSERT);
+
+            if (return_value != EBPF_SUCCESS) {
+                return return_value;
+            }
+
+            return_value = ebpf_hash_table_find(
+                _ebpf_state_thread_table, (const uint8_t*)&current_thread_id, (uint8_t**)&local_entry);
+        }
+    } else {
+        uint32_t current_cpu = ebpf_get_current_cpu();
+        if (current_cpu >= _ebpf_state_cpu_table_size) {
+            return EBPF_OPERATION_NOT_SUPPORTED;
+        }
+        local_entry = _ebpf_state_cpu_table + current_cpu;
+    }
+    *entry = local_entry;
+    return EBPF_SUCCESS;
+}
+
+ebpf_result_t
+ebpf_state_store(size_t index, uintptr_t value)
+{
+    ebpf_state_entry_t* entry = NULL;
+    ebpf_result_t return_value;
+
+    return_value = _ebpf_state_get_entry(&entry);
+    if (return_value == EBPF_SUCCESS) {
+        entry->state[index] = value;
+    }
+    return return_value;
+}
+
+ebpf_result_t
+ebpf_state_load(size_t index, _Out_ uintptr_t* value)
+{
+    ebpf_state_entry_t* entry = NULL;
+    ebpf_result_t return_value;
+
+    return_value = _ebpf_state_get_entry(&entry);
+    if (return_value == EBPF_SUCCESS) {
+        *value = entry->state[index];
+    }
+    return return_value;
+}

libs/platform/ebpf_state.h

@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include "ebpf_platform.h"
+
+#ifdef __cplusplus
+extern "C"
+{
+#endif
+
+    /**
+     * @brief Initialize the eBPF state tracking module.
+     *
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_NO_MEMORY Unable to allocate resources for this
+     *  operation.
+     */
+    ebpf_result_t
+    ebpf_state_initiate();
+
+    /**
+     * @brief Uninitialize the eBPF state tracking module.
+     *
+     */
+    void
+    ebpf_state_terminate();
+
+    /**
+     * @brief Allocate a new index in the state tracker.
+     *
+     * @param[out] new_index Pointer to memory that contains the index on success.
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_NO_MEMORY Unable to allocate resources for this
+     *  operation.
+     */
+    ebpf_result_t
+    ebpf_state_allocate_index(_Out_ size_t* new_index);
+
+    /**
+     * @brief Store a value in the state tracker.
+     *
+     * @param[in] index Assigned for storing state.
+     * @param[in] value Value to be stored.
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_NO_MEMORY Unable to allocate resources for this
+     *  operation.
+     */
+    ebpf_result_t
+    ebpf_state_store(size_t index, uintptr_t value);
+
+    /**
+     * @brief Load a value in the state tracker.
+     *
+     * @param[in] index Assigned for storing state.
+     * @param[out] value Value to be loaded.
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_NO_MEMORY Unable to allocate resources for this
+     *  operation.
+     */
+    ebpf_result_t
+    ebpf_state_load(size_t index, _Out_ uintptr_t* value);
+
+#ifdef __cplusplus
+}
+#endif

libs/platform/kernel/platform_kernel.vcxproj

@@ -45,6 +45,7 @@
     <ClCompile Include="..\ebpf_pinning_table.c" />
     <ClCompile Include="..\ebpf_program_types.c" />
     <ClCompile Include="..\ebpf_serialize.c" />
+    <ClCompile Include="..\ebpf_state.c" />
     <ClCompile Include="..\ebpf_trampoline.c" />
     <ClCompile Include="ebpf_extension_kernel.c" />
     <ClCompile Include="ebpf_handle_kernel.c" />
@@ -57,6 +58,7 @@
     <ClInclude Include="..\ebpf_pinning_table.h" />
     <ClInclude Include="..\ebpf_platform.h" />
     <ClInclude Include="..\ebpf_serialize.h" />
+    <ClInclude Include="..\ebpf_state.h" />
     <ClInclude Include="framework.h" />
     <ClInclude Include="stdbool.h" />
     <ClInclude Include="stdint.h" />

libs/platform/kernel/platform_kernel.vcxproj.filters

@@ -49,6 +49,9 @@
     <ClCompile Include="..\ebpf_serialize.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\ebpf_state.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="..\ebpf_epoch.h">
@@ -78,6 +81,9 @@
     <ClInclude Include="..\ebpf_serialize.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\ebpf_state.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <None Include="..\ebpf_program_types.acf">

libs/platform/unit/platform_unit_test.cpp

@@ -19,6 +19,7 @@
 #include "ebpf_program_types.h"
 #include "ebpf_serialize.h"
 #include "ebpf_xdp_program_data.h"
+#include "ebpf_state.h"
 
 class _test_helper
 {
@@ -631,3 +632,22 @@ TEST_CASE("serialize_program_info_test", "[platform]")
 
     free(buffer);
 }
+
+TEST_CASE("state_test", "[state]")
+{
+    size_t allocated_index_1 = 0;
+    size_t allocated_index_2 = 0;
+    struct
+    {
+        uint32_t some_value;
+    } foo;
+    uintptr_t retreived_value = 0;
+    REQUIRE(ebpf_state_initiate() == EBPF_SUCCESS);
+    REQUIRE(ebpf_state_allocate_index(&allocated_index_1) == EBPF_SUCCESS);
+    REQUIRE(ebpf_state_allocate_index(&allocated_index_2) == EBPF_SUCCESS);
+    REQUIRE(allocated_index_2 != allocated_index_1);
+    REQUIRE(ebpf_state_store(allocated_index_1, reinterpret_cast<uintptr_t>(&foo)) == EBPF_SUCCESS);
+    REQUIRE(ebpf_state_load(allocated_index_1, &retreived_value) == EBPF_SUCCESS);
+    REQUIRE(retreived_value == reinterpret_cast<uintptr_t>(&foo));
+    ebpf_state_terminate();
+}

libs/platform/user/platform_user.vcxproj

@@ -29,6 +29,7 @@
     <ClCompile Include="..\ebpf_pinning_table.c" />
     <ClCompile Include="..\ebpf_program_types.c" />
     <ClCompile Include="..\ebpf_serialize.c" />
+    <ClCompile Include="..\ebpf_state.c" />
     <ClCompile Include="..\ebpf_trampoline.c" />
     <ClCompile Include="ebpf_extension_user.c" />
     <ClCompile Include="ebpf_handle_user.c" />
@@ -38,6 +39,7 @@
     <ClInclude Include="..\ebpf_object.h" />
     <ClInclude Include="..\ebpf_pinning_table.h" />
     <ClInclude Include="..\ebpf_platform.h" />
+    <ClInclude Include="..\ebpf_state.h" />
     <ClInclude Include="framework.h" />
   </ItemGroup>
   <ItemGroup>

libs/platform/user/platform_user.vcxproj.filters

@@ -52,6 +52,9 @@
     <ClCompile Include="..\ebpf_serialize.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\ebpf_state.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <Midl Include="..\ebpf_program_types.idl">
@@ -82,5 +85,8 @@
     <ClInclude Include="framework.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\ebpf_state.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
 </Project>

libs/service/api_service.cpp

@@ -16,18 +16,23 @@ extern "C"
 }
 #include "Verifier.h"
 #include "verifier_service.h"
+#include "windows_platform.hpp"
 
 #define MAX_CODE_SIZE_IN_BYTES (32 * 1024) // 32 KB
 
 static ebpf_result_t
 _build_helper_id_to_address_map(
-    ebpf_handle_t program_handle, ebpf_code_buffer_t& byte_code, std::vector<uint64_t>& helper_addresses)
+    ebpf_handle_t program_handle,
+    ebpf_code_buffer_t& byte_code,
+    std::vector<uint64_t>& helper_addresses,
+    uint32_t& unwind_index)
 {
     // Note:
     // eBPF supports helper IDs in the range [1, MAXUINT32]
     // uBPF jitter only supports helper IDs in the range [0,63]
     // Build a table to map [1, MAXUINT32] -> [0,63]
     std::map<uint32_t, uint32_t> helper_id_mapping;
+    unwind_index = MAXUINT32;
 
     ebpf_inst* instructions = reinterpret_cast<ebpf_inst*>(byte_code.data());
     for (size_t index = 0; index < byte_code.size() / sizeof(ebpf_inst); index++) {
@@ -86,6 +91,12 @@ _build_helper_id_to_address_map(
         }
         instruction.imm = helper_id_mapping[instruction.imm];
     }
+    for (auto& [old_helper_id, new_helper_id] : helper_id_mapping) {
+        if (get_helper_prototype_windows(old_helper_id).return_type != EBPF_RETURN_TYPE_INTEGER_OR_NO_RETURN_IF_SUCCEED)
+            continue;
+        unwind_index = new_helper_id;
+        break;
+    }
 
     return EBPF_SUCCESS;
 }
@@ -335,7 +346,8 @@ ebpf_verify_and_load_program(
         }
 
         std::vector<uint64_t> helper_id_adddress;
-        result = _build_helper_id_to_address_map(program_handle, byte_code_buffer, helper_id_adddress);
+        uint32_t unwind_index;
+        result = _build_helper_id_to_address_map(program_handle, byte_code_buffer, helper_id_adddress, unwind_index);
         if (result != EBPF_SUCCESS)
             goto Exit;
 
@@ -358,6 +370,9 @@ ebpf_verify_and_load_program(
                 }
             }
 
+            if (unwind_index != MAXUINT32)
+                ubpf_set_unwind_function_index(vm, unwind_index);
+
             ubpf_set_error_print(
                 vm, reinterpret_cast<int (*)(FILE * stream, const char* format, ...)>(log_function_address));
 

tests/sample/tail_call.c

@@ -10,12 +10,23 @@
 __attribute__((section("maps"), used)) struct bpf_map map = {
     sizeof(struct bpf_map), BPF_MAP_TYPE_PROG_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1};
 
+__attribute__((section("maps"), used)) struct bpf_map canary = {
+    sizeof(struct bpf_map), BPF_MAP_TYPE_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1};
+
 __attribute__((section("xdp_prog"), used)) int
 caller(struct xdp_md* ctx)
 {
+    uint32_t key = 0;
+    uint32_t* value;
+
     bpf_tail_call(ctx, &map, 0);
 
     // If we get to here it means bpf_tail_call failed.
+    value = bpf_map_lookup_elem(&canary, &key);
+    if (value) {
+        *value = 1;
+    }
+
     return 6;
 }
 

tests/sample/tail_call_bad.c

@@ -10,12 +10,23 @@
 __attribute__((section("maps"), used)) struct bpf_map map = {
     sizeof(struct bpf_map), BPF_MAP_TYPE_PROG_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1};
 
+__attribute__((section("maps"), used)) struct bpf_map canary = {
+    sizeof(struct bpf_map), BPF_MAP_TYPE_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1};
+
 __attribute__((section("xdp_prog"), used)) int
 caller(struct xdp_md* ctx)
 {
+    uint32_t key = 0;
+    uint32_t* value;
+
     // This should fail since the index is past the end of the array.
     long error = bpf_tail_call(ctx, &map, 1);
 
+    value = bpf_map_lookup_elem(&canary, &key);
+    if (value) {
+        *value = 1;
+    }
+
     return (int)error;
 }
 

tests/unit/libbpf_test.cpp

@@ -286,6 +286,8 @@ _ebpf_test_tail_call(_In_z_ const char* filename, int expected_result)
 
     struct bpf_map* map = bpf_map__next(nullptr, object);
     REQUIRE(map != nullptr);
+    struct bpf_map* canary_map = bpf_map__next(map, object);
+    REQUIRE(canary_map != nullptr);
 
     int callee_fd = bpf_program__fd(callee);
     REQUIRE(callee_fd >= 0);
@@ -293,6 +295,9 @@ _ebpf_test_tail_call(_In_z_ const char* filename, int expected_result)
     int map_fd = bpf_map__fd(map);
     REQUIRE(map_fd >= 0);
 
+    int canary_map_fd = bpf_map__fd(canary_map);
+    REQUIRE(canary_map_fd >= 0);
+
     // First do some negative tests.
     int index = 1;
     error = bpf_map_update_elem(map_fd, (uint8_t*)&index, (uint8_t*)&callee_fd, 0);
@@ -317,6 +322,19 @@ _ebpf_test_tail_call(_In_z_ const char* filename, int expected_result)
     REQUIRE(hook.fire(&ctx, &result) == EBPF_SUCCESS);
     REQUIRE(result == expected_result);
 
+    uint32_t key = 0;
+    uint32_t value = 0;
+    error = bpf_map_lookup_elem(canary_map_fd, &key, &value);
+    REQUIRE(error == 0);
+
+    // Is bpf_tail_call expected to work?
+    // Verify stack unwind occured.
+    if (expected_result >= 0) {
+        REQUIRE(value == 0);
+    } else {
+        REQUIRE(value != 0);
+    }
+
     result = bpf_link__destroy(link);
     REQUIRE(result == 0);
     bpf_object__close(object);
@@ -325,8 +343,7 @@ _ebpf_test_tail_call(_In_z_ const char* filename, int expected_result)
 TEST_CASE("good tail_call", "[libbpf]")
 {
     // Verify that 42 is returned, which is done by the callee.
-    // TODO(issue #344): change the 6 below to 42 once tail call is done correctly.
-    _ebpf_test_tail_call("tail_call.o", 6);
+    _ebpf_test_tail_call("tail_call.o", 42);
 }
 
 TEST_CASE("bad tail_call", "[libbpf]") { _ebpf_test_tail_call("tail_call_bad.o", -EBPF_INVALID_ARGUMENT); }