From 9c2df33238e44e1373fb8830c0a6b4fc7b14de9a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Em=C4=ABls?= <emils@mullvad.net>
Date: Fri, 20 Dec 2024 10:36:16 +0100
Subject: [PATCH] fixup! Use DAITAv2 on iOS

---
 .../ios_tcp_connection.rs                     |  1 -
 mullvad-ios/src/ephemeral_peer_proxy/mod.rs   | 40 +++++++++++++------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/mullvad-ios/src/ephemeral_peer_proxy/ios_tcp_connection.rs b/mullvad-ios/src/ephemeral_peer_proxy/ios_tcp_connection.rs
index d7f2260acd83..1be9a18bf339 100644
--- a/mullvad-ios/src/ephemeral_peer_proxy/ios_tcp_connection.rs
+++ b/mullvad-ios/src/ephemeral_peer_proxy/ios_tcp_connection.rs
@@ -70,7 +70,6 @@ impl WgTcpConnectionFunctions {
     }
 }
 
-
 #[derive(Clone)]
 pub struct IosTcpProvider {
     tunnel_handle: i32,
diff --git a/mullvad-ios/src/ephemeral_peer_proxy/mod.rs b/mullvad-ios/src/ephemeral_peer_proxy/mod.rs
index 805bd2ef8c6f..eac4e2eeb3b9 100644
--- a/mullvad-ios/src/ephemeral_peer_proxy/mod.rs
+++ b/mullvad-ios/src/ephemeral_peer_proxy/mod.rs
@@ -17,6 +17,10 @@ pub struct PacketTunnelBridge {
 
 impl PacketTunnelBridge {
     fn fail_exchange(self) {
+        // # Safety
+        // Call is safe as long as the `packet_tunnel` pointer is valid. Since a valid instance of
+        // `PacketTunnelBridge` requires the packet tunnel pointer to be valid, it is assumed this
+        // call is safe.
         unsafe {
             swift_ephemeral_peer_ready(self.packet_tunnel, ptr::null(), ptr::null(), ptr::null())
         };
@@ -38,7 +42,11 @@ impl PacketTunnelBridge {
             .as_ref()
             .map(|params| params as *const _)
             .unwrap_or(ptr::null());
-
+        // # Safety
+        // The `packet_tunnel` pointer must be valid, much like the call in `fail_exchange`, but
+        // since the other arguments here are non-null, these pointers (`preshared_ptr`,
+        // `ephmerela_ptr` and `daita_ptr`) have to be valid too. Since they point to local
+        // variables or are null, the pointer values will be valid for the lifetime of the call.
         unsafe {
             swift_ephemeral_peer_ready(self.packet_tunnel, preshared_ptr, ephemeral_ptr, daita_ptr)
         };
@@ -77,7 +85,7 @@ impl DaitaParameters {
 
 impl Drop for DaitaParameters {
     fn drop(&mut self) {
-        // Safety:
+        // # Safety
         // `machines` pointer must be a valid pointer to a CString. This can be achieved by
         // ensuring that `DaitaParameters` are constructed via `DaitaParameters::new` and the
         // `machines` pointer is never written to.
@@ -86,11 +94,15 @@ impl Drop for DaitaParameters {
 }
 
 extern "C" {
-    /// Called when the preshared post quantum key is ready,
-    /// or when a Daita peer has been successfully requested.
-    /// `raw_preshared_key` will be NULL if:
-    /// - The post quantum key negotiation failed
-    /// - A Daita peer has been requested without enabling post quantum keys.
+    /// To be called when ephemeral peer exchange has finished. All parameters except
+    /// `raw_packet_tunnel` are optional.
+    ///
+    /// # Safety:
+    /// If the key exchange failed, all pointers except `raw_packet_tunnel` must be null. If the
+    /// key exchange was successful, `raw_ephemeral_private_key` must be a valid pointer to 32
+    /// bytes for the lifetime of this call. If PQ was enabled, `raw_preshared_key` must be a valid
+    /// pointer to 32 bytes for the lifetime of this call. If DAITA was requested, the
+    /// `daita_prameters` must point to a valid instance of `DaitaParameters`.
     pub fn swift_ephemeral_peer_ready(
         raw_packet_tunnel: *const c_void,
         raw_preshared_key: *const u8,
@@ -131,11 +143,11 @@ pub unsafe extern "C" fn drop_ephemeral_peer_exchange_token(
 /// Entry point for requesting ephemeral peers on iOS.
 /// The TCP connection must be created to go through the tunnel.
 /// # Safety
-/// `public_key` and `ephemeral_key` must be valid respective `PublicKey` and `PrivateKey` types.
-/// They will not be valid after this function is called, and thus must be copied here.
-/// `packet_tunnel` must be valid pointers to a packet tunnel, the packet tunnel pointer must
-/// outlive the ephemeral peer exchange. `cancel_token` should be owned by the caller of this
-/// function.
+/// `public_key` and `ephemeral_key` must be valid respective `PublicKey` and `PrivateKey` types,
+/// specifically, they must be valid pointers to 32 bytes. They will not be valid after this
+/// function is called, and thus must be copied here. `packet_tunnel` must be valid pointers to a
+/// packet tunnel, the packet tunnel pointer must outlive the ephemeral peer exchange.
+/// `cancel_token` should be owned by the caller of this function.
 #[no_mangle]
 pub unsafe extern "C" fn request_ephemeral_peer(
     public_key: *const u8,
@@ -150,7 +162,11 @@ pub unsafe extern "C" fn request_ephemeral_peer(
             .init();
     });
 
+    /// # Safety
+    /// `public_key` pointer must be a valid pointer to 32 unsigned bytes.
     let pub_key: [u8; 32] = unsafe { ptr::read(public_key as *const [u8; 32]) };
+    /// # Safety
+    /// `ephemeral_key` pointer must be a valid pointer to 32 unsigned bytes.
     let eph_key: [u8; 32] = unsafe { ptr::read(ephemeral_key as *const [u8; 32]) };
 
     let handle = match crate::mullvad_ios_runtime() {