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 use std::cmp;
  14 use std::env;
  15 use std::fs::{self, File, Permissions};
  16 use std::io::prelude::*;
  17 use std::io::{self, BufReader};
  18 use std::net::{TcpListener, TcpStream};
  19 use std::os::unix::prelude::*;
  20 use std::path::{Path, PathBuf};
  21 use std::process::{Command, Stdio};
  22 use std::str;
  23 use std::sync::atomic::{AtomicUsize, Ordering};
  24 use std::sync::{Arc, Mutex};
  25 use std::thread;
  26
  27 macro_rules! t {
  28     ($e:expr) => (match $e {
  29         Ok(e) => e,
  30         Err(e) => panic!("{} failed with {}", stringify!($e), e),
  31     })
  32 }
  33
  34 static TEST: AtomicUsize = AtomicUsize::new(0);
  35
  36 struct Config {
  37     pub remote: bool,
  38     pub verbose: bool,
  39 }
  40
  41 impl Config {
  42     pub fn default() -> Config {
  43         Config {
  44             remote: false,
  45             verbose: false,
  46         }
  47     }
  48
  49     pub fn parse_args() -> Config {
  50         let mut config = Config::default();
  51
  52         let args = env::args().skip(1);
  53         for argument in args {
  54             match &argument[..] {
  55                 "remote" => {
  56                     config.remote = true;
  57                 },
  58                 "verbose" | "-v" => {
  59                     config.verbose = true;
  60                 }
  61                 arg => panic!("unknown argument: {}", arg),
  62             }
  63         }
  64
  65         config
  66     }
  67 }
  68
  69 fn main() {
  70     println!("starting test server");
  71
  72     let config = Config::parse_args();
  73
  74     let bind_addr = if cfg!(target_os = "android") || config.remote {
  75         "0.0.0.0:12345"
  76     } else {
  77         "10.0.2.15:12345"
  78     };
  79
  80     let (listener, work) = if cfg!(target_os = "android") {
  81         (t!(TcpListener::bind(bind_addr)), "/data/tmp/work")
  82     } else {
  83         (t!(TcpListener::bind(bind_addr)), "/tmp/work")
  84     };
  85     println!("listening!");
  86
  87     let work = Path::new(work);
  88     t!(fs::create_dir_all(work));
  89
  90     let lock = Arc::new(Mutex::new(()));
  91
  92     for socket in listener.incoming() {
  93         let mut socket = t!(socket);
  94         let mut buf = [0; 4];
  95         if socket.read_exact(&mut buf).is_err() {
  96             continue
  97         }
  98         if &buf[..] == b"ping" {
  99             t!(socket.write_all(b"pong"));
 100         } else if &buf[..] == b"push" {
 101             handle_push(socket, work);
 102         } else if &buf[..] == b"run " {
 103             let lock = lock.clone();
 104             thread::spawn(move || handle_run(socket, work, &lock));
 105         } else {
 106             panic!("unknown command {:?}", buf);
 107         }
 108     }
 109 }
 110
 111 fn handle_push(socket: TcpStream, work: &Path) {
 112     let mut reader = BufReader::new(socket);
 113     recv(&work, &mut reader);
 114
 115     let mut socket = reader.into_inner();
 116     t!(socket.write_all(b"ack "));
 117 }
 118
 119 struct RemoveOnDrop<'a> {
 120     inner: &'a Path,
 121 }
 122
 123 impl Drop for RemoveOnDrop<'_> {
 124     fn drop(&mut self) {
 125         t!(fs::remove_dir_all(self.inner));
 126     }
 127 }
 128
 129 fn handle_run(socket: TcpStream, work: &Path, lock: &Mutex<()>) {
 130     let mut arg = Vec::new();
 131     let mut reader = BufReader::new(socket);
 132
 133     // Allocate ourselves a directory that we'll delete when we're done to save
 134     // space.
 135     let n = TEST.fetch_add(1, Ordering::SeqCst);
 136     let path = work.join(format!("test{}", n));
 137     t!(fs::create_dir(&path));
 138     let _a = RemoveOnDrop { inner: &path };
 139
 140     // First up we'll get a list of arguments delimited with 0 bytes. An empty
 141     // argument means that we're done.
 142     let mut args = Vec::new();
 143     while t!(reader.read_until(0, &mut arg)) > 1 {
 144         args.push(t!(str::from_utf8(&arg[..arg.len() - 1])).to_string());
 145         arg.truncate(0);
 146     }
 147
 148     // Next we'll get a bunch of env vars in pairs delimited by 0s as well
 149     let mut env = Vec::new();
 150     arg.truncate(0);
 151     while t!(reader.read_until(0, &mut arg)) > 1 {
 152         let key_len = arg.len() - 1;
 153         let val_len = t!(reader.read_until(0, &mut arg)) - 1;
 154         {
 155             let key = &arg[..key_len];
 156             let val = &arg[key_len + 1..][..val_len];
 157             let key = t!(str::from_utf8(key)).to_string();
 158             let val = t!(str::from_utf8(val)).to_string();
 159             env.push((key, val));
 160         }
 161         arg.truncate(0);
 162     }
 163
 164     // The section of code from here down to where we drop the lock is going to
 165     // be a critical section for us. On Linux you can't execute a file which is
 166     // open somewhere for writing, as you'll receive the error "text file busy".
 167     // Now here we never have the text file open for writing when we spawn it,
 168     // so why do we still need a critical section?
 169     //
 170     // Process spawning first involves a `fork` on Unix, which clones all file
 171     // descriptors into the child process. This means that it's possible for us
 172     // to open the file for writing (as we're downloading it), then some other
 173     // thread forks, then we close the file and try to exec. At that point the
 174     // other thread created a child process with the file open for writing, and
 175     // we attempt to execute it, so we get an error.
 176     //
 177     // This race is resolve by ensuring that only one thread can write the file
 178     // and spawn a child process at once. Kinda an unfortunate solution, but we
 179     // don't have many other choices with this sort of setup!
 180     //
 181     // In any case the lock is acquired here, before we start writing any files.
 182     // It's then dropped just after we spawn the child. That way we don't lock
 183     // the execution of the child, just the creation of its files.
 184     let lock = lock.lock();
 185
 186     // Next there's a list of dynamic libraries preceded by their filenames.
 187     while t!(reader.fill_buf())[0] != 0 {
 188         recv(&path, &mut reader);
 189     }
 190     assert_eq!(t!(reader.read(&mut [0])), 1);
 191
 192     // Finally we'll get the binary. The other end will tell us how big the
 193     // binary is and then we'll download it all to the exe path we calculated
 194     // earlier.
 195     let exe = recv(&path, &mut reader);
 196
 197     let mut cmd = Command::new(&exe);
 198     for arg in args {
 199         cmd.arg(arg);
 200     }
 201     for (k, v) in env {
 202         cmd.env(k, v);
 203     }
 204
 205     // Support libraries were uploaded to `work` earlier, so make sure that's
 206     // in `LD_LIBRARY_PATH`. Also include our own current dir which may have
 207     // had some libs uploaded.
 208     cmd.env("LD_LIBRARY_PATH",
 209             format!("{}:{}", work.display(), path.display()));
 210
 211     // Spawn the child and ferry over stdout/stderr to the socket in a framed
 212     // fashion (poor man's style)
 213     let mut child = t!(cmd.stdin(Stdio::null())
 214                           .stdout(Stdio::piped())
 215                           .stderr(Stdio::piped())
 216                           .spawn());
 217     drop(lock);
 218     let mut stdout = child.stdout.take().unwrap();
 219     let mut stderr = child.stderr.take().unwrap();
 220     let socket = Arc::new(Mutex::new(reader.into_inner()));
 221     let socket2 = socket.clone();
 222     let thread = thread::spawn(move || my_copy(&mut stdout, 0, &*socket2));
 223     my_copy(&mut stderr, 1, &*socket);
 224     thread.join().unwrap();
 225
 226     // Finally send over the exit status.
 227     let status = t!(child.wait());
 228     let (which, code) = match status.code() {
 229         Some(n) => (0, n),
 230         None => (1, status.signal().unwrap()),
 231     };
 232     t!(socket.lock().unwrap().write_all(&[
 233         which,
 234         (code >> 24) as u8,
 235         (code >> 16) as u8,
 236         (code >>  8) as u8,
 237         (code >>  0) as u8,
 238     ]));
 239 }
 240
 241 fn recv<B: BufRead>(dir: &Path, io: &mut B) -> PathBuf {
 242     let mut filename = Vec::new();
 243     t!(io.read_until(0, &mut filename));
 244
 245     // We've got some tests with *really* long names. We try to name the test
 246     // executable the same on the target as it is on the host to aid with
 247     // debugging, but the targets we're emulating are often more restrictive
 248     // than the hosts as well.
 249     //
 250     // To ensure we can run a maximum number of tests without modifications we
 251     // just arbitrarily truncate the filename to 50 bytes. That should
 252     // hopefully allow us to still identify what's running while staying under
 253     // the filesystem limits.
 254     let len = cmp::min(filename.len() - 1, 50);
 255     let dst = dir.join(t!(str::from_utf8(&filename[..len])));
 256     let amt = read_u32(io) as u64;
 257     t!(io::copy(&mut io.take(amt),
 258                 &mut t!(File::create(&dst))));
 259     t!(fs::set_permissions(&dst, Permissions::from_mode(0o755)));
 260     dst
 261 }
 262
 263 fn my_copy(src: &mut dyn Read, which: u8, dst: &Mutex<dyn Write>) {
 264     let mut b = [0; 1024];
 265     loop {
 266         let n = t!(src.read(&mut b));
 267         let mut dst = dst.lock().unwrap();
 268         t!(dst.write_all(&[
 269             which,
 270             (n >> 24) as u8,
 271             (n >> 16) as u8,
 272             (n >>  8) as u8,
 273             (n >>  0) as u8,
 274         ]));
 275         if n > 0 {
 276             t!(dst.write_all(&b[..n]));
 277         } else {
 278             break
 279         }
 280     }
 281 }
 282
 283 fn read_u32(r: &mut dyn Read) -> u32 {
 284     let mut len = [0; 4];
 285     t!(r.read_exact(&mut len));
 286     ((len[0] as u32) << 24) |
 287     ((len[1] as u32) << 16) |
 288     ((len[2] as u32) <<  8) |
 289     ((len[3] as u32) <<  0)
 290 }