Refactor CLI to use TunnelState from mullvad-types

6 files changed, 85 insertions, 182 deletions

diff --git a/mullvad-cli/src/cmds/connect.rs b/mullvad-cli/src/cmds/connect.rs
index 3a879fc6a8..0f470d3d2a 100644
--- a/mullvad-cli/src/cmds/connect.rs
+++ b/mullvad-cli/src/cmds/connect.rs
@@ -1,6 +1,6 @@
 use crate::{format, new_rpc_client, state, Command, Error, Result};
 use futures::StreamExt;
-use mullvad_management_interface::types::tunnel_state::State;
+use mullvad_types::states::TunnelState;
 
 pub struct Connect;
 
@@ -35,9 +35,9 @@ impl Command for Connect {
                 while let Some(state) = receiver.next().await {
                     let state = state?;
                     format::print_state(&state, false);
-                    match state.state.unwrap() {
-                        State::Connected(_) => return Ok(()),
-                        State::Error(_) => return Err(Error::CommandFailed("connect")),
+                    match state {
+                        TunnelState::Connected { .. } => return Ok(()),
+                        TunnelState::Error(_) => return Err(Error::CommandFailed("connect")),
                         _ => {}
                     }
                 }
diff --git a/mullvad-cli/src/cmds/disconnect.rs b/mullvad-cli/src/cmds/disconnect.rs
index 2864001656..4ea5722fe9 100644
--- a/mullvad-cli/src/cmds/disconnect.rs
+++ b/mullvad-cli/src/cmds/disconnect.rs
@@ -1,6 +1,5 @@
 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;
 
@@ -35,7 +34,7 @@ impl Command for Disconnect {
                 while let Some(state) = receiver.next().await {
                     let state = state?;
                     format::print_state(&state, false);
-                    if let Disconnected(_) = state.state.unwrap() {
+                    if state.is_disconnected() {
                         return Ok(());
                     }
                 }
diff --git a/mullvad-cli/src/cmds/reconnect.rs b/mullvad-cli/src/cmds/reconnect.rs
index 53ca666482..0a39d9f33d 100644
--- a/mullvad-cli/src/cmds/reconnect.rs
+++ b/mullvad-cli/src/cmds/reconnect.rs
@@ -1,6 +1,6 @@
 use crate::{format, new_rpc_client, state, Command, Error, Result};
 use futures::StreamExt;
-use mullvad_management_interface::types::tunnel_state::State;
+use mullvad_types::states::TunnelState;
 
 pub struct Reconnect;
 
@@ -35,9 +35,9 @@ impl Command for Reconnect {
                 while let Some(state) = receiver.next().await {
                     let state = state?;
                     format::print_state(&state, false);
-                    match state.state.unwrap() {
-                        State::Connected(_) => return Ok(()),
-                        State::Error(_) => return Err(Error::CommandFailed("reconnect")),
+                    match state {
+                        TunnelState::Connected { .. } => return Ok(()),
+                        TunnelState::Error { .. } => return Err(Error::CommandFailed("reconnect")),
                         _ => {}
                     }
                 }
diff --git a/mullvad-cli/src/cmds/status.rs b/mullvad-cli/src/cmds/status.rs
index 71f31f5fef..7d1ad09327 100644
--- a/mullvad-cli/src/cmds/status.rs
+++ b/mullvad-cli/src/cmds/status.rs
@@ -2,6 +2,7 @@ use crate::{format, new_rpc_client, Command, Error, Result};
 use mullvad_management_interface::{
     types::daemon_event::Event as EventType, ManagementServiceClient,
 };
+use mullvad_types::states::TunnelState;
 
 pub struct Status;
 
@@ -45,6 +46,7 @@ impl Command for Status {
         if debug {
             println!("Tunnel state: {:#?}", state);
         } else {
+            let state = TunnelState::try_from(state).expect("invalid tunnel state");
             format::print_state(&state, verbose);
         }
 
@@ -58,15 +60,17 @@ impl Command for Status {
             while let Some(event) = events.message().await? {
                 match event.event.unwrap() {
                     EventType::TunnelState(new_state) => {
+                        let new_state =
+                            TunnelState::try_from(new_state).expect("invalid tunnel state");
+
                         if debug {
                             println!("New tunnel state: {:#?}", new_state);
                         } else {
                             format::print_state(&new_state, verbose);
                         }
 
-                        use mullvad_management_interface::types::tunnel_state::State::*;
-                        match new_state.state.unwrap() {
-                            Connected(..) | Disconnected(..) => {
+                        match new_state {
+                            TunnelState::Connected { .. } | TunnelState::Disconnected => {
                                 if show_full_location {
                                     print_location(&mut rpc).await?;
                                 }
diff --git a/mullvad-cli/src/format.rs b/mullvad-cli/src/format.rs
index 5a2fd8b654..693f5c9c7e 100644
--- a/mullvad-cli/src/format.rs
+++ b/mullvad-cli/src/format.rs
@@ -1,89 +1,91 @@
-use mullvad_management_interface::types::{
-    error_state::{
-        firewall_policy_error::ErrorType as FirewallPolicyErrorType, Cause as ErrorStateCause,
-        FirewallPolicyError, GenerationError,
-    },
-    tunnel_state,
-    tunnel_state::State::*,
-    ErrorState, ObfuscationType, ProxyType, TransportProtocol, TunnelState, TunnelStateRelayInfo,
-    TunnelType,
+use mullvad_types::{location::GeoIpLocation, states::TunnelState};
+use talpid_types::{
+    net::{Endpoint, TunnelEndpoint},
+    tunnel::ErrorState,
 };
-use mullvad_types::auth_failed::AuthFailed;
-use std::borrow::Cow;
 
 pub fn print_state(state: &TunnelState, verbose: bool) {
-    match state.state.as_ref().unwrap() {
-        Error(error) => print_error_state(error.error_state.as_ref().unwrap()),
-        Connected(tunnel_state::Connected { relay_info }) => {
+    use TunnelState::*;
+
+    match state {
+        Error(error) => print_error_state(error),
+        Connected { endpoint, location } => {
             println!(
                 "Connected to {}",
-                format_relay_connection(relay_info.as_ref().unwrap(), verbose)
+                format_relay_connection(endpoint, location.as_ref(), verbose)
             );
         }
-        Connecting(tunnel_state::Connecting { relay_info }) => {
+        Connecting { endpoint, location } => {
             let ellipsis = if !verbose { "..." } else { "" };
             println!(
                 "Connecting to {}{ellipsis}",
-                format_relay_connection(relay_info.as_ref().unwrap(), verbose)
+                format_relay_connection(endpoint, location.as_ref(), verbose)
             );
         }
-        Disconnected(_) => println!("Disconnected"),
+        Disconnected => println!("Disconnected"),
         Disconnecting(_) => println!("Disconnecting..."),
     }
 }
 
-fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) -> String {
-    let endpoint = relay_info.tunnel_endpoint.as_ref().unwrap();
-    let location = &relay_info.location.as_ref();
-
+fn format_relay_connection(
+    endpoint: &TunnelEndpoint,
+    location: Option<&GeoIpLocation>,
+    verbose: bool,
+) -> String {
     let prefix_separator = if verbose { "\n\t" } else { " " };
     let mut obfuscator_overlaps = false;
 
     let exit_endpoint = {
-        let mut address = Cow::Borrowed(endpoint.address.as_str());
-        let mut protocol = endpoint.protocol;
-        if let Some(obfuscator) = endpoint.obfuscation.as_ref() {
+        let mut exit_endpoint = &endpoint.endpoint;
+        if let Some(obfuscator) = &endpoint.obfuscation {
             if location
                 .map(|l| l.hostname == l.obfuscator_hostname)
                 .unwrap_or(false)
             {
                 obfuscator_overlaps = true;
-                address = Cow::Owned(format!("{}:{}", obfuscator.address, obfuscator.port));
-                protocol = obfuscator.protocol;
+                exit_endpoint = &obfuscator.endpoint;
             }
         };
 
         let exit = format_endpoint(
-            location.map(|l| l.hostname.as_str()),
-            protocol,
-            &*address,
+            location.and_then(|l| l.hostname.as_deref()),
+            exit_endpoint,
             verbose,
         );
-        if let Some(location) = location {
-            format!("{exit} in {}, {}", &location.city, &location.country)
-        } else {
-            exit
+        match location {
+            Some(GeoIpLocation {
+                country,
+                city: Some(city),
+                ..
+            }) => {
+                format!("{exit} in {}, {}", city, country)
+            }
+            Some(GeoIpLocation {
+                country,
+                city: None,
+                ..
+            }) => {
+                format!("{exit} in {}", country)
+            }
+            None => exit,
         }
     };
 
     let first_hop = endpoint.entry_endpoint.as_ref().map(|entry| {
-        let mut address = entry.address.as_str();
-        let mut protocol = entry.protocol;
-        if let Some(obfuscator) = endpoint.obfuscation.as_ref() {
-            obfuscator_overlaps = true;
+        let mut entry_endpoint = entry;
+        if let Some(obfuscator) = &endpoint.obfuscation {
             if location
-                .map(|l| l.hostname == l.obfuscator_hostname)
+                .map(|l| l.entry_hostname == l.obfuscator_hostname)
                 .unwrap_or(false)
             {
-                address = &obfuscator.address;
-                protocol = obfuscator.protocol;
+                obfuscator_overlaps = true;
+                entry_endpoint = &obfuscator.endpoint;
             }
         };
 
         let endpoint = format_endpoint(
-            location.map(|l| l.entry_hostname.as_str()),
-            protocol,
-            address,
+            location.and_then(|l| l.entry_hostname.as_deref()),
+            entry_endpoint,
             verbose,
         );
         format!("{prefix_separator}via {endpoint}")
@@ -92,9 +94,8 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
     let obfuscator = endpoint.obfuscation.as_ref().map(|obfuscator| {
         if !obfuscator_overlaps {
             let endpoint_str = format_endpoint(
-                location.map(|l| l.obfuscator_hostname.as_str()),
-                obfuscator.protocol,
-                obfuscator.address.as_str(),
+                location.and_then(|l| l.obfuscator_hostname.as_deref()),
+                &obfuscator.endpoint,
                 verbose,
             );
             format!("{prefix_separator}obfuscated via {endpoint_str}")
@@ -105,22 +106,15 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
 
     let bridge = endpoint.proxy.as_ref().map(|proxy| {
         let proxy_endpoint = format_endpoint(
-            location.map(|l| l.bridge_hostname.as_str()),
-            proxy.protocol,
-            proxy.address.as_str(),
+            location.and_then(|l| l.bridge_hostname.as_deref()),
+            &proxy.endpoint,
             verbose,
         );
 
         format!("{prefix_separator}via {proxy_endpoint}")
     });
     let tunnel_type = if verbose {
-        let tunnel = match TunnelType::from_i32(endpoint.tunnel_type).expect("invalid tunnel type")
-        {
-            TunnelType::Wireguard => "WireGuard",
-            TunnelType::Openvpn => "OpenVPN",
-        };
-
-        format!("\nTunnel type: {tunnel}")
+        format!("\nTunnel type: {}", endpoint.tunnel_type)
     } else {
         String::new()
     };
@@ -135,16 +129,11 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
     let mut bridge_type = String::new();
     let mut obfuscator_type = String::new();
     if verbose {
-        if let Some(bridge) = endpoint.proxy.as_ref() {
-            let bridge = match ProxyType::from_i32(bridge.proxy_type).expect("invalid proxy type") {
-                ProxyType::Shadowsocks => "Shadowsocks",
-                ProxyType::Custom => "custom bridge",
-            };
-            bridge_type = format!("\nBridge type: {}", bridge);
+        if let Some(bridge) = &endpoint.proxy {
+            bridge_type = format!("\nBridge type: {}", bridge.proxy_type);
         }
-        if let Some(obfuscator) = endpoint.obfuscation.as_ref() {
-            let obfuscation = convert_obfuscator_type(obfuscator.obfuscation_type);
-            obfuscator_type = format!("\nObfuscator: {obfuscation}");
+        if let Some(obfuscator) = &endpoint.obfuscation {
+            obfuscator_type = format!("\nObfuscator: {}", obfuscator.obfuscation_type);
         }
     }
 
@@ -156,118 +145,27 @@ fn format_relay_connection(relay_info: &TunnelStateRelayInfo, verbose: bool) ->
     )
 }
 
-fn convert_obfuscator_type(obfuscator: i32) -> &'static str {
-    match ObfuscationType::from_i32(obfuscator).expect("invalid obfuscator type") {
-        ObfuscationType::Udp2tcp => "Udp2Tcp",
-    }
-}
-
-fn format_endpoint(
-    hostname: Option<&str>,
-    protocol_enum: i32,
-    addr: &str,
-    verbose: bool,
-) -> String {
-    let protocol = format_protocol(
-        TransportProtocol::from_i32(protocol_enum).expect("invalid transport protocol"),
-    );
-
+fn format_endpoint(hostname: Option<&str>, endpoint: &Endpoint, verbose: bool) -> String {
     match (hostname, verbose) {
-        (Some(hostname), true) => format!("{hostname} ({addr}/{protocol})"),
-        (None, true) => format!("{addr}/{protocol}"),
+        (Some(hostname), true) => format!("{hostname} ({endpoint})"),
+        (None, true) => endpoint.to_string(),
         (Some(hostname), false) => hostname.to_string(),
-        (None, false) => addr.to_string(),
+        (None, false) => endpoint.address.to_string(),
     }
 }
 
 fn print_error_state(error_state: &ErrorState) {
-    if error_state.blocking_error.is_some() {
+    if error_state.block_failure().is_some() {
         eprintln!("Mullvad daemon failed to setup firewall rules!");
         eprintln!("Daemon 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())
-            );
-        }
+    match error_state.cause() {
         #[cfg(target_os = "linux")]
-        Some(ErrorStateCause::SetFirewallPolicyError) => {
-            println!("Blocked: {}", error_state_to_string(error_state));
+        cause @ talpid_types::tunnel::ErrorStateCause::SetFirewallPolicyError(_) => {
+            println!("Blocked: {}", cause);
             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",
-        #[cfg(target_os = "android")]
-        VpnPermissionDenied => "The Android VPN permission was denied when creating the tunnel",
-        #[cfg(target_os = "windows")]
-        SplitTunnelError => "The split tunneling module reported an error",
-        #[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",
+        cause => println!("Blocked: {}", cause),
     }
 }
diff --git a/mullvad-cli/src/state.rs b/mullvad-cli/src/state.rs
index f237689a3b..7b3dfdc955 100644
--- a/mullvad-cli/src/state.rs
+++ b/mullvad-cli/src/state.rs
@@ -4,9 +4,9 @@ use futures::{
     SinkExt,
 };
 use mullvad_management_interface::{
-    types::{daemon_event::Event as EventType, TunnelState},
-    ManagementServiceClient,
+    types::daemon_event::Event as EventType, ManagementServiceClient,
 };
+use mullvad_types::states::TunnelState;
 
 // 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.
@@ -19,7 +19,9 @@ pub fn state_listen(mut rpc: ManagementServiceClient) -> Receiver<Result<TunnelS
                 loop {
                     let forward = match events.message().await {
                         Ok(Some(event)) => match event.event.unwrap() {
-                            EventType::TunnelState(new_state) => Ok(new_state),
+                            EventType::TunnelState(new_state) => {
+                                Ok(TunnelState::try_from(new_state).expect("invalid tunnel state"))
+                            }
                             _ => continue,
                         },
                         Ok(None) => break,