Merge branch 'pq-quickfix-mtu' into main

5 files changed, 159 insertions, 13 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1fe3de44c8..83184591d9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,6 +50,9 @@ Line wrap the file at 100 chars.                                              Th
 - Fix duplicated notifications in some situations.
 - Fix notification setting being inverted. When non-important notifications were disabled it instead
   disabled important ones and showed non-important ones.
+- Work around issues with PQ and multihop caused by fragmentation in the tunnel. The workaround
+  doesn't fix fragmentation issues in general but prevents the PSK exchange packets from being
+  fragmented by setting an explicit maximum segment size.
 
 #### Android
 - Fix adaptive app icon which previously had a displaced nose and some other oddities.
diff --git a/Cargo.lock b/Cargo.lock
index ce7e86ffe7..f1d8e294a3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3359,6 +3359,7 @@ name = "talpid-tunnel-config-client"
 version = "0.0.0"
 dependencies = [
  "classic-mceliece-rust",
+ "libc",
  "log",
  "pqc_kyber",
  "prost",
@@ -3368,6 +3369,7 @@ dependencies = [
  "tonic",
  "tonic-build",
  "tower",
+ "windows-sys 0.45.0",
  "zeroize",
 ]
 
diff --git a/talpid-tunnel-config-client/Cargo.toml b/talpid-tunnel-config-client/Cargo.toml
index 62542da683..0f39c4d1ea 100644
--- a/talpid-tunnel-config-client/Cargo.toml
+++ b/talpid-tunnel-config-client/Cargo.toml
@@ -18,6 +18,12 @@ tokio = "1"
 classic-mceliece-rust = { version = "2.0.0", features = ["mceliece460896f", "zeroize"] }
 pqc_kyber = { version = "0.4.0", features = ["std", "kyber1024", "zeroize"] }
 zeroize = "1.5.7"
+libc = "0.2"
+
+[target.'cfg(target_os="windows")'.dependencies]
+windows-sys = { version = "0.45.0", features = [
+     "Win32_Networking_WinSock"
+]}
 
 [dev-dependencies]
 tokio = { version = "1", features = ["rt-multi-thread"] }
diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs
index b795f6b814..f826eefe4f 100644
--- a/talpid-tunnel-config-client/src/lib.rs
+++ b/talpid-tunnel-config-client/src/lib.rs
@@ -1,6 +1,11 @@
-use std::{fmt, net::IpAddr};
+use std::{
+    fmt,
+    net::{IpAddr, SocketAddr},
+};
 use talpid_types::net::wireguard::{PresharedKey, PublicKey};
-use tonic::transport::Channel;
+use tokio::net::TcpSocket;
+use tonic::transport::{Channel, Endpoint};
+use tower::service_fn;
 use zeroize::Zeroize;
 
 mod classic_mceliece;
@@ -11,6 +16,20 @@ mod proto {
     tonic::include_proto!("tunnel_config");
 }
 
+use libc::setsockopt;
+
+#[cfg(not(target_os = "windows"))]
+mod sys {
+    pub use libc::{socklen_t, IPPROTO_TCP, TCP_MAXSEG};
+    pub use std::os::fd::{AsRawFd, RawFd};
+}
+#[cfg(target_os = "windows")]
+mod sys {
+    pub use std::os::windows::io::{AsRawSocket, RawSocket};
+    pub use windows_sys::Win32::Networking::WinSock::{IPPROTO_IP, IPPROTO_TCP, IP_USER_MTU};
+}
+use sys::*;
+
 #[derive(Debug)]
 pub enum Error {
     GrpcConnectError(tonic::transport::Error),
@@ -72,10 +91,19 @@ pub async fn push_pq_key(
     wg_pubkey: PublicKey,
     wg_psk_pubkey: PublicKey,
 ) -> Result<PresharedKey, Error> {
+    push_pq_key_with_opts(service_address, wg_pubkey, wg_psk_pubkey, None).await
+}
+
+pub async fn push_pq_key_with_opts(
+    service_address: IpAddr,
+    wg_pubkey: PublicKey,
+    wg_psk_pubkey: PublicKey,
+    mtu: Option<u16>,
+) -> Result<PresharedKey, Error> {
     let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
     let kyber_keypair = kyber::keypair(&mut rand::thread_rng());
 
-    let mut client = new_client(service_address).await?;
+    let mut client = new_client(service_address, mtu).await?;
     let response = client
         .psk_exchange_v1(proto::PskRequestV1 {
             wg_pubkey: wg_pubkey.as_bytes().to_vec(),
@@ -111,8 +139,9 @@ pub async fn push_pq_key(
         let mut shared_secret = classic_mceliece::decapsulate(&cme_kem_secret, cme_ciphertext)?;
         xor_assign(&mut psk_data, shared_secret.as_array());
 
-        // This should happen automatically due to `SharedSecret` implementing ZeroizeOnDrop. But doing it explicitly
-        // provides a stronger guarantee that it's not accidentally removed.
+        // This should happen automatically due to `SharedSecret` implementing ZeroizeOnDrop. But
+        // doing it explicitly provides a stronger guarantee that it's not accidentally
+        // removed.
         shared_secret.zeroize();
     }
     // Decapsulate Kyber and mix into PSK
@@ -137,8 +166,83 @@ fn xor_assign(dst: &mut [u8; 32], src: &[u8; 32]) {
     }
 }
 
-async fn new_client(addr: IpAddr) -> Result<RelayConfigService, Error> {
-    RelayConfigService::connect(format!("tcp://{addr}:{CONFIG_SERVICE_PORT}"))
+async fn new_client(addr: IpAddr, mtu: Option<u16>) -> Result<RelayConfigService, Error> {
+    let endpoint = Endpoint::from_static("tcp://0.0.0.0:0");
+
+    let conn = endpoint
+        .connect_with_connector(service_fn(move |_| async move {
+            let sock = TcpSocket::new_v4()?;
+
+            if let Some(mtu) = mtu {
+                #[cfg(target_os = "windows")]
+                try_set_tcp_sock_mtu(sock.as_raw_socket(), mtu);
+
+                #[cfg(not(target_os = "windows"))]
+                try_set_tcp_sock_mtu(&addr, sock.as_raw_fd(), mtu);
+            }
+
+            sock.connect(SocketAddr::new(addr, CONFIG_SERVICE_PORT))
+                .await
+        }))
         .await
-        .map_err(Error::GrpcConnectError)
+        .map_err(Error::GrpcConnectError)?;
+
+    Ok(RelayConfigService::new(conn))
+}
+
+#[cfg(windows)]
+fn try_set_tcp_sock_mtu(sock: RawSocket, mtu: u16) {
+    let mtu = u32::from(mtu);
+    log::debug!("Config client socket MTU: {mtu}");
+
+    let raw_sock = usize::try_from(sock).unwrap();
+
+    let result = unsafe {
+        setsockopt(
+            raw_sock,
+            IPPROTO_IP as i32,
+            IP_USER_MTU as i32,
+            &mtu as *const _ as _,
+            std::ffi::c_int::try_from(std::mem::size_of_val(&mtu)).unwrap(),
+        )
+    };
+    if result != 0 {
+        log::error!(
+            "Failed to set user MTU on config client socket: {}",
+            std::io::Error::last_os_error()
+        );
+    }
+}
+
+#[cfg(not(windows))]
+fn try_set_tcp_sock_mtu(dest: &IpAddr, sock: RawFd, mut mtu: u16) {
+    const IPV4_HEADER_SIZE: u16 = 20;
+    const IPV6_HEADER_SIZE: u16 = 40;
+    const MAX_TCP_HEADER_SIZE: u16 = 60;
+
+    if dest.is_ipv4() {
+        mtu = mtu.saturating_sub(IPV4_HEADER_SIZE);
+    } else {
+        mtu = mtu.saturating_sub(IPV6_HEADER_SIZE);
+    }
+
+    let mss = u32::from(mtu.saturating_sub(MAX_TCP_HEADER_SIZE));
+
+    log::debug!("Config client socket MSS: {mss}");
+
+    let result = unsafe {
+        setsockopt(
+            sock,
+            IPPROTO_TCP,
+            TCP_MAXSEG,
+            &mss as *const _ as _,
+            socklen_t::try_from(std::mem::size_of_val(&mss)).unwrap(),
+        )
+    };
+    if result != 0 {
+        log::error!(
+            "Failed to set MSS on config client socket: {}",
+            std::io::Error::last_os_error()
+        );
+    }
 }
diff --git a/talpid-wireguard/src/lib.rs b/talpid-wireguard/src/lib.rs
index 85b4c5a6df..3d542641b7 100644
--- a/talpid-wireguard/src/lib.rs
+++ b/talpid-wireguard/src/lib.rs
@@ -121,8 +121,8 @@ pub struct WireguardMonitor {
     obfuscator: Arc<AsyncMutex<Option<ObfuscatorHandle>>>,
 }
 
-const INITIAL_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(4);
-const MAX_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(15);
+const INITIAL_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(8);
+const MAX_PSK_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(48);
 const PSK_EXCHANGE_TIMEOUT_MULTIPLIER: u32 = 2;
 
 /// Simple wrapper that automatically cancels the future which runs an obfuscator.
@@ -473,8 +473,9 @@ impl WireguardMonitor {
         let metadata = Self::tunnel_metadata(iface_name, config);
         (on_event)(TunnelEvent::InterfaceUp(metadata, allowed_traffic.clone())).await;
 
+        let mtu = Self::get_exit_psk_socket_mtu(config);
         let exit_psk =
-            Self::perform_psk_negotiation(retry_attempt, config, wg_psk_privkey.public_key())
+            Self::perform_psk_negotiation(retry_attempt, config, wg_psk_privkey.public_key(), mtu)
                 .await?;
 
         log::debug!("Successfully exchanged PSK with exit peer");
@@ -507,6 +508,7 @@ impl WireguardMonitor {
                     retry_attempt,
                     &entry_config,
                     wg_psk_privkey.public_key(),
+                    None,
                 )
                 .await?,
             );
@@ -638,6 +640,7 @@ impl WireguardMonitor {
         retry_attempt: u32,
         config: &Config,
         wg_psk_pubkey: PublicKey,
+        mss: Option<u16>,
     ) -> std::result::Result<PresharedKey, CloseMsg> {
         log::debug!("Performing PQ-safe PSK exchange");
 
@@ -649,10 +652,11 @@ impl WireguardMonitor {
 
         let psk = tokio::time::timeout(
             timeout,
-            talpid_tunnel_config_client::push_pq_key(
-                IpAddr::V4(config.ipv4_gateway),
+            talpid_tunnel_config_client::push_pq_key_with_opts(
+                Self::get_tunnel_config_client_addr(config),
                 config.tunnel.private_key.public_key(),
                 wg_psk_pubkey,
+                mss,
             ),
         )
         .await
@@ -666,6 +670,33 @@ impl WireguardMonitor {
         Ok(psk)
     }
 
+    /// TODO: This hack makes sure that the max segment size is small enough to prevent
+    /// fragmentation. This is an issue when using multihop, because the MTU is too small to
+    /// accommodate the headers of the inner tunnel, and lowering it does no good.
+    //
+    /// When we have fixed fragmentation issues for multihop, this workaround can be removed.
+    ///
+    /// Note that TCP imposes max and min bounds on the specified size as well.
+    fn get_exit_psk_socket_mtu(config: &Config) -> Option<u16> {
+        // macOS is averse(?) to large values for MSS, so we just use the minimum
+        const MIN_IPV4_MTU: u16 = 576;
+        const MIN_IPV6_MTU: u16 = 1280;
+
+        if config.peers.len() == 1 {
+            return None;
+        }
+
+        if Self::get_tunnel_config_client_addr(config).is_ipv4() {
+            Some(MIN_IPV4_MTU)
+        } else {
+            Some(MIN_IPV6_MTU)
+        }
+    }
+
+    fn get_tunnel_config_client_addr(config: &Config) -> IpAddr {
+        IpAddr::V4(config.ipv4_gateway)
+    }
+
     #[allow(unused_variables)]
     fn open_tunnel(
         runtime: tokio::runtime::Handle,