]> git.lizzy.rs Git - rust.git/blob - src/tools/remote-test-server/src/main.rs
Auto merge of #103137 - dtolnay:readdir, r=Mark-Simulacrum
[rust.git] / src / tools / remote-test-server / src / main.rs
1 //! This is a small server which is intended to run inside of an emulator or
2 //! on a remote test device. This server pairs with the `remote-test-client`
3 //! program in this repository. The `remote-test-client` connects to this
4 //! server over a TCP socket and performs work such as:
5 //!
6 //! 1. Pushing shared libraries to the server
7 //! 2. Running tests through the server
8 //!
9 //! The server supports running tests concurrently and also supports tests
10 //! themselves having support libraries. All data over the TCP sockets is in a
11 //! basically custom format suiting our needs.
12
13 #[cfg(not(windows))]
14 use std::fs::Permissions;
15 use std::net::SocketAddr;
16 #[cfg(not(windows))]
17 use std::os::unix::prelude::*;
18
19 use std::cmp;
20 use std::env;
21 use std::fs::{self, File};
22 use std::io::prelude::*;
23 use std::io::{self, BufReader};
24 use std::net::{TcpListener, TcpStream};
25 use std::path::{Path, PathBuf};
26 use std::process::{Command, ExitStatus, Stdio};
27 use std::str;
28 use std::sync::atomic::{AtomicUsize, Ordering};
29 use std::sync::{Arc, Mutex};
30 use std::thread;
31
32 macro_rules! t {
33     ($e:expr) => {
34         match $e {
35             Ok(e) => e,
36             Err(e) => panic!("{} failed with {}", stringify!($e), e),
37         }
38     };
39 }
40
41 static TEST: AtomicUsize = AtomicUsize::new(0);
42
43 #[derive(Copy, Clone)]
44 struct Config {
45     verbose: bool,
46     sequential: bool,
47     bind: SocketAddr,
48 }
49
50 impl Config {
51     pub fn default() -> Config {
52         Config {
53             verbose: false,
54             sequential: false,
55             bind: if cfg!(target_os = "android") || cfg!(windows) {
56                 ([0, 0, 0, 0], 12345).into()
57             } else {
58                 ([10, 0, 2, 15], 12345).into()
59             },
60         }
61     }
62
63     pub fn parse_args() -> Config {
64         let mut config = Config::default();
65
66         let args = env::args().skip(1);
67         let mut next_is_bind = false;
68         for argument in args {
69             match &argument[..] {
70                 bind if next_is_bind => {
71                     config.bind = t!(bind.parse());
72                     next_is_bind = false;
73                 }
74                 "--bind" => next_is_bind = true,
75                 "--sequential" => config.sequential = true,
76                 "--verbose" | "-v" => config.verbose = true,
77                 "--help" | "-h" => {
78                     show_help();
79                     std::process::exit(0);
80                 }
81                 arg => panic!("unknown argument: {}, use `--help` for known arguments", arg),
82             }
83         }
84         if next_is_bind {
85             panic!("missing value for --bind");
86         }
87
88         config
89     }
90 }
91
92 fn show_help() {
93     eprintln!(
94         r#"Usage:
95
96 {} [OPTIONS]
97
98 OPTIONS:
99     --bind <IP>:<PORT>   Specify IP address and port to listen for requests, e.g. "0.0.0.0:12345"
100     --sequential         Run only one test at a time
101     -v, --verbose        Show status messages
102     -h, --help           Show this help screen
103 "#,
104         std::env::args().next().unwrap()
105     );
106 }
107
108 fn print_verbose(s: &str, conf: Config) {
109     if conf.verbose {
110         println!("{}", s);
111     }
112 }
113
114 fn main() {
115     let config = Config::parse_args();
116     println!("starting test server");
117
118     let listener = t!(TcpListener::bind(config.bind));
119     let (work, tmp): (PathBuf, PathBuf) = if cfg!(target_os = "android") {
120         ("/data/local/tmp/work".into(), "/data/local/tmp/work/tmp".into())
121     } else {
122         let mut work_dir = env::temp_dir();
123         work_dir.push("work");
124         let mut tmp_dir = work_dir.clone();
125         tmp_dir.push("tmp");
126         (work_dir, tmp_dir)
127     };
128     println!("listening on {}!", config.bind);
129
130     t!(fs::create_dir_all(&work));
131     t!(fs::create_dir_all(&tmp));
132
133     let lock = Arc::new(Mutex::new(()));
134
135     for socket in listener.incoming() {
136         let mut socket = t!(socket);
137         let mut buf = [0; 4];
138         if socket.read_exact(&mut buf).is_err() {
139             continue;
140         }
141         if &buf[..] == b"ping" {
142             print_verbose("Received ping", config);
143             t!(socket.write_all(b"pong"));
144         } else if &buf[..] == b"push" {
145             handle_push(socket, &work, config);
146         } else if &buf[..] == b"run " {
147             let lock = lock.clone();
148             let work = work.clone();
149             let tmp = tmp.clone();
150             let f = move || handle_run(socket, &work, &tmp, &lock, config);
151             if config.sequential {
152                 f();
153             } else {
154                 thread::spawn(f);
155             }
156         } else {
157             panic!("unknown command {:?}", buf);
158         }
159     }
160 }
161
162 fn handle_push(socket: TcpStream, work: &Path, config: Config) {
163     let mut reader = BufReader::new(socket);
164     let dst = recv(&work, &mut reader);
165     print_verbose(&format!("push {:#?}", dst), config);
166
167     let mut socket = reader.into_inner();
168     t!(socket.write_all(b"ack "));
169 }
170
171 struct RemoveOnDrop<'a> {
172     inner: &'a Path,
173 }
174
175 impl Drop for RemoveOnDrop<'_> {
176     fn drop(&mut self) {
177         t!(fs::remove_dir_all(self.inner));
178     }
179 }
180
181 fn handle_run(socket: TcpStream, work: &Path, tmp: &Path, lock: &Mutex<()>, config: Config) {
182     let mut arg = Vec::new();
183     let mut reader = BufReader::new(socket);
184
185     // Allocate ourselves a directory that we'll delete when we're done to save
186     // space.
187     let n = TEST.fetch_add(1, Ordering::SeqCst);
188     let path = work.join(format!("test{}", n));
189     t!(fs::create_dir(&path));
190     let _a = RemoveOnDrop { inner: &path };
191
192     // First up we'll get a list of arguments delimited with 0 bytes. An empty
193     // argument means that we're done.
194     let mut args = Vec::new();
195     while t!(reader.read_until(0, &mut arg)) > 1 {
196         args.push(t!(str::from_utf8(&arg[..arg.len() - 1])).to_string());
197         arg.truncate(0);
198     }
199
200     // Next we'll get a bunch of env vars in pairs delimited by 0s as well
201     let mut env = Vec::new();
202     arg.truncate(0);
203     while t!(reader.read_until(0, &mut arg)) > 1 {
204         let key_len = arg.len() - 1;
205         let val_len = t!(reader.read_until(0, &mut arg)) - 1;
206         {
207             let key = &arg[..key_len];
208             let val = &arg[key_len + 1..][..val_len];
209             let key = t!(str::from_utf8(key)).to_string();
210             let val = t!(str::from_utf8(val)).to_string();
211             env.push((key, val));
212         }
213         arg.truncate(0);
214     }
215
216     // The section of code from here down to where we drop the lock is going to
217     // be a critical section for us. On Linux you can't execute a file which is
218     // open somewhere for writing, as you'll receive the error "text file busy".
219     // Now here we never have the text file open for writing when we spawn it,
220     // so why do we still need a critical section?
221     //
222     // Process spawning first involves a `fork` on Unix, which clones all file
223     // descriptors into the child process. This means that it's possible for us
224     // to open the file for writing (as we're downloading it), then some other
225     // thread forks, then we close the file and try to exec. At that point the
226     // other thread created a child process with the file open for writing, and
227     // we attempt to execute it, so we get an error.
228     //
229     // This race is resolve by ensuring that only one thread can write the file
230     // and spawn a child process at once. Kinda an unfortunate solution, but we
231     // don't have many other choices with this sort of setup!
232     //
233     // In any case the lock is acquired here, before we start writing any files.
234     // It's then dropped just after we spawn the child. That way we don't lock
235     // the execution of the child, just the creation of its files.
236     let lock = lock.lock();
237
238     // Next there's a list of dynamic libraries preceded by their filenames.
239     while t!(reader.fill_buf())[0] != 0 {
240         recv(&path, &mut reader);
241     }
242     assert_eq!(t!(reader.read(&mut [0])), 1);
243
244     // Finally we'll get the binary. The other end will tell us how big the
245     // binary is and then we'll download it all to the exe path we calculated
246     // earlier.
247     let exe = recv(&path, &mut reader);
248     print_verbose(&format!("run {:#?}", exe), config);
249
250     let mut cmd = Command::new(&exe);
251     cmd.args(args);
252     cmd.envs(env);
253
254     // On windows, libraries are just searched in the executable directory,
255     // system directories, PWD, and PATH, in that order. PATH is the only one
256     // we can change for this.
257     let library_path = if cfg!(windows) { "PATH" } else { "LD_LIBRARY_PATH" };
258
259     // Support libraries were uploaded to `work` earlier, so make sure that's
260     // in `LD_LIBRARY_PATH`. Also include our own current dir which may have
261     // had some libs uploaded.
262     let mut paths = vec![work.to_owned(), path.clone()];
263     if let Some(library_path) = env::var_os(library_path) {
264         paths.extend(env::split_paths(&library_path));
265     }
266     cmd.env(library_path, env::join_paths(paths).unwrap());
267
268     // Some tests assume RUST_TEST_TMPDIR exists
269     cmd.env("RUST_TEST_TMPDIR", tmp.to_owned());
270
271     // Spawn the child and ferry over stdout/stderr to the socket in a framed
272     // fashion (poor man's style)
273     let mut child =
274         t!(cmd.stdin(Stdio::null()).stdout(Stdio::piped()).stderr(Stdio::piped()).spawn());
275     drop(lock);
276     let mut stdout = child.stdout.take().unwrap();
277     let mut stderr = child.stderr.take().unwrap();
278     let socket = Arc::new(Mutex::new(reader.into_inner()));
279     let socket2 = socket.clone();
280     let thread = thread::spawn(move || my_copy(&mut stdout, 0, &*socket2));
281     my_copy(&mut stderr, 1, &*socket);
282     thread.join().unwrap();
283
284     // Finally send over the exit status.
285     let status = t!(child.wait());
286
287     let (which, code) = get_status_code(&status);
288
289     t!(socket.lock().unwrap().write_all(&[
290         which,
291         (code >> 24) as u8,
292         (code >> 16) as u8,
293         (code >> 8) as u8,
294         (code >> 0) as u8,
295     ]));
296 }
297
298 #[cfg(not(windows))]
299 fn get_status_code(status: &ExitStatus) -> (u8, i32) {
300     match status.code() {
301         Some(n) => (0, n),
302         None => (1, status.signal().unwrap()),
303     }
304 }
305
306 #[cfg(windows)]
307 fn get_status_code(status: &ExitStatus) -> (u8, i32) {
308     (0, status.code().unwrap())
309 }
310
311 fn recv<B: BufRead>(dir: &Path, io: &mut B) -> PathBuf {
312     let mut filename = Vec::new();
313     t!(io.read_until(0, &mut filename));
314
315     // We've got some tests with *really* long names. We try to name the test
316     // executable the same on the target as it is on the host to aid with
317     // debugging, but the targets we're emulating are often more restrictive
318     // than the hosts as well.
319     //
320     // To ensure we can run a maximum number of tests without modifications we
321     // just arbitrarily truncate the filename to 50 bytes. That should
322     // hopefully allow us to still identify what's running while staying under
323     // the filesystem limits.
324     let len = cmp::min(filename.len() - 1, 50);
325     let dst = dir.join(t!(str::from_utf8(&filename[..len])));
326     let amt = read_u32(io) as u64;
327     t!(io::copy(&mut io.take(amt), &mut t!(File::create(&dst))));
328     set_permissions(&dst);
329     dst
330 }
331
332 #[cfg(not(windows))]
333 fn set_permissions(path: &Path) {
334     t!(fs::set_permissions(&path, Permissions::from_mode(0o755)));
335 }
336 #[cfg(windows)]
337 fn set_permissions(_path: &Path) {}
338
339 fn my_copy(src: &mut dyn Read, which: u8, dst: &Mutex<dyn Write>) {
340     let mut b = [0; 1024];
341     loop {
342         let n = t!(src.read(&mut b));
343         let mut dst = dst.lock().unwrap();
344         t!(dst.write_all(&[
345             which,
346             (n >> 24) as u8,
347             (n >> 16) as u8,
348             (n >> 8) as u8,
349             (n >> 0) as u8,
350         ]));
351         if n > 0 {
352             t!(dst.write_all(&b[..n]));
353         } else {
354             break;
355         }
356     }
357 }
358
359 fn read_u32(r: &mut dyn Read) -> u32 {
360     let mut len = [0; 4];
361     t!(r.read_exact(&mut len));
362     ((len[0] as u32) << 24)
363         | ((len[1] as u32) << 16)
364         | ((len[2] as u32) << 8)
365         | ((len[3] as u32) << 0)
366 }