4 files changed, 31 insertions, 12 deletions

diff --git a/src/device.c b/src/device.c
index b6cc32a..1e3cc83 100644
--- a/src/device.c
+++ b/src/device.c
@@ -112,6 +112,7 @@ struct mux_device
         uint16_t next_sport;
         unsigned char *pktbuf;
         uint32_t pktlen;
+        void *preflight_cb_data;
 };
 
 static struct collection device_list;
@@ -648,6 +649,7 @@ int device_add(struct usb_device *usbdev)
         dev->next_sport = 1;
         dev->pktbuf = malloc(DEV_MRU);
         dev->pktlen = 0;
+        dev->preflight_cb_data = NULL;
         struct version_header vh;
         vh.major = htonl(1);
         vh.minor = htonl(0);
@@ -674,6 +676,9 @@ void device_remove(struct usb_device *usbdev)
                                 client_device_remove(dev->id);
                                 collection_free(&dev->connections);
                         }
+                        if (dev->preflight_cb_data) {
+                                preflight_device_remove_cb(dev->preflight_cb_data);
+                        }
                         collection_remove(&device_list, dev);
                         free(dev->pktbuf);
                         free(dev);
@@ -690,7 +695,17 @@ void device_set_visible(int device_id)
                         dev->visible = 1;
                         break;
                 }
-        } ENDFOREACH    
+        } ENDFOREACH
+}
+
+void device_set_preflight_cb_data(int device_id, void* data)
+{
+        FOREACH(struct mux_device *dev, &device_list) {
+                if(dev->id == device_id) {
+                        dev->preflight_cb_data = data;
+                        break;
+                }
+        } ENDFOREACH
 }
 
 int device_get_count(int include_hidden)
diff --git a/src/device.h b/src/device.h
index d7c82ca..95d470e 100644
--- a/src/device.h
+++ b/src/device.h
@@ -41,6 +41,7 @@ void device_client_process(int device_id, struct mux_client *client, short event
 void device_abort_connect(int device_id, struct mux_client *client);
 
 void device_set_visible(int device_id);
+void device_set_preflight_cb_data(int device_id, void* data);
 
 int device_get_count(int include_hidden);
 int device_get_list(int include_hidden, struct device_info *p);
diff --git a/src/preflight.c b/src/preflight.c
index 283c6d9..b011344 100644
--- a/src/preflight.c
+++ b/src/preflight.c
@@ -30,11 +30,14 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 
 #include <sys/time.h>
 
+#ifdef HAVE_LIBIMOBILEDEVICE
 #include <libimobiledevice/libimobiledevice.h>
 #include <libimobiledevice/lockdown.h>
 #include <libimobiledevice/notification_proxy.h>
+#endif
 
 #include "preflight.h"
+#include "device.h"
 #include "client.h"
 #include "conf.h"
 #include "log.h"
@@ -65,15 +68,12 @@ static void lockdownd_set_untrusted_host_buid(lockdownd_client_t lockdown)
         free(system_buid);
 }
 
-static void idevice_callback(const idevice_event_t* event, void* userdata)
+void preflight_device_remove_cb(void *data)
 {
-        struct cb_data *cbdata = (struct cb_data*)userdata;
-        idevice_t dev = cbdata->dev;
-        struct idevice_private *_dev = (struct idevice_private*)dev;
-
-        if (event->event == IDEVICE_DEVICE_REMOVE && !strcmp(_dev->udid, event->udid)) {
-                cbdata->is_device_connected = 0;
-        }
+        if (!data)
+                return;
+        struct cb_data *cbdata = (struct cb_data*)data;
+        cbdata->is_device_connected = 0;
 }
 
 static void np_callback(const char* notification, void* userdata)
@@ -246,7 +246,7 @@ retry:
                 cbdata.is_device_connected = 1;
 
                 np_set_notify_callback(np, np_callback, (void*)&cbdata);
-                idevice_event_subscribe(idevice_callback, (void*)&cbdata);
+                device_set_preflight_cb_data(info->id, (void*)&cbdata);
 
                 const char* spec[] = {
                         "com.apple.mobile.lockdown.request_pair",
@@ -267,8 +267,6 @@ retry:
                 }
                 usbmuxd_log(LL_INFO, "%s: Finished waiting for notification from device %s, is_device_connected %d", __func__, _dev->udid, cbdata.is_device_connected);
 
-                idevice_event_unsubscribe();
-
                 if (cbdata.np) {
                         np_client_free(cbdata.np);
                 }
@@ -326,6 +324,10 @@ leave:
 
         return NULL;
 }
+#else
+void preflight_device_remove_cb(void *data)
+{
+}
 #endif
 
 void preflight_worker_device_add(struct device_info* info)
diff --git a/src/preflight.h b/src/preflight.h
index 62349d1..4c8aa60 100644
--- a/src/preflight.h
+++ b/src/preflight.h
@@ -26,6 +26,7 @@ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 extern void userpref_get_system_buid(char **systembuid);
 extern void userpref_device_record_get_host_id(const char *udid, char **host_id);
 
+void preflight_device_remove_cb(void *data);
 void preflight_worker_device_add(struct device_info* info);
 
 #endif