Pass program handle in GET_PROGRAM_INFO message. (#463)

* fix

* PR feedback.

* Apply suggestions from code review

Co-authored-by: Dave Thaler <dthaler@microsoft.com>

Co-authored-by: Dave Thaler <dthaler@microsoft.com>

ebpfsvc/rpc_api.cpp

@@ -21,6 +21,9 @@ ebpf_server_verify_and_load_program(
         return EBPF_INVALID_ARGUMENT;
     }
 
+    // Set the handle of program being verified in thread-local storage.
+    set_program_under_verification(info->program_handle);
+
     result = ebpf_verify_and_load_program(
         &info->program_type,
         info->program_handle,

libs/api/Verifier.cpp

@@ -35,6 +35,11 @@ typedef struct _section_offset_to_map
     string map_name;
 } section_offset_to_map_t;
 
+struct _thread_local_storage_cache
+{
+    ~_thread_local_storage_cache() { ebpf_clear_thread_local_storage(); }
+};
+
 static void
 _get_program_and_map_names(
     _In_ string& path,
@@ -329,6 +334,8 @@ ebpf_api_elf_enumerate_sections(
     ebpf_verifier_options_t verifier_options{false, false, false, false, true};
     const ebpf_platform_t* platform = &g_ebpf_platform_windows;
     std::ostringstream str;
+    struct _thread_local_storage_cache tls_cache;
+
     try {
         auto raw_programs = read_elf(file, section ? std::string(section) : std::string(), &verifier_options, platform);
         tlv_sequence sequence;
@@ -380,6 +387,8 @@ ebpf_api_elf_disassemble_section(
     const ebpf_platform_t* platform = &g_ebpf_platform_windows;
     std::ostringstream error;
     std::ostringstream output;
+    struct _thread_local_storage_cache tls_cache;
+
     try {
         auto raw_programs = read_elf(file, section, &verifier_options, platform);
         raw_program raw_program = raw_programs.back();
@@ -414,8 +423,8 @@ ebpf_api_elf_verify_section(
     ebpf_api_verifier_stats_t* stats)
 {
     std::ostringstream error;
-
     std::ostringstream output;
+    struct _thread_local_storage_cache tls_cache;
 
     *report = nullptr;
     *error_message = nullptr;

libs/api_common/api_common.cpp

@@ -13,6 +13,8 @@
 
 #include "ebpf_verifier_wrapper.hpp"
 
+#include "map_descriptors.hpp"
+
 thread_local static const ebpf_program_type_t* _global_program_type = nullptr;
 thread_local static const ebpf_attach_type_t* _global_attach_type = nullptr;
 
@@ -100,3 +102,11 @@ get_global_attach_type()
 {
     return _global_attach_type;
 }
+
+void
+ebpf_clear_thread_local_storage() noexcept
+{
+    clear_map_descriptors();
+    clear_program_info_cache();
+    set_program_under_verification(ebpf_handle_invalid);
+}

libs/api_common/api_common.hpp

@@ -156,3 +156,17 @@ get_global_program_type();
 
 const ebpf_attach_type_t*
 get_global_attach_type();
+
+/**
+ * @brief Save handle to program being verified in thread-local storage.
+ *
+ * @param[in] program Handle to program being verified.
+ */
+void
+set_program_under_verification(ebpf_handle_t program);
+
+/**
+ * @brief Clear thread-local storage used for storing data needed for program verification.
+ */
+void
+ebpf_clear_thread_local_storage() noexcept;

libs/api_common/windows_helpers.cpp

@@ -40,19 +40,32 @@ static thread_local std::map<GUID, ebpf_program_info_ptr_t, guid_compare> _progr
 
 static thread_local std::map<GUID, ebpf_helper::ebpf_memory_ptr, guid_compare> _static_program_info_cache;
 
+static thread_local ebpf_handle_t _program_under_verification = ebpf_handle_invalid;
+
+void
+set_program_under_verification(ebpf_handle_t program)
+{
+    _program_under_verification = program;
+}
+
 void
 clear_program_info_cache()
 {
     _program_info_cache.clear();
 }
 
+#define GET_PROGRAM_INFO_REPLY_BUFFER_SIZE 2048
+
 ebpf_result_t
 get_program_info_data(ebpf_program_type_t program_type, _Outptr_ ebpf_program_info_t** program_info)
 {
-    ebpf_protocol_buffer_t reply_buffer(1024);
+    ebpf_protocol_buffer_t reply_buffer(GET_PROGRAM_INFO_REPLY_BUFFER_SIZE);
     size_t required_buffer_length;
     ebpf_operation_get_program_info_request_t request{
-        sizeof(request), ebpf_operation_id_t::EBPF_OPERATION_GET_PROGRAM_INFO, program_type};
+        sizeof(request),
+        ebpf_operation_id_t::EBPF_OPERATION_GET_PROGRAM_INFO,
+        program_type,
+        reinterpret_cast<uint64_t>(_program_under_verification)};
 
     *program_info = nullptr;
 

libs/execution_context/ebpf_core.c

@@ -12,6 +12,8 @@
 #include "ebpf_serialize.h"
 #include "ebpf_state.h"
 
+const ebpf_handle_t ebpf_handle_invalid = (ebpf_handle_t)-1;
+
 GUID ebpf_general_helper_function_interface_id = {/* 8d2a1d3f-9ce6-473d-b48e-17aa5c5581fe */
                                                   0x8d2a1d3f,
                                                   0x9ce6,
@@ -835,15 +837,21 @@ _ebpf_core_protocol_get_program_info(
     size_t serialization_buffer_size;
     size_t required_length;
 
-    retval = ebpf_program_create(&program);
-    if (retval != EBPF_SUCCESS)
-        goto Done;
-
     program_parameters.program_type = request->program_type;
 
-    retval = ebpf_program_initialize(program, &program_parameters);
-    if (retval != EBPF_SUCCESS)
-        goto Done;
+    if (request->program_handle == ebpf_handle_invalid) {
+        retval = ebpf_program_create(&program);
+        if (retval != EBPF_SUCCESS)
+            goto Done;
+        retval = ebpf_program_initialize(program, &program_parameters);
+        if (retval != EBPF_SUCCESS)
+            goto Done;
+    } else {
+        retval =
+            ebpf_reference_object_by_handle(request->program_handle, EBPF_OBJECT_PROGRAM, (ebpf_object_t**)&program);
+        if (retval != EBPF_SUCCESS)
+            goto Done;
+    }
 
     retval = ebpf_program_get_program_info(program, &program_info);
     if (retval != EBPF_SUCCESS)

libs/execution_context/ebpf_protocol.h

@@ -278,6 +278,7 @@ typedef struct _ebpf_operation_get_program_info_request
 {
     struct _ebpf_operation_header header;
     ebpf_program_type_t program_type;
+    uint64_t program_handle;
 } ebpf_operation_get_program_info_request_t;
 
 typedef struct _ebpf_operation_get_program_info_reply

libs/service/api_service.cpp

@@ -234,13 +234,6 @@ _query_and_cache_map_descriptors(
     return EBPF_SUCCESS;
 }
 
-void
-ebpf_clear_thread_local_storage() noexcept
-{
-    clear_map_descriptors();
-    clear_program_info_cache();
-}
-
 ebpf_result_t
 ebpf_verify_program(
     const GUID* program_type,

libs/service/api_service.h

@@ -38,7 +38,4 @@ uint32_t
 ebpf_service_initialize() noexcept;
 
 void
-ebpf_service_cleanup() noexcept;
-
-void
-ebpf_clear_thread_local_storage() noexcept;
+ebpf_service_cleanup() noexcept;

tests/end_to_end/end_to_end.cpp

@@ -486,6 +486,7 @@ TEST_CASE("verify section", "[end_to_end]")
     const char* error_message = nullptr;
     const char* report = nullptr;
     uint32_t result;
+    program_info_provider_t xdp_program_info(EBPF_PROGRAM_TYPE_XDP);
 
     ebpf_api_verifier_stats_t stats;
     REQUIRE((