Bug 1525126 - [geckodriver] Wait for Marionette handshake rather than TCP connection. r=jgraham,webdriver-reviewers,nalexander
authorHenrik Skupin <mail@hskupin.info>
Wed, 25 Sep 2019 23:00:10 +0000
changeset 495005 d1f3680d88b34ffedd1e01dd52766f99d607cc7e
parent 495004 cee49076fdd8bc1821ac798ba393d6efe552f2fb
child 495006 e20ffbef6a295e9f9b1302f2dd3b9ca92bbad2b0
push id114131
push userdluca@mozilla.com
push dateThu, 26 Sep 2019 09:47:34 +0000
treeherdermozilla-inbound@1dc1a755079a [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersjgraham, webdriver-reviewers, nalexander
bugs1525126
milestone71.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1525126 - [geckodriver] Wait for Marionette handshake rather than TCP connection. r=jgraham,webdriver-reviewers,nalexander This allows to port-forward using adb. Connecting to such ports always establishes a TCP connection to the ADB daemon, but not all such connections are to bound ports. Waiting for the Marionette handshake ensures the connection is real, and not just partially forwarded. Depends on D44895 Differential Revision: https://phabricator.services.mozilla.com/D44896
testing/geckodriver/src/marionette.rs
--- a/testing/geckodriver/src/marionette.rs
+++ b/testing/geckodriver/src/marionette.rs
@@ -13,17 +13,16 @@ use marionette_rs::webdriver::{
     Selector as MarionetteSelector, WindowRect as MarionetteWindowRect,
 };
 use mozprofile::preferences::Pref;
 use mozprofile::profile::Profile;
 use mozrunner::runner::{FirefoxProcess, FirefoxRunner, Runner, RunnerProcess};
 use serde::de::{self, Deserialize, Deserializer};
 use serde::ser::{Serialize, Serializer};
 use serde_json::{self, Map, Value};
-use std::error::Error;
 use std::io::prelude::*;
 use std::io::Error as IoError;
 use std::io::ErrorKind;
 use std::io::Result as IoResult;
 use std::net::{TcpListener, TcpStream};
 use std::path::PathBuf;
 use std::sync::Mutex;
 use std::thread;
@@ -1178,16 +1177,17 @@ impl MarionetteConnection {
         let now = time::Instant::now();
 
         debug!(
             "Waiting {}s to connect to browser on {}:{}",
             timeout.as_secs(),
             self.host,
             self.port
         );
+
         loop {
             // immediately abort connection attempts if process disappears
             if let Some(ref mut runner) = *browser {
                 let exit_status = match runner.try_wait() {
                     Ok(Some(status)) => Some(
                         status
                             .code()
                             .map(|c| c.to_string())
@@ -1199,70 +1199,70 @@ impl MarionetteConnection {
                 if let Some(s) = exit_status {
                     return Err(WebDriverError::new(
                         ErrorStatus::UnknownError,
                         format!("Process unexpectedly closed with status {}", s),
                     ));
                 }
             }
 
-            match TcpStream::connect((&self.host[..], self.port)) {
-                Ok(stream) => {
+            let try_connect = || -> WebDriverResult<(TcpStream, MarionetteHandshake)> {
+                let mut stream = TcpStream::connect((&self.host[..], self.port))?;
+                let data = MarionetteConnection::handshake(&mut stream)?;
+
+                Ok((stream, data))
+            };
+
+            match try_connect() {
+                Ok((stream, data)) => {
+                    debug!(
+                        "Connection to Marionette established on {}:{}.",
+                        self.host, self.port,
+                    );
+
                     self.stream = Some(stream);
+                    self.session.application_type = Some(data.application_type);
+                    self.session.protocol = Some(data.protocol);
                     break;
                 }
                 Err(e) => {
                     if now.elapsed() < timeout {
                         thread::sleep(poll_interval);
                     } else {
                         return Err(WebDriverError::new(
-                            ErrorStatus::UnknownError,
-                            e.description().to_owned(),
+                            ErrorStatus::Timeout,
+                            e.to_string(),
                         ));
                     }
                 }
             }
         }
 
-        debug!(
-            "Connection established on {}:{}. Waiting for Marionette handshake",
-            self.host, self.port,
-        );
-
-        let data = self.handshake()?;
-        self.session.application_type = Some(data.application_type);
-        self.session.protocol = Some(data.protocol);
-
-        debug!("Connected to Marionette");
         Ok(())
     }
 
-    fn handshake(&mut self) -> WebDriverResult<MarionetteHandshake> {
-        let resp = (match self.stream.as_mut().unwrap().read_timeout() {
+    fn handshake(stream: &mut TcpStream) -> WebDriverResult<MarionetteHandshake> {
+        let resp = (match stream.read_timeout() {
             Ok(timeout) => {
                 // If platform supports changing the read timeout of the stream,
                 // use a short one only for the handshake with Marionette.
-                self.stream
-                    .as_mut()
-                    .unwrap()
-                    .set_read_timeout(Some(time::Duration::from_secs(10)))
+                stream
+                    .set_read_timeout(Some(time::Duration::from_millis(100)))
                     .ok();
-                let data = self.read_resp();
-                self.stream.as_mut().unwrap().set_read_timeout(timeout).ok();
+                let data = MarionetteConnection::read_resp(stream);
+                stream.set_read_timeout(timeout).ok();
 
                 data
             }
-            _ => self.read_resp(),
+            _ => MarionetteConnection::read_resp(stream),
         })
-        .or_else(|e| {
-            Err(WebDriverError::new(
-                ErrorStatus::UnknownError,
-                format!("Socket timeout reading Marionette handshake data: {}", e),
-            ))
-        })?;
+        .map_err(|e| WebDriverError::new(
+            ErrorStatus::UnknownError,
+            format!("Socket timeout reading Marionette handshake data: {}", e),
+        ))?;
 
         let data = serde_json::from_str::<MarionetteHandshake>(&resp)?;
 
         if data.application_type != "gecko" {
             return Err(WebDriverError::new(
                 ErrorStatus::UnknownError,
                 format!("Unrecognized application type {}", data.application_type),
             ));
@@ -1292,55 +1292,55 @@ impl MarionetteConnection {
         let enc_cmd = MarionetteCommand::from_webdriver_message(id, capabilities, msg)?;
         let resp_data = self.send(enc_cmd)?;
         let data: MarionetteResponse = serde_json::from_str(&resp_data)?;
 
         self.session.response(msg, data)
     }
 
     fn send(&mut self, data: String) -> WebDriverResult<String> {
-        match self.stream {
+        let stream = match self.stream {
             Some(ref mut stream) => {
                 if stream.write(&*data.as_bytes()).is_err() {
                     let mut err = WebDriverError::new(
                         ErrorStatus::UnknownError,
-                        "Failed to write response to stream",
+                        "Failed to write request to stream",
                     );
                     err.delete_session = true;
                     return Err(err);
                 }
+
+                stream
             }
             None => {
                 let mut err = WebDriverError::new(
                     ErrorStatus::UnknownError,
                     "Tried to write before opening stream",
                 );
                 err.delete_session = true;
                 return Err(err);
             }
-        }
+        };
 
-        match self.read_resp() {
+        match MarionetteConnection::read_resp(stream) {
             Ok(resp) => Ok(resp),
             Err(_) => {
                 let mut err = WebDriverError::new(
                     ErrorStatus::UnknownError,
                     "Failed to decode response from marionette",
                 );
                 err.delete_session = true;
                 Err(err)
             }
         }
     }
 
-    fn read_resp(&mut self) -> IoResult<String> {
+    fn read_resp(stream: &mut TcpStream) -> IoResult<String> {
         let mut bytes = 0usize;
 
-        // TODO(jgraham): Check before we unwrap?
-        let stream = self.stream.as_mut().unwrap();
         loop {
             let buf = &mut [0 as u8];
             let num_read = stream.read(buf)?;
             let byte = match num_read {
                 0 => {
                     return Err(IoError::new(
                         ErrorKind::Other,
                         "EOF reading marionette message",