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