Update RPCs to handle device states and event variants

7 files changed, 182 insertions, 102 deletions

diff --git a/mullvad-cli/src/cmds/account.rs b/mullvad-cli/src/cmds/account.rs
index d42a46388c..c6c723c5ce 100644
--- a/mullvad-cli/src/cmds/account.rs
+++ b/mullvad-cli/src/cmds/account.rs
@@ -7,7 +7,8 @@ use mullvad_management_interface::{
 use mullvad_types::{account::AccountToken, device::Device};
 use std::io::{self, Write};
 
-const NOT_LOGGED_IN_ERROR: &str = "Not logged in to any account";
+const NOT_LOGGED_IN_MESSAGE: &str = "Not logged in on any account";
+const REVOKED_MESSAGE: &str = "The current device has been revoked";
 const DEVICE_NOT_FOUND_ERROR: &str = "There is no such device";
 const INVALID_ACCOUNT_ERROR: &str = "The account does not exist";
 const TOO_MANY_DEVICES_ERROR: &str =
@@ -138,37 +139,45 @@ impl Account {
 
         let _ = rpc.update_device(()).await;
 
-        let device = rpc
+        let state = rpc
             .get_device(())
             .await
-            .map_err(|error| match error.code() {
-                Code::NotFound => Error::Other(NOT_LOGGED_IN_ERROR),
-                _other => map_device_error(error),
-            })?
+            .map_err(map_device_error)?
             .into_inner();
-        if !device.account_token.is_empty() {
-            println!("Mullvad account: {}", device.account_token);
-            let inner_device = Device::try_from(device.device.unwrap()).unwrap();
-            println!("Device name    : {}", inner_device.pretty_name());
-            if verbose {
-                println!("Device id      : {}", inner_device.id);
-                println!("Device pubkey  : {}", inner_device.pubkey);
-                for port in inner_device.ports {
-                    println!("Device port    : {}", port);
+
+        use types::device_state::State;
+
+        match State::from_i32(state.state).unwrap() {
+            State::LoggedIn => {
+                let device = state.device.expect("Device must be provided if logged in");
+                println!("Mullvad account: {}", device.account_token);
+                let inner_device = Device::try_from(device.device.unwrap()).unwrap();
+                println!("Device name    : {}", inner_device.pretty_name());
+                if verbose {
+                    println!("Device id      : {}", inner_device.id);
+                    println!("Device pubkey  : {}", inner_device.pubkey);
+                    for port in inner_device.ports {
+                        println!("Device port    : {}", port);
+                    }
                 }
+                let expiry = rpc
+                    .get_account_data(device.account_token)
+                    .await
+                    .map_err(|error| Error::RpcFailedExt("Failed to fetch account data", error))?
+                    .into_inner();
+                println!(
+                    "Expires at     : {}",
+                    Self::format_expiry(&expiry.expiry.unwrap())
+                );
+            }
+            State::LoggedOut => {
+                println!("{}", NOT_LOGGED_IN_MESSAGE);
+            }
+            State::Revoked => {
+                println!("{}", REVOKED_MESSAGE);
             }
-            let expiry = rpc
-                .get_account_data(device.account_token)
-                .await
-                .map_err(|error| Error::RpcFailedExt("Failed to fetch account data", error))?
-                .into_inner();
-            println!(
-                "Expires at     : {}",
-                Self::format_expiry(&expiry.expiry.unwrap())
-            );
-        } else {
-            println!("No account configured");
         }
+
         Ok(())
     }
 
@@ -241,17 +250,15 @@ impl Account {
         match matches.value_of("account").map(str::to_string) {
             Some(token) => Ok(token),
             None => {
-                let device = rpc
+                let state = rpc
                     .get_device(())
                     .await
-                    .map_err(|error| match error.code() {
-                        mullvad_management_interface::Code::NotFound => {
-                            Error::Other("Log in or specify an account")
-                        }
-                        _ => Error::RpcFailedExt("Failed to obtain device", error),
-                    })?
+                    .map_err(|error| Error::RpcFailedExt("Failed to obtain device", error))?
                     .into_inner();
-                Ok(device.account_token)
+                if state.state != types::device_state::State::LoggedIn as i32 {
+                    return Err(Error::Other("Log in or specify an account"));
+                }
+                Ok(state.device.unwrap().account_token)
             }
         }
     }
diff --git a/mullvad-daemon/src/device/mod.rs b/mullvad-daemon/src/device/mod.rs
index 505923523e..84fa1afa50 100644
--- a/mullvad-daemon/src/device/mod.rs
+++ b/mullvad-daemon/src/device/mod.rs
@@ -7,7 +7,10 @@ use futures::{
 use mullvad_api::rest;
 use mullvad_types::{
     account::AccountToken,
-    device::{AccountAndDevice, Device, DeviceEvent, DeviceId, DeviceName, DevicePort},
+    device::{
+        AccountAndDevice, Device, DeviceEvent, DeviceEventCause, DeviceId, DeviceName, DevicePort,
+        DeviceState,
+    },
     wireguard::{self, RotationInterval, WireguardData},
 };
 use std::{
@@ -121,6 +124,16 @@ impl PrivateDeviceState {
     }
 }
 
+impl From<PrivateDeviceState> for DeviceState {
+    fn from(state: PrivateDeviceState) -> Self {
+        match state {
+            PrivateDeviceState::LoggedIn(dev) => DeviceState::LoggedIn(AccountAndDevice::from(dev)),
+            PrivateDeviceState::LoggedOut => DeviceState::LoggedOut,
+            PrivateDeviceState::Revoked => DeviceState::Revoked,
+        }
+    }
+}
+
 /// Same as [PrivateDevice] but also contains the associated account token.
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub struct PrivateAccountAndDevice {
@@ -190,32 +203,30 @@ impl From<PrivateDevice> for Device {
 
 #[derive(Clone)]
 pub(crate) enum PrivateDeviceEvent {
+    /// Logged in on a new device.
+    Login(PrivateAccountAndDevice),
     /// The device was removed due to user (or daemon) action.
     Logout,
-    /// Logged in to a new device.
-    Login(PrivateAccountAndDevice),
+    /// Device was removed because it was not found remotely.
+    Revoked,
     /// The device was updated remotely, but not its key.
     Updated(PrivateAccountAndDevice),
     /// The key was rotated.
     RotatedKey(PrivateAccountAndDevice),
-    /// Device was removed because it was not found remotely.
-    Revoked,
 }
 
 impl From<PrivateDeviceEvent> for DeviceEvent {
     fn from(event: PrivateDeviceEvent) -> DeviceEvent {
-        match event {
-            PrivateDeviceEvent::Logout => DeviceEvent::revoke(false),
-            PrivateDeviceEvent::Login(config) => {
-                DeviceEvent::from_device(AccountAndDevice::from(config), false)
-            }
-            PrivateDeviceEvent::Updated(config) => {
-                DeviceEvent::from_device(AccountAndDevice::from(config), true)
-            }
-            PrivateDeviceEvent::RotatedKey(config) => {
-                DeviceEvent::from_device(AccountAndDevice::from(config), false)
-            }
-            PrivateDeviceEvent::Revoked => DeviceEvent::revoke(true),
+        let cause = match event {
+            PrivateDeviceEvent::Login(_) => DeviceEventCause::LoggedIn,
+            PrivateDeviceEvent::Logout => DeviceEventCause::LoggedOut,
+            PrivateDeviceEvent::Revoked => DeviceEventCause::Revoked,
+            PrivateDeviceEvent::Updated(_) => DeviceEventCause::Updated,
+            PrivateDeviceEvent::RotatedKey(_) => DeviceEventCause::RotatedKey,
+        };
+        DeviceEvent {
+            cause,
+            new_state: DeviceState::from(event.state()),
         }
     }
 }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index a7eee39a70..5015580b63 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -40,7 +40,7 @@ use mullvad_relay_selector::{
 };
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
-    device::{AccountAndDevice, Device, DeviceEvent, DeviceId, RemoveDeviceEvent},
+    device::{Device, DeviceEvent, DeviceEventCause, DeviceId, DeviceState, RemoveDeviceEvent},
     location::GeoIpLocation,
     relay_constraints::{BridgeSettings, BridgeState, ObfuscationSettings, RelaySettingsUpdate},
     relay_list::RelayList,
@@ -194,8 +194,8 @@ pub enum DaemonCommand {
     LoginAccount(ResponseTx<(), Error>, AccountToken),
     /// Log out of the current account and remove the device, if they exist.
     LogoutAccount(ResponseTx<(), Error>),
-    /// Return the current device configuration, if there is one.
-    GetDevice(ResponseTx<Option<AccountAndDevice>, Error>),
+    /// Return the current device configuration.
+    GetDevice(ResponseTx<DeviceState, Error>),
     /// Update/check the current device, if there is one.
     UpdateDevice(ResponseTx<(), Error>),
     /// Return all the devices for a given account token.
@@ -1087,7 +1087,10 @@ where
                     error.display_chain_with_msg("Failed to move over account from old settings")
                 );
                 // Synthesize a logout event.
-                event_listener.notify_device_event(DeviceEvent::revoke(false));
+                event_listener.notify_device_event(DeviceEvent {
+                    cause: DeviceEventCause::LoggedOut,
+                    new_state: DeviceState::LoggedOut,
+                });
             }
         });
     }
@@ -1326,17 +1329,16 @@ where
         });
     }
 
-    async fn on_get_device(&mut self, tx: ResponseTx<Option<AccountAndDevice>, Error>) {
+    async fn on_get_device(&mut self, tx: ResponseTx<DeviceState, Error>) {
         let account_manager = self.account_manager.clone();
         tokio::spawn(async move {
             Self::oneshot_send(
                 tx,
-                Ok(account_manager
+                account_manager
                     .data()
                     .await
-                    .map(|s| s.into_device())
-                    .unwrap_or(None)
-                    .map(AccountAndDevice::from)),
+                    .map_err(|_| Error::NoAccountToken)
+                    .map(DeviceState::from),
                 "get_device response",
             );
         });
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 0e581055a9..078fb52d97 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -511,16 +511,12 @@ impl ManagementService for ManagementServiceImpl {
     }
 
     // Device management
-    async fn get_device(&self, _: Request<()>) -> ServiceResult<types::DeviceConfig> {
+    async fn get_device(&self, _: Request<()>) -> ServiceResult<types::DeviceState> {
         log::debug!("get_device");
         let (tx, rx) = oneshot::channel();
         self.send_command_to_daemon(DaemonCommand::GetDevice(tx))?;
-        let device = self
-            .wait_for_result(rx)
-            .await?
-            .map_err(map_daemon_error)?
-            .ok_or(Status::new(Code::NotFound, "no device is set"))?;
-        Ok(Response::new(types::DeviceConfig::from(device)))
+        let device = self.wait_for_result(rx).await?.map_err(map_daemon_error)?;
+        Ok(Response::new(types::DeviceState::from(device)))
     }
 
     async fn update_device(&self, _: Request<()>) -> ServiceResult<()> {
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index c14d621cac..42a73ca7bd 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -56,7 +56,7 @@ service ManagementService {
 	rpc SubmitVoucher(google.protobuf.StringValue) returns (VoucherSubmission) {}
 
 	// Device management
-	rpc GetDevice(google.protobuf.Empty) returns (DeviceConfig) {}
+	rpc GetDevice(google.protobuf.Empty) returns (DeviceState) {}
 	rpc UpdateDevice(google.protobuf.Empty) returns (google.protobuf.Empty) {}
 	rpc ListDevices(google.protobuf.StringValue) returns (DeviceList) {}
 	rpc RemoveDevice(DeviceRemoval) returns (google.protobuf.Empty) {}
@@ -570,7 +570,7 @@ message RelayList {
 	repeated RelayListCountry countries = 1;
 }
 
-message DeviceConfig {
+message AccountAndDevice {
 	string account_token = 1;
 	Device device = 2;
 }
@@ -595,9 +595,26 @@ message DeviceRemoval {
 	string device_id = 2;
 }
 
+message DeviceState {
+	enum State {
+		LOGGED_IN = 0;
+		LOGGED_OUT = 1;
+		REVOKED = 2;
+	}
+	State state = 1;
+	AccountAndDevice device = 2;
+}
+
 message DeviceEvent {
-	DeviceConfig device = 1;
-	bool remote = 2;
+	enum Cause {
+		LOGGED_IN = 0;
+		LOGGED_OUT = 1;
+		REVOKED = 2;
+		UPDATED = 3;
+		ROTATED_KEY = 4;
+	}
+	Cause cause = 1;
+	DeviceState new_state = 2;
 }
 
 message RemoveDeviceEvent {
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index 2b45fbe840..4785f7221d 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -227,14 +227,47 @@ impl From<mullvad_types::device::DevicePort> for DevicePort {
     }
 }
 
+impl From<mullvad_types::device::DeviceState> for DeviceState {
+    fn from(state: mullvad_types::device::DeviceState) -> Self {
+        DeviceState {
+            state: device_state::State::from(&state) as i32,
+            device: state.into_device().map(|device| AccountAndDevice {
+                account_token: device.account_token,
+                device: Some(Device::from(device.device)),
+            }),
+        }
+    }
+}
+
+impl From<&mullvad_types::device::DeviceState> for device_state::State {
+    fn from(state: &mullvad_types::device::DeviceState) -> Self {
+        use mullvad_types::device::DeviceState as MullvadState;
+        match state {
+            MullvadState::LoggedIn(_) => device_state::State::LoggedIn,
+            MullvadState::LoggedOut => device_state::State::LoggedOut,
+            MullvadState::Revoked => device_state::State::Revoked,
+        }
+    }
+}
+
 impl From<mullvad_types::device::DeviceEvent> for DeviceEvent {
     fn from(event: mullvad_types::device::DeviceEvent) -> Self {
         DeviceEvent {
-            device: event.device.map(|config| DeviceConfig {
-                account_token: config.account_token,
-                device: Some(Device::from(config.device)),
-            }),
-            remote: event.remote,
+            cause: device_event::Cause::from(event.cause) as i32,
+            new_state: Some(DeviceState::from(event.new_state)),
+        }
+    }
+}
+
+impl From<mullvad_types::device::DeviceEventCause> for device_event::Cause {
+    fn from(cause: mullvad_types::device::DeviceEventCause) -> Self {
+        use mullvad_types::device::DeviceEventCause as MullvadEvent;
+        match cause {
+            MullvadEvent::LoggedIn => device_event::Cause::LoggedIn,
+            MullvadEvent::LoggedOut => device_event::Cause::LoggedOut,
+            MullvadEvent::Revoked => device_event::Cause::Revoked,
+            MullvadEvent::Updated => device_event::Cause::Updated,
+            MullvadEvent::RotatedKey => device_event::Cause::RotatedKey,
         }
     }
 }
@@ -249,9 +282,9 @@ impl From<mullvad_types::device::RemoveDeviceEvent> for RemoveDeviceEvent {
     }
 }
 
-impl From<mullvad_types::device::AccountAndDevice> for DeviceConfig {
+impl From<mullvad_types::device::AccountAndDevice> for AccountAndDevice {
     fn from(device: mullvad_types::device::AccountAndDevice) -> Self {
-        DeviceConfig {
+        AccountAndDevice {
             account_token: device.account_token,
             device: Some(Device::from(device.device)),
         }
diff --git a/mullvad-types/src/device.rs b/mullvad-types/src/device.rs
index 93836c674a..b39690f603 100644
--- a/mullvad-types/src/device.rs
+++ b/mullvad-types/src/device.rs
@@ -61,6 +61,25 @@ impl fmt::Display for DevicePort {
     }
 }
 
+/// Contains a device state.
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+#[serde(rename_all = "snake_case")]
+#[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
+pub enum DeviceState {
+    LoggedIn(AccountAndDevice),
+    LoggedOut,
+    Revoked,
+}
+
+impl DeviceState {
+    pub fn into_device(self) -> Option<AccountAndDevice> {
+        match self {
+            DeviceState::LoggedIn(dev) => Some(dev),
+            _ => None,
+        }
+    }
+}
+
 /// A [Device] and its associated account token.
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 #[cfg_attr(target_os = "android", derive(IntoJava))]
@@ -79,35 +98,30 @@ impl AccountAndDevice {
     }
 }
 
-/// Emitted when logging in or out of an account, or when the device changes.
+/// Reason why a [DeviceEvent] was emitted.
 #[derive(Clone, Debug)]
 #[cfg_attr(target_os = "android", derive(IntoJava))]
 #[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
-pub struct DeviceEvent {
-    /// Device that was affected.
-    pub device: Option<AccountAndDevice>,
-    /// Indicates whether the change was initiated remotely or by the daemon.
-    pub remote: bool,
+pub enum DeviceEventCause {
+    /// Logged in on a new device.
+    LoggedIn,
+    /// The device was removed due to user (or daemon) action.
+    LoggedOut,
+    /// Device was removed because it was not found remotely.
+    Revoked,
+    /// The device was updated, but not its key.
+    Updated,
+    /// The key was rotated.
+    RotatedKey,
 }
 
-impl DeviceEvent {
-    pub fn new(device: Option<AccountAndDevice>, remote: bool) -> DeviceEvent {
-        DeviceEvent { device, remote }
-    }
-
-    pub fn from_device(device: AccountAndDevice, remote: bool) -> DeviceEvent {
-        DeviceEvent {
-            device: Some(device),
-            remote,
-        }
-    }
-
-    pub fn revoke(remote: bool) -> Self {
-        Self {
-            device: None,
-            remote,
-        }
-    }
+/// Emitted when logging in or out of an account, or when the device changes.
+#[derive(Clone, Debug)]
+#[cfg_attr(target_os = "android", derive(IntoJava))]
+#[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
+pub struct DeviceEvent {
+    pub cause: DeviceEventCause,
+    pub new_state: DeviceState,
 }
 
 /// Emitted when a device is removed using the `RemoveDevice` RPC.