1 files changed, 36 insertions, 26 deletions

diff --git a/serverd.c b/serverd.c
index 6f21530..b99e4ea 100644
--- a/serverd.c
+++ b/serverd.c
@@ -15,10 +15,42 @@ int main(void) {
 	struct icmpd *server = icmpd_create_server(-1, 0);
 	struct icmpd *conn = NULL;
 
+	int server_fd = icmpd_get_select_fd(server);
+	int conn_fd = -1;
+
 	while (true) {
-		// printf("Iter\n");
+		fd_set inset;
+		FD_ZERO(&inset);
+		FD_SET(0, &inset);
+		FD_SET(server_fd, &inset);
+		if (conn) FD_SET(conn_fd, &inset);
+		int nfds = maxi(0, maxi(server_fd, conn_fd)) + 1;
+
+		int ret = select(nfds, &inset, NULL, NULL, NULL);
+		if (ret < 0) {
+			if (errno == EINTR) continue;
+			perror("select");
+			return 1;
+		}
+
+		if (ret == 0) continue;  // timeout
+
+		if (FD_ISSET(0, &inset)) {
+			char *line = NULL;
+			size_t linelen = 0;
+			errno = 0;
+			ssize_t nr = getline(&line, &linelen, stdin);
+			if (nr < 0) {
+				if (errno == 0) break;  // EOF
+				perror("getline");
+				exit(1);
+			}
+			if (nr > 0) nr--;
+			icmpd_send(conn, line, nr);
+			free(line);
+		}
 
-		if (icmpd_peek(server)) {
+		if (FD_ISSET(server_fd, &inset) && icmpd_peek(server)) {
 			struct icmpd_received msg = icmpd_recv(server);
 			printf("Server recv: %zu\n", msg.length);
 			xxd(msg.data, msg.length);
@@ -32,6 +64,7 @@ int main(void) {
 			} else {
 				conn = icmpd_create_server(msg.id, msg.source_addr);
 				icmpd_server_set_outstanding(conn, msg.seqnum);
+				conn_fd = icmpd_get_select_fd(conn);
 
 				// icmpd_send(conn, "dankje", 6);
 				// icmpd_destroy(conn);
@@ -40,34 +73,11 @@ int main(void) {
 			free(msg.data);
 		}
 
-		if (icmpd_peek(conn)) {
+		if (conn && FD_ISSET(conn_fd, &inset) && icmpd_peek(conn)) {
 			struct icmpd_received msg = icmpd_recv(conn);
 			printf("Connection recv: %zu\n", msg.length);
 			xxd(msg.data, msg.length);
 			free(msg.data);
 		}
-
-		fd_set inset;
-		FD_ZERO(&inset);
-		FD_SET(0, &inset);
-		struct timeval tv;
-		tv.tv_sec = 1;
-		tv.tv_usec = 0;
-		int ret = select(1, &inset, NULL, NULL, &tv);
-		if (ret < 0) {
-			if (errno == EINTR) continue;
-			perror("select");
-			return 1;
-		}
-		if (ret == 0) continue;  // timeout
-		if (FD_ISSET(0, &inset)) {
-			char *line = NULL;
-			size_t linelen = 0;
-			ssize_t nr = getline(&line, &linelen, stdin);
-			assert(nr >= 0);
-			if (nr > 0) nr--;
-			icmpd_send(conn, line, nr);
-			free(line);
-		}
 	}
 }