Improve CLI handling by permitting either std::cin or file for progra…

…m data (#513) Signed-off-by: Alan Jowett <[email protected]> Co-authored-by: Alan Jowett <[email protected]>
iovisor · Jun 26, 2024 · adc563e · adc563e
1 parent a0f7abb
commit adc563e
Show file tree

Hide file tree

Showing 11 changed files with 89 additions and 25 deletions.
diff --git 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
@@ -4,6 +4,8 @@
 #include <cassert>
 #include <cstdint>
 #include <cstring>
+#include <fstream>
+#include <iostream>
 #include <vector>
 #include <string>
 #include <sstream>
@@ -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<std::string> 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
@@ -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
@@ -39,18 +39,20 @@ external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie)
 
 int main(int argc, char **argv)
 {
-    std::vector<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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
@@ -37,18 +37,20 @@ external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie)
 
 int main(int argc, char **argv)
 {
-    std::vector<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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
@@ -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<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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
@@ -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<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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
@@ -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<ebpf_inst> 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
@@ -62,16 +62,18 @@ test_helpers_validater(unsigned int idx, const struct ubpf_vm* vm)
 int
 main(int argc, char** argv)
 {
-    std::vector<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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
@@ -39,6 +39,7 @@ int main(int argc, char **argv)
 {
     std::vector<std::string> 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<ubpf_vm, decltype(&ubpf_destroy)> 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