From adc563ebaf0cb23555854d09f20142508de7d13d Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 25 Jun 2024 20:39:26 -0700 Subject: [PATCH] Improve CLI handling by permitting either std::cin or file for program data (#513) Signed-off-by: Alan Jowett Co-authored-by: Alan Jowett --- custom_tests/CMakeLists.txt | 2 +- custom_tests/srcs/ubpf_custom_test_support.cc | 41 +++++++++++++++++++ custom_tests/srcs/ubpf_custom_test_support.h | 11 +++++ ...t_external_dispatcher_context_overwrite.cc | 8 ++-- ...test_external_dispatcher_simple_context.cc | 8 ++-- .../srcs/ubpf_test_external_stack_contents.cc | 8 ++-- custom_tests/srcs/ubpf_test_frame_pointer.cc | 8 ++-- .../ubpf_test_jit_unexpected_instruction.cc | 9 ++-- .../srcs/ubpf_test_update_dispatcher.cc | 8 ++-- custom_tests/srcs/ubpf_test_update_helpers.cc | 7 +++- custom_tests/windows/test.cmd | 4 -- 11 files changed, 89 insertions(+), 25 deletions(-) delete mode 100644 custom_tests/windows/test.cmd diff --git a/custom_tests/CMakeLists.txt b/custom_tests/CMakeLists.txt index 612ef8df5..d0184720d 100644 --- a/custom_tests/CMakeLists.txt +++ b/custom_tests/CMakeLists.txt @@ -57,7 +57,7 @@ foreach(test_file ${test_descr_files}) string(REPLACE "/" "\\" test_exe ${test_exe}) add_test( NAME ${test_name}-Custom - COMMAND ${CMAKE_SOURCE_DIR}/custom_tests/windows/test.cmd ${potential_input_file} ${test_exe} + COMMAND ${test_exe} --program ${potential_input_file} ) else() add_test( diff --git a/custom_tests/srcs/ubpf_custom_test_support.cc b/custom_tests/srcs/ubpf_custom_test_support.cc index cdc2368e0..be31a87c9 100644 --- a/custom_tests/srcs/ubpf_custom_test_support.cc +++ b/custom_tests/srcs/ubpf_custom_test_support.cc @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -96,3 +98,42 @@ bool ubpf_setup_custom_test(ubpf_vm_up &vm, free(error_s); return true; } + +bool get_program_string(int argc, char **argv, std::string &program_string, std::string &error) +{ + std::vector args(argv, argv + argc); + + if (args.size() == 1) + { + std::string line; + while (std::getline(std::cin, line)) + { + program_string += line; + } + } + else if (args.size() == 3 && args[1] == "--program") + { + std::ifstream file(args[2]); + if (file.is_open()) + { + std::string line; + while (std::getline(file, line)) + { + program_string += line; + } + file.close(); + } + else + { + error = "Failed to open file: " + args[2]; + return false; + } + } + else + { + error = "Usage: " + args[0] + " [program]"; + return false; + } + + return true; +} diff --git a/custom_tests/srcs/ubpf_custom_test_support.h b/custom_tests/srcs/ubpf_custom_test_support.h index 8957c5c60..a6d6b48e3 100644 --- a/custom_tests/srcs/ubpf_custom_test_support.h +++ b/custom_tests/srcs/ubpf_custom_test_support.h @@ -56,3 +56,14 @@ bool ubpf_setup_custom_test(ubpf_vm_up &vm, ubpf_jit_fn &jit_fn, std::string &error); +/** + * @brief Get the program string object from the command line arguments or stdin. + * + * @param[in] argc The number of command line arguments. + * @param[in] argv The command line arguments. + * @param[out] program_string The program string to return. + * @param[out] error The error message to return if there is a problem. + * @return true If the program string was successfully obtained. + * @return false If there was a problem obtaining the program string. + */ +bool get_program_string(int argc, char **argv, std::string &program_string, std::string &error); diff --git a/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc b/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc index b99d0448f..ca97d399a 100644 --- a/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc +++ b/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc @@ -39,18 +39,20 @@ external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie) int main(int argc, char **argv) { - std::vector args(argv, argv + argc); std::string program_string{}; + std::string error{}; ubpf_jit_fn jit_fn; uint64_t memory{0x123456789}; // The program modifies (eBPF) r0 (see test description) and then invokes // a helper function that will be invoked through the external // dispatcher. - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc b/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc index e3fa6e65e..377363645 100644 --- a/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc +++ b/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc @@ -37,18 +37,20 @@ external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie) int main(int argc, char **argv) { - std::vector args(argv, argv + argc); std::string program_string{}; + std::string error{}; ubpf_jit_fn jit_fn; uint64_t memory{0x123456789}; // The test program invokes an external function (which is invoked via the registered // external helper dispatcher). The result of that external function is given as the // result of the eBPF's program execution. Therefore, ... - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/srcs/ubpf_test_external_stack_contents.cc b/custom_tests/srcs/ubpf_test_external_stack_contents.cc index 23a930f60..145487177 100644 --- a/custom_tests/srcs/ubpf_test_external_stack_contents.cc +++ b/custom_tests/srcs/ubpf_test_external_stack_contents.cc @@ -27,11 +27,14 @@ stack_usage_calculator(const struct ubpf_vm* vm, uint16_t pc, void* cookie) int main(int argc, char** argv) { - std::vector args(argv, argv + argc); std::string program_string{}; + std::string error{}; ubpf_jit_fn jit_fn; - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } const size_t stack_size{32}; uint8_t expected_result[] = { @@ -41,7 +44,6 @@ main(int argc, char** argv) bool success = true; std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/srcs/ubpf_test_frame_pointer.cc b/custom_tests/srcs/ubpf_test_frame_pointer.cc index 7c19c0412..bfeec7246 100644 --- a/custom_tests/srcs/ubpf_test_frame_pointer.cc +++ b/custom_tests/srcs/ubpf_test_frame_pointer.cc @@ -36,11 +36,14 @@ overwrite_stack_usage_calculator(const struct ubpf_vm* vm, uint16_t pc, void* co int main(int argc, char** argv) { - std::vector args(argv, argv + argc); std::string program_string{}; + std::string error{}; ubpf_jit_fn jit_fn; - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } uint64_t no_overwrite_interp_result = 0; uint64_t no_overwrite_jit_result = 0; @@ -50,7 +53,6 @@ main(int argc, char** argv) { std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc index 5b819cc31..e4d9523e9 100644 --- a/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc +++ b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc @@ -13,20 +13,23 @@ extern "C" #include "ubpf_custom_test_support.h" -int main() +int main(int argc, char** argv) { std::string expected_error{"Failed to load program: unknown opcode 0x8f at PC 0" }; ubpf_jit_fn jit_fn; std::string program_string; + std::string error{}; // The program's first instruction contains an invalid opcode. Attempting to load this // program should elicit an error alerting the user to an unknown opcode (see above). - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } std::vector program = bytes_to_ebpf_inst(base16_decode(program_string)); ubpf_vm_up vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, diff --git a/custom_tests/srcs/ubpf_test_update_dispatcher.cc b/custom_tests/srcs/ubpf_test_update_dispatcher.cc index 5124122f2..f28eb3596 100644 --- a/custom_tests/srcs/ubpf_test_update_dispatcher.cc +++ b/custom_tests/srcs/ubpf_test_update_dispatcher.cc @@ -62,16 +62,18 @@ test_helpers_validater(unsigned int idx, const struct ubpf_vm* vm) int main(int argc, char** argv) { - std::vector args(argv, argv + argc); std::string program_string; + std::string error{}; std::string memory_string; - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } ubpf_jit_fn jit_fn; uint64_t memory{0x123456789}; std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/srcs/ubpf_test_update_helpers.cc b/custom_tests/srcs/ubpf_test_update_helpers.cc index 8566e42d2..c3d36be08 100644 --- a/custom_tests/srcs/ubpf_test_update_helpers.cc +++ b/custom_tests/srcs/ubpf_test_update_helpers.cc @@ -39,6 +39,7 @@ int main(int argc, char **argv) { std::vector args(argv, argv + argc); std::string program_string{}; + std::string error{}; bool success{true}; const char memfrob_testcase_name[] = "memfrob"; @@ -99,13 +100,15 @@ int main(int argc, char **argv) } }; - std::getline(std::cin, program_string); + if (!get_program_string(argc, argv, program_string, error)) { + std::cerr << error << std::endl; + return 1; + } for (auto testcase : test_cases) { ubpf_jit_fn jit_fn; uint64_t memory{0x123456789}; std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; if (!ubpf_setup_custom_test( vm, program_string, diff --git a/custom_tests/windows/test.cmd b/custom_tests/windows/test.cmd deleted file mode 100644 index dc905f14b..000000000 --- a/custom_tests/windows/test.cmd +++ /dev/null @@ -1,4 +0,0 @@ -rem Copyright (c) Microsoft Corporation -rem SPDX-License-Identifier: Apache-2.0 - -type %1 | %2