Set up peers for entry and exit endpoints

4 files changed, 21 insertions, 3 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index d44ce55153..6d92566bac 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -24,6 +24,7 @@ use futures::{
     future::{abortable, AbortHandle, Future},
     SinkExt, StreamExt,
 };
+use ipnetwork::IpNetwork;
 use log::{debug, error, info, warn};
 use mullvad_rpc::AccountsProxy;
 use mullvad_types::{
@@ -1058,10 +1059,23 @@ where
                 .into())
             }
             MullvadEndpoint::Wireguard {
-                peer,
+                mut peer,
                 ipv4_gateway,
                 ipv6_gateway,
             } => {
+                let exit_peer = match self.settings.get_relay_settings() {
+                    RelaySettings::Normal(ref relay_constraints) => self
+                        .relay_selector
+                        .get_tunnel_exit_endpoint(relay_constraints)
+                        .and_then(|(_relay, mullvad_endpoint)| match mullvad_endpoint {
+                            MullvadEndpoint::Wireguard { peer, .. } => Some(peer),
+                            _ => None,
+                        }),
+                    _ => None,
+                };
+                if let Some(ref exit) = exit_peer {
+                    peer.allowed_ips = vec![IpNetwork::from(exit.endpoint.ip())];
+                }
                 let wg_data = self
                     .account_history
                     .get(&account_token)
@@ -1080,6 +1094,7 @@ where
                     connection: wireguard::ConnectionConfig {
                         tunnel,
                         peer,
+                        additional_peers: exit_peer.map(|peer| vec![peer]).unwrap_or(vec![]),
                         ipv4_gateway,
                         ipv6_gateway: Some(ipv6_gateway),
                     },
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index b844c96fe7..07df610045 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -773,6 +773,7 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                                     ))?
                                     .into(),
                             },
+                            additional_peers: vec![],
                             ipv4_gateway,
                             ipv6_gateway,
                         })
diff --git a/talpid-core/src/tunnel/wireguard/config.rs b/talpid-core/src/tunnel/wireguard/config.rs
index eddc51b25b..a7b291438e 100644
--- a/talpid-core/src/tunnel/wireguard/config.rs
+++ b/talpid-core/src/tunnel/wireguard/config.rs
@@ -47,10 +47,11 @@ impl Config {
     /// Constructs a Config from parameters
     pub fn from_parameters(params: &wireguard::TunnelParameters) -> Result<Config, Error> {
         let tunnel = params.connection.tunnel.clone();
-        let peer = vec![params.connection.peer.clone()];
+        let mut peers = vec![params.connection.peer.clone()];
+        peers.append(&mut params.connection.additional_peers.clone());
         Self::new(
             tunnel,
-            peer,
+            peers,
             &params.connection,
             &params.options,
             &params.generic_options,
diff --git a/talpid-types/src/net/wireguard.rs b/talpid-types/src/net/wireguard.rs
index 0fd8dfb843..a3cc75c88c 100644
--- a/talpid-types/src/net/wireguard.rs
+++ b/talpid-types/src/net/wireguard.rs
@@ -25,6 +25,7 @@ pub struct TunnelParameters {
 pub struct ConnectionConfig {
     pub tunnel: TunnelConfig,
     pub peer: PeerConfig,
+    pub additional_peers: Vec<PeerConfig>,
     /// Gateway used by the tunnel (a private address).
     pub ipv4_gateway: Ipv4Addr,
     pub ipv6_gateway: Option<Ipv6Addr>,