Implement initial PQ PSK exchange for wireguard-nt

4 files changed, 87 insertions, 3 deletions

diff --git a/talpid-core/src/tunnel/mod.rs b/talpid-core/src/tunnel/mod.rs
index ea7ec35133..8820f26c72 100644
--- a/talpid-core/src/tunnel/mod.rs
+++ b/talpid-core/src/tunnel/mod.rs
@@ -198,6 +198,7 @@ impl TunnelMonitor {
         let monitor = wireguard::WireguardMonitor::start(
             runtime,
             config,
+            Some(params.connection.peer.public_key.clone()),
             log.as_deref(),
             resource_dir,
             on_event,
diff --git a/talpid-core/src/tunnel/wireguard/connectivity_check.rs b/talpid-core/src/tunnel/wireguard/connectivity_check.rs
index ec2af873a0..bcb28d7c17 100644
--- a/talpid-core/src/tunnel/wireguard/connectivity_check.rs
+++ b/talpid-core/src/tunnel/wireguard/connectivity_check.rs
@@ -391,12 +391,16 @@ impl ConnState {
 
 #[cfg(test)]
 mod test {
+    use futures::Future;
+
     use super::*;
     use crate::tunnel::wireguard::{
+        config::Config,
         stats::{self, Stats},
         TunnelError,
     };
     use std::{
+        pin::Pin,
         sync::{
             atomic::{AtomicBool, Ordering},
             Arc, Mutex,
@@ -598,6 +602,13 @@ mod test {
         fn get_tunnel_stats(&self) -> Result<stats::StatsMap, TunnelError> {
             (self.on_get_stats)()
         }
+
+        fn set_config(
+            &self,
+            _config: Config,
+        ) -> Pin<Box<dyn Future<Output = std::result::Result<(), TunnelError>> + Send>> {
+            Box::pin(async { Ok(()) })
+        }
     }
 
     fn mock_monitor(
diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 80c8e2ae8e..dc2300b0ab 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -25,7 +25,10 @@ use std::{
 };
 #[cfg(windows)]
 use talpid_types::BoxedError;
-use talpid_types::{net::obfuscation::ObfuscatorConfig, ErrorExt};
+use talpid_types::{
+    net::{obfuscation::ObfuscatorConfig, wireguard::PublicKey},
+    ErrorExt,
+};
 use tunnel_obfuscation::{
     create_obfuscator, Error as ObfuscationError, Settings as ObfuscationSettings, Udp2TcpSettings,
 };
@@ -73,6 +76,10 @@ pub enum Error {
     #[error(display = "Connectivity monitor failed")]
     ConnectivityMonitorError(#[error(source)] connectivity_check::Error),
 
+    /// Failed to negotiate PQ PSK
+    #[error(display = "Failed to negotiate PQ PSK")]
+    PskNegotiationError(talpid_relay_config_client::Error),
+
     /// Failed to set up IP interfaces.
     #[cfg(windows)]
     #[error(display = "Failed to set up IP interfaces")]
@@ -188,6 +195,7 @@ impl WireguardMonitor {
     >(
         runtime: tokio::runtime::Handle,
         mut config: Config,
+        psk_negotiation: Option<PublicKey>,
         log_path: Option<&Path>,
         resource_dir: &Path,
         on_event: F,
@@ -237,6 +245,7 @@ impl WireguardMonitor {
         .map_err(Error::ConnectivityMonitorError)?;
 
         let metadata = Self::tunnel_metadata(&iface_name, &config);
+        let tunnel = monitor.tunnel.clone();
 
         let tunnel_fut = async move {
             #[cfg(windows)]
@@ -280,6 +289,35 @@ impl WireguardMonitor {
                 .map_err(Error::SetupRoutingError)
                 .map_err(CloseMsg::SetupError)?;
 
+            if let Some(pubkey) = psk_negotiation {
+                // TODO: add timeout
+                let (private_key, psk) = talpid_relay_config_client::push_pq_key(
+                    IpAddr::V4(config.ipv4_gateway),
+                    config.tunnel.private_key.public_key(),
+                )
+                .await
+                .map_err(Error::PskNegotiationError)
+                .map_err(CloseMsg::SetupError)?;
+
+                config.tunnel.private_key = private_key;
+
+                for peer in &mut config.peers {
+                    if pubkey == peer.public_key {
+                        peer.psk = Some(psk);
+                        break;
+                    }
+                }
+
+                log::trace!("Ephemeral pubkey: {}", config.tunnel.private_key.public_key());
+
+                if let Some(tunnel) = &*tunnel.lock().unwrap() {
+                    tunnel
+                        .set_config(&config)
+                        .map_err(Error::TunnelError)
+                        .map_err(CloseMsg::SetupError)?;
+                }
+            }
+
             let mut connectivity_monitor = tokio::task::spawn_blocking(move || {
                 match connectivity_monitor.establish_connectivity(retry_attempt) {
                     Ok(true) => Ok(connectivity_monitor),
@@ -594,6 +632,9 @@ pub(crate) trait Tunnel: Send {
     fn get_interface_name(&self) -> String;
     fn stop(self: Box<Self>) -> std::result::Result<(), TunnelError>;
     fn get_tunnel_stats(&self) -> std::result::Result<stats::StatsMap, TunnelError>;
+    fn set_config(&self, _config: &Config) -> std::result::Result<(), TunnelError> {
+        unimplemented!()
+    }
 }
 
 /// Errors to be returned from WireGuard implementations, namely implementers of the Tunnel trait
@@ -630,6 +671,10 @@ pub enum TunnelError {
     #[error(display = "Failed to get config of WireGuard tunnel")]
     GetConfigError,
 
+    /// Failed to set WireGuard tunnel config on device
+    #[error(display = "Failed to set config of WireGuard tunnel")]
+    SetConfigError,
+
     /// Failed to duplicate tunnel file descriptor for wireguard-go
     #[cfg(any(target_os = "linux", target_os = "macos", target_os = "android"))]
     #[error(display = "Failed to duplicate tunnel file descriptor for wireguard-go")]
diff --git a/talpid-core/src/tunnel/wireguard/wireguard_nt.rs b/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
index 21c9b705ce..31222d2ce8 100644
--- a/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
+++ b/talpid-core/src/tunnel/wireguard/wireguard_nt.rs
@@ -833,11 +833,20 @@ fn serialize_config(config: &Config) -> Result<Vec<MaybeUninit<u8>>> {
     buffer.extend(windows::as_uninit_byte_slice(&header));
 
     for peer in &config.peers {
+        let flags = if peer.psk.is_some() {
+            WgPeerFlag::HAS_PRESHARED_KEY | WgPeerFlag::HAS_PUBLIC_KEY | WgPeerFlag::HAS_ENDPOINT
+        } else {
+            WgPeerFlag::HAS_PUBLIC_KEY | WgPeerFlag::HAS_ENDPOINT
+        };
         let wg_peer = WgPeer {
-            flags: WgPeerFlag::HAS_PUBLIC_KEY | WgPeerFlag::HAS_ENDPOINT,
+            flags,
             reserved: 0,
             public_key: peer.public_key.as_bytes().clone(),
-            preshared_key: [0u8; WIREGUARD_KEY_LENGTH],
+            preshared_key: peer
+                .psk
+                .as_ref()
+                .map(|psk| psk.as_bytes().clone())
+                .unwrap_or([0u8; WIREGUARD_KEY_LENGTH]),
             persistent_keepalive: 0,
             endpoint: windows::inet_sockaddr_from_socketaddr(peer.endpoint).into(),
             tx_bytes: 0,
@@ -976,6 +985,24 @@ impl Tunnel for WgNtTunnel {
         self.stop_tunnel();
         Ok(())
     }
+
+    fn set_config(&self, config: &Config) -> std::result::Result<(), super::TunnelError> {
+        if let Some(ref device) = &*self.device.lock().unwrap() {
+            if let Err(error) = device.set_config(&config) {
+                log::error!(
+                    "{}",
+                    error.display_chain_with_msg("Failed to set wg-nt tunnel config")
+                );
+                Err(super::TunnelError::SetConfigError)
+            } else {
+                Ok(())
+            }
+        } else {
+            Err(super::TunnelError::StatsError(
+                super::stats::Error::NoTunnelDevice,
+            ))
+        }
+    }
 }
 
 #[cfg(test)]