1 files changed, 74 insertions, 150 deletions

diff --git a/tools/iproxy.c b/tools/iproxy.c
index 8b57e94..8321143 100644
--- a/tools/iproxy.c
+++ b/tools/iproxy.c
@@ -54,169 +54,61 @@ typedef unsigned int socklen_t;
 #endif
 
 static int debug_level = 0;
-static uint16_t listen_port = 0;
-static uint16_t device_port = 0;
-static char* device_udid = NULL;
-static enum usbmux_lookup_options lookup_opts = 0;
 
 struct client_data {
 	int fd;
 	int sfd;
-	volatile int stop_ctos;
-	volatile int stop_stoc;
+	char* udid;
+	enum usbmux_lookup_options lookup_opts;
+	uint16_t device_port;
 };
 
-static void *run_stoc_loop(void *arg)
-{
-	struct client_data *cdata = (struct client_data*)arg;
-	int recv_len;
-	int sent;
-	char buffer[131072];
-
-	printf("%s: fd = %d\n", __func__, cdata->fd);
-
-	while (!cdata->stop_stoc && cdata->fd > 0 && cdata->sfd > 0) {
-		recv_len = socket_receive_timeout(cdata->sfd, buffer, sizeof(buffer), 0, 5000);
-		if (recv_len <= 0) {
-			if (recv_len == 0 || recv_len == -ETIMEDOUT) {
-				// try again
-				continue;
-			} else {
-				fprintf(stderr, "recv failed: %s\n", strerror(-recv_len));
-				break;
-			}
-		} else {
-			// send to socket
-			sent = socket_send(cdata->fd, buffer, recv_len);
-			if (sent < recv_len) {
-				if (sent <= 0) {
-					fprintf(stderr, "send failed: %s\n", strerror(errno));
-					break;
-				} else {
-					fprintf(stderr, "only sent %d from %d bytes\n", sent, recv_len);
-				}
-			}
-		}
-	}
-
-	socket_close(cdata->fd);
-
-	cdata->fd = -1;
-	cdata->stop_ctos = 1;
-
-	return NULL;
-}
-
-static void *run_ctos_loop(void *arg)
-{
-	struct client_data *cdata = (struct client_data*)arg;
-	int recv_len;
-	int sent;
-	char buffer[131072];
-#ifdef WIN32
-	HANDLE stoc = NULL;
-#else
-	pthread_t stoc;
-#endif
-
-	printf("%s: fd = %d\n", __func__, cdata->fd);
-
-	cdata->stop_stoc = 0;
-#ifdef WIN32
-	stoc = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)run_stoc_loop, cdata, 0, NULL);
-#else
-	pthread_create(&stoc, NULL, run_stoc_loop, cdata);
-#endif
-
-	while (!cdata->stop_ctos && cdata->fd>0 && cdata->sfd>0) {
-		recv_len = socket_receive_timeout(cdata->fd, buffer, sizeof(buffer), 0, 5000);
-		if (recv_len <= 0) {
-			if (recv_len == 0 || recv_len == -ETIMEDOUT) {
-				// try again
-				continue;
-			} else {
-				fprintf(stderr, "recv failed: %s\n", strerror(-recv_len));
-				break;
-			}
-		} else {
-			// send to local socket
-			sent = socket_send(cdata->sfd, buffer, recv_len);
-			if (sent < recv_len) {
-				if (sent <= 0) {
-					fprintf(stderr, "send failed: %s\n", strerror(errno));
-					break;
-				} else {
-					fprintf(stderr, "only sent %d from %d bytes\n", sent, recv_len);
-				}
-			}
-		}
-	}
-
-	socket_close(cdata->fd);
-
-	cdata->fd = -1;
-	cdata->stop_stoc = 1;
-
-#ifdef WIN32
-	WaitForSingleObject(stoc, INFINITE);
-#else
-	pthread_join(stoc, NULL);
-#endif
-
-	return NULL;
+#define CDATA_FREE(x) if (x) { \
+	if (x->fd > 0) socket_close(x->fd); \
+	if (x->sfd > 0) socket_close(x->sfd); \
+	free(x->udid); \
+	free(x); \
 }
 
 static void *acceptor_thread(void *arg)
 {
-	struct client_data *cdata;
+	char buffer[32768];
+	struct client_data *cdata = (struct client_data*)arg;
 	usbmuxd_device_info_t *dev_list = NULL;
 	usbmuxd_device_info_t *dev = NULL;
 	usbmuxd_device_info_t muxdev;
-#ifdef WIN32
-	HANDLE ctos = NULL;
-#else
-	pthread_t ctos;
-#endif
 	int count;
 
-	if (!arg) {
+	if (!cdata) {
 		fprintf(stderr, "invalid client_data provided!\n");
 		return NULL;
 	}
 
-	cdata = (struct client_data*)arg;
-
-	if (device_udid) {
-		if (usbmuxd_get_device(device_udid, &muxdev, lookup_opts) > 0) {
+	if (cdata->udid) {
+		if (usbmuxd_get_device(cdata->udid, &muxdev, cdata->lookup_opts) > 0) {
 			dev = &muxdev;
 		}
 	} else {
 		if ((count = usbmuxd_get_device_list(&dev_list)) < 0) {
 			printf("Connecting to usbmuxd failed, terminating.\n");
 			free(dev_list);
-			if (cdata->fd > 0) {
-				socket_close(cdata->fd);
-			}
-			free(cdata);
+			CDATA_FREE(cdata);
 			return NULL;
 		}
 
 		if (dev_list == NULL || dev_list[0].handle == 0) {
 			printf("No connected device found, terminating.\n");
 			free(dev_list);
-			if (cdata->fd > 0) {
-				socket_close(cdata->fd);
-			}
-			free(cdata);
+			CDATA_FREE(cdata);
 			return NULL;
 		}
 
 		int i;
 		for (i = 0; i < count; i++) {
-			if (dev_list[i].conn_type == CONNECTION_TYPE_USB && (lookup_opts & DEVICE_LOOKUP_USBMUX)) {
+			if (dev_list[i].conn_type == CONNECTION_TYPE_USB && (cdata->lookup_opts & DEVICE_LOOKUP_USBMUX)) {
 				dev = &(dev_list[i]);
 				break;
-			} else if (dev_list[i].conn_type == CONNECTION_TYPE_NETWORK && (lookup_opts & DEVICE_LOOKUP_NETWORK)) {
+			} else if (dev_list[i].conn_type == CONNECTION_TYPE_NETWORK && (cdata->lookup_opts & DEVICE_LOOKUP_NETWORK)) {
 				dev = &(dev_list[i]);
 				break;
 			}
@@ -226,10 +118,7 @@ static void *acceptor_thread(void *arg)
 	if (dev == NULL || dev->handle == 0) {
 		printf("No connected/matching device found, disconnecting client.\n");
 		free(dev_list);
-		if (cdata->fd > 0) {
-			socket_close(cdata->fd);
-		}
-		free(cdata);
+		CDATA_FREE(cdata);
 		return NULL;
 	}
 
@@ -248,13 +137,13 @@ static void *acceptor_thread(void *arg)
 			memcpy(&saddr->sa_data[0], (char*)dev->conn_data+2, 26);
 #else
 			fprintf(stderr, "ERROR: Got an IPv6 address but this system doesn't support IPv6\n");
-			free(cdata);
+			CDATA_FREE(cdata);
 			return NULL;
 #endif
 		}
 		else {
 			fprintf(stderr, "Unsupported address family 0x%02x\n", ((char*)dev->conn_data)[1]);
-			free(cdata);
+			CDATA_FREE(cdata);
 			return NULL;
 		}
 		char addrtxt[48];
@@ -262,35 +151,61 @@ static void *acceptor_thread(void *arg)
 		if (!socket_addr_to_string(saddr, addrtxt, sizeof(addrtxt))) {
 			fprintf(stderr, "Failed to convert network address: %d (%s)\n", errno, strerror(errno));
 		}
-		fprintf(stdout, "Requesting connecion to NETWORK device %s (serial: %s), port %d\n", addrtxt, dev->udid, device_port);
-		cdata->sfd = socket_connect_addr(saddr, device_port);
+		fprintf(stdout, "Requesting connecion to NETWORK device %s (serial: %s), port %d\n", addrtxt, dev->udid, cdata->device_port);
+		cdata->sfd = socket_connect_addr(saddr, cdata->device_port);
 	} else if (dev->conn_type == CONNECTION_TYPE_USB) {
-		fprintf(stdout, "Requesting connecion to USB device handle %d (serial: %s), port %d\n", dev->handle, dev->udid, device_port);
+		fprintf(stdout, "Requesting connecion to USB device handle %d (serial: %s), port %d\n", dev->handle, dev->udid, cdata->device_port);
 
-		cdata->sfd = usbmuxd_connect(dev->handle, device_port);
+		cdata->sfd = usbmuxd_connect(dev->handle, cdata->device_port);
 	}
 	free(dev_list);
 	if (cdata->sfd < 0) {
 		fprintf(stderr, "Error connecting to device: %s\n", strerror(errno));
 	} else {
-		cdata->stop_ctos = 0;
+		fd_set fds;
+		FD_ZERO(&fds);
+		FD_SET(cdata->fd, &fds);
+		FD_SET(cdata->sfd, &fds);
 
-#ifdef WIN32
-		ctos = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)run_ctos_loop, cdata, 0, NULL);
-		WaitForSingleObject(ctos, INFINITE);
-#else
-		pthread_create(&ctos, NULL, run_ctos_loop, cdata);
-		pthread_join(ctos, NULL);
-#endif
+		while (1) {
+			fd_set read_fds = fds;
+			int ret_sel = select(cdata->sfd+1, &read_fds, NULL, NULL, NULL);
+			if (ret_sel < 0) {
+				perror("select");
+				break;
+			}
+			if (FD_ISSET(cdata->fd, &read_fds)) {
+				int r = socket_receive_timeout(cdata->fd, buffer, sizeof(buffer), 0, 100);
+				if (r <= 0) {
+					break;
+				}
+				int sent = 0;
+				while (sent < r) {
+					int s = socket_send(cdata->sfd, buffer+sent, r-sent);
+					if (s <= 0) {
+						break;
+					}
+					sent += s;
+				}
+			}
+			if (FD_ISSET(cdata->sfd, &read_fds)) {
+				int r = socket_receive_timeout(cdata->sfd, buffer, sizeof(buffer), 0, 100);
+				if (r <= 0) {
+					break;
+				}
+				int sent = 0;
+				while (sent < r) {
+					int s = socket_send(cdata->fd, buffer+sent, r-sent);
+					if (s <= 0) {
+						break;
+					}
+					sent += s;
+				}
+			}
+		}
 	}
 
-	if (cdata->fd > 0) {
-		socket_close(cdata->fd);
-	}
-	if (cdata->sfd > 0) {
-		socket_close(cdata->sfd);
-	}
-	free(cdata);
+	CDATA_FREE(cdata);
 
 	return NULL;
 }
@@ -317,6 +232,11 @@ static void print_usage(int argc, char **argv, int is_error)
 int main(int argc, char **argv)
 {
 	int mysock = -1;
+	char* device_udid = NULL;
+	uint16_t listen_port = 0;
+	uint16_t device_port = 0;
+	enum usbmux_lookup_options lookup_opts = 0;
+
 	const struct option longopts[] = {
 		{ "debug", no_argument, NULL, 'd' },
 		{ "help", no_argument, NULL, 'h' },
@@ -413,6 +333,10 @@ int main(int argc, char **argv)
 					return -1;
 				}
 				cdata->fd = c_sock;
+				cdata->sfd = -1;
+				cdata->udid = strdup(device_udid);
+				cdata->lookup_opts = lookup_opts;
+				cdata->device_port = device_port;
 #ifdef WIN32
 				acceptor = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)acceptor_thread, cdata, 0, NULL);
 				CloseHandle(acceptor);