Skip to content

Commit

Permalink
Remove GetCurrentLocation. Make the daemon send two tunnel state
Browse files Browse the repository at this point in the history 
updates, one with out IP being empty, and another with
it being filled when am.i.mullvad.net responds.

Update CLI for this change. Other front ends are left out.
  • Loading branch information
@Serock3
Serock3 committed Dec 1, 2023
1 parent 63da800 commit a374f6a
Show file tree
Hide file tree
Showing 9 changed files with 94 additions and 134 deletions.
42 changes: 18 additions & 24 deletions mullvad-cli/src/cmds/status.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use anyhow::Result;
use clap::{Args, Subcommand};
use futures::StreamExt;
use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
use mullvad_types::{device::DeviceState, states::TunnelState};
use mullvad_types::{device::DeviceState, location::GeoIpLocation, states::TunnelState};

use crate::format;

Expand Down Expand Up @@ -38,13 +38,21 @@ impl Status {
format::print_state(&new_state, args.verbose);
}

match new_state {
TunnelState::Connected { .. } | TunnelState::Disconnected => {
if args.location {
print_location(&mut rpc).await?;
}
let location = match new_state {
TunnelState::Connected {
location,
endpoint: _,
} => location,
TunnelState::Disconnected(location) => location,
_ => return Ok(()), /* No location data is available during
* connecting/disconnecting/error states */
};
if args.location {
if let Some(location) = location {
print_location(&location);
} else {
eprintln!("No location available")
}
_ => {}
}
}
DaemonEvent::Settings(settings) => {
Expand Down Expand Up @@ -88,30 +96,17 @@ pub async fn handle(cmd: Option<Status>, args: StatusArgs) -> Result<()> {
if args.debug {
println!("Tunnel state: {state:#?}");
} else {
// TODO: respect location arg?
format::print_state(&state, args.verbose);
}

if args.location {
print_location(&mut rpc).await?;
}

if cmd == Some(Status::Listen) {
Status::listen(rpc, args).await?;
}
Ok(())
}

async fn print_location(rpc: &mut MullvadProxyClient) -> Result<()> {
let location = match rpc.get_current_location().await {
Ok(location) => location,
Err(error) => match &error {
mullvad_management_interface::Error::NoLocationData => {
println!("Location data unavailable");
return Ok(());
}
_ => return Err(error.into()),
},
};
pub fn print_location(location: &GeoIpLocation) {
if let Some(ipv4) = location.ipv4 {
println!("IPv4: {ipv4}");
}
Expand All @@ -123,7 +118,6 @@ async fn print_location(rpc: &mut MullvadProxyClient) -> Result<()> {
"Position: {:.5}°N, {:.5}°W",
location.latitude, location.longitude
);
Ok(())
}

fn print_account_loggedout(state: &TunnelState, device: &DeviceState) {
Expand All @@ -137,6 +131,6 @@ fn print_account_loggedout(state: &TunnelState, device: &DeviceState) {
DeviceState::LoggedIn(_) => (),
}
}
TunnelState::Disconnected | TunnelState::Disconnecting(_) => (),
TunnelState::Disconnected(_) | TunnelState::Disconnecting(_) => (),
}
}
9 changes: 8 additions & 1 deletion mullvad-cli/src/format.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ use talpid_types::{
tunnel::ErrorState,
};

use crate::cmds::status::print_location;

#[macro_export]
macro_rules! print_option {
($value:expr $(,)?) => {{
Expand Down Expand Up @@ -37,7 +39,12 @@ pub fn print_state(state: &TunnelState, verbose: bool) {
format_relay_connection(endpoint, location.as_ref(), verbose)
);
}
Disconnected => println!("Disconnected"),
Disconnected(disconnected_location) => {
println!("Disconnected");
if let Some(disconnected_location) = disconnected_location {
print_location(disconnected_location);
}
}
Disconnecting(_) => println!("Disconnecting..."),
}
}
Expand Down
13 changes: 3 additions & 10 deletions mullvad-daemon/src/geoip.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,13 +44,13 @@ const LOCATION_RETRY_STRATEGY: Jittered<ExponentialBackoff> =
Jittered::jitter(ExponentialBackoff::new(Duration::from_secs(1), 4));

/// Fetch the current `GeoIpLocation` with retrys
pub async fn get_geo_location(
pub async fn get_geo_location_with_retry(
rest_service: RequestServiceHandle,
use_ipv6: bool,
api_handle: ApiAvailabilityHandle,
) -> Option<GeoIpLocation> {
) -> Result<GeoIpLocation, Error> {
log::debug!("Fetching GeoIpLocation");
match retry_future(
retry_future(
move || send_location_request(rest_service.clone(), use_ipv6),
move |result| match result {
Err(error) if error.is_network_error() => !api_handle.get_state().is_offline(),
Expand All @@ -59,13 +59,6 @@ pub async fn get_geo_location(
LOCATION_RETRY_STRATEGY,
)
.await
{
Ok(loc) => Some(loc),
Err(e) => {
log::warn!("Unable to fetch GeoIP location: {}", e.display_chain());
None
}
}
}

async fn send_location_request(
Expand Down
107 changes: 48 additions & 59 deletions mullvad-daemon/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ mod tunnel;
pub mod version;
mod version_check;

use crate::{geoip::get_geo_location, target_state::PersistentTargetState};
use crate::target_state::PersistentTargetState;
use device::{AccountEvent, PrivateAccountAndDevice, PrivateDeviceEvent};
use futures::{
channel::{mpsc, oneshot},
Expand Down Expand Up @@ -200,8 +200,6 @@ pub enum DaemonCommand {
Reconnect(oneshot::Sender<bool>),
/// Request the current state.
GetState(oneshot::Sender<TunnelState>),
/// Get the current geographical location.
GetCurrentLocation(oneshot::Sender<Option<GeoIpLocation>>),
CreateNewAccount(ResponseTx<String, Error>),
/// Request the metadata for an account.
GetAccountData(
Expand Down Expand Up @@ -417,7 +415,7 @@ impl DaemonExecutionState {
match self {
Running => {
match tunnel_state {
TunnelState::Disconnected => mem::replace(self, Finished),
TunnelState::Disconnected(_) => mem::replace(self, Finished),
_ => mem::replace(self, Exiting),
};
}
Expand Down Expand Up @@ -614,6 +612,7 @@ pub struct Daemon<L: EventListener> {
tunnel_state_machine_handle: TunnelStateMachineHandle,
#[cfg(target_os = "windows")]
volume_update_tx: mpsc::UnboundedSender<()>,
location_abort_handle: Option<tokio::task::AbortHandle>,
}

impl<L> Daemon<L>
Expand Down Expand Up @@ -847,7 +846,7 @@ where
relay_list_updater.update().await;

let daemon = Daemon {
tunnel_state: TunnelState::Disconnected,
tunnel_state: TunnelState::Disconnected(None),
target_state,
state: DaemonExecutionState::Running,
#[cfg(target_os = "linux")]
Expand All @@ -873,6 +872,7 @@ where
tunnel_state_machine_handle,
#[cfg(target_os = "windows")]
volume_update_tx,
location_abort_handle: None,
};

api_availability.unsuspend();
Expand Down Expand Up @@ -971,20 +971,33 @@ where
&mut self,
tunnel_state_transition: TunnelStateTransition,
) {
// Abort any ongoing calls to `self.notify_tunnel_state_when_ip_arrives` from
// previous tunnel state transitions to avoid sending an outdated information.
if let Some(handle) = self.location_abort_handle.take() {
handle.abort();
}
self.reset_rpc_sockets_on_tunnel_state_transition(&tunnel_state_transition);
self.device_checker
.handle_state_transition(&tunnel_state_transition);

let tunnel_state = match tunnel_state_transition {
TunnelStateTransition::Disconnected => TunnelState::Disconnected,
TunnelStateTransition::Disconnected => {
self.notify_tunnel_state_when_ip_arrives(TunnelState::Disconnected)
.await;
TunnelState::Disconnected(self.parameters_generator.get_last_location().await)
}

TunnelStateTransition::Connecting(endpoint) => TunnelState::Connecting {
endpoint,
location: self.parameters_generator.get_last_location().await,
},
TunnelStateTransition::Connected(endpoint) => TunnelState::Connected {
endpoint,
location: self.parameters_generator.get_last_location().await,
},
TunnelStateTransition::Connected(endpoint) => {
let make_tunnel_state = |location| TunnelState::Connected { endpoint, location };
self.notify_tunnel_state_when_ip_arrives(make_tunnel_state.clone())
.await;

make_tunnel_state(self.parameters_generator.get_last_location().await)
}
TunnelStateTransition::Disconnecting(after_disconnect) => {
TunnelState::Disconnecting(after_disconnect)
}
Expand All @@ -1001,7 +1014,7 @@ where
log::debug!("New tunnel state: {:?}", tunnel_state);

match tunnel_state {
TunnelState::Disconnected => {
TunnelState::Disconnected(_) => {
self.api_handle.availability.reset_inactivity_timer();
}
_ => {
Expand All @@ -1010,7 +1023,7 @@ where
}

match &tunnel_state {
TunnelState::Disconnected => self.state.disconnected(),
TunnelState::Disconnected(_) => self.state.disconnected(),
TunnelState::Connecting { .. } => {
log::debug!("Settings: {}", self.settings.summary());
}
Expand Down Expand Up @@ -1040,6 +1053,29 @@ where
self.event_listener.notify_new_state(tunnel_state);
}

/// Get the out IP from am.i.mullvad.net. When it arrives, send another
/// TunnelState with the IP filled in.
async fn notify_tunnel_state_when_ip_arrives(
&mut self,
make_tunnel_state: impl FnOnce(Option<GeoIpLocation>) -> TunnelState + Send + 'static,
) {
let rest_service = self.api_runtime.rest_handle().await;
let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
let api_handle = self.api_handle.availability.clone();
let even_listener = self.event_listener.clone();

self.location_abort_handle = Some(
tokio::spawn(async move {
let location =
geoip::get_geo_location_with_retry(rest_service, use_ipv6, api_handle)
.await
.ok();
even_listener.notify_new_state(make_tunnel_state(location));
})
.abort_handle(),
);
}

fn reset_rpc_sockets_on_tunnel_state_transition(
&mut self,
tunnel_state_transition: &TunnelStateTransition,
Expand Down Expand Up @@ -1091,7 +1127,6 @@ where
SetTargetState(tx, state) => self.on_set_target_state(tx, state).await,
Reconnect(tx) => self.on_reconnect(tx),
GetState(tx) => self.on_get_state(tx),
GetCurrentLocation(tx) => self.on_get_current_location(tx).await,
CreateNewAccount(tx) => self.on_create_new_account(tx),
GetAccountData(tx, account_token) => self.on_get_account_data(tx, account_token),
GetWwwAuthToken(tx) => self.on_get_www_auth_token(tx).await,
Expand Down Expand Up @@ -1342,52 +1377,6 @@ where
Self::oneshot_send(tx, performing_post_upgrade, "performing post upgrade");
}

async fn on_get_current_location(&mut self, tx: oneshot::Sender<Option<GeoIpLocation>>) {
use self::TunnelState::*;

match &self.tunnel_state {
Disconnected => self.update_and_send_geo_location(tx, None).await,
Connecting { location, .. } => {
Self::oneshot_send(tx, location.clone(), "current location")
}
Disconnecting(..) => Self::oneshot_send(
tx,
self.parameters_generator.get_last_location().await,
"current location",
),
Connected { location, .. } => {
self.update_and_send_geo_location(tx, location.clone())
.await
}
// We are not online at all at this stage so no location data is available.
Error(_) => Self::oneshot_send(tx, None, "current location"),
};
}

/// Fetch the current `GeoIpLocation` and send it on the return channel,
/// in a non-blocking fashion. Optionally give a chached previous location.
async fn update_and_send_geo_location(
&mut self,
tx: oneshot::Sender<Option<GeoIpLocation>>,
current_relay_location: Option<GeoIpLocation>,
) {
let rest_service = self.api_runtime.rest_handle().await;
let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
let api_handle = self.api_handle.availability.clone();
tokio::spawn(async move {
let new_location = get_geo_location(rest_service, use_ipv6, api_handle).await;
Self::oneshot_send(
tx,
new_location.map(|fetched_location| GeoIpLocation {
ipv4: fetched_location.ipv4,
ipv6: fetched_location.ipv6,
..current_relay_location.unwrap_or(fetched_location)
}),
"current location",
);
});
}

fn on_create_new_account(&mut self, tx: ResponseTx<String, Error>) {
let account_manager = self.account_manager.clone();
tokio::spawn(async move {
Expand Down
11 changes: 0 additions & 11 deletions mullvad-daemon/src/management_interface.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,17 +191,6 @@ impl ManagementService for ManagementServiceImpl {
.map(|relays| Response::new(types::RelayList::from(relays)))
}

async fn get_current_location(&self, _: Request<()>) -> ServiceResult<types::GeoIpLocation> {
log::debug!("get_current_location");
let (tx, rx) = oneshot::channel();
self.send_command_to_daemon(DaemonCommand::GetCurrentLocation(tx))?;
let result = self.wait_for_result(rx).await?;
match result {
Some(geoip) => Ok(Response::new(types::GeoIpLocation::from(geoip))),
None => Err(Status::not_found("no location was found")),
}
}

async fn set_bridge_settings(
&self,
request: Request<types::BridgeSettings>,
Expand Down
3 changes: 1 addition & 2 deletions mullvad-management-interface/proto/management_interface.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ service ManagementService {
// Relays and tunnel constraints
rpc UpdateRelayLocations(google.protobuf.Empty) returns (google.protobuf.Empty) {}
rpc GetRelayLocations(google.protobuf.Empty) returns (RelayList) {}
rpc GetCurrentLocation(google.protobuf.Empty) returns (GeoIpLocation) {}
rpc SetRelaySettings(RelaySettings) returns (google.protobuf.Empty) {}
rpc SetBridgeSettings(BridgeSettings) returns (google.protobuf.Empty) {}
rpc SetBridgeState(BridgeState) returns (google.protobuf.Empty) {}
Expand Down Expand Up @@ -174,7 +173,7 @@ message ErrorState {
}

message TunnelState {
message Disconnected {}
message Disconnected { GeoIpLocation disconnected_location = 1; }
message Connecting { TunnelStateRelayInfo relay_info = 1; }
message Connected { TunnelStateRelayInfo relay_info = 1; }
message Disconnecting { AfterDisconnect after_disconnect = 1; }
Expand Down
Loading

0 comments on commit a374f6a

Please sign in to comment.