diff --git a/tests/sample/ext/app/sample_ext_app.cpp b/tests/sample/ext/app/sample_ext_app.cpp index b89351ce0..795d73a1d 100644 --- a/tests/sample/ext/app/sample_ext_app.cpp +++ b/tests/sample/ext/app/sample_ext_app.cpp @@ -85,6 +85,24 @@ struct _sample_extension_helper nullptr) == TRUE); } + void + invoke_batch(std::vector& input_buffer, std::vector& output_buffer) + { + uint32_t count_of_bytes_returned; + + // Issue IOCTL. + REQUIRE( + ::DeviceIoControl( + device_handle, + IOCTL_SAMPLE_EBPF_EXT_CTL_RUN_BATCH, + input_buffer.data(), + static_cast(input_buffer.size()), + output_buffer.data(), + static_cast(output_buffer.size()), + (unsigned long*)&count_of_bytes_returned, + nullptr) == TRUE); + } + private: HANDLE device_handle; }; @@ -118,6 +136,41 @@ sample_ebpf_ext_test(_In_ const struct bpf_object* object) REQUIRE(memcmp(output_buffer.data(), expected_output, strlen(expected_output)) == 0); } +void +sample_ebpf_ext_test_batch(_In_ const struct bpf_object* object) +{ + struct bpf_map* map = nullptr; + fd_t map_fd; + const char* strings[] = {"rainy", "sunny"}; + std::vector> map_entry_buffers(EBPF_COUNT_OF(strings), std::vector(32)); + const char* input_string = "rainy rainy rainy rainy rainy"; + size_t input_string_length = strlen(input_string); + std::vector input_buffer(EBPF_OFFSET_OF(sample_ebpf_ext_batch_run_request_t, data) + input_string_length); + const char* expected_output = "sunny sunny sunny sunny rainy"; + std::vector output_buffer(256); + + sample_ebpf_ext_batch_run_request_t* request = (sample_ebpf_ext_batch_run_request_t*)input_buffer.data(); + request->count = 4; + memcpy(request->data, input_string, input_string_length); + sample_ebpf_ext_batch_run_reply_t* reply = (sample_ebpf_ext_batch_run_reply_t*)output_buffer.data(); + _sample_extension_helper extension; + + // Get map and insert data. + map = bpf_object__find_map_by_name(object, "test_map"); + REQUIRE(map != nullptr); + map_fd = bpf_map__fd(map); + REQUIRE(map_fd > 0); + + for (uint32_t key = 0; key < EBPF_COUNT_OF(strings); key++) { + std::copy(strings[key], strings[key] + strlen(strings[key]), map_entry_buffers[key].begin()); + REQUIRE(bpf_map_update_elem(map_fd, &key, map_entry_buffers[key].data(), EBPF_ANY) == EBPF_SUCCESS); + } + + extension.invoke_batch(input_buffer, output_buffer); + + REQUIRE(memcmp(reply->data, expected_output, strlen(expected_output)) == 0); +} + #if !defined(CONFIG_BPF_JIT_DISABLED) TEST_CASE("jit_test", "[sample_ext_test]") { @@ -146,6 +199,32 @@ TEST_CASE("interpret_test", "[sample_ext_test]") } #endif +TEST_CASE("native_test", "[sample_ext_test]") +{ + struct bpf_object* object = nullptr; + hook_helper_t hook(EBPF_ATTACH_TYPE_SAMPLE); + program_load_attach_helper_t _helper( + "test_sample_ebpf.sys", BPF_PROG_TYPE_SAMPLE, "test_program_entry", EBPF_EXECUTION_ANY, nullptr, 0, hook); + + object = _helper.get_object(); + + sample_ebpf_ext_test(object); +} + +#if !defined(CONFIG_BPF_JIT_DISABLED) +TEST_CASE("batch_test", "[sample_ext_test]") +{ + struct bpf_object* object = nullptr; + hook_helper_t hook(EBPF_ATTACH_TYPE_SAMPLE); + program_load_attach_helper_t _helper( + "test_sample_ebpf.o", BPF_PROG_TYPE_SAMPLE, "test_program_entry", EBPF_EXECUTION_ANY, nullptr, 0, hook); + + object = _helper.get_object(); + + sample_ebpf_ext_test_batch(object); +} +#endif + void utility_helpers_test(ebpf_execution_type_t execution_type) { diff --git a/tests/sample/ext/drv/sample_ext.c b/tests/sample/ext/drv/sample_ext.c index 264bc6516..9219c45ec 100644 --- a/tests/sample/ext/drv/sample_ext.c +++ b/tests/sample/ext/drv/sample_ext.c @@ -217,6 +217,9 @@ typedef struct _sample_ebpf_extension_hook_client const void* client_binding_context; const ebpf_extension_data_t* client_data; ebpf_program_invoke_function_t invoke_program; + ebpf_program_batch_begin_invoke_function_t begin_batch_program_invoke; + ebpf_program_batch_end_invoke_function_t end_batch_program_invoke; + ebpf_program_batch_invoke_function_t batch_program_invoke; } sample_ebpf_extension_hook_client_t; /** @@ -421,6 +424,9 @@ _sample_ebpf_extension_hook_provider_attach_client( goto Exit; } hook_client->invoke_program = client_dispatch_table->ebpf_program_invoke_function; + hook_client->batch_program_invoke = client_dispatch_table->ebpf_program_batch_invoke_function; + hook_client->begin_batch_program_invoke = client_dispatch_table->ebpf_program_batch_begin_invoke_function; + hook_client->end_batch_program_invoke = client_dispatch_table->ebpf_program_batch_end_invoke_function; local_provider_context->attached_client = hook_client; @@ -526,6 +532,73 @@ Exit: return return_value; } +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_begin_program(_Inout_ ebpf_execution_context_state_t* state) +{ + ebpf_result_t return_value = EBPF_SUCCESS; + + sample_ebpf_extension_hook_provider_t* hook_provider_context = &_sample_ebpf_extension_hook_provider_context; + + sample_ebpf_extension_hook_client_t* hook_client = hook_provider_context->attached_client; + + if (hook_client == NULL) { + return_value = EBPF_FAILED; + goto Exit; + } + ebpf_program_batch_begin_invoke_function_t batch_begin_function = hook_client->begin_batch_program_invoke; + const void* client_binding_context = hook_client->client_binding_context; + + return_value = batch_begin_function(client_binding_context, sizeof(ebpf_execution_context_state_t), state); + +Exit: + return return_value; +} + +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_program( + _Inout_ sample_program_context_t* context, _In_ const ebpf_execution_context_state_t* state, _Out_ uint32_t* result) +{ + ebpf_result_t return_value = EBPF_SUCCESS; + + sample_ebpf_extension_hook_provider_t* hook_provider_context = &_sample_ebpf_extension_hook_provider_context; + + sample_ebpf_extension_hook_client_t* hook_client = hook_provider_context->attached_client; + + if (hook_client == NULL) { + return_value = EBPF_FAILED; + goto Exit; + } + ebpf_program_batch_invoke_function_t batch_invoke_program = hook_client->batch_program_invoke; + const void* client_binding_context = hook_client->client_binding_context; + + return_value = batch_invoke_program(client_binding_context, context, result, state); + +Exit: + return return_value; +} + +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_end_program(_Inout_ ebpf_execution_context_state_t* state) +{ + ebpf_result_t return_value = EBPF_SUCCESS; + + sample_ebpf_extension_hook_provider_t* hook_provider_context = &_sample_ebpf_extension_hook_provider_context; + + sample_ebpf_extension_hook_client_t* hook_client = hook_provider_context->attached_client; + + if (hook_client == NULL) { + return_value = EBPF_FAILED; + goto Exit; + } + ebpf_program_batch_end_invoke_function_t batch_end_function = hook_client->end_batch_program_invoke; + const void* client_binding_context = hook_client->client_binding_context; + + return_value = batch_end_function(client_binding_context, state); + +Exit: + return return_value; +} + _Must_inspect_result_ ebpf_result_t sample_ebpf_extension_profile_program( _Inout_ sample_ebpf_ext_profile_request_t* request, diff --git a/tests/sample/ext/drv/sample_ext.h b/tests/sample/ext/drv/sample_ext.h index 6d9361938..0c0199299 100644 --- a/tests/sample/ext/drv/sample_ext.h +++ b/tests/sample/ext/drv/sample_ext.h @@ -71,3 +71,41 @@ sample_ebpf_extension_profile_program( _Inout_ sample_ebpf_ext_profile_request_t* request, size_t request_length, _Inout_ sample_ebpf_ext_profile_reply_t* reply); + +/** + * @brief Invoke batch begin function. + * + * @param[in, out] state Pointer to the execution context state. + * + * @retval EBPF_SUCCESS Operation succeeded. + * @retval EBPF_OPERATION_NOT_SUPPORTED Operation not supported. + */ +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_begin_program(_Inout_ ebpf_execution_context_state_t* state); + +/** + * @brief Batch invoke eBPF program attached to a hook provider instance. + * + * @param[in] context Pointer to eBPF program context. + * @param[in] state Pointer to the execution context state. + * @param[out] result Result returned by eBPF program at the end of execution. + * + * @retval EBPF_SUCCESS Operation succeeded. + * @retval EBPF_OPERATION_NOT_SUPPORTED Operation not supported. + */ +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_program( + _Inout_ sample_program_context_t* context, + _In_ const ebpf_execution_context_state_t* state, + _Out_ uint32_t* result); + +/** + * @brief Invoke batch end function. + * + * @param[in, out] state Pointer to the execution context state. + * + * @retval EBPF_SUCCESS Operation succeeded. + * @retval EBPF_OPERATION_NOT_SUPPORTED Operation not supported. + */ +_Must_inspect_result_ ebpf_result_t +sample_ebpf_extension_invoke_batch_end_program(_Inout_ ebpf_execution_context_state_t* state); diff --git a/tests/sample/ext/drv/sample_ext_drv.c b/tests/sample/ext/drv/sample_ext_drv.c index 97292b22b..478315aea 100644 --- a/tests/sample/ext/drv/sample_ext_drv.c +++ b/tests/sample/ext/drv/sample_ext_drv.c @@ -291,6 +291,100 @@ _sample_ebpf_ext_driver_io_device_control( goto Done; } break; + case IOCTL_SAMPLE_EBPF_EXT_CTL_RUN_BATCH: + sample_ebpf_ext_batch_run_request_t* batch_run_request = NULL; + sample_ebpf_ext_batch_run_reply_t* batch_run_reply = NULL; + ebpf_execution_context_state_t context_state = {0}; + if (input_buffer_length != 0) { + // Retrieve the input buffer associated with the request object. + status = WdfRequestRetrieveInputBuffer( + request, // Request object. + input_buffer_length, // Length of input buffer. + &input_buffer, // Pointer to buffer. + &actual_input_length // Length of buffer. + ); + + if (!NT_SUCCESS(status)) { + KdPrintEx( + (DPFLTR_IHVDRIVER_ID, + DPFLTR_INFO_LEVEL, + "%s: Input buffer failure %d\n", + SAMPLE_EBPF_EXT_NAME_A, + status)); + goto Done; + } + + if (input_buffer == NULL) { + status = STATUS_INVALID_PARAMETER; + goto Done; + } + + size_t minimum_request_size = sizeof(sample_ebpf_ext_batch_run_request_t); + size_t minimum_reply_size; + + if (actual_input_length < minimum_request_size) { + status = STATUS_INVALID_PARAMETER; + goto Done; + } + + minimum_reply_size = actual_input_length; + + // Be aware: Input and output buffer point to the same memory. + if (minimum_reply_size > 0) { + // Retrieve output buffer associated with the request object. + status = WdfRequestRetrieveOutputBuffer( + request, output_buffer_length, &output_buffer, &actual_output_length); + if (!NT_SUCCESS(status)) { + KdPrintEx( + (DPFLTR_IHVDRIVER_ID, + DPFLTR_INFO_LEVEL, + "%s: Output buffer failure %d\n", + SAMPLE_EBPF_EXT_NAME_A, + status)); + goto Done; + } + if (output_buffer == NULL) { + status = STATUS_INVALID_PARAMETER; + goto Done; + } + + if (actual_output_length < minimum_reply_size) { + status = STATUS_BUFFER_TOO_SMALL; + goto Done; + } + } + + batch_run_request = (sample_ebpf_ext_batch_run_request_t*)input_buffer; + batch_run_reply = (sample_ebpf_ext_batch_run_reply_t*)output_buffer; + + result = sample_ebpf_extension_invoke_batch_begin_program(&context_state); + if (result != EBPF_SUCCESS) { + status = STATUS_UNSUCCESSFUL; + goto Done; + } + + program_context.data_start = batch_run_request->data; + program_context.data_end = (uint8_t*)batch_run_request + input_buffer_length; + + // Invoke the eBPF program. Pass the output buffer as program context data. + for (uint32_t i = 0; i < batch_run_request->count; i++) { + result = sample_ebpf_extension_invoke_batch_program(&program_context, &context_state, &program_result); + if (result != EBPF_SUCCESS) { + status = STATUS_UNSUCCESSFUL; + break; + } + } + + result = sample_ebpf_extension_invoke_batch_end_program(&context_state); + if (result != EBPF_SUCCESS) { + status = STATUS_UNSUCCESSFUL; + goto Done; + } + } else { + status = STATUS_INVALID_PARAMETER; + goto Done; + } + break; case IOCTL_SAMPLE_EBPF_EXT_CTL_PROFILE: { size_t minimum_request_size = sizeof(sample_ebpf_ext_profile_request_t); size_t minimum_reply_size = sizeof(sample_ebpf_ext_profile_reply_t); diff --git a/tests/sample/ext/inc/sample_ext_ioctls.h b/tests/sample/ext/inc/sample_ext_ioctls.h index bdc9b78fa..b61bcb835 100644 --- a/tests/sample/ext/inc/sample_ext_ioctls.h +++ b/tests/sample/ext/inc/sample_ext_ioctls.h @@ -22,6 +22,7 @@ typedef enum _sample_ebpf_ext_control_code { SAMPLE_EBPF_EXT_CONTROL_RUN, + SAMPLE_EBPF_EXT_CONTROL_RUN_BATCH, SAMPLE_EBPF_EXT_CONTROL_PROFILE, } sample_ebpf_ext_control_code_t; @@ -42,7 +43,23 @@ typedef struct _sample_ebpf_ext_profile_reply uint64_t duration; } sample_ebpf_ext_profile_reply_t; +typedef struct _sample_ebpf_ext_batch_run_request +{ + uint32_t count; + uint8_t data[1]; +} sample_ebpf_ext_batch_run_request_t; + +typedef struct _sample_ebpf_ext_batch_run_reply +{ + uint32_t status; + uint8_t data[1]; +} sample_ebpf_ext_batch_run_reply_t; + +#define SAMPLE_EBPF_PROGRAM_BATCH_INVOCATION_COUNT 10 + #define IOCTL_SAMPLE_EBPF_EXT_CTL_RUN \ CTL_CODE(FILE_DEVICE_NETWORK, SAMPLE_EBPF_EXT_CONTROL_RUN, METHOD_BUFFERED, FILE_ANY_ACCESS) +#define IOCTL_SAMPLE_EBPF_EXT_CTL_RUN_BATCH \ + CTL_CODE(FILE_DEVICE_NETWORK, SAMPLE_EBPF_EXT_CONTROL_RUN_BATCH, METHOD_BUFFERED, FILE_ANY_ACCESS) #define IOCTL_SAMPLE_EBPF_EXT_CTL_PROFILE \ CTL_CODE(FILE_DEVICE_NETWORK, SAMPLE_EBPF_EXT_CONTROL_PROFILE, METHOD_BUFFERED, FILE_ANY_ACCESS)