Remove hidden assumptions from WireGuard config

5 files changed, 85 insertions, 91 deletions

diff --git a/talpid-wireguard/src/config.rs b/talpid-wireguard/src/config.rs
index fa4b7e078e..0e462102b2 100644
--- a/talpid-wireguard/src/config.rs
+++ b/talpid-wireguard/src/config.rs
@@ -10,8 +10,10 @@ use talpid_types::net::{obfuscation::ObfuscatorConfig, wireguard, GenericTunnelO
 pub struct Config {
     /// Contains tunnel endpoint specific config
     pub tunnel: wireguard::TunnelConfig,
-    /// List of peer configurations
-    pub peers: Vec<wireguard::PeerConfig>,
+    /// Entry peer
+    pub entry_peer: wireguard::PeerConfig,
+    /// Multihop exit peer
+    pub exit_peer: Option<wireguard::PeerConfig>,
     /// IPv4 gateway
     pub ipv4_gateway: Ipv4Addr,
     /// IPv6 gateway
@@ -46,54 +48,28 @@ pub enum Error {
     /// Peer has no valid IPs
     #[error(display = "Supplied peer has no valid IPs")]
     InvalidPeerIpError,
-
-    /// Parameters don't contain any peers
-    #[error(display = "No peers supplied")]
-    NoPeersSuppliedError,
 }
 
 impl Config {
     /// Constructs a Config from parameters
     pub fn from_parameters(params: &wireguard::TunnelParameters) -> Result<Config, Error> {
-        let tunnel = params.connection.tunnel.clone();
-        let mut peers = vec![params.connection.peer.clone()];
-        if let Some(exit_peer) = &params.connection.exit_peer {
-            peers.push(exit_peer.clone());
-        }
         Self::new(
-            tunnel,
-            peers,
             &params.connection,
             &params.options,
             &params.generic_options,
-            params.obfuscation.clone(),
+            &params.obfuscation,
         )
     }
 
     /// Constructs a new Config struct
-    pub fn new(
-        mut tunnel: wireguard::TunnelConfig,
-        mut peers: Vec<wireguard::PeerConfig>,
-        connection_config: &wireguard::ConnectionConfig,
+    fn new(
+        connection: &wireguard::ConnectionConfig,
         wg_options: &wireguard::TunnelOptions,
         generic_options: &GenericTunnelOptions,
-        obfuscator_config: Option<ObfuscatorConfig>,
+        obfuscator_config: &Option<ObfuscatorConfig>,
     ) -> Result<Config, Error> {
-        if peers.is_empty() {
-            return Err(Error::NoPeersSuppliedError);
-        }
+        let mut tunnel = connection.tunnel.clone();
         let mtu = wg_options.mtu.unwrap_or(DEFAULT_MTU);
-        for peer in &mut peers {
-            peer.allowed_ips = peer
-                .allowed_ips
-                .iter()
-                .cloned()
-                .filter(|ip| ip.is_ipv4() || generic_options.enable_ipv6)
-                .collect();
-            if peer.allowed_ips.is_empty() {
-                return Err(Error::InvalidPeerIpError);
-            }
-        }
 
         if tunnel.addresses.is_empty() {
             return Err(Error::InvalidTunnelIpError);
@@ -102,24 +78,33 @@ impl Config {
             .addresses
             .retain(|ip| ip.is_ipv4() || generic_options.enable_ipv6);
 
-        let ipv6_gateway = if generic_options.enable_ipv6 {
-            connection_config.ipv6_gateway
-        } else {
-            None
-        };
+        let ipv6_gateway = connection
+            .ipv6_gateway
+            .filter(|_opt| generic_options.enable_ipv6);
 
-        Ok(Config {
+        let mut config = Config {
             tunnel,
-            peers,
-            ipv4_gateway: connection_config.ipv4_gateway,
+            entry_peer: connection.peer.clone(),
+            exit_peer: connection.exit_peer.clone(),
+            ipv4_gateway: connection.ipv4_gateway,
             ipv6_gateway,
             mtu,
             #[cfg(target_os = "linux")]
-            fwmark: connection_config.fwmark,
+            fwmark: connection.fwmark,
             #[cfg(target_os = "linux")]
             enable_ipv6: generic_options.enable_ipv6,
-            obfuscator_config,
-        })
+            obfuscator_config: obfuscator_config.to_owned(),
+        };
+
+        for peer in config.peers_mut() {
+            peer.allowed_ips
+                .retain(|ip| ip.is_ipv4() || generic_options.enable_ipv6);
+            if peer.allowed_ips.is_empty() {
+                return Err(Error::InvalidPeerIpError);
+            }
+        }
+
+        Ok(config)
     }
 
     /// Returns a CString with the appropriate config for WireGuard-go
@@ -138,7 +123,7 @@ impl Config {
 
         wg_conf.add("replace_peers", "true");
 
-        for peer in &self.peers {
+        for peer in self.peers() {
             wg_conf
                 .add("public_key", peer.public_key.as_bytes().as_ref())
                 .add("endpoint", peer.endpoint.to_string().as_str())
@@ -154,6 +139,35 @@ impl Config {
         let bytes = wg_conf.into_config();
         CString::new(bytes).expect("null bytes inside config")
     }
+
+    /// Return whether the config connects to an exit peer from another remote peer.
+    pub fn is_multihop(&self) -> bool {
+        self.exit_peer.is_some()
+    }
+
+    /// Return the exit peer. `exit_peer` if it is set, otherwise `entry_peer`.
+    pub fn exit_peer_mut(&mut self) -> &mut wireguard::PeerConfig {
+        if let Some(ref mut peer) = self.exit_peer {
+            return peer;
+        }
+        &mut self.entry_peer
+    }
+
+    /// Return an iterator over all peers.
+    pub fn peers(&self) -> impl Iterator<Item = &wireguard::PeerConfig> {
+        self.exit_peer
+            .as_ref()
+            .into_iter()
+            .chain(std::iter::once(&self.entry_peer))
+    }
+
+    /// Return a mutable iterator over all peers.
+    pub fn peers_mut(&mut self) -> impl Iterator<Item = &mut wireguard::PeerConfig> {
+        self.exit_peer
+            .as_mut()
+            .into_iter()
+            .chain(std::iter::once(&mut self.entry_peer))
+    }
 }
 
 enum ConfValue<'a> {
diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index 9ca230b652..27f7b3d9a5 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -93,10 +93,6 @@ pub enum Error {
     #[error(display = "Failed to negotiate PQ PSK")]
     PskNegotiationError(#[error(source)] talpid_tunnel_config_client::Error),
 
-    /// Too many peers in the config
-    #[error(display = "There are too many peers in the tunnel config")]
-    TooManyPeers,
-
     /// Failed to set up IP interfaces.
     #[cfg(windows)]
     #[error(display = "Failed to set up IP interfaces")]
@@ -171,10 +167,6 @@ async fn maybe_create_obfuscator(
     config: &mut Config,
     close_msg_sender: sync_mpsc::Sender<CloseMsg>,
 ) -> Result<Option<ObfuscatorHandle>> {
-    // There are one or two peers.
-    // The first one is always the entry relay.
-    let first_peer = config.peers.get_mut(0).expect("missing peer");
-
     if let Some(ref obfuscator_config) = config.obfuscator_config {
         match obfuscator_config {
             ObfuscatorConfig::Udp2Tcp { endpoint } => {
@@ -190,7 +182,7 @@ async fn maybe_create_obfuscator(
                 let endpoint = obfuscator.endpoint();
 
                 log::trace!("Patching first WireGuard peer to become {:?}", endpoint);
-                first_peer.endpoint = endpoint;
+                config.entry_peer.endpoint = endpoint;
 
                 #[cfg(target_os = "android")]
                 let remote_socket_fd = obfuscator.remote_socket_fd();
@@ -238,8 +230,7 @@ impl WireguardMonitor {
     ) -> Result<WireguardMonitor> {
         let on_event = args.on_event.clone();
 
-        let endpoint_addrs: Vec<IpAddr> =
-            config.peers.iter().map(|peer| peer.endpoint.ip()).collect();
+        let endpoint_addrs: Vec<IpAddr> = config.peers().map(|peer| peer.endpoint.ip()).collect();
 
         let (close_obfs_sender, close_obfs_listener) = sync_mpsc::channel();
         let obfuscator = args.runtime.block_on(maybe_create_obfuscator(
@@ -256,7 +247,7 @@ impl WireguardMonitor {
         // properly so fragmentation does not happen.
         let init_tunnel_config = if cfg!(target_os = "macos") {
             let mut init_tunnel_config = config.clone();
-            if psk_negotiation && config.peers.len() > 1 {
+            if psk_negotiation && config.is_multihop() {
                 const MH_PQ_HANDSHAKE_MTU: u16 = 1280;
                 init_tunnel_config.mtu = MH_PQ_HANDSHAKE_MTU;
             }
@@ -457,13 +448,16 @@ impl WireguardMonitor {
             talpid_tunnel_config_client::CONFIG_SERVICE_PORT,
             TransportProtocol::Tcp,
         );
-        let allowed_traffic = if config.peers.len() > 1 {
+        let allowed_traffic = if config.is_multihop() {
             // NOTE: We need to let traffic meant for the exit IP through the firewall. This
             // should not allow any non-PQ traffic to leak since you can only reach the
             // exit peer with these rules and not the broader internet.
             AllowedTunnelTraffic::Two(
                 allowed_traffic,
-                Endpoint::from_socket_address(config.peers[1].endpoint, TransportProtocol::Udp),
+                Endpoint::from_socket_address(
+                    config.exit_peer_mut().endpoint,
+                    TransportProtocol::Udp,
+                ),
             )
         } else {
             AllowedTunnelTraffic::One(allowed_traffic)
@@ -478,18 +472,11 @@ impl WireguardMonitor {
 
         log::debug!("Successfully exchanged PSK with exit peer");
 
-        let mut entry_psk = None;
-
-        if config.peers.len() > 1 {
-            if config.peers.len() != 2 {
-                return Err(CloseMsg::TooManyPeers);
-            }
+        if config.is_multihop() {
             // Set up tunnel to lead to entry
             let mut entry_tun_config = config.clone();
             entry_tun_config
-                .peers
-                .get_mut(0)
-                .expect("entry peer not found")
+                .entry_peer
                 .allowed_ips
                 .push(IpNetwork::new(IpAddr::V4(config.ipv4_gateway), 32).unwrap());
 
@@ -503,7 +490,7 @@ impl WireguardMonitor {
                 &tun_provider,
             )
             .await?;
-            entry_psk = Some(
+            let entry_psk = Some(
                 Self::perform_psk_negotiation(
                     retry_attempt,
                     &entry_config,
@@ -513,18 +500,13 @@ impl WireguardMonitor {
                 .await?,
             );
             log::debug!("Successfully exchanged PSK with entry peer");
+
+            config.entry_peer.psk = entry_psk;
         }
 
-        // Set new priv key and psks
+        config.exit_peer_mut().psk = Some(exit_psk);
+
         config.tunnel.private_key = wg_psk_privkey;
-        if let Some(entry_psk) = entry_psk {
-            // The first peer is the entry peer and there is guaranteed to be a second peer
-            // which is the exit
-            config.peers.get_mut(0).expect("entry peer not found").psk = Some(entry_psk);
-            config.peers.get_mut(1).expect("exit peer not found").psk = Some(exit_psk);
-        } else {
-            config.peers.get_mut(0).expect("peer not found").psk = Some(exit_psk);
-        }
 
         *config = Self::reconfigure_tunnel(
             tunnel,
@@ -596,7 +578,7 @@ impl WireguardMonitor {
             let gateway_net_v6 = config
                 .ipv6_gateway
                 .map(|net| ipnetwork::IpNetwork::from(IpAddr::from(net)));
-            for peer in &mut patched_config.peers {
+            for peer in patched_config.peers_mut() {
                 peer.allowed_ips = peer
                     .allowed_ips
                     .iter()
@@ -700,7 +682,7 @@ impl WireguardMonitor {
         const MIN_IPV4_MTU: u16 = 576;
         const MIN_IPV6_MTU: u16 = 1280;
 
-        if config.peers.len() == 1 {
+        if !config.is_multihop() {
             return None;
         }
 
@@ -800,7 +782,6 @@ impl WireguardMonitor {
             Ok(CloseMsg::Stop) | Ok(CloseMsg::ObfuscatorExpired) => Ok(()),
             Ok(CloseMsg::SetupError(error)) => Err(error),
             Ok(CloseMsg::ObfuscatorFailed(error)) => Err(error),
-            Ok(CloseMsg::TooManyPeers) => Err(Error::TooManyPeers),
             Err(_) => Ok(()),
         };
 
@@ -925,7 +906,7 @@ impl WireguardMonitor {
 
     #[cfg(target_os = "linux")]
     fn apply_route_mtu_for_multihop(route: RequiredRoute, config: &Config) -> RequiredRoute {
-        if config.peers.len() == 1 {
+        if !config.is_multihop() {
             route
         } else {
             // Set route MTU by subtracting the WireGuard overhead from the tunnel MTU.
@@ -948,8 +929,7 @@ impl WireguardMonitor {
     /// Return routes for all allowed IPs.
     fn get_tunnel_destinations(config: &Config) -> impl Iterator<Item = ipnetwork::IpNetwork> + '_ {
         config
-            .peers
-            .iter()
+            .peers()
             .flat_map(|peer| peer.allowed_ips.iter())
             .cloned()
     }
@@ -989,7 +969,6 @@ enum CloseMsg {
     SetupError(Error),
     ObfuscatorExpired,
     ObfuscatorFailed(Error),
-    TooManyPeers,
 }
 
 pub(crate) trait Tunnel: Send {
diff --git a/talpid-wireguard/src/wireguard_kernel/nm_tunnel.rs b/talpid-wireguard/src/wireguard_kernel/nm_tunnel.rs
index 3f2661a4dc..7b5966b9e4 100644
--- a/talpid-wireguard/src/wireguard_kernel/nm_tunnel.rs
+++ b/talpid-wireguard/src/wireguard_kernel/nm_tunnel.rs
@@ -130,7 +130,7 @@ fn convert_config_to_dbus(config: &Config) -> DeviceConfig {
     );
     wireguard_config.insert("private-key-flags".into(), Variant(Box::new(0x0u32)));
 
-    for peer in config.peers.iter() {
+    for peer in config.peers() {
         let mut peer_config: VariantMap = HashMap::new();
         let allowed_ips = peer
             .allowed_ips
diff --git a/talpid-wireguard/src/wireguard_kernel/wg_message.rs b/talpid-wireguard/src/wireguard_kernel/wg_message.rs
index 7ed972c3ea..4dc84ed503 100644
--- a/talpid-wireguard/src/wireguard_kernel/wg_message.rs
+++ b/talpid-wireguard/src/wireguard_kernel/wg_message.rs
@@ -78,7 +78,7 @@ impl DeviceMessage {
     pub fn reset_config(message_type: u16, interface_index: u32, config: &Config) -> DeviceMessage {
         let mut peers = vec![];
 
-        for peer in config.peers.iter() {
+        for peer in config.peers() {
             let peer_endpoint = InetAddr::from_std(&peer.endpoint);
             let allowed_ips = peer.allowed_ips.iter().map(From::from).collect();
             let mut peer_nlas = vec![
diff --git a/talpid-wireguard/src/wireguard_nt.rs b/talpid-wireguard/src/wireguard_nt.rs
index 588d5a7f82..0a9cc15219 100644
--- a/talpid-wireguard/src/wireguard_nt.rs
+++ b/talpid-wireguard/src/wireguard_nt.rs
@@ -811,12 +811,12 @@ fn serialize_config(config: &Config) -> Result<Vec<MaybeUninit<u8>>> {
         listen_port: 0,
         private_key: config.tunnel.private_key.to_bytes(),
         public_key: [0u8; WIREGUARD_KEY_LENGTH],
-        peers_count: u32::try_from(config.peers.len()).unwrap(),
+        peers_count: u32::try_from(config.peers().count()).unwrap(),
     };
 
     buffer.extend(as_uninit_byte_slice(&header));
 
-    for peer in &config.peers {
+    for peer in config.peers() {
         let flags = if peer.psk.is_some() {
             WgPeerFlag::HAS_PRESHARED_KEY | WgPeerFlag::HAS_PUBLIC_KEY | WgPeerFlag::HAS_ENDPOINT
         } else {
@@ -1005,12 +1005,13 @@ mod tests {
             private_key: WG_PRIVATE_KEY.clone(),
             addresses: vec![],
         },
-        peers: vec![wireguard::PeerConfig {
+        entry_peer: wireguard::PeerConfig {
             public_key: WG_PUBLIC_KEY.clone(),
             allowed_ips: vec!["1.3.3.0/24".parse().unwrap()],
             endpoint: "1.2.3.4:1234".parse().unwrap(),
             psk: None,
-        }],
+        },
+        exit_peer: None,
         ipv4_gateway: "0.0.0.0".parse().unwrap(),
         ipv6_gateway: None,
         mtu: 0,