1// vim: ts=2 sw=2 et
2
3#include <format>
4#include <iostream>
5#include <map>
6#include <string>
7#include <vector>
8#include <stdlib.h>
9#include <unistd.h>
10
11class Handler {
12  private:
13    const std::string kPytestName = "pytest";
14    const std::string kCleanupSuffix = ":cleanup";
15    const std::string kPythonPathEnv = "PYTHONPATH";
16    const std::string kAtfVar = "_ATF_VAR_";
17  public:
18    // Test listing requested
19    bool flag_list = false;
20    // Output debug data (will break listing)
21    bool flag_debug = false;
22    // Cleanup for the test requested
23    bool flag_cleanup = false;
24    // Test source directory (provided by ATF)
25    std::string src_dir;
26    // Path to write test status to (provided by ATF)
27    std::string dst_file;
28    // Path to add to PYTHONPATH (provided by the schebang args)
29    std::string python_path;
30    // Path to the script (provided by the schebang wrapper)
31    std::string script_path;
32    // Name of the test to run (provided by ATF)
33    std::string test_name;
34    // kv pairs (provided by ATF)
35    std::map<std::string,std::string> kv_map;
36    // our binary name
37    std::string binary_name;
38
39    static std::vector<std::string> ToVector(int argc, char **argv) {
40      std::vector<std::string> ret;
41
42      for (int i = 0; i < argc; i++) {
43        ret.emplace_back(std::string(argv[i]));
44      }
45      return ret;
46    }
47
48    static void PrintVector(std::string prefix, const std::vector<std::string> &vec) {
49      std::cerr << prefix << ": ";
50      for (auto &val: vec) {
51        std::cerr << "'" << val << "' ";
52      }
53      std::cerr << std::endl;
54    }
55
56    void Usage(std::string msg, bool exit_with_error) {
57      std::cerr << binary_name << ": ERROR: " << msg << "." << std::endl;
58      std::cerr << binary_name << ": See atf-test-program(1) for usage details." << std::endl;
59      exit(exit_with_error != 0);
60    }
61
62    // Parse args received from the OS. There can be multiple valid options:
63    // * with schebang args (#!/binary -P/path):
64    // atf_wrap '-P /path' /path/to/script -l
65    // * without schebang args
66    // atf_wrap /path/to/script -l
67    // Running test:
68    // atf_wrap '-P /path' /path/to/script -r /path1 -s /path2 -vk1=v1 testname
69    void Parse(int argc, char **argv) {
70      if (flag_debug) {
71        PrintVector("IN", ToVector(argc, argv));
72      }
73      // getopt() skips the first argument (as it is typically binary name)
74      // it is possible to have either '-P\s*/path' followed by the script name
75      // or just the script name. Parse kernel-provided arg manually and adjust
76      // array to make getopt work
77
78      binary_name = std::string(argv[0]);
79      argc--; argv++;
80      // parse -P\s*path from the kernel.
81      if (argc > 0 && !strncmp(argv[0], "-P", 2)) {
82        char *path = &argv[0][2];
83        while (*path == ' ')
84          path++;
85        python_path = std::string(path);
86        argc--; argv++;
87      }
88
89      // The next argument is a script name. Copy and keep argc/argv the same
90      // Show usage for empty args
91      if (argc == 0) {
92        Usage("Must provide a test case name", true);
93      }
94      script_path = std::string(argv[0]);
95
96      int c;
97      while ((c = getopt(argc, argv, "lr:s:v:")) != -1) {
98        switch (c) {
99          case 'l':
100            flag_list = true;
101            break;
102          case 's':
103            src_dir = std::string(optarg);
104            break;
105          case 'r':
106            dst_file = std::string(optarg);
107            break;
108          case 'v':
109            {
110              std::string kv = std::string(optarg);
111              size_t splitter = kv.find("=");
112              if (splitter == std::string::npos) {
113                Usage("Unknown variable: " + kv, true);
114              }
115              kv_map[kv.substr(0, splitter)] = kv.substr(splitter + 1);
116            }
117            break;
118          default:
119            Usage("Unknown option -" + std::string(1, static_cast<char>(c)), true);
120        }
121      }
122      argc -= optind;
123      argv += optind;
124
125      if (flag_list) {
126        return;
127      }
128      // There should be just one argument with the test name
129      if (argc != 1) {
130        Usage("Must provide a test case name", true);
131      }
132      test_name = std::string(argv[0]);
133      if (test_name.size() > kCleanupSuffix.size() &&
134          std::equal(kCleanupSuffix.rbegin(), kCleanupSuffix.rend(), test_name.rbegin())) {
135        test_name = test_name.substr(0, test_name.size() - kCleanupSuffix.size());
136        flag_cleanup = true;
137      }
138    }
139
140    std::vector<std::string> BuildArgs() {
141      std::vector<std::string> args = {"pytest", "-vv", "-p",
142        "no:cacheprovider", "-s", "--atf"};
143
144      args.push_back("--confcutdir=" + python_path);
145
146      if (flag_list) {
147        args.push_back("--co");
148        args.push_back(script_path);
149        return args;
150      }
151      if (flag_cleanup) {
152        args.push_back("--atf-cleanup");
153      }
154      // workaround pytest parser bug:
155      // https://github.com/pytest-dev/pytest/issues/3097
156      // use '--arg=value' format instead of '--arg value' for all
157      // path-like options
158      if (!src_dir.empty()) {
159        args.push_back("--atf-source-dir=" + src_dir);
160      }
161      if (!dst_file.empty()) {
162        args.push_back("--atf-file=" + dst_file);
163      }
164      // Create nodeid from the test path &name
165      args.push_back(script_path + "::" + test_name);
166      return args;
167    }
168
169    void SetPythonPath() {
170      if (!python_path.empty()) {
171        char *env_path = getenv(kPythonPathEnv.c_str());
172        if (env_path != nullptr) {
173          python_path = python_path + ":" + std::string(env_path);
174        }
175        setenv(kPythonPathEnv.c_str(), python_path.c_str(), 1);
176      }
177    }
178
179    void SetEnv() {
180      SetPythonPath();
181
182      // Pass ATF kv pairs as env variables to avoid dealing with
183      // pytest parser
184      for (auto [k, v]: kv_map) {
185        setenv((kAtfVar + k).c_str(), v.c_str(), 1);
186      }
187    }
188
189    bool Run(std::string binary, std::vector<std::string> args) {
190      if (flag_debug) {
191        PrintVector("OUT", args);
192      }
193      // allocate array with final NULL
194      char **arr = new char*[args.size() + 1]();
195      for (unsigned long i = 0; i < args.size(); i++) {
196        // work around 'char *const *'
197        arr[i] = strdup(args[i].c_str());
198      }
199      return execvp(binary.c_str(), arr) == 0;
200    }
201
202    void ReportError() {
203      if (flag_list) {
204        std::cout << "Content-Type: application/X-atf-tp; version=\"1\"";
205        std::cout << std::endl << std::endl;
206        std::cout << "ident: __test_cases_list_"<< kPytestName << "_binary_" <<
207          "not_found__" << std::endl;
208      } else {
209        std::cout << "execvp(" << kPytestName << ") failed: " <<
210          std::strerror(errno) << std::endl;
211      }
212    }
213
214    int Process() {
215      SetEnv();
216      if (!Run(kPytestName, BuildArgs())) {
217        ReportError();
218      }
219      return 0;
220    }
221};
222
223
224int main(int argc, char **argv) {
225  Handler handler;
226
227  handler.Parse(argc, argv);
228  return handler.Process();
229}
230