From 715e777f72d3ed9d08fa85c165db956152ac25a1 Mon Sep 17 00:00:00 2001
From: Jackson Coxson <jkcoxson@gmail.com>
Date: Mon, 24 Mar 2025 19:09:24 -0600
Subject: [PATCH] Don't accidentially free provider when connecting to
 installation proxy

---
 ffi/examples/list_apps.c      |  6 +----
 ffi/src/installation_proxy.rs | 44 ++++++++++++++++++++++++++---------
 ffi/src/provider.rs           |  1 +
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/ffi/examples/list_apps.c b/ffi/examples/list_apps.c
index 7b72deb..7fe1308 100644
--- a/ffi/examples/list_apps.c
+++ b/ffi/examples/list_apps.c
@@ -26,7 +26,6 @@ int main() {
     fprintf(stderr, "Failed to read pairing file: %d\n", err);
     return 1;
   }
-  printf("Read the pairing file");
 
   // Create TCP provider
   TcpProviderHandle *provider = NULL;
@@ -62,11 +61,9 @@ int main() {
     return 1;
   }
 
-  // Cast the result to plist_t array (assuming the binding returns plist_t)
+  // Cast the result to plist_t array
   plist_t *app_list = (plist_t *)apps;
 
-  // Iterate through apps (this is a simplified example - you'd need proper
-  // plist handling)
   printf("Found %zu apps:\n", apps_len);
   for (size_t i = 0; i < apps_len; i++) {
     plist_t app = app_list[i];
@@ -82,7 +79,6 @@ int main() {
   }
 
   // Cleanup
-  // Note: You'd need to properly free the plist array here
   installation_proxy_client_free(client);
   tcp_provider_free(provider);
 
diff --git a/ffi/src/installation_proxy.rs b/ffi/src/installation_proxy.rs
index 8973a50..9ea104c 100644
--- a/ffi/src/installation_proxy.rs
+++ b/ffi/src/installation_proxy.rs
@@ -31,16 +31,24 @@ pub unsafe extern "C" fn installation_proxy_connect_tcp(
     provider: *mut TcpProviderHandle,
     client: *mut *mut InstallationProxyClientHandle,
 ) -> IdeviceErrorCode {
-    if provider.is_null() {
-        log::error!("Provider is null");
+    if provider.is_null() || client.is_null() {
+        log::error!("Null pointer provided");
         return IdeviceErrorCode::InvalidArg;
     }
-    let provider = unsafe { Box::from_raw(provider) }.0;
 
     let res: Result<InstallationProxyClient, IdeviceError> = RUNTIME.block_on(async move {
-        let res = InstallationProxyClient::connect(&provider).await;
-        std::mem::forget(provider);
-        res
+        // Take ownership of the provider (without immediately dropping it)
+        let provider_box = unsafe { Box::from_raw(provider) };
+
+        // Get a reference to the inner value
+        let provider_ref = &provider_box.0;
+
+        // Connect using the reference
+        let result = InstallationProxyClient::connect(provider_ref).await;
+
+        // Explicitly keep the provider_box alive until after connect completes
+        std::mem::forget(provider_box);
+        result
     });
 
     match res {
@@ -49,7 +57,12 @@ pub unsafe extern "C" fn installation_proxy_connect_tcp(
             unsafe { *client = Box::into_raw(boxed) };
             IdeviceErrorCode::IdeviceSuccess
         }
-        Err(e) => e.into(),
+        Err(e) => {
+            // If connection failed, the provider_box was already forgotten,
+            // so we need to reconstruct it to avoid leak
+            let _ = unsafe { Box::from_raw(provider) };
+            e.into()
+        }
     }
 }
 
@@ -74,12 +87,20 @@ pub unsafe extern "C" fn installation_proxy_connect_usbmuxd(
         log::error!("Provider is null");
         return IdeviceErrorCode::InvalidArg;
     }
-    let provider = unsafe { Box::from_raw(provider) }.0;
 
     let res: Result<InstallationProxyClient, IdeviceError> = RUNTIME.block_on(async move {
-        let res = InstallationProxyClient::connect(&provider).await;
-        std::mem::forget(provider);
-        res
+        // Take ownership of the provider (without immediately dropping it)
+        let provider_box = unsafe { Box::from_raw(provider) };
+
+        // Get a reference to the inner value
+        let provider_ref = &provider_box.0;
+
+        // Connect using the reference
+        let result = InstallationProxyClient::connect(provider_ref).await;
+
+        // Explicitly keep the provider_box alive until after connect completes
+        std::mem::forget(provider_box);
+        result
     });
 
     match res {
@@ -211,6 +232,7 @@ pub unsafe extern "C" fn installation_proxy_client_free(
     handle: *mut InstallationProxyClientHandle,
 ) {
     if !handle.is_null() {
+        log::debug!("Freeing installation_proxy_client");
         let _ = unsafe { Box::from_raw(handle) };
     }
 }
diff --git a/ffi/src/provider.rs b/ffi/src/provider.rs
index 5b3dcd7..19dacac 100644
--- a/ffi/src/provider.rs
+++ b/ffi/src/provider.rs
@@ -69,6 +69,7 @@ pub unsafe extern "C" fn idevice_tcp_provider_new(
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn tcp_provider_free(provider: *mut TcpProviderHandle) {
     if !provider.is_null() {
+        log::debug!("Freeing TCP provider");
         unsafe { drop(Box::from_raw(provider)) };
     }
 }