libusbmuxd: handle remove messages in usbmuxd_get_device_list

1 files changed, 49 insertions, 26 deletions

diff --git a/libusbmuxd/libusbmuxd.c b/libusbmuxd/libusbmuxd.c
index 861eefc..172da80 100644
--- a/libusbmuxd/libusbmuxd.c
+++ b/libusbmuxd/libusbmuxd.c
@@ -657,12 +657,15 @@ int usbmuxd_get_device_list(usbmuxd_device_info_t **device_list)
         int sfd;
         int listen_success = 0;
         uint32_t res;
+        struct collection tmpdevs;
         usbmuxd_device_info_t *newlist = NULL;
         struct usbmuxd_header hdr;
-        struct usbmuxd_device_record *dev_info;
+        struct usbmuxd_device_record *dev;
         int dev_cnt = 0;
         void *payload = NULL;
 
+        *device_list = NULL;
+
 #ifdef HAVE_PLIST
 retry:
 #endif
@@ -698,32 +701,42 @@ retry:
                 return -1;
         }
 
-        *device_list = NULL;
+        collection_init(&tmpdevs);
+
         // receive device list
         while (1) {
                 if (receive_packet(sfd, &hdr, &payload, 1000) > 0) {
                         if (hdr.message == MESSAGE_DEVICE_ADD) {
-                                dev_info = payload;
-                                newlist = (usbmuxd_device_info_t *) realloc(*device_list, sizeof(usbmuxd_device_info_t) * (dev_cnt + 1));
-                                if (newlist) {
-                                        newlist[dev_cnt].handle =
-                                                (int) dev_info->device_id;
-                                        newlist[dev_cnt].product_id =
-                                                dev_info->product_id;
-                                        memset(newlist[dev_cnt].uuid, '\0',
-                                                   sizeof(newlist[dev_cnt].uuid));
-                                        memcpy(newlist[dev_cnt].uuid,
-                                                   dev_info->serial_number,
-                                                   sizeof(newlist[dev_cnt].uuid));
-                                        *device_list = newlist;
-                                        dev_cnt++;
-                                } else {
-                                        fprintf(stderr,
-                                                "%s: ERROR: out of memory when trying to realloc!\n",
-                                                __func__);
-                                        if (payload)
-                                                free(payload);
-                                        break;
+                                dev = payload;
+                                usbmuxd_device_info_t *devinfo = (usbmuxd_device_info_t*)malloc(sizeof(usbmuxd_device_info_t));
+                                if (!devinfo) {
+                                        fprintf(stderr, "%s: Out of memory!\n", __func__);
+                                        free(payload);
+                                        return -1;
+                                }
+
+                                devinfo->handle = dev->device_id;
+                                devinfo->product_id = dev->product_id;
+                                memset(devinfo->uuid, '\0', sizeof(devinfo->uuid));
+                                memcpy(devinfo->uuid, dev->serial_number, sizeof(devinfo->uuid));
+
+                                collection_add(&tmpdevs, devinfo);
+
+                        } else if (hdr.message == MESSAGE_DEVICE_REMOVE) {
+                                uint32_t handle;
+                                usbmuxd_device_info_t *devinfo = NULL;
+
+                                memcpy(&handle, payload, sizeof(uint32_t));
+
+                                FOREACH(usbmuxd_device_info_t *di, &tmpdevs) {
+                                        if (di && di->handle == handle) {
+                                                devinfo = di;
+                                                break;
+                                        }
+                                } ENDFOREACH
+                                if (devinfo) {
+                                        collection_remove(&tmpdevs, devinfo);
+                                        free(devinfo);
                                 }
                         } else {
                                 fprintf(stderr, "%s: Unexpected message %d\n", __func__, hdr.message);
@@ -740,9 +753,19 @@ retry:
         // explicitly close connection
         close_socket(sfd);
 
-        // terminating zero record
-        newlist = (usbmuxd_device_info_t*) realloc(*device_list, sizeof(usbmuxd_device_info_t) * (dev_cnt + 1));
-        memset(newlist + dev_cnt, 0, sizeof(usbmuxd_device_info_t));
+        // create copy of device info entries from collection
+        newlist = (usbmuxd_device_info_t*)malloc(sizeof(usbmuxd_device_info_t) * (collection_count(&tmpdevs) + 1));
+        dev_cnt = 0;
+        FOREACH(usbmuxd_device_info_t *di, &tmpdevs) {
+                if (di) {
+                        memcpy(&newlist[dev_cnt], di, sizeof(usbmuxd_device_info_t));
+                        free(di);
+                        dev_cnt++;
+                }
+        } ENDFOREACH
+        collection_free(&tmpdevs);
+
+        memset(&newlist[dev_cnt], 0, sizeof(usbmuxd_device_info_t));
         *device_list = newlist;
 
         return dev_cnt;