libusbmuxd: add generic receive_packet function and use it

1 files changed, 117 insertions, 91 deletions

diff --git a/libusbmuxd/libusbmuxd.c b/libusbmuxd/libusbmuxd.c
index 9d54a88..db8d1d4 100644
--- a/libusbmuxd/libusbmuxd.c
+++ b/libusbmuxd/libusbmuxd.c
@@ -74,39 +74,85 @@ static int connect_usbmuxd_socket()
 #endif
 }
 
+static int receive_packet(int sfd, struct usbmuxd_header *header, void **payload, int timeout)
+{
+	int recv_len;
+	struct usbmuxd_header hdr;
+	char *payload_loc = NULL;
+
+	header->length = 0;
+	header->version = 0;
+	header->message = 0;
+	header->tag = 0;
+
+	recv_len = recv_buf_timeout(sfd, &hdr, sizeof(hdr), 0, timeout);
+	if (recv_len < 0) {
+		return recv_len;
+	} else if (recv_len < sizeof(hdr)) {
+		return recv_len;
+	}
+
+	uint32_t payload_size = hdr.length - sizeof(hdr);
+	if (payload_size > 0) {
+		payload_loc = (char*)malloc(payload_size);
+		if (recv_buf_timeout(sfd, payload_loc, payload_size, 0, 5000) != payload_size) {
+			fprintf(stderr, "%s: Error receiving payload of size %d\n", __func__, payload_size);
+			free(payload_loc);
+			return -EBADMSG;
+		}
+	}
+
+	*payload = payload_loc;
+
+	memcpy(header, &hdr, sizeof(hdr));
+
+	return hdr.length;
+}
+
 /**
  * Retrieves the result code to a previously sent request.
  */
 static int usbmuxd_get_result(int sfd, uint32_t tag, uint32_t * result)
 {
-	struct usbmuxd_result_msg res;
+	struct usbmuxd_header hdr;
 	int recv_len;
+	uint32_t *res = NULL;
 
 	if (!result) {
 		return -EINVAL;
 	}
+	*result = -1;
 
-	if ((recv_len = recv_buf(sfd, &res, sizeof(res))) <= 0) {
-		perror("recv");
+	if ((recv_len = receive_packet(sfd, &hdr, (void**)&res, 5000)) < 0) {
+		fprintf(stderr, "%s: Error receiving packet: %d\n", __func__, errno);
+		if (res)
+			free(res);
 		return -errno;
-	} else {
-		if ((recv_len == sizeof(res))
-			&& (res.header.length == (uint32_t) recv_len)
-			&& (res.header.version == USBMUXD_PROTOCOL_VERSION)
-			&& (res.header.message == MESSAGE_RESULT)
-			) {
-			*result = res.result;
-			if (res.header.tag == tag) {
-				return 1;
-			} else {
-				return 0;
-			}
-		}
+	}
+	if (recv_len < sizeof(hdr)) {
+		fprintf(stderr, "%s: Received packet is too small!\n", __func__);
+		if (res)
+			free(res);
+		return -EPROTO;
 	}
 
-	return -1;
+	if (hdr.message == MESSAGE_RESULT) {
+		int ret = 0;
+		if (res && (hdr.tag == tag)) {
+			memcpy(result, res, sizeof(uint32_t));
+			ret = 1;
+		}
+		if (res)
+			free(res);
+		return ret;
+	}
+	fprintf(stderr, "%s: Unexpected message of type %d received!\n", __func__, hdr.message);
+	if (res)
+		free(res);
+	return -EPROTO;
 }
 
+
 /**
  * Generates an event, i.e. calls the callback function.
  * A reference to a populated usbmuxd_event_t with information about the event
@@ -163,8 +209,8 @@ static int usbmuxd_listen()
 		return -1;
 	}
 	if (usbmuxd_get_result(sfd, req.header.tag, &res) && (res != 0)) {
-		fprintf(stderr, "%s: ERROR: did not get OK\n", __func__);
 		close(sfd);
+		fprintf(stderr, "%s: ERROR: did not get OK but %d\n", __func__, res);
 		return -1;
 	}
 
@@ -177,12 +223,11 @@ static int usbmuxd_listen()
  */
 int get_next_event(int sfd, usbmuxd_event_cb_t callback, void *user_data)
 {
-	int recv_len;
 	struct usbmuxd_header hdr;
+	void *payload = NULL;
 
 	/* block until we receive something */
-	recv_len = recv_buf_timeout(sfd, &hdr, sizeof(hdr), 0, 0);
-	if (recv_len < 0) {
+	if (receive_packet(sfd, &hdr, &payload, 0) < 0) {
 		// when then usbmuxd connection fails,
 		// generate remove events for every device that
 		// is still present so applications know about it
@@ -190,57 +235,48 @@ int get_next_event(int sfd, usbmuxd_event_cb_t callback, void *user_data)
 			generate_event(callback, dev, UE_DEVICE_REMOVE, user_data);
 			collection_remove(&devices, dev);
 		} ENDFOREACH
-		return recv_len;
-	} else if (recv_len == sizeof(hdr)) {
-		if (hdr.message == MESSAGE_DEVICE_ADD) {
-			struct usbmuxd_device_record dev;
-			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__);
-				return -1;
-			}
+		return -EIO;
+	}
 
-			if (hdr.length != sizeof(struct usbmuxd_header)+sizeof(struct usbmuxd_device_record)) {
-				fprintf(stderr, "%s: WARNING: unexpected packet size %d for MESSAGE_DEVICE_ADD (expected %d)!\n", __func__, hdr.length, (int)(sizeof(struct usbmuxd_header)+sizeof(struct usbmuxd_device_record)));
-			}
-			recv_len =  recv_buf_timeout(sfd, &dev, hdr.length - sizeof(struct usbmuxd_header), 0, 5000);
-			if (recv_len != (hdr.length - sizeof(struct usbmuxd_header))) {
-				fprintf(stderr, "%s: ERROR: Could not receive packet\n", __func__);
-				return recv_len;
-			}
+	if ((hdr.length > sizeof(hdr)) && !payload) {
+		fprintf(stderr, "%s: Invalid packet received, payload is missing!\n", __func__);
+		return -EBADMSG;
+	}
 
-			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));
+	if (hdr.message == MESSAGE_DEVICE_ADD) {
+		struct usbmuxd_device_record *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;
+		}
 
-			collection_add(&devices, devinfo);
-			generate_event(callback, devinfo, UE_DEVICE_ADD, user_data);
-		} else if (hdr.message == MESSAGE_DEVICE_REMOVE) {
-			uint32_t handle;
-			usbmuxd_device_info_t *dev;
+		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));
 
-			if (hdr.length != sizeof(struct usbmuxd_header)+sizeof(uint32_t)) {
-				fprintf(stderr, "%s: WARNING: unexpected packet size %d for MESSAGE_DEVICE_REMOVE (expected %d)!\n", __func__, hdr.length, (int)(sizeof(struct usbmuxd_header)+sizeof(uint32_t)));
-			}
-			recv_len = recv_buf_timeout(sfd, &handle, sizeof(uint32_t), 0, 5000);
-			if (recv_len != sizeof(uint32_t)) {
-				fprintf(stderr, "%s: ERROR: Could not receive packet\n", __func__);
-				return recv_len;
-			}
+		collection_add(&devices, devinfo);
+		generate_event(callback, devinfo, UE_DEVICE_ADD, user_data);
+	} else if (hdr.message == MESSAGE_DEVICE_REMOVE) {
+		uint32_t handle;
+		usbmuxd_device_info_t *devinfo;
 
-			dev = devices_find(handle);
-			if (!dev) {
-				fprintf(stderr, "%s: WARNING: got device remove message for handle %d, but couldn't find the corresponding handle in the device list. This event will be ignored.\n", __func__, handle);
-			} else {
-				generate_event(callback, dev, UE_DEVICE_REMOVE, user_data);
-				collection_remove(&devices, dev);
-			}
+		memcpy(&handle, payload, sizeof(uint32_t));
+
+		devinfo = devices_find(handle);
+		if (!devinfo) {
+			fprintf(stderr, "%s: WARNING: got device remove message for handle %d, but couldn't find the corresponding handle in the device list. This event will be ignored.\n", __func__, handle);
 		} else {
-			fprintf(stderr, "%s: Unknown message type %d length %d\n", __func__, hdr.message, hdr.length);
+			generate_event(callback, devinfo, UE_DEVICE_REMOVE, user_data);
+			collection_remove(&devices, devinfo);
 		}
 	} else {
-		fprintf(stderr, "%s: ERROR: incomplete packet received!\n", __func__);
+		fprintf(stderr, "%s: Unexpected message type %d length %d received!\n", __func__, hdr.message, hdr.length);
+	}
+	if (payload) {
+		free(payload);
 	}
 	return 0;
 }
@@ -311,11 +347,11 @@ int usbmuxd_get_device_list(usbmuxd_device_info_t **device_list)
 	int sfd;
 	int listen_success = 0;
 	uint32_t res;
-	int recv_len;
 	usbmuxd_device_info_t *newlist = NULL;
 	struct usbmuxd_header hdr;
-	struct usbmuxd_device_record dev_info;
+	struct usbmuxd_device_record *dev_info;
 	int dev_cnt = 0;
+	void *payload = NULL;
 
 	sfd = connect_usbmuxd_socket();
 	if (sfd < 0) {
@@ -336,10 +372,10 @@ int usbmuxd_get_device_list(usbmuxd_device_info_t **device_list)
 		if (usbmuxd_get_result(sfd, s_req.header.tag, &res) && (res == 0)) {
 			listen_success = 1;
 		} else {
+			close(sfd);
 			fprintf(stderr,
 					"%s: Did not get response to scan request (with result=0)...\n",
 					__func__);
-			close(sfd);
 			return res;
 		}
 	}
@@ -352,45 +388,35 @@ int usbmuxd_get_device_list(usbmuxd_device_info_t **device_list)
 	*device_list = NULL;
 	// receive device list
 	while (1) {
-		if (recv_buf_timeout(sfd, &hdr, sizeof(hdr), 0, 1000) == sizeof(hdr)) {
-			if ((hdr.length < 48) || (hdr.length > sizeof(hdr)+sizeof(dev_info))) {
-				// invalid packet size received!
-				fprintf(stderr,
-						"%s: Invalid packet size (%d) received when expecting a device info record.\n",
-						__func__, hdr.length);
-				break;
-			}
-
-			recv_len = recv_buf(sfd, &dev_info, hdr.length - sizeof(hdr));
-			if (recv_len <= 0) {
-				fprintf(stderr,
-						"%s: Error when receiving device info record\n",
-						__func__);
-				break;
-			} else if ((uint32_t) recv_len < hdr.length - sizeof(hdr)) {
-				fprintf(stderr,
-						"%s: received less data than specified in header!\n", __func__);
-			} else {
+		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;
+						(int) dev_info->device_id;
 					newlist[dev_cnt].product_id =
-						dev_info.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,
+						   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__);
+						"%s: ERROR: out of memory when trying to realloc!\n",
+						__func__);
+					if (payload)
+						free(payload);
 					break;
 				}
+			} else {
+				fprintf(stderr, "%s: Unexpected message %d\n", __func__, hdr.message);
 			}
+			if (payload)
+				free(payload);
 		} else {
 			// we _should_ have all of them now.
 			// or perhaps an error occured.