Skip to content

Commit

Permalink
Merge branch 'fix-vendored-mig-types'
Browse files Browse the repository at this point in the history
  • Loading branch information
dlon committed Mar 23, 2022
2 parents b1910fc + 2c02d17 commit e53c455
Show file tree
Hide file tree
Showing 5 changed files with 149 additions and 120 deletions.
4 changes: 2 additions & 2 deletions mullvad-daemon/src/device.rs
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,7 @@ impl AccountManager {
Ok(())
}

async fn logout_inner(&mut self) -> tokio::task::JoinHandle<()> {
fn logout_inner(&mut self) -> tokio::task::JoinHandle<()> {
let prev_data = self.data.take();
let service = self.device_service.clone();

Expand Down Expand Up @@ -392,7 +392,7 @@ impl AccountManager {
{
// Remove the existing device if its ID differs. Otherwise, only update
// the data.
self.logout_inner().await;
self.logout_inner();
}

self.data = data.cloned();
Expand Down
29 changes: 18 additions & 11 deletions mullvad-daemon/src/migrations/account_history.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,19 @@
use super::{Error, Result};
use mullvad_types::{account::AccountToken, wireguard::WireguardData};
use mullvad_types::account::AccountToken;
use regex::Regex;
use serde::Deserialize;
use std::path::Path;
use talpid_types::ErrorExt;
use tokio::{
fs,
io::{self, AsyncReadExt, AsyncSeekExt, AsyncWriteExt},
};

// ======================================================
// Section for vendoring types.

// ======================================================

const ACCOUNT_HISTORY_FILE: &str = "account-history.json";

lazy_static::lazy_static! {
Expand Down Expand Up @@ -77,7 +83,7 @@ fn migrate_formats_inner(
settings: &mut serde_json::Value,
) -> Result<AccountToken> {
if let Some((token, wg_data)) = try_format_v2(account_bytes) {
settings["wireguard"] = serde_json::json!(wg_data);
settings["wireguard"] = wg_data;
Ok(token)
} else if let Some(token) = try_format_v1(account_bytes) {
Ok(token)
Expand All @@ -93,19 +99,20 @@ fn is_format_v3(bytes: &[u8]) -> bool {
}
}

fn try_format_v2(bytes: &[u8]) -> Option<(AccountToken, Option<WireguardData>)> {
#[derive(Serialize, Deserialize, Clone, Debug)]
fn try_format_v2(bytes: &[u8]) -> Option<(AccountToken, serde_json::Value)> {
#[derive(Deserialize, Clone)]
pub struct AccountEntry {
pub account: AccountToken,
pub wireguard: Option<WireguardData>,
pub wireguard: serde_json::Value,
}
serde_json::from_slice(bytes)
.map(|entries: Vec<AccountEntry>| {
.ok()
.and_then(|entries: Vec<AccountEntry>| {
entries
.first()
.map(|entry| (entry.account.clone(), entry.wireguard.clone()))
.into_iter()
.next()
.map(|entry| (entry.account, entry.wireguard))
})
.unwrap_or(None)
}

fn try_format_v1(bytes: &[u8]) -> Option<AccountToken> {
Expand All @@ -114,8 +121,8 @@ fn try_format_v1(bytes: &[u8]) -> Option<AccountToken> {
accounts: Vec<AccountToken>,
}
serde_json::from_slice(bytes)
.map(|old_format: OldFormat| old_format.accounts.first().cloned())
.unwrap_or(None)
.ok()
.and_then(|old_format: OldFormat| old_format.accounts.into_iter().next())
}

#[cfg(test)]
Expand Down
97 changes: 97 additions & 0 deletions mullvad-daemon/src/migrations/device.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
//! Generates a `device.json` from a WireGuard key and account token by matching them against
//! devices returned by the API and sending the `DeviceMigrationEvent` event to the daemon.
//! The account token and private key may be lost if it fails, but this should not be not
//! critical since the account history also contains the token.
//!
//! This module is allowed to import a number of types, unlike other migration modules, as it
//! does not modify any files directly and may safely fail.
use super::v5::MigrationData;
use crate::{
device::{self, DeviceService},
DaemonEventSender, InternalDaemonEvent,
};
use mullvad_types::{account::AccountToken, device::DeviceData, wireguard::WireguardData};
use talpid_core::mpsc::Sender;
use talpid_types::ErrorExt;

pub(crate) fn generate_device(
migration_data: MigrationData,
rest_handle: mullvad_api::rest::MullvadRestHandle,
daemon_tx: DaemonEventSender,
) {
tokio::spawn(async move {
let wg_data: Option<WireguardData> = migration_data.wg_data.and_then(|data| {
serde_json::from_value(data)
.map(Some)
.unwrap_or_else(|error| {
log::error!(
"{}",
error.display_chain_with_msg("Failed to parse WireGuard data")
);
None
})
});

let api_handle = rest_handle.availability.clone();
let service = DeviceService::new(rest_handle, api_handle);
let result = match (migration_data.token, wg_data) {
(token, Some(wg_data)) => {
log::info!("Creating a new device cache from previous settings");
cache_from_wireguard_key(service, token, wg_data).await
}
(token, None) => {
log::info!("Generating a new device for the account");
cache_from_account(service, token).await
}
};
if let Ok(data) = result {
let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(data));
}
});
}

async fn cache_from_wireguard_key(
service: DeviceService,
token: AccountToken,
wg_data: WireguardData,
) -> Result<DeviceData, device::Error> {
let devices = service
.list_devices_with_backoff(token.clone())
.await
.map_err(|error| {
log::error!(
"{}",
error.display_chain_with_msg("Failed to enumerate devices for account")
);
error
})?;

for device in devices.into_iter() {
if device.pubkey == wg_data.private_key.public_key() {
return Ok(DeviceData {
token,
device,
wg_data,
});
}
}
log::info!("The existing WireGuard key is not valid; generating a new device");
cache_from_account(service, token).await
}

async fn cache_from_account(
service: DeviceService,
token: AccountToken,
) -> Result<DeviceData, device::Error> {
service
.generate_for_account_with_backoff(token)
.await
.map_err(|error| {
log::error!(
"{}",
error.display_chain_with_msg("Failed to generate new device for account")
);
error
})
}
9 changes: 7 additions & 2 deletions mullvad-daemon/src/migrations/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
//! Migration modules may NOT import and use structs that may
//! change. Because then a later change to the current code can break
//! old migrations. The only items a settings migration module may import
//! are anything from `std`, `jnix` and the following:
//! are anything from `std`, `jnix`, `serde` and the following:
//!
//! ```ignore
//! use super::{Error, Result};
Expand Down Expand Up @@ -38,6 +38,7 @@ use tokio::{
};

mod account_history;
mod device;
mod v1;
mod v2;
mod v3;
Expand Down Expand Up @@ -123,7 +124,11 @@ pub(crate) async fn migrate_all(
account_history::migrate_location(cache_dir, settings_dir).await;
account_history::migrate_formats(settings_dir, &mut settings).await?;

v5::migrate(&mut settings, rest_handle, daemon_tx).await?;
let migration_data = v5::migrate(&mut settings).await?;

if let Some(migration_data) = migration_data {
device::generate_device(migration_data, rest_handle, daemon_tx);
}

if settings == old_settings {
// Nothing changed
Expand Down
130 changes: 25 additions & 105 deletions mullvad-daemon/src/migrations/v5.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,19 @@
use super::{Error, Result};
use crate::{device::DeviceService, DaemonEventSender, InternalDaemonEvent};
use mullvad_types::{
account::AccountToken, device::DeviceData, settings::SettingsVersion, wireguard::WireguardData,
};
use talpid_core::mpsc::Sender;
use talpid_types::ErrorExt;
use mullvad_types::settings::SettingsVersion;

// ======================================================
// Section for vendoring types and values that
// this settings version depend on. See `mod.rs`.

pub type AccountToken = String;

// ======================================================

pub(crate) struct MigrationData {
pub token: AccountToken,
pub wg_data: Option<serde_json::Value>,
}

/// This is an open ended migration. There is no v6 yet!
/// The migrations performed by this function are still backwards compatible.
/// The JSON coming out of this migration can be read by any v5 compatible daemon.
Expand All @@ -29,42 +31,8 @@ use talpid_types::ErrorExt;
/// It is also no longer valid to have `entry_location` set to null. So remove the field if it
/// is null in order to make it default back to the default location.
///
/// This also removes the account token and WireGuard key from the settings, looks up the
/// corresponding device, and eventually stores them in `device.json` instead. This is done by
/// sending the `DeviceMigrationEvent` event to the daemon. Because this is fallible, it can
/// result in the account token and private key being lost. This should not be not critical since
/// the account token is also stored in the account history.
pub(crate) async fn migrate(
settings: &mut serde_json::Value,
rest_handle: mullvad_api::rest::MullvadRestHandle,
daemon_tx: DaemonEventSender,
) -> Result<()> {
let migration_data = migrate_inner(settings).await?;

if let Some(migration_data) = migration_data {
let api_handle = rest_handle.availability.clone();
let service = DeviceService::new(rest_handle, api_handle);
match (migration_data.token, migration_data.wg_data) {
(token, Some(wg_data)) => {
log::info!("Creating a new device cache from previous settings");
tokio::spawn(cache_from_wireguard_key(daemon_tx, service, token, wg_data));
}
(token, None) => {
log::info!("Generating a new device for the account");
tokio::spawn(cache_from_account(daemon_tx, service, token));
}
}
}

Ok(())
}

struct MigrationData {
token: AccountToken,
wg_data: Option<WireguardData>,
}

async fn migrate_inner(settings: &mut serde_json::Value) -> Result<Option<MigrationData>> {
/// This also removes the account token and WireGuard key from the settings.
pub(crate) async fn migrate(settings: &mut serde_json::Value) -> Result<Option<MigrationData>> {
if !version_matches(settings) {
return Ok(None);
}
Expand Down Expand Up @@ -94,25 +62,24 @@ async fn migrate_inner(settings: &mut serde_json::Value) -> Result<Option<Migrat
if let Some(token) = settings.get("account_token").filter(|t| !t.is_null()) {
let token: AccountToken =
serde_json::from_value(token.clone()).map_err(Error::ParseError)?;
let mig_data = if let Some(wg_data) = settings.get("wireguard").filter(|wg| !wg.is_null()) {
let wg_data: WireguardData =
serde_json::from_value(wg_data.clone()).map_err(Error::ParseError)?;
Ok(Some(MigrationData {
token,
wg_data: Some(wg_data),
}))
} else {
Ok(Some(MigrationData {
token,
wg_data: None,
}))
};
let migration_data =
if let Some(wg_data) = settings.get("wireguard").filter(|wg| !wg.is_null()) {
Ok(Some(MigrationData {
token,
wg_data: Some(wg_data.clone()),
}))
} else {
Ok(Some(MigrationData {
token,
wg_data: None,
}))
};

let settings_map = settings.as_object_mut().ok_or(Error::NoMatchingVersion)?;
settings_map.remove("account_token");
settings_map.remove("wireguard");

return mig_data;
return migration_data;
}

// Note: Not incrementing the version number yet, since this migration is still open
Expand All @@ -129,56 +96,9 @@ fn version_matches(settings: &mut serde_json::Value) -> bool {
.unwrap_or(false)
}

async fn cache_from_wireguard_key(
daemon_tx: DaemonEventSender,
service: DeviceService,
token: AccountToken,
wg_data: WireguardData,
) {
let devices = match service.list_devices_with_backoff(token.clone()).await {
Ok(devices) => devices,
Err(error) => {
log::error!(
"{}",
error.display_chain_with_msg("Failed to enumerate devices for account")
);
return;
}
};

for device in devices.into_iter() {
if device.pubkey == wg_data.private_key.public_key() {
let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(DeviceData {
token,
device,
wg_data,
}));
return;
}
}
log::info!("The existing WireGuard key is not valid; generating a new device");
cache_from_account(daemon_tx, service, token).await;
}

async fn cache_from_account(
daemon_tx: DaemonEventSender,
service: DeviceService,
token: AccountToken,
) {
match service.generate_for_account_with_backoff(token).await {
Ok(device_data) => {
let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(device_data));
}
Err(error) => log::error!(
"{}",
error.display_chain_with_msg("Failed to generate new device for account")
),
}
}

#[cfg(test)]
mod test {
use super::{migrate_inner, version_matches};
use super::{migrate, version_matches};
use serde_json;

pub const V5_SETTINGS_V1: &str = r#"
Expand Down Expand Up @@ -324,7 +244,7 @@ mod test {
let mut old_settings = serde_json::from_str(V5_SETTINGS_V1).unwrap();

assert!(version_matches(&mut old_settings));
migrate_inner(&mut old_settings).await.unwrap();
migrate(&mut old_settings).await.unwrap();
let new_settings: serde_json::Value = serde_json::from_str(V5_SETTINGS_V2).unwrap();

assert_eq!(&old_settings, &new_settings);
Expand Down

0 comments on commit e53c455

Please sign in to comment.