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 
11 class 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 
224 int main(int argc, char **argv) {
225   Handler handler;
226 
227   handler.Parse(argc, argv);
228   return handler.Process();
229 }
230