Merge branch 'linux-lower-entry-multihop-mtu' into main

9 files changed, 153 insertions, 94 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3ce9628a1a..7be653e994 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,10 @@ Line wrap the file at 100 chars.                                              Th
 #### Linux
 - Rename interface name from `wg-mullvad` to `wg0-mullvad`.
 
+### Fixed
+#### Linux
+- Prevent fragmentation when multihop is enabled by setting a default route MTU.
+
 
 ## [2023.6-beta1] - 2023-11-23
 ### Added
diff --git a/talpid-routing/Cargo.toml b/talpid-routing/Cargo.toml
index e5cfb9d9f9..31545fe705 100644
--- a/talpid-routing/Cargo.toml
+++ b/talpid-routing/Cargo.toml
@@ -23,7 +23,7 @@ talpid-types = { path = "../talpid-types" }
 libc = "0.2"
 once_cell = { workspace = true }
 rtnetlink = "0.11"
-netlink-packet-route = "0.13"
+netlink-packet-route = { version = "0.13", features = ["rich_nlas"] }
 netlink-sys = "0.8.3"
 
 [target.'cfg(target_os = "macos")'.dependencies]
diff --git a/talpid-routing/src/lib.rs b/talpid-routing/src/lib.rs
index dd5fd3a761..f1bf28d1a2 100644
--- a/talpid-routing/src/lib.rs
+++ b/talpid-routing/src/lib.rs
@@ -38,6 +38,8 @@ pub struct Route {
     metric: Option<u32>,
     #[cfg(target_os = "linux")]
     table_id: u32,
+    #[cfg(target_os = "linux")]
+    mtu: Option<u32>,
 }
 
 impl Route {
@@ -49,6 +51,8 @@ impl Route {
             metric: None,
             #[cfg(target_os = "linux")]
             table_id: u32::from(RT_TABLE_MAIN),
+            #[cfg(target_os = "linux")]
+            mtu: None,
         }
     }
 
@@ -72,6 +76,10 @@ impl fmt::Display for Route {
         }
         #[cfg(target_os = "linux")]
         write!(f, " table {}", self.table_id)?;
+        #[cfg(target_os = "linux")]
+        if let Some(mtu) = self.mtu {
+            write!(f, " mtu {mtu}")?;
+        }
         Ok(())
     }
 }
@@ -87,6 +95,9 @@ pub struct RequiredRoute {
     /// Specifies whether the route should be added to the main routing table or not.
     #[cfg(target_os = "linux")]
     main_table: bool,
+    /// Specifies route MTU
+    #[cfg(target_os = "linux")]
+    mtu: Option<u16>,
 }
 
 impl RequiredRoute {
@@ -97,6 +108,8 @@ impl RequiredRoute {
             prefix,
             #[cfg(target_os = "linux")]
             main_table: true,
+            #[cfg(target_os = "linux")]
+            mtu: None,
         }
     }
 
@@ -106,6 +119,13 @@ impl RequiredRoute {
         self.main_table = main_table;
         self
     }
+
+    /// Set route MTU to the given value.
+    #[cfg(target_os = "linux")]
+    pub fn mtu(mut self, mtu: u16) -> Self {
+        self.mtu = Some(mtu);
+        self
+    }
 }
 
 /// A NetNode represents a network node - either a real one or a symbolic default one.
diff --git a/talpid-routing/src/unix/linux.rs b/talpid-routing/src/unix/linux.rs
index a2cf19d3b4..600ddd6794 100644
--- a/talpid-routing/src/unix/linux.rs
+++ b/talpid-routing/src/unix/linux.rs
@@ -20,7 +20,7 @@ use libc::{AF_INET, AF_INET6};
 use netlink_packet_route::{
     constants::{ARPHRD_LOOPBACK, FIB_RULE_INVERT, FR_ACT_TO_TBL, NLM_F_REQUEST},
     link::{nlas::Nla as LinkNla, LinkMessage},
-    route::{nlas::Nla as RouteNla, RouteHeader, RouteMessage},
+    route::{nlas::Nla as RouteNla, Metrics, RouteHeader, RouteMessage},
     rtnl::{
         constants::{
             RTN_UNSPEC, RTPROT_UNSPEC, RT_SCOPE_LINK, RT_SCOPE_UNIVERSE, RT_TABLE_COMPAT,
@@ -293,7 +293,9 @@ impl RouteManagerImpl {
                     } else {
                         self.table_id
                     };
-                    required_normal_routes.insert(Route::new(node, route.prefix).table(table));
+                    let mut new_route = Route::new(node, route.prefix).table(table);
+                    new_route.mtu = route.mtu.map(u32::from);
+                    required_normal_routes.insert(new_route);
                 }
             }
         }
@@ -450,12 +452,13 @@ impl RouteManagerImpl {
             destination_length,
         )
         .map_err(Error::InvalidNetworkPrefix)?;
+
         let mut node_addr = None;
         let mut device = None;
         let mut metric = None;
         let mut gateway: Option<IpAddr> = None;
-
         let mut table_id = u32::from(msg.header.table);
+        let mut route_mtu = None;
 
         for nla in msg.nlas.iter() {
             match nla {
@@ -501,6 +504,10 @@ impl RouteManagerImpl {
                 RouteNla::Table(id) => {
                     table_id = *id;
                 }
+
+                RouteNla::Metrics(Metrics::Mtu(mtu)) => {
+                    route_mtu = Some(*mtu);
+                }
                 _ => continue,
             }
         }
@@ -519,6 +526,7 @@ impl RouteManagerImpl {
             prefix,
             metric,
             table_id,
+            mtu: route_mtu,
         }))
     }
 
@@ -700,6 +708,11 @@ impl RouteManagerImpl {
             add_message.nlas.push(RouteNla::Priority(metric));
         }
 
+        // Set route MTU
+        if let Some(mtu) = route.mtu {
+            add_message.nlas.push(RouteNla::Metrics(Metrics::Mtu(mtu)));
+        }
+
         // Need to modify the request in place to set the correct flags to be able to replace any
         // existing routes - self.handle.route().add_v4().execute() sets the NLM_F_EXCL flag which
         // will make the request fail if a route with the same destination already exists.
@@ -743,6 +756,7 @@ impl RouteManagerImpl {
     async fn get_mtu_for_route(&self, ip: IpAddr) -> Result<u16> {
         // RECURSION_LIMIT controls how many times we recurse to find the device name by looking up
         // an IP with `get_destination_route`.
+        // TODO: Check route MTU first
         const RECURSION_LIMIT: usize = 10;
         const STANDARD_MTU: u16 = 1500;
         let mut attempted_ip = ip;
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 c95a5d371b..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(()),
         };
 
@@ -880,6 +861,10 @@ impl WireguardMonitor {
 
         let (node_v4, node_v6) = Self::get_tunnel_nodes(iface_name, config);
 
+        #[cfg(target_os = "linux")]
+        let gateway_routes =
+            gateway_routes.map(|route| Self::apply_route_mtu_for_multihop(route, config));
+
         let routes = gateway_routes.chain(
             Self::get_tunnel_destinations(config)
                 .filter(|allowed_ip| allowed_ip.prefix() != 0)
@@ -916,13 +901,35 @@ impl WireguardMonitor {
 
         #[cfg(target_os = "linux")]
         iter.map(|route| route.use_main_table(false))
+            .map(|route| Self::apply_route_mtu_for_multihop(route, config))
+    }
+
+    #[cfg(target_os = "linux")]
+    fn apply_route_mtu_for_multihop(route: RequiredRoute, config: &Config) -> RequiredRoute {
+        if !config.is_multihop() {
+            route
+        } else {
+            // Set route MTU by subtracting the WireGuard overhead from the tunnel MTU.
+            // NOTE: Somewhat incorrect since it doesn't account for packet padding/alignment?
+            // TODO: Move consts to shared location
+            const IPV4_HEADER_SIZE: u16 = 20;
+            const IPV6_HEADER_SIZE: u16 = 40;
+            const WIREGUARD_HEADER_SIZE: u16 = 40;
+
+            let ip_overhead = match route.prefix.is_ipv4() {
+                true => IPV4_HEADER_SIZE,
+                false => IPV6_HEADER_SIZE,
+            };
+            let mtu = config.mtu - ip_overhead - WIREGUARD_HEADER_SIZE;
+
+            route.mtu(mtu)
+        }
     }
 
     /// 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()
     }
@@ -962,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,