Add --wait argument to connect command

3 files changed, 76 insertions, 5 deletions

diff --git a/mullvad-cli/src/cmds/connect.rs b/mullvad-cli/src/cmds/connect.rs
index 6fc7072580..c93ed8c2e2 100644
--- a/mullvad-cli/src/cmds/connect.rs
+++ b/mullvad-cli/src/cmds/connect.rs
@@ -1,5 +1,6 @@
-use crate::{new_rpc_client, Command, Result};
-use talpid_types::ErrorExt;
+use crate::{format, new_rpc_client, state, Command, Error, Result};
+use futures::StreamExt;
+use mullvad_management_interface::types::tunnel_state::State;
 
 pub struct Connect;
 
@@ -12,13 +13,38 @@ impl Command for Connect {
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
         clap::SubCommand::with_name(self.name())
             .about("Command the client to start establishing a VPN tunnel")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until connected before exiting"),
+            )
     }
 
-    async fn run(&self, _: &clap::ArgMatches<'_>) -> Result<()> {
+    async fn run(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        if let Err(e) = rpc.connect_tunnel(()).await {
-            eprintln!("{}", e.display_chain());
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.connect_tunnel(()).await?.into_inner() {
+            if let Some(mut receiver) = receiver_option {
+                while let Some(state) = receiver.next().await {
+                    let state = state?;
+                    format::print_state(&state);
+                    match state.state.unwrap() {
+                        State::Connected(_) => return Ok(()),
+                        State::Error(_) => return Err(Error::CommandFailed("connect")),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
         }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/main.rs b/mullvad-cli/src/main.rs
index 206e6386a9..17019c0ab2 100644
--- a/mullvad-cli/src/main.rs
+++ b/mullvad-cli/src/main.rs
@@ -10,6 +10,7 @@ pub use mullvad_management_interface::{self, new_rpc_client};
 mod cmds;
 mod format;
 mod location;
+mod state;
 
 pub const BIN_NAME: &str = "mullvad";
 pub const PRODUCT_VERSION: &str = include_str!(concat!(env!("OUT_DIR"), "/product-version.txt"));
@@ -33,6 +34,9 @@ pub enum Error {
 
     #[error(display = "Command failed: {}", _0)]
     CommandFailed(&'static str),
+
+    #[error(display = "Failed to listen for status updates")]
+    StatusListenerFailed,
 }
 
 #[tokio::main]
diff --git a/mullvad-cli/src/state.rs b/mullvad-cli/src/state.rs
new file mode 100644
index 0000000000..9890218545
--- /dev/null
+++ b/mullvad-cli/src/state.rs
@@ -0,0 +1,41 @@
+use crate::{Error, Result};
+use futures::{
+    channel::{mpsc, mpsc::Receiver},
+    SinkExt,
+};
+use mullvad_management_interface::{
+    types::{daemon_event::Event as EventType, TunnelState},
+    ManagementServiceClient,
+};
+
+// Spawns a new task that listens for tunnel state changes and forwards it through the returned
+// channel. Panics if called from outside of the Tokio runtime.
+pub fn state_listen(mut rpc: ManagementServiceClient) -> Receiver<Result<TunnelState>> {
+    let (mut sender, receiver) = mpsc::channel::<Result<TunnelState>>(1);
+    tokio::spawn(async move {
+        match rpc.events_listen(()).await {
+            Ok(events) => {
+                let mut events = events.into_inner();
+                loop {
+                    let forward = match events.message().await {
+                        Ok(Some(event)) => match event.event.unwrap() {
+                            EventType::TunnelState(new_state) => Ok(new_state),
+                            _ => continue,
+                        },
+                        Ok(None) => break,
+                        Err(status) => Err(Error::GrpcClientError(status)),
+                    };
+
+                    if let Err(_) = sender.send(forward).await {
+                        break;
+                    }
+                }
+            }
+            Err(status) => {
+                let _ = sender.send(Err(Error::GrpcClientError(status))).await;
+            }
+        }
+    });
+
+    receiver
+}