invalid prog data tests. (microsoft#3718)

* invalid prog data tests.

changes.

* analysis_assume.

* PR feedback.

---------

Co-authored-by: Anurag Saxena <43585259+saxena-anurag@users.noreply.github.com>

libs/execution_context/ebpf_core.c

@@ -144,7 +144,7 @@ static const void* _ebpf_general_helpers[] = {
 static const ebpf_helper_function_addresses_t _ebpf_global_helper_function_dispatch_table = {
     EBPF_HELPER_FUNCTION_ADDRESSES_HEADER, EBPF_COUNT_OF(_ebpf_general_helpers), (uint64_t*)_ebpf_general_helpers};
 static const ebpf_program_data_t _ebpf_global_helper_function_program_data = {
-    EBPF_PROGRAM_DATA_HEADER, &_ebpf_global_helper_program_info, &_ebpf_global_helper_function_dispatch_table};
+    EBPF_PROGRAM_DATA_HEADER, &_ebpf_global_helper_program_info, NULL, &_ebpf_global_helper_function_dispatch_table};
 
 static NPI_PROVIDER_ATTACH_CLIENT_FN _ebpf_general_helper_function_provider_attach_client;
 static NPI_PROVIDER_DETACH_CLIENT_FN _ebpf_general_helper_function_provider_detach_client;
@@ -256,8 +256,8 @@ ebpf_core_initiate()
         goto Done;
     }
 
-    _ebpf_global_helper_program_info.count_of_program_type_specific_helpers = ebpf_core_helper_functions_count;
-    _ebpf_global_helper_program_info.program_type_specific_helper_prototype = ebpf_core_helper_function_prototype;
+    _ebpf_global_helper_program_info.count_of_global_helpers = ebpf_core_helper_functions_count;
+    _ebpf_global_helper_program_info.global_helper_prototype = ebpf_core_helper_function_prototype;
 
     status = NmrRegisterProvider(
         &_ebpf_global_helper_function_provider_characteristics, NULL, &_ebpf_global_helper_function_nmr_binding_handle);

libs/execution_context/ebpf_program.c

@@ -346,8 +346,7 @@ _ebpf_program_general_program_information_attach_provider(
 
     program->general_helper_program_data = program_data;
 
-    program->global_helper_function_count =
-        program_data->program_type_specific_helper_function_addresses->helper_function_count;
+    program->global_helper_function_count = program_data->global_helper_function_addresses->helper_function_count;
 
     ebpf_lock_unlock(&program->lock, state);
     lock_held = false;
@@ -1589,13 +1588,17 @@ _Success_(return == true)
 {
     bool found = false;
     const ebpf_program_data_t* program_data = program->extension_program_data;
+    ebpf_assert_assume(program_data != NULL);
     const ebpf_program_data_t* general_program_data = program->general_helper_program_data;
+    ebpf_assert_assume(general_program_data != NULL);
 
     if (helper_function_id < EBPF_MAX_GENERAL_HELPER_FUNCTION) {
         // First check the global helper function table of the program type for any helper functions that are
         // overwritten.
+        ebpf_assert_assume(program_data->program_info != NULL);
         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) {
+                ebpf_assert_assume(program_data->global_helper_function_addresses != NULL);
                 address->address = program_data->global_helper_function_addresses->helper_function_address[index];
                 address->implicit_context = program_data->program_info->global_helper_prototype[index].implicit_context;
                 found = true;
@@ -1605,18 +1608,18 @@ _Success_(return == true)
 
         if (!found) {
             // Next check the general helper function table of the general program type.
-            for (size_t index = 0; index < general_program_data->program_info->count_of_program_type_specific_helpers;
-                 index++) {
-                if (general_program_data->program_info->program_type_specific_helper_prototype[index].helper_id ==
+            ebpf_assert_assume(general_program_data->program_info != NULL);
+            for (size_t index = 0; index < general_program_data->program_info->count_of_global_helpers; index++) {
+                if (general_program_data->program_info->global_helper_prototype[index].helper_id ==
                     helper_function_id) {
-                    address->address = general_program_data->program_type_specific_helper_function_addresses
-                                           ->helper_function_address[index];
+                    ebpf_assert_assume(general_program_data->global_helper_function_addresses != NULL);
+                    address->address =
+                        general_program_data->global_helper_function_addresses->helper_function_address[index];
                     // In case of helper functions that do not have any default / general implementation, the function
                     // address will be NULL.
                     if (address->address != 0) {
                         address->implicit_context =
-                            general_program_data->program_info->program_type_specific_helper_prototype[index]
-                                .implicit_context;
+                            general_program_data->program_info->global_helper_prototype[index].implicit_context;
                         found = true;
                     } else {
                         EBPF_LOG_MESSAGE_UINT64(
@@ -1631,9 +1634,11 @@ _Success_(return == true)
         }
     } else {
         // Check the program type specific helper function table of the program type.
+        ebpf_assert_assume(program_data->program_info != NULL);
         for (size_t index = 0; index < program_data->program_info->count_of_program_type_specific_helpers; index++) {
             if (program_data->program_info->program_type_specific_helper_prototype[index].helper_id ==
                 helper_function_id) {
+                ebpf_assert_assume(program_data->program_type_specific_helper_function_addresses != NULL);
                 address->address =
                     program_data->program_type_specific_helper_function_addresses->helper_function_address[index];
                 address->implicit_context =
@@ -1929,6 +1934,7 @@ ebpf_program_get_program_info(_In_ const ebpf_program_t* program, _Outptr_ ebpf_
     }
     provider_data_referenced = true;
     program_data = program->extension_program_data;
+    ebpf_assert_assume(program_data != NULL);
 
     if (!program->general_helper_program_data) {
         EBPF_LOG_MESSAGE_GUID(
@@ -1940,11 +1946,12 @@ ebpf_program_get_program_info(_In_ const ebpf_program_t* program, _Outptr_ ebpf_
         goto Exit;
     }
     general_helper_program_data = program->general_helper_program_data;
+    ebpf_assert_assume(general_helper_program_data != NULL);
 
     total_count_of_helpers = program_data->program_info->count_of_program_type_specific_helpers +
-                             general_helper_program_data->program_info->count_of_program_type_specific_helpers;
+                             general_helper_program_data->program_info->count_of_global_helpers;
     if ((total_count_of_helpers < program_data->program_info->count_of_program_type_specific_helpers) ||
-        (total_count_of_helpers < general_helper_program_data->program_info->count_of_program_type_specific_helpers)) {
+        (total_count_of_helpers < general_helper_program_data->program_info->count_of_global_helpers)) {
         result = EBPF_ARITHMETIC_OVERFLOW;
         goto Exit;
     }
@@ -1976,12 +1983,10 @@ ebpf_program_get_program_info(_In_ const ebpf_program_t* program, _Outptr_ ebpf_
                 program_data->program_info->program_type_specific_helper_prototype[index];
         }
 
-        for (uint32_t index = 0;
-             index < general_helper_program_data->program_info->count_of_program_type_specific_helpers;
-             index++) {
+        for (uint32_t index = 0; index < general_helper_program_data->program_info->count_of_global_helpers; index++) {
             __analysis_assume(helper_index < total_count_of_helpers);
             helper_prototype[helper_index++] =
-                general_helper_program_data->program_info->program_type_specific_helper_prototype[index];
+                general_helper_program_data->program_info->global_helper_prototype[index];
         }
     }
 
@@ -2198,14 +2203,14 @@ _IRQL_requires_max_(PASSIVE_LEVEL) static ebpf_result_t _ebpf_program_compute_pr
         }
 
         // Copy global helpers to helper_id_to_index.
-        for (uint32_t index = 0; index < general_program_info->count_of_program_type_specific_helpers; index++) {
-            uint32_t helper_id = general_program_info->program_type_specific_helper_prototype[index].helper_id;
+        for (uint32_t index = 0; index < general_program_info->count_of_global_helpers; index++) {
+            uint32_t helper_id = general_program_info->global_helper_prototype[index].helper_id;
             if (!_ebpf_contains_helper_id(actual_helper_ids, count_of_actual_helper_ids, helper_id)) {
                 continue;
             }
             helper_id_to_index[helper_function_index].helper_id = helper_id;
             helper_id_to_index[helper_function_index].helper_function_prototype =
-                &general_program_info->program_type_specific_helper_prototype[index];
+                &general_program_info->global_helper_prototype[index];
             helper_function_index++;
         }
 

libs/execution_context/unit/execution_context_unit_test.cpp

@@ -1937,6 +1937,191 @@ TEST_CASE("EBPF_OPERATION_LOAD_NATIVE_MODULE short header", "[execution_context]
             completion) == EBPF_ARITHMETIC_OVERFLOW);
 }
 
+#define EBPF_PROGRAM_TYPE_TEST_GUID                                                    \
+    {                                                                                  \
+        0x8ee1b757, 0xc0b2, 0x4c84, { 0xac, 0x07, 0x0c, 0x76, 0x29, 0x8f, 0x1d, 0xc9 } \
+    }
+
+void
+test_register_provider(
+    _In_ const NPI_PROVIDER_CHARACTERISTICS* provider_characteristics, bool expected_to_succeed = false)
+{
+    class provider_deregister_helper
+    {
+      public:
+        void
+        operator()(HANDLE handle)
+        {
+            NmrDeregisterProvider(handle);
+        }
+    };
+
+    typedef std::unique_ptr<void, provider_deregister_helper> provider_ptr;
+    provider_ptr provider;
+    HANDLE nmr_provider_handle;
+
+    REQUIRE(NmrRegisterProvider(provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS);
+    provider.reset(nmr_provider_handle);
+
+    ebpf_program_type_t EBPF_PROGRAM_TYPE_TEST = EBPF_PROGRAM_TYPE_TEST_GUID;
+    const cxplat_utf8_string_t program_name{(uint8_t*)("foo"), 3};
+    const cxplat_utf8_string_t section_name{(uint8_t*)("bar"), 3};
+    const ebpf_program_parameters_t program_parameters{
+        EBPF_PROGRAM_TYPE_TEST, EBPF_ATTACH_TYPE_SAMPLE, program_name, section_name};
+    program_ptr program;
+    {
+        ebpf_program_t* local_program = nullptr;
+        REQUIRE(
+            ebpf_program_create(&program_parameters, &local_program) ==
+            (expected_to_succeed ? EBPF_SUCCESS : EBPF_EXTENSION_FAILED_TO_LOAD));
+        program.reset(local_program);
+        if (expected_to_succeed) {
+            helper_function_address_t addresses[1] = {};
+            uint32_t helper_function_ids[] = {EBPF_MAX_GENERAL_HELPER_FUNCTION + 1};
+            REQUIRE(
+                ebpf_program_set_helper_function_ids(
+                    program.get(), EBPF_COUNT_OF(helper_function_ids), helper_function_ids) == EBPF_SUCCESS);
+            REQUIRE(
+                ebpf_program_get_helper_function_addresses(
+                    program.get(), EBPF_COUNT_OF(helper_function_ids), addresses) == EBPF_SUCCESS);
+            REQUIRE(addresses[0].address != 0);
+        }
+    }
+}
+
+TEST_CASE("INVALID_PROGRAM_DATA", "[execution_context][negative]")
+{
+    _ebpf_core_initializer core;
+    core.initialize();
+
+    ebpf_context_descriptor_t _test_context_descriptor = {sizeof(ebpf_context_descriptor_t), -1, -1, -1};
+
+    const uint32_t _test_prog_type = 1000;
+    ebpf_program_type_descriptor_t _test_program_type_descriptor = {
+        EBPF_PROGRAM_TYPE_DESCRIPTOR_HEADER,
+        "test_program_type",
+        &_test_context_descriptor,
+        EBPF_PROGRAM_TYPE_TEST_GUID,
+        _test_prog_type,
+        0};
+
+    ebpf_helper_function_prototype_t _test_helper_function_prototype[] = {
+        {EBPF_HELPER_FUNCTION_PROTOTYPE_HEADER,
+         EBPF_MAX_GENERAL_HELPER_FUNCTION + 1,
+         "test_helper_function",
+         EBPF_RETURN_TYPE_INTEGER,
+         {EBPF_ARGUMENT_TYPE_DONTCARE}}};
+
+    ebpf_program_info_t _test_program_info = {
+        EBPF_PROGRAM_INFORMATION_HEADER,
+        &_test_program_type_descriptor,
+        EBPF_COUNT_OF(_test_helper_function_prototype),
+        _test_helper_function_prototype};
+
+    auto provider_function1 = []() { return (ebpf_result_t)TEST_FUNCTION_RETURN; };
+    ebpf_result_t (*function_pointer1)() = provider_function1;
+    const void* helper_functions[] = {(void*)function_pointer1};
+    ebpf_helper_function_addresses_t helper_function_addresses = {
+        EBPF_HELPER_FUNCTION_ADDRESSES_HEADER, EBPF_COUNT_OF(helper_functions), (uint64_t*)helper_functions};
+
+    ebpf_program_data_t _test_program_data = {
+        EBPF_PROGRAM_DATA_HEADER, &_test_program_info, &helper_function_addresses};
+
+    auto provider_attach_client_callback =
+        [](HANDLE, void*, const NPI_REGISTRATION_INSTANCE*, void*, const void*, void**, const void**) -> NTSTATUS {
+        return STATUS_SUCCESS;
+    };
+    auto provider_detach_client_callback = [](void*) -> NTSTATUS { return STATUS_SUCCESS; };
+
+    NPI_MODULEID module_id = {sizeof(NPI_MODULEID), MIT_GUID, EBPF_PROGRAM_TYPE_TEST_GUID};
+
+    NPI_PROVIDER_CHARACTERISTICS provider_characteristics{
+        0,
+        sizeof(NPI_PROVIDER_CHARACTERISTICS),
+        provider_attach_client_callback,
+        provider_detach_client_callback,
+        nullptr,
+        {
+            0,
+            sizeof(NPI_REGISTRATION_INSTANCE),
+            &EBPF_PROGRAM_INFO_EXTENSION_IID,
+            &module_id,
+            0,
+            &_test_program_data,
+        },
+    };
+
+    // Register the provider with valid program data. The test is expected to succeed.
+    bool expected_to_succeed = true;
+    test_register_provider(&provider_characteristics, expected_to_succeed);
+
+    // In the next tests, exactly one field of the program data will be invalidated. The tests are expected to fail.
+
+    // Set bad version for program data header.
+    _test_program_data.header.version = 0;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_data.header.version = EBPF_PROGRAM_DATA_CURRENT_VERSION;
+
+    // Set bigger size than supported for the program data header.
+    _test_program_data.header.size = EBPF_PROGRAM_DATA_CURRENT_VERSION_SIZE + 1;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_data.header.size = EBPF_PROGRAM_DATA_CURRENT_VERSION_SIZE;
+
+    // Remove the program data from the provider characteristics struct.
+    provider_characteristics.ProviderRegistrationInstance.NpiSpecificCharacteristics = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    provider_characteristics.ProviderRegistrationInstance.NpiSpecificCharacteristics = &_test_program_data;
+
+    // Remove the address from the helper function.
+    helper_function_addresses.helper_function_address = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    helper_function_addresses.helper_function_address = (uint64_t*)helper_functions;
+
+    // Remove the program info struct from the program data.
+    _test_program_data.program_info = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_data.program_info = &_test_program_info;
+
+    // Remove the program type descriptor from the program info.
+    _test_program_info.program_type_descriptor = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_info.program_type_descriptor = &_test_program_type_descriptor;
+
+    // Remove name to the program type descriptor.
+    _test_program_type_descriptor.name = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_type_descriptor.name = "test_program_type";
+
+    // Remove the context descriptor from the program type descriptor.
+    _test_program_type_descriptor.context_descriptor = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_type_descriptor.context_descriptor = &_test_context_descriptor;
+
+    // Invalidate the context descriptor (size = 0).
+    _test_context_descriptor.size = 0;
+    test_register_provider(&provider_characteristics);
+    // Fix up the context descriptor.
+    _test_context_descriptor.size = sizeof(ebpf_context_descriptor_t);
+
+    // Remove the helper function prototype from the program info.
+    _test_program_info.program_type_specific_helper_prototype = nullptr;
+    test_register_provider(&provider_characteristics);
+    // Restore.
+    _test_program_info.program_type_specific_helper_prototype = _test_helper_function_prototype;
+
+    // Invalidate the helper function prototype by removing the name of the helper function.
+    _test_helper_function_prototype[0].name = nullptr;
+    test_register_provider(&provider_characteristics);
+}
+
 // TODO: Add more native module loading IOCTL negative tests.
 // https://github.com/microsoft/ebpf-for-windows/issues/1139
 // EBPF_OPERATION_LOAD_NATIVE_MODULE

libs/shared/ebpf_shared_framework.h

@@ -26,6 +26,10 @@ CXPLAT_EXTERN_C_BEGIN
 #define EBPF_SYMBOLIC_DEVICE_NAME L"\\GLOBAL??\\EbpfIoDevice"
 #define EBPF_DEVICE_WIN32_NAME L"\\\\.\\EbpfIoDevice"
 
+#define ebpf_assert_assume(x) \
+    ebpf_assert(x);           \
+    __analysis_assume(x);
+
 // Macro locally suppresses "Unreferenced variable" warning, which in 'Release' builds is treated as an error.
 #define ebpf_assert_success(x)                                     \
     _Pragma("warning(push)") _Pragma("warning(disable : 4189)") do \

libs/shared/shared_common.c

@@ -127,6 +127,10 @@ ebpf_validate_helper_function_prototype_array(
     _In_reads_(count) const ebpf_helper_function_prototype_t* helper_prototype_array, uint32_t count)
 {
     if (count > 0) {
+        // The helper_prototype_array cannot be NULL.
+        if (helper_prototype_array == NULL) {
+            return false;
+        }
         // Use "total_size" to calculate the actual size of the ebpf_helper_function_prototype_t struct.
         size_t helper_prototype_size = helper_prototype_array[0].header.total_size;
         for (uint32_t i = 0; i < count; i++) {