Merge branch 'add-wait-flag-to-cli' into master

7 files changed, 308 insertions, 181 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/cmds/disconnect.rs b/mullvad-cli/src/cmds/disconnect.rs
index f976b96fab..96d5843d20 100644
--- a/mullvad-cli/src/cmds/disconnect.rs
+++ b/mullvad-cli/src/cmds/disconnect.rs
@@ -1,4 +1,6 @@
-use crate::{new_rpc_client, Command, Result};
+use crate::{format, new_rpc_client, state, Command, Error, Result};
+use futures::StreamExt;
+use mullvad_management_interface::types::tunnel_state::State::Disconnected;
 
 pub struct Disconnect;
 
@@ -11,11 +13,37 @@ impl Command for Disconnect {
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
         clap::SubCommand::with_name(self.name())
             .about("Command the client to disconnect the VPN tunnel")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until disconnected before exiting"),
+            )
     }
 
-    async fn run(&self, _: &clap::ArgMatches<'_>) -> Result<()> {
+    async fn run(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        rpc.disconnect_tunnel(()).await?;
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.disconnect_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() {
+                        Disconnected(_) => return Ok(()),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
+        }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/cmds/reconnect.rs b/mullvad-cli/src/cmds/reconnect.rs
index ecc0b089c0..87de63144f 100644
--- a/mullvad-cli/src/cmds/reconnect.rs
+++ b/mullvad-cli/src/cmds/reconnect.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 Reconnect;
 
@@ -10,14 +11,40 @@ impl Command for Reconnect {
     }
 
     fn clap_subcommand(&self) -> clap::App<'static, 'static> {
-        clap::SubCommand::with_name(self.name()).about("Command the client to reconnect")
+        clap::SubCommand::with_name(self.name())
+            .about("Command the client to reconnect")
+            .arg(
+                clap::Arg::with_name("wait")
+                    .long("wait")
+                    .short("w")
+                    .help("Wait until reconnected 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.reconnect_tunnel(()).await {
-            eprintln!("{}", e.display_chain());
+
+        let receiver_option = if matches.is_present("wait") {
+            Some(state::state_listen(rpc.clone()))
+        } else {
+            None
+        };
+
+        if rpc.reconnect_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("reconnect")),
+                        _ => {}
+                    }
+                }
+                return Err(Error::StatusListenerFailed);
+            }
         }
+
         Ok(())
     }
 }
diff --git a/mullvad-cli/src/cmds/status.rs b/mullvad-cli/src/cmds/status.rs
index b0c4d40ad0..76a4ed2e67 100644
--- a/mullvad-cli/src/cmds/status.rs
+++ b/mullvad-cli/src/cmds/status.rs
@@ -1,17 +1,7 @@
-use crate::{format::print_keygen_event, new_rpc_client, Command, Error, Result};
+use crate::{format, format::print_keygen_event, new_rpc_client, Command, Error, Result};
 use mullvad_management_interface::{
-    types::{
-        daemon_event::Event as EventType,
-        error_state::{
-            firewall_policy_error::ErrorType as FirewallPolicyErrorType, Cause as ErrorStateCause,
-            FirewallPolicyError, GenerationError,
-        },
-        ErrorState, ProxyType, TransportProtocol, TunnelEndpoint, TunnelState, TunnelType,
-    },
-    ManagementServiceClient,
+    types::daemon_event::Event as EventType, ManagementServiceClient,
 };
-use mullvad_types::auth_failed::AuthFailed;
-use std::fmt::Write;
 
 pub struct Status;
 
@@ -45,7 +35,7 @@ impl Command for Status {
         let mut rpc = new_rpc_client().await?;
         let state = rpc.get_tunnel_state(()).await?.into_inner();
 
-        print_state(&state);
+        format::print_state(&state);
         if matches.is_present("location") {
             print_location(&mut rpc).await?;
         }
@@ -58,7 +48,7 @@ impl Command for Status {
             while let Some(event) = events.message().await? {
                 match event.event.unwrap() {
                     EventType::TunnelState(new_state) => {
-                        print_state(&new_state);
+                        format::print_state(&new_state);
                         use mullvad_management_interface::types::tunnel_state::State::*;
                         match new_state.state.unwrap() {
                             Connected(..) | Disconnected(..) => {
@@ -98,151 +88,6 @@ impl Command for Status {
     }
 }
 
-fn print_state(state: &TunnelState) {
-    use mullvad_management_interface::types::{tunnel_state, tunnel_state::State::*};
-
-    print!("Tunnel status: ");
-    match state.state.as_ref().unwrap() {
-        Error(error) => print_error_state(error.error_state.as_ref().unwrap()),
-        Connected(tunnel_state::Connected { relay_info }) => {
-            let endpoint = relay_info
-                .as_ref()
-                .unwrap()
-                .tunnel_endpoint
-                .as_ref()
-                .unwrap();
-            println!("Connected to {}", format_endpoint(&endpoint));
-        }
-        Connecting(tunnel_state::Connecting { relay_info }) => {
-            let endpoint = relay_info
-                .as_ref()
-                .unwrap()
-                .tunnel_endpoint
-                .as_ref()
-                .unwrap();
-            println!("Connecting to {}...", format_endpoint(&endpoint));
-        }
-        Disconnected(_) => println!("Disconnected"),
-        Disconnecting(_) => println!("Disconnecting..."),
-    }
-}
-
-fn format_endpoint(endpoint: &TunnelEndpoint) -> String {
-    let mut out = format!(
-        "{} {} over {}",
-        match TunnelType::from_i32(endpoint.tunnel_type).expect("unknown tunnel protocol") {
-            TunnelType::Wireguard => "WireGuard",
-            TunnelType::Openvpn => "OpenVPN",
-        },
-        endpoint.address,
-        format_protocol(
-            TransportProtocol::from_i32(endpoint.protocol).expect("unknown transport protocol")
-        ),
-    );
-
-    if let Some(ref proxy) = endpoint.proxy {
-        write!(
-            &mut out,
-            " via {} {} over {}",
-            match ProxyType::from_i32(proxy.proxy_type).expect("unknown proxy type") {
-                ProxyType::Shadowsocks => "Shadowsocks",
-                ProxyType::Custom => "custom bridge",
-            },
-            proxy.address,
-            format_protocol(
-                TransportProtocol::from_i32(proxy.protocol).expect("unknown transport protocol")
-            ),
-        )
-        .unwrap();
-    }
-
-    out
-}
-
-fn print_error_state(error_state: &ErrorState) {
-    if error_state.blocking_error.is_some() {
-        eprintln!("Mullvad daemon failed to setup firewall rules!");
-        eprintln!("Deamon cannot block traffic from flowing, non-local traffic will leak");
-    }
-
-    match ErrorStateCause::from_i32(error_state.cause) {
-        Some(ErrorStateCause::AuthFailed) => {
-            println!(
-                "Blocked: {}",
-                AuthFailed::from(error_state.auth_fail_reason.as_ref())
-            );
-        }
-        #[cfg(target_os = "linux")]
-        Some(ErrorStateCause::SetFirewallPolicyError) => {
-            println!("Blocked: {}", error_state_to_string(error_state));
-            println!("Your kernel might be terribly out of date or missing nftables");
-        }
-        _ => println!("Blocked: {}", error_state_to_string(error_state)),
-    }
-}
-
-fn error_state_to_string(error_state: &ErrorState) -> String {
-    use ErrorStateCause::*;
-
-    let error_str = match ErrorStateCause::from_i32(error_state.cause).expect("unknown error cause")
-    {
-        AuthFailed => {
-            return if error_state.auth_fail_reason.is_empty() {
-                "Authentication with remote server failed".to_string()
-            } else {
-                format!(
-                    "Authentication with remote server failed: {}",
-                    error_state.auth_fail_reason
-                )
-            };
-        }
-        Ipv6Unavailable => "Failed to configure IPv6 because it's disabled in the platform",
-        SetFirewallPolicyError => {
-            return policy_error_to_string(error_state.policy_error.as_ref().unwrap())
-        }
-        SetDnsError => "Failed to set system DNS server",
-        StartTunnelError => "Failed to start connection to remote server",
-        TunnelParameterError => {
-            return format!(
-                "Failure to generate tunnel parameters: {}",
-                tunnel_parameter_error_to_string(error_state.parameter_error)
-            );
-        }
-        IsOffline => "This device is offline, no tunnels can be established",
-        TapAdapterProblem => "A problem with the TAP adapter has been detected",
-        #[cfg(target_os = "android")]
-        VpnPermissionDenied => "The Android VPN permission was denied when creating the tunnel",
-        #[cfg(not(target_os = "android"))]
-        _ => unreachable!("unknown error cause"),
-    };
-
-    error_str.to_string()
-}
-
-fn tunnel_parameter_error_to_string(parameter_error: i32) -> &'static str {
-    match GenerationError::from_i32(parameter_error).expect("unknown generation error") {
-        GenerationError::NoMatchingRelay => "Failure to select a matching tunnel relay",
-        GenerationError::NoMatchingBridgeRelay => "Failure to select a matching bridge relay",
-        GenerationError::NoWireguardKey => "No wireguard key available",
-        GenerationError::CustomTunnelHostResolutionError => {
-            "Can't resolve hostname for custom tunnel host"
-        }
-    }
-}
-
-fn policy_error_to_string(policy_error: &FirewallPolicyError) -> String {
-    let cause = match FirewallPolicyErrorType::from_i32(policy_error.r#type)
-        .expect("unknown policy error")
-    {
-        FirewallPolicyErrorType::Generic => return "Failed to set firewall policy".to_string(),
-        FirewallPolicyErrorType::Locked => format!(
-            "An application prevented the firewall policy from being set: {} (pid {})",
-            policy_error.lock_name, policy_error.lock_pid
-        ),
-    };
-    format!("Failed to set firewall policy: {}", cause)
-}
-
 async fn print_location(rpc: &mut ManagementServiceClient) -> Result<()> {
     let location = rpc.get_current_location(()).await;
     let location = match location {
@@ -278,10 +123,3 @@ async fn print_location(rpc: &mut ManagementServiceClient) -> Result<()> {
     );
     Ok(())
 }
-
-fn format_protocol(protocol: TransportProtocol) -> &'static str {
-    match protocol {
-        TransportProtocol::Udp => "UDP",
-        TransportProtocol::Tcp => "TCP",
-    }
-}
diff --git a/mullvad-cli/src/format.rs b/mullvad-cli/src/format.rs
index 4f69d05317..92f6076e66 100644
--- a/mullvad-cli/src/format.rs
+++ b/mullvad-cli/src/format.rs
@@ -1,4 +1,14 @@
-use mullvad_management_interface::types::KeygenEvent;
+use mullvad_management_interface::types::{
+    error_state::{
+        firewall_policy_error::ErrorType as FirewallPolicyErrorType, Cause as ErrorStateCause,
+        FirewallPolicyError, GenerationError,
+    },
+    tunnel_state,
+    tunnel_state::State::*,
+    ErrorState, KeygenEvent, ProxyType, TransportProtocol, TunnelEndpoint, TunnelState, TunnelType,
+};
+use mullvad_types::auth_failed::AuthFailed;
+use std::fmt::Write;
 
 pub fn print_keygen_event(key_event: &KeygenEvent) {
     use mullvad_management_interface::types::keygen_event::KeygenEvent as EventType;
@@ -18,3 +28,153 @@ pub fn print_keygen_event(key_event: &KeygenEvent) {
         }
     }
 }
+
+pub fn print_state(state: &TunnelState) {
+    print!("Tunnel status: ");
+    match state.state.as_ref().unwrap() {
+        Error(error) => print_error_state(error.error_state.as_ref().unwrap()),
+        Connected(tunnel_state::Connected { relay_info }) => {
+            let endpoint = relay_info
+                .as_ref()
+                .unwrap()
+                .tunnel_endpoint
+                .as_ref()
+                .unwrap();
+            println!("Connected to {}", format_endpoint(&endpoint));
+        }
+        Connecting(tunnel_state::Connecting { relay_info }) => {
+            let endpoint = relay_info
+                .as_ref()
+                .unwrap()
+                .tunnel_endpoint
+                .as_ref()
+                .unwrap();
+            println!("Connecting to {}...", format_endpoint(&endpoint));
+        }
+        Disconnected(_) => println!("Disconnected"),
+        Disconnecting(_) => println!("Disconnecting..."),
+    }
+}
+
+fn format_endpoint(endpoint: &TunnelEndpoint) -> String {
+    let mut out = format!(
+        "{} {} over {}",
+        match TunnelType::from_i32(endpoint.tunnel_type).expect("unknown tunnel protocol") {
+            TunnelType::Wireguard => "WireGuard",
+            TunnelType::Openvpn => "OpenVPN",
+        },
+        endpoint.address,
+        format_protocol(
+            TransportProtocol::from_i32(endpoint.protocol).expect("unknown transport protocol")
+        ),
+    );
+
+    if let Some(ref proxy) = endpoint.proxy {
+        write!(
+            &mut out,
+            " via {} {} over {}",
+            match ProxyType::from_i32(proxy.proxy_type).expect("unknown proxy type") {
+                ProxyType::Shadowsocks => "Shadowsocks",
+                ProxyType::Custom => "custom bridge",
+            },
+            proxy.address,
+            format_protocol(
+                TransportProtocol::from_i32(proxy.protocol).expect("unknown transport protocol")
+            ),
+        )
+        .unwrap();
+    }
+
+    out
+}
+
+fn print_error_state(error_state: &ErrorState) {
+    if error_state.blocking_error.is_some() {
+        eprintln!("Mullvad daemon failed to setup firewall rules!");
+        eprintln!("Deamon cannot block traffic from flowing, non-local traffic will leak");
+    }
+
+    match ErrorStateCause::from_i32(error_state.cause) {
+        Some(ErrorStateCause::AuthFailed) => {
+            println!(
+                "Blocked: {}",
+                AuthFailed::from(error_state.auth_fail_reason.as_ref())
+            );
+        }
+        #[cfg(target_os = "linux")]
+        Some(ErrorStateCause::SetFirewallPolicyError) => {
+            println!("Blocked: {}", error_state_to_string(error_state));
+            println!("Your kernel might be terribly out of date or missing nftables");
+        }
+        _ => println!("Blocked: {}", error_state_to_string(error_state)),
+    }
+}
+
+fn error_state_to_string(error_state: &ErrorState) -> String {
+    use ErrorStateCause::*;
+
+    let error_str = match ErrorStateCause::from_i32(error_state.cause).expect("unknown error cause")
+    {
+        AuthFailed => {
+            return if error_state.auth_fail_reason.is_empty() {
+                "Authentication with remote server failed".to_string()
+            } else {
+                format!(
+                    "Authentication with remote server failed: {}",
+                    error_state.auth_fail_reason
+                )
+            };
+        }
+        Ipv6Unavailable => "Failed to configure IPv6 because it's disabled in the platform",
+        SetFirewallPolicyError => {
+            return policy_error_to_string(error_state.policy_error.as_ref().unwrap())
+        }
+        SetDnsError => "Failed to set system DNS server",
+        StartTunnelError => "Failed to start connection to remote server",
+        TunnelParameterError => {
+            return format!(
+                "Failure to generate tunnel parameters: {}",
+                tunnel_parameter_error_to_string(error_state.parameter_error)
+            );
+        }
+        IsOffline => "This device is offline, no tunnels can be established",
+        TapAdapterProblem => "A problem with the TAP adapter has been detected",
+        #[cfg(target_os = "android")]
+        VpnPermissionDenied => "The Android VPN permission was denied when creating the tunnel",
+        #[cfg(not(target_os = "android"))]
+        _ => unreachable!("unknown error cause"),
+    };
+
+    error_str.to_string()
+}
+
+fn tunnel_parameter_error_to_string(parameter_error: i32) -> &'static str {
+    match GenerationError::from_i32(parameter_error).expect("unknown generation error") {
+        GenerationError::NoMatchingRelay => "Failure to select a matching tunnel relay",
+        GenerationError::NoMatchingBridgeRelay => "Failure to select a matching bridge relay",
+        GenerationError::NoWireguardKey => "No wireguard key available",
+        GenerationError::CustomTunnelHostResolutionError => {
+            "Can't resolve hostname for custom tunnel host"
+        }
+    }
+}
+
+fn policy_error_to_string(policy_error: &FirewallPolicyError) -> String {
+    let cause = match FirewallPolicyErrorType::from_i32(policy_error.r#type)
+        .expect("unknown policy error")
+    {
+        FirewallPolicyErrorType::Generic => return "Failed to set firewall policy".to_string(),
+        FirewallPolicyErrorType::Locked => format!(
+            "An application prevented the firewall policy from being set: {} (pid {})",
+            policy_error.lock_name, policy_error.lock_pid
+        ),
+    };
+    format!("Failed to set firewall policy: {}", cause)
+}
+
+fn format_protocol(protocol: TransportProtocol) -> &'static str {
+    match protocol {
+        TransportProtocol::Udp => "UDP",
+        TransportProtocol::Tcp => "TCP",
+    }
+}
diff --git a/mullvad-cli/src/main.rs b/mullvad-cli/src/main.rs
index 356a50ce69..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"));
@@ -30,6 +31,12 @@ pub enum Error {
     /// The given command is not correct in some way
     #[error(display = "Invalid command: {}", _0)]
     InvalidCommand(&'static str),
+
+    #[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
+}