Add PQ support for multihop, and allow listing two endpoints in the

tunnel

4 files changed, 99 insertions, 44 deletions

diff --git a/talpid-core/src/firewall/linux.rs b/talpid-core/src/firewall/linux.rs
index 14900afb51..ced8390693 100644
--- a/talpid-core/src/firewall/linux.rs
+++ b/talpid-core/src/firewall/linux.rs
@@ -576,9 +576,13 @@ impl<'a> PolicyBatch<'a> {
                             self.add_allow_tunnel_rules(&tunnel.interface)?;
                         }
                         AllowedTunnelTraffic::None => (),
-                        AllowedTunnelTraffic::Only(endpoint) => {
+                        AllowedTunnelTraffic::One(endpoint) => {
                             self.add_allow_in_tunnel_endpoint_rules(&tunnel.interface, endpoint)?;
                         }
+                        AllowedTunnelTraffic::Two(endpoint1, endpoint2) => {
+                            self.add_allow_in_tunnel_endpoint_rules(&tunnel.interface, endpoint1)?;
+                            self.add_allow_in_tunnel_endpoint_rules(&tunnel.interface, endpoint2)?;
+                        }
                     }
                     if *allow_lan {
                         self.add_block_cve_2019_14899(tunnel);
diff --git a/talpid-core/src/firewall/macos.rs b/talpid-core/src/firewall/macos.rs
index 6f59020990..6ca12d0924 100644
--- a/talpid-core/src/firewall/macos.rs
+++ b/talpid-core/src/firewall/macos.rs
@@ -130,7 +130,7 @@ impl Firewall {
 
                 if let Some(tunnel) = tunnel {
                     rules.extend(
-                        self.get_allow_tunnel_rule(&tunnel.interface, allowed_tunnel_traffic)?
+                        self.get_allow_tunnel_rules(&tunnel.interface, allowed_tunnel_traffic)?
                             .into_iter(),
                     );
                 }
@@ -159,7 +159,7 @@ impl Firewall {
                 rules.append(&mut self.get_block_dns_rules()?);
 
                 rules.extend(
-                    self.get_allow_tunnel_rule(
+                    self.get_allow_tunnel_rules(
                         tunnel.interface.as_str(),
                         &AllowedTunnelTraffic::All,
                     )?
@@ -330,28 +330,55 @@ impl Firewall {
         Ok(vec![block_tcp_dns_rule, block_udp_dns_rule])
     }
 
-    fn get_allow_tunnel_rule(
+    fn base_rule(
         &self,
+        action: FilterRuleAction,
         tunnel_interface: &str,
-        allowed_traffic: &AllowedTunnelTraffic,
-    ) -> Result<Option<pfctl::FilterRule>> {
-        let mut rule_builder = self.create_rule_builder(FilterRuleAction::Pass);
-        let mut base_rule = rule_builder
+    ) -> pfctl::FilterRuleBuilder {
+        let mut rule_builder = self.create_rule_builder(action);
+        rule_builder
             .quick(true)
             .interface(tunnel_interface)
             .keep_state(pfctl::StatePolicy::Keep)
             .tcp_flags(Self::get_tcp_flags());
-        match allowed_traffic {
-            AllowedTunnelTraffic::Only(endpoint) => {
+        rule_builder
+    }
+
+    fn get_allow_tunnel_rules(
+        &self,
+        tunnel_interface: &str,
+        allowed_traffic: &AllowedTunnelTraffic,
+    ) -> Result<Vec<pfctl::FilterRule>> {
+        Ok(match allowed_traffic {
+            AllowedTunnelTraffic::One(endpoint) => {
+                let mut base_rule = &mut self.base_rule(FilterRuleAction::Pass, tunnel_interface);
                 let pfctl_proto = as_pfctl_proto(endpoint.protocol);
                 base_rule = base_rule.to(endpoint.address).proto(pfctl_proto);
+                vec![base_rule.build()?]
+            }
+            AllowedTunnelTraffic::Two(endpoint1, endpoint2) => {
+                let mut rules = Vec::with_capacity(2);
+
+                let mut base_rule = &mut self.base_rule(FilterRuleAction::Pass, tunnel_interface);
+                let pfctl_proto = as_pfctl_proto(endpoint1.protocol);
+                base_rule = base_rule.to(endpoint1.address).proto(pfctl_proto);
+                rules.push(base_rule.build()?);
+
+                let mut base_rule = &mut self.base_rule(FilterRuleAction::Pass, tunnel_interface);
+                let pfctl_proto = as_pfctl_proto(endpoint2.protocol);
+                base_rule = base_rule.to(endpoint2.address).proto(pfctl_proto);
+                rules.push(base_rule.build()?);
+
+                rules
+            }
+            AllowedTunnelTraffic::All => {
+                let base_rule = &mut self.base_rule(FilterRuleAction::Pass, tunnel_interface);
+                vec![base_rule.build()?]
             }
-            AllowedTunnelTraffic::All => {}
             AllowedTunnelTraffic::None => {
-                return Ok(None);
+                vec![]
             }
-        };
-        Ok(Some(base_rule.build()?))
+        })
     }
 
     fn get_allow_loopback_rules(&self) -> Result<Vec<pfctl::FilterRule>> {
diff --git a/talpid-core/src/firewall/windows.rs b/talpid-core/src/firewall/windows.rs
index e2854e0197..b015db2385 100644
--- a/talpid-core/src/firewall/windows.rs
+++ b/talpid-core/src/firewall/windows.rs
@@ -168,27 +168,52 @@ impl Firewall {
             ptr::null()
         };
 
-        let allowed_tun_ip;
-        let allowed_tunnel_endpoint =
-            if let AllowedTunnelTraffic::Only(endpoint) = allowed_tunnel_traffic {
-                allowed_tun_ip = widestring_ip(endpoint.address.ip());
-                Some(WinFwEndpoint {
-                    ip: allowed_tun_ip.as_ptr(),
-                    port: endpoint.address.port(),
-                    protocol: WinFwProt::from(endpoint.protocol),
-                })
-            } else {
-                None
-            };
+        // SAFETY: `allowed_tun_ips`, `entry_endpoint` and `exit_endpoint` must not be dropped until
+        // `WinFw_ApplyPolicyConnecting` has returned.
+        let mut allowed_tun_ips = [WideCString::new(), WideCString::new()];
+        let (entry_endpoint, exit_endpoint) = match allowed_tunnel_traffic {
+            AllowedTunnelTraffic::One(endpoint) => {
+                allowed_tun_ips[0] = widestring_ip(endpoint.address.ip());
+                (
+                    Some(WinFwEndpoint {
+                        ip: allowed_tun_ips[0].as_ptr(),
+                        port: endpoint.address.port(),
+                        protocol: WinFwProt::from(endpoint.protocol),
+                    }),
+                    None,
+                )
+            }
+            AllowedTunnelTraffic::Two(endpoint1, endpoint2) => {
+                allowed_tun_ips[0] = widestring_ip(endpoint1.address.ip());
+                let entry_endpoint = Some(WinFwEndpoint {
+                    ip: allowed_tun_ips[0].as_ptr(),
+                    port: endpoint1.address.port(),
+                    protocol: WinFwProt::from(endpoint1.protocol),
+                });
+                allowed_tun_ips[1] = widestring_ip(endpoint2.address.ip());
+                let exit_endpoint = Some(WinFwEndpoint {
+                    ip: allowed_tun_ips[1].as_ptr(),
+                    port: endpoint2.address.port(),
+                    protocol: WinFwProt::from(endpoint2.protocol),
+                });
+                (entry_endpoint, exit_endpoint)
+            }
+            AllowedTunnelTraffic::None | AllowedTunnelTraffic::All => (None, None),
+        };
+
         let allowed_tunnel_traffic = WinFwAllowedTunnelTraffic {
             type_: WinFwAllowedTunnelTrafficType::from(allowed_tunnel_traffic),
-            endpoint: allowed_tunnel_endpoint
+            entry_endpoint: entry_endpoint
+                .as_ref()
+                .map(|ep| ep as *const _)
+                .unwrap_or(ptr::null()),
+            exit_endpoint: exit_endpoint
                 .as_ref()
                 .map(|ep| ep as *const _)
                 .unwrap_or(ptr::null()),
         };
 
-        unsafe {
+        let res = unsafe {
             WinFw_ApplyPolicyConnecting(
                 winfw_settings,
                 &winfw_relay,
@@ -199,7 +224,14 @@ impl Firewall {
             )
             .into_result()
             .map_err(Error::ApplyingConnectingPolicy)
-        }
+        };
+        // SAFETY: All of these hold stack allocated memory which is pointed to by
+        // `allowed_tunnel_traffic` and must remain allocated until `WinFw_ApplyPolicyConnecting`
+        // has returned.
+        drop(allowed_tun_ips);
+        drop(entry_endpoint);
+        drop(exit_endpoint);
+        res
     }
 
     fn set_connected_state(
@@ -439,7 +471,8 @@ mod winfw {
     #[repr(C)]
     pub struct WinFwAllowedTunnelTraffic {
         pub type_: WinFwAllowedTunnelTrafficType,
-        pub endpoint: *const WinFwEndpoint,
+        pub entry_endpoint: *const WinFwEndpoint,
+        pub exit_endpoint: *const WinFwEndpoint,
     }
 
     #[repr(u8)]
@@ -447,7 +480,8 @@ mod winfw {
     pub enum WinFwAllowedTunnelTrafficType {
         None,
         All,
-        Only,
+        One,
+        Two,
     }
 
     impl From<&AllowedTunnelTraffic> for WinFwAllowedTunnelTrafficType {
@@ -455,7 +489,8 @@ mod winfw {
             match traffic {
                 AllowedTunnelTraffic::None => WinFwAllowedTunnelTrafficType::None,
                 AllowedTunnelTraffic::All => WinFwAllowedTunnelTrafficType::All,
-                AllowedTunnelTraffic::Only(..) => WinFwAllowedTunnelTrafficType::Only,
+                AllowedTunnelTraffic::One(..) => WinFwAllowedTunnelTrafficType::One,
+                AllowedTunnelTraffic::Two(..) => WinFwAllowedTunnelTrafficType::Two,
             }
         }
     }
diff --git a/talpid-core/src/tunnel/mod.rs b/talpid-core/src/tunnel/mod.rs
index 29f3f9c746..15422f0b94 100644
--- a/talpid-core/src/tunnel/mod.rs
+++ b/talpid-core/src/tunnel/mod.rs
@@ -143,18 +143,7 @@ impl TunnelMonitor {
         let config = talpid_wireguard::config::Config::from_parameters(params)?;
         let monitor = talpid_wireguard::WireguardMonitor::start(
             config,
-            if params.options.quantum_resistant {
-                Some(
-                    params
-                        .connection
-                        .exit_peer
-                        .as_ref()
-                        .map(|peer| peer.public_key.clone())
-                        .unwrap_or_else(|| params.connection.peer.public_key.clone()),
-                )
-            } else {
-                None
-            },
+            params.options.quantum_resistant,
             log.as_deref(),
             args,
         )?;