Add multihop state separate from entry location

5 files changed, 40 insertions, 36 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index c9e7bef6f6..11e36e0d98 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -571,6 +571,7 @@ impl Relay {
         }
         if let Some(entry) = matches.values_of("entry location") {
             wireguard_constraints.entry_location = parse_entry_location_constraint(entry);
+            wireguard_constraints.use_multihop = wireguard_constraints.entry_location.is_some();
         }
 
         self.update_constraints(types::RelaySettingsUpdate {
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 4eee17ac23..48b46dca9a 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -58,7 +58,8 @@ const WIREGUARD_EXIT_CONSTRAINTS: WireguardConstraints = WireguardConstraints {
         port: Constraint::Only(DEFAULT_WIREGUARD_PORT),
     }),
     ip_version: Constraint::Only(IpVersion::V4),
-    entry_location: None,
+    use_multihop: false,
+    entry_location: Constraint::Any,
 };
 const WIREGUARD_TCP_PORTS: [(u16, u16); 3] = [(80, 80), (443, 443), (5001, 5001)];
 
@@ -244,12 +245,13 @@ impl RelaySelector {
         wg_key_exists: bool,
     ) -> Result<(Relay, Option<Relay>, MullvadEndpoint), Error> {
         let mut exit_relay_constraints = relay_constraints.clone();
-        let wg_entry_is_subset = if let Some(entry_location) =
-            exit_relay_constraints.wireguard_constraints.entry_location
-        {
+        let wg_entry_is_subset = if exit_relay_constraints.wireguard_constraints.use_multihop {
+            let use_multihop = exit_relay_constraints.wireguard_constraints.use_multihop;
+            let entry_location = exit_relay_constraints.wireguard_constraints.entry_location;
             let is_subset = entry_location.is_subset(&exit_relay_constraints.location);
             exit_relay_constraints.wireguard_constraints = WireguardConstraints {
-                entry_location: Some(entry_location),
+                use_multihop,
+                entry_location,
                 ..WIREGUARD_EXIT_CONSTRAINTS
             };
             is_subset
@@ -257,16 +259,12 @@ impl RelaySelector {
             false
         };
 
-        let entry_endpoint = if wg_entry_is_subset
-            && relay_constraints
-                .wireguard_constraints
-                .entry_location
-                .is_some()
-        {
-            self.select_entry_endpoint(None, &relay_constraints, retry_attempt)
-        } else {
-            None
-        };
+        let entry_endpoint =
+            if wg_entry_is_subset && relay_constraints.wireguard_constraints.use_multihop {
+                self.select_entry_endpoint(None, &relay_constraints, retry_attempt)
+            } else {
+                None
+            };
 
         let (exit_relay, mut endpoint) = self.get_tunnel_exit_endpoint(
             &exit_relay_constraints,
@@ -283,12 +281,7 @@ impl RelaySelector {
         )?;
 
         let mut entry_endpoint = entry_endpoint.or_else(|| {
-            if !wg_entry_is_subset
-                && relay_constraints
-                    .wireguard_constraints
-                    .entry_location
-                    .is_some()
-            {
+            if !wg_entry_is_subset && relay_constraints.wireguard_constraints.use_multihop {
                 if let MullvadEndpoint::Wireguard { peer, .. } = &endpoint {
                     self.select_entry_endpoint(Some(peer), &relay_constraints, retry_attempt)
                 } else {
@@ -308,11 +301,7 @@ impl RelaySelector {
                     entry_relay.hostname, addr_in
                 );
                 return Ok((exit_relay, Some(entry_relay), entry_endpoint));
-            } else if relay_constraints
-                .wireguard_constraints
-                .entry_location
-                .is_some()
-            {
+            } else if relay_constraints.wireguard_constraints.use_multihop {
                 return Err(Error::NoRelay);
             }
         }
@@ -450,10 +439,13 @@ impl RelaySelector {
         relay_constraints: &RelayConstraints,
         retry_attempt: u32,
     ) -> Option<(Relay, MullvadEndpoint)> {
+        if !relay_constraints.wireguard_constraints.use_multihop {
+            return None;
+        }
         let entry_location = relay_constraints
             .wireguard_constraints
             .entry_location
-            .clone()?;
+            .clone();
         let entry_constraints = RelayConstraints {
             location: entry_location,
             tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
@@ -1393,14 +1385,15 @@ mod test {
             ..RelayConstraints::default()
         };
 
-        relay_constraints.wireguard_constraints.entry_location = Some(Constraint::Only(location1));
+        relay_constraints.wireguard_constraints.use_multihop = true;
+        relay_constraints.wireguard_constraints.entry_location = Constraint::Only(location1);
 
         // The same host cannot be used for entry and exit
         assert!(relay_selector
             .get_tunnel_endpoint(&relay_constraints, BridgeState::Off, 0, true)
             .is_err());
 
-        relay_constraints.wireguard_constraints.entry_location = Some(Constraint::Only(location2));
+        relay_constraints.wireguard_constraints.entry_location = Constraint::Only(location2);
 
         // If the entry and exit differ, this should succeed
         assert!(relay_selector
@@ -1427,8 +1420,9 @@ mod test {
             ..RelayConstraints::default()
         };
 
+        relay_constraints.wireguard_constraints.use_multihop = true;
         relay_constraints.wireguard_constraints.entry_location =
-            Some(Constraint::Only(location_specific.clone()));
+            Constraint::Only(location_specific.clone());
 
         // The exit must not equal the entry
         let (exit_relay, _entry_relay, _exit_endpoint) = relay_selector
@@ -1439,7 +1433,7 @@ mod test {
 
         relay_constraints.location = Constraint::Only(location_specific);
         relay_constraints.wireguard_constraints.entry_location =
-            Some(Constraint::Only(location_general));
+            Constraint::Only(location_general);
 
         // The entry must not equal the exit
         let (exit_relay, _entry_relay, exit_endpoint) = relay_selector
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index ded1445fb8..a259b2f51a 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -336,7 +336,8 @@ message IpVersionConstraint {
 message WireguardConstraints {
 	TransportPort port = 1;
 	IpVersionConstraint ip_version = 2;
-	RelayLocation entry_location = 3;
+	bool use_multihop = 3;
+	RelayLocation entry_location = 4;
 }
 
 message CustomRelaySettings {
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index 05e41ceee7..1504abee39 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -505,9 +505,11 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for RelaySettings {
                             .option()
                             .map(IpVersion::from)
                             .map(IpVersionConstraint::from),
+                        use_multihop: constraints.wireguard_constraints.use_multihop,
                         entry_location: constraints
                             .wireguard_constraints
                             .entry_location
+                            .option()
                             .map(RelayLocation::from),
                     }),
 
@@ -715,10 +717,12 @@ impl TryFrom<&WireguardConstraints> for mullvad_types::relay_constraints::Wiregu
         Ok(mullvad_constraints::WireguardConstraints {
             port: Constraint::from(wireguard_transport_port),
             ip_version: Constraint::from(ip_version),
+            use_multihop: constraints.use_multihop,
             entry_location: constraints
                 .entry_location
                 .clone()
-                .map(Constraint::<mullvad_types::relay_constraints::LocationConstraint>::from),
+                .map(Constraint::<mullvad_types::relay_constraints::LocationConstraint>::from)
+                .unwrap_or(Constraint::Any),
         })
     }
 }
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index dac8c5e2a8..2c599bd6b0 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -494,7 +494,8 @@ impl Match<OpenVpnEndpointData> for OpenVpnConstraints {
 pub struct WireguardConstraints {
     pub port: Constraint<TransportPort>,
     pub ip_version: Constraint<IpVersion>,
-    pub entry_location: Option<Constraint<LocationConstraint>>,
+    pub use_multihop: bool,
+    pub entry_location: Constraint<LocationConstraint>,
 }
 
 impl fmt::Display for WireguardConstraints {
@@ -514,8 +515,11 @@ impl fmt::Display for WireguardConstraints {
             Constraint::Any => write!(f, "IPv4 or IPv6")?,
             Constraint::Only(protocol) => write!(f, "{}", protocol)?,
         }
-        if let Some(Constraint::Only(ref entry)) = self.entry_location {
-            write!(f, " (via {})", entry)
+        if self.use_multihop {
+            match &self.entry_location {
+                Constraint::Any => write!(f, " (via any location)"),
+                Constraint::Only(location) => write!(f, " (via {})", location),
+            }
         } else {
             Ok(())
         }