1 files changed, 53 insertions, 5 deletions

diff --git a/ssh/client.c b/ssh/client.c
index a763c68..fbf4c1d 100644
--- a/ssh/client.c
+++ b/ssh/client.c
@@ -448,9 +448,7 @@ static void handle_event(struct state *state, const struct tomsg_event event) {
 	}
 }
 
-static bool hostkey_checker(const unsigned char *hash, size_t length, void *stdinbuf_) {
-	struct readbuffer *stdinbuf = stdinbuf_;
-
+static bool hostkey_checker(const unsigned char *hash, size_t length, struct readbuffer *stdinbuf) {
 	const char *fingerprint = tomsg_print_hash(hash, length);
 
 	printf("Server host key fingerprint: %s\n", fingerprint);
@@ -459,6 +457,57 @@ static bool hostkey_checker(const unsigned char *hash, size_t length, void *stdi
 		"member that you trust and is already connected to the server?");
 }
 
+static enum tomsg_retval connect_server(
+	const char *hostname, int port,
+	struct readbuffer *stdinbuf,
+	struct tomsg_client **clientp
+) {
+	struct tomsg_async_connect *async;
+	enum tomsg_retval ret = tomsg_async_connect(hostname, port, &async);
+	if (ret != TOMSG_OK) return ret;
+
+	const int fd = tomsg_async_connect_poll_fd(async);
+
+	struct pollfd pfd;
+	pfd.fd = fd;
+	pfd.events = POLLIN;
+
+	while (true) {
+		pfd.revents = 0;
+		int pollret = poll(&pfd, 1, -1);
+		if (pollret < 0) {
+			perror("poll");
+			// Can't really close the async connector?
+			exit(1);
+		}
+
+		if (pfd.revents & (POLLIN | POLLHUP | POLLERR)) {
+			struct tomsg_async_connect_event event;
+			ret = tomsg_async_connect_next_event(async, &event);
+			if (ret == TOMSG_ERR_AGAIN) continue;
+			if (ret != TOMSG_OK) {
+				fprintf(stderr, "next_event returned %d\n", ret);
+				return ret;
+			}
+
+			switch (event.type) {
+				case TOMSG_AC_HOSTKEY:
+					ret = tomsg_async_connect_accept(async,
+							hostkey_checker(event.key.hostkey, event.key.length, stdinbuf));
+					if (ret != TOMSG_OK) {
+						fprintf(stderr, "connect_accept returned %d\n", ret);
+						return ret;
+					}
+					break;
+
+				case TOMSG_AC_SUCCESS:
+					*clientp = event.client;
+					return TOMSG_OK;
+			}
+		}
+	}
+}
+
 int main(int argc, char **argv) {
 	if (argc != 2) {
 		printf("Usage: %s <hostname:port>\n", argv[0]);
@@ -477,8 +526,7 @@ int main(int argc, char **argv) {
 	printf("Connecting...\n");
 
 	struct tomsg_client *client;
-	enum tomsg_retval ret = tomsg_connect(
-			hostname, port, hostkey_checker, &stdinbuf, &client);
+	enum tomsg_retval ret = connect_server(hostname, port, &stdinbuf, &client);
 	if (ret != TOMSG_OK) {
 		printf("Could not connect: %s\n", tomsg_strerror(ret));
 		return 1;