Reconnect wip

tmate-decoder.c

@@ -114,19 +114,29 @@ out:
 	free(cmd_str);
 }
 
-static void handle_set_env(__unused struct tmate_session *session,
+static void maybe_save_reconnection_data(struct tmate_session *session,
+				const char *name, const char *value)
+{
+	if (!strcmp(name, "tmate_reconnection_data")) {
+		free(session->reconnection_data);
+		session->reconnection_data = xstrdup(value);
+	}
+}
+
+static void handle_set_env(struct tmate_session *session,
 			   struct tmate_unpacker *uk)
 {
 	char *name = unpack_string(uk);
 	char *value = unpack_string(uk);
 
 	tmate_set_env(name, value);
+	maybe_save_reconnection_data(session, name, value);
 
 	free(name);
 	free(value);
 }
 
-static void handle_ready(__unused struct tmate_session *session,
+static void handle_ready(struct tmate_session *session,
 			 __unused struct tmate_unpacker *uk)
 {
 	session->tmate_env_ready = 1;

tmate-encoder.c

@@ -236,3 +236,126 @@ void tmate_write_fin(void)
 	pack(array, 1);
 	pack(int, TMATE_OUT_FIN);
 }
+
+static void do_snapshot(unsigned int max_history_lines,
+			struct window_pane *pane)
+{
+	struct screen *screen;
+	struct grid *grid;
+	struct grid_line *line;
+	struct grid_cell gc;
+	unsigned int line_i, i;
+	unsigned int max_lines;
+	size_t str_len;
+
+	screen = &pane->base;
+	grid = screen->grid;
+
+	pack(array, 4);
+	pack(int, pane->id);
+
+	pack(array, 2);
+	pack(int, screen->cx);
+	pack(int, screen->cy);
+
+	pack(unsigned_int, screen->mode);
+
+	max_lines = max_history_lines + grid->sy;
+
+#define grid_num_lines(grid) (grid->hsize + grid->sy)
+
+	if (grid_num_lines(grid) > max_lines)
+		line_i = grid_num_lines(grid) - max_lines;
+	else
+		line_i = 0;
+
+	pack(array, grid_num_lines(grid) - line_i);
+	for (; line_i < grid_num_lines(grid); line_i++) {
+		line = &grid->linedata[line_i];
+
+		pack(array, 2);
+		str_len = 0;
+		for (i = 0; i < line->cellsize; i++) {
+			grid_get_cell(grid, i, line_i, &gc);
+			str_len += gc.data.size;
+		}
+
+		pack(str, str_len);
+		for (i = 0; i < line->cellsize; i++) {
+			grid_get_cell(grid, i, line_i, &gc);
+			pack(str_body, gc.data.data, gc.data.size);
+		}
+
+		pack(array, line->cellsize);
+		for (i = 0; i < line->cellsize; i++) {
+			grid_get_cell(grid, i, line_i, &gc);
+			pack(unsigned_int, ((gc.flags << 24) |
+					    (gc.attr  << 16) |
+					    (gc.bg    << 8)  |
+					     gc.fg        ));
+		}
+	}
+}
+
+static void tmate_send_session_snapshot(unsigned int max_history_lines)
+{
+	struct session *s;
+	struct winlink *wl;
+	struct window *w;
+	struct window_pane *pane;
+	int num_panes;
+
+	pack(array, 2);
+	pack(int, TMATE_OUT_SNAPSHOT);
+
+	s = RB_MIN(sessions, &sessions);
+	if (!s)
+		tmate_fatal("no session?");
+
+	num_panes = 0;
+	RB_FOREACH(wl, winlinks, &s->windows) {
+		w = wl->window;
+		if (!w)
+			continue;
+
+		TAILQ_FOREACH(pane, &w->panes, entry)
+			num_panes++;
+	}
+
+	pack(array, num_panes);
+	RB_FOREACH(wl, winlinks, &s->windows) {
+		w = wl->window;
+		if (!w)
+			continue;
+
+		TAILQ_FOREACH(pane, &w->panes, entry)
+			do_snapshot(max_history_lines, pane);
+	}
+}
+
+static void tmate_send_reconnection_data(struct tmate_session *session)
+{
+	if (!session->reconnection_data)
+		return;
+
+	pack(array, 2);
+	pack(int, TMATE_OUT_RECONNECT);
+	pack(string, session->reconnection_data);
+}
+
+#define RECONNECTION_MAX_HISTORY_LINE 300
+
+void tmate_send_reconnection_state(struct tmate_session *session)
+{
+	/* Start with a fresh encoder */
+	tmate_encoder_destroy(&session->encoder);
+	tmate_encoder_init(&session->encoder, NULL, session);
+
+	tmate_write_header();
+	tmate_send_reconnection_data(session);
+	/* TODO send all option variables */
+	tmate_write_ready();
+
+	tmate_sync_layout();
+	tmate_send_session_snapshot(RECONNECTION_MAX_HISTORY_LINE);
+}

tmate-msgpack.c

@@ -65,13 +65,22 @@ void tmate_encoder_init(struct tmate_encoder *encoder,
 	encoder->ev_active = false;
 }
 
+void tmate_encoder_destroy(struct tmate_encoder *encoder)
+{
+	/* encoder->pk doesn't need any cleanup */
+	evbuffer_free(encoder->buffer);
+	event_del(&encoder->ev_buffer);
+	memset(encoder, 0, sizeof(*encoder));
+}
+
 void tmate_encoder_set_ready_callback(struct tmate_encoder *encoder,
 				      tmate_encoder_write_cb *callback,
 				      void *userdata)
 {
 	encoder->ready_callback = callback;
 	encoder->userdata = userdata;
-	encoder->ready_callback(encoder->userdata, encoder->buffer);
+	if (encoder->ready_callback)
+		encoder->ready_callback(encoder->userdata, encoder->buffer);
 }
 
 void tmate_decoder_error(void)
@@ -178,6 +187,12 @@ void tmate_decoder_init(struct tmate_decoder *decoder, tmate_decoder_reader *rea
 	decoder->userdata = userdata;
 }
 
+void tmate_decoder_destroy(struct tmate_decoder *decoder)
+{
+	msgpack_unpacker_destroy(&decoder->unpacker);
+	memset(decoder, 0, sizeof(*decoder));
+}
+
 void tmate_decoder_get_buffer(struct tmate_decoder *decoder,
 			      char **buf, size_t *len)
 {

tmate-protocol.h

@@ -29,6 +29,7 @@ enum tmate_control_in_msg_types {
 	TMATE_CTL_PANE_KEYS,
 	TMATE_CTL_RESIZE,
 	TMATE_CTL_EXEC_RESPONSE,
+	TMATE_CTL_RENAME_SESSION,
 };
 
 /*
@@ -37,6 +38,7 @@ enum tmate_control_in_msg_types {
 [TMATE_CTL_PANE_KEYS, int: pane_id, string: keys]
 [TMATE_CTL_RESIZE, int: sx, int: sy] // sx == -1: no clients
 [TMATE_CTL_EXEC_RESPONSE, int: exit_code, string: message]
+[TMATE_CTL_RENAME_SESSION, string: stoken, string: stoken_ro]
 */
 
 enum tmate_daemon_out_msg_types {
@@ -50,6 +52,8 @@ enum tmate_daemon_out_msg_types {
 	TMATE_OUT_WRITE_COPY_MODE,
 	TMATE_OUT_FIN,
 	TMATE_OUT_READY,
+	TMATE_OUT_RECONNECT,
+	TMATE_OUT_SNAPSHOT,
 };
 
 /*

tmate-session.c

@@ -11,7 +11,8 @@
 
 #include "tmate.h"
 
-#define TMATE_DNS_RETRY_TIMEOUT 10
+#define TMATE_DNS_RETRY_TIMEOUT 2
+#define TMATE_RECONNECT_RETRY_TIMEOUT 2
 
 struct tmate_session tmate_session;
 
@@ -129,7 +130,8 @@ void tmate_session_init(struct event_base *base)
 
 void tmate_session_start(void)
 {
-	/* We split init and start because:
+	/*
+	 * We split init and start because:
 	 * - We need to process the tmux config file during the connection as
 	 *   we are setting up the tmate identity.
 	 * - While we are parsing the config file, we need to be able to
@@ -137,3 +139,43 @@ void tmate_session_start(void)
 	 */
 	lookup_and_connect();
 }
+
+static void on_reconnect_retry(__unused evutil_socket_t fd, __unused short what, void *arg)
+{
+	struct tmate_session *session = arg;
+
+	if (session->last_server_ip) {
+		/*
+		 * We have a previous server ip. Let's try that again first,
+		 * but then connect to any server if it fails again.
+		 */
+		(void)tmate_ssh_client_alloc(&tmate_session, session->last_server_ip);
+		free(session->last_server_ip);
+		session->last_server_ip = NULL;
+	} else {
+		lookup_and_connect();
+	}
+}
+
+void tmate_reconnect_session(struct tmate_session *session)
+{
+	/*
+	 * We no longer have an SSH connection. Time to reconnect.
+	 * We'll reuse some of the session information if we can,
+	 * and we'll try to reconnect to the same server if possible,
+	 * to avoid an SSH connection string change.
+	 */
+	struct timeval tv = { .tv_sec = TMATE_RECONNECT_RETRY_TIMEOUT, .tv_usec = 0 };
+
+	evtimer_assign(&session->ev_connection_retry, session->ev_base,
+		       on_reconnect_retry, session);
+	evtimer_add(&session->ev_connection_retry, &tv);
+
+	tmate_status_message("Reconnecting...");
+
+	/*
+	 * This says that we'll need to send a snapshot of the current state.
+	 * Until we have persisted logs...
+	 */
+	session->reconnected = true;
+}

tmate-ssh-client.c

@@ -12,7 +12,7 @@ static void __on_ssh_client_event(evutil_socket_t fd, short what, void *arg);
 
 static void printflike(2, 3) kill_ssh_client(struct tmate_ssh_client *client,
 					     const char *fmt, ...);
-static void printflike(2, 3) reconnect_ssh_client(struct tmate_ssh_client *client,
+static void printflike(2, 3) kill_ssh_client(struct tmate_ssh_client *client,
 						  const char *fmt, ...);
 
 static void read_channel(struct tmate_ssh_client *client)
@@ -25,8 +25,8 @@ static void read_channel(struct tmate_ssh_client *client)
 		tmate_decoder_get_buffer(decoder, &buf, &len);
 		len = ssh_channel_read_nonblocking(client->channel, buf, len, 0);
 		if (len < 0) {
-			reconnect_ssh_client(client, "Error reading from channel: %s",
+			kill_ssh_client(client, "Error reading from channel: %s",
-					     ssh_get_error(client->session));
+					ssh_get_error(client->session));
 			break;
 		}
 
@@ -61,8 +61,8 @@ static void on_encoder_write(void *userdata, struct evbuffer *buffer)
 
 		written = ssh_channel_write(client->channel, buf, len);
 		if (written < 0) {
-			reconnect_ssh_client(client, "Error writing to channel: %s",
+			kill_ssh_client(client, "Error writing to channel: %s",
-					     ssh_get_error(client->session));
+					ssh_get_error(client->session));
 			break;
 		}
 
@@ -245,8 +245,8 @@ static void on_ssh_client_event(struct tmate_ssh_client *client)
 			init_conn_fd(client);
 			return;
 		case SSH_ERROR:
-			reconnect_ssh_client(client, "Error connecting: %s",
+			kill_ssh_client(client, "Error connecting: %s",
-					     ssh_get_error(session));
+					ssh_get_error(session));
 			return;
 		case SSH_OK:
 			init_conn_fd(client);
@@ -315,19 +315,20 @@ static void on_ssh_client_event(struct tmate_ssh_client *client)
 		case SSH_AUTH_PARTIAL:
 		case SSH_AUTH_INFO:
 		case SSH_AUTH_DENIED:
-			if (client->tmate_session->need_passphrase)
+			if (client->tmate_session->need_passphrase) {
 				request_passphrase(client);
-			else
+			} else {
 				kill_ssh_client(client, "SSH keys not found."
 				" Run 'ssh-keygen' to create keys and try again.");
+				return;
+			}
 
 			if (client->tried_passphrase)
 				tmate_status_message("Can't load SSH key."
 				" Try typing passphrase again in case of typo. ctrl-c to abort.");
 			return;
 		case SSH_AUTH_ERROR:
-			reconnect_ssh_client(client, "Auth error: %s",
+			kill_ssh_client(client, "Auth error: %s", ssh_get_error(session));
-					     ssh_get_error(session));
 			return;
 		case SSH_AUTH_SUCCESS:
 			tmate_debug("Auth successful");
@@ -340,8 +341,8 @@ static void on_ssh_client_event(struct tmate_ssh_client *client)
 		case SSH_AGAIN:
 			return;
 		case SSH_ERROR:
-			reconnect_ssh_client(client, "Error opening channel: %s",
+			kill_ssh_client(client, "Error opening channel: %s",
-					     ssh_get_error(session));
+					ssh_get_error(session));
 			return;
 		case SSH_OK:
 			tmate_debug("Session opened, initalizing tmate");
@@ -354,8 +355,8 @@ static void on_ssh_client_event(struct tmate_ssh_client *client)
 		case SSH_AGAIN:
 			return;
 		case SSH_ERROR:
-			reconnect_ssh_client(client, "Error initializing tmate: %s",
+			kill_ssh_client(client, "Error initializing tmate: %s",
-					     ssh_get_error(session));
+					ssh_get_error(session));
 			return;
 		case SSH_OK:
 			tmate_debug("Ready");
@@ -365,19 +366,22 @@ static void on_ssh_client_event(struct tmate_ssh_client *client)
 
 			client->state = SSH_READY;
 
+			if (client->tmate_session->reconnected)
+				tmate_send_reconnection_state(client->tmate_session);
+
 			tmate_encoder_set_ready_callback(&client->tmate_session->encoder,
 							 on_encoder_write, client);
 			tmate_decoder_init(&client->tmate_session->decoder,
 					   on_decoder_read, client);
+
+			free(client->tmate_session->last_server_ip);
+			client->tmate_session->last_server_ip = xstrdup(client->server_ip);
+
 			/* fall through */
 		}
 
 	case SSH_READY:
 		read_channel(client);
-		if (!ssh_is_connected(session)) {
-			reconnect_ssh_client(client, "Disconnected");
-			return;
-		}
 	}
 }
 
@@ -386,19 +390,37 @@ static void __on_ssh_client_event(__unused evutil_socket_t fd, __unused short wh
 	on_ssh_client_event(arg);
 }
 
-static void __kill_ssh_client(struct tmate_ssh_client *client,
+static void kill_ssh_client(struct tmate_ssh_client *client,
-			      const char *fmt, va_list va)
+			    const char *fmt, ...)
 {
-	if (fmt && TAILQ_EMPTY(&client->tmate_session->clients))
+	bool last_client;
-		__tmate_status_message(fmt, va);
+	va_list ap;
-	else
+
-		tmate_debug("Disconnecting %s", client->server_ip);
+	TAILQ_REMOVE(&client->tmate_session->clients, client, node);
+	last_client = TAILQ_EMPTY(&client->tmate_session->clients);
+
+	if (fmt && last_client) {
+		va_start(ap, fmt);
+		__tmate_status_message(fmt, ap);
+		va_end(ap);
+	}
+
+	tmate_debug("SSH client killed (%s)", client->server_ip);
 
 	if (client->has_init_conn_fd) {
 		event_del(&client->ev_ssh);
 		client->has_init_conn_fd = false;
 	}
 
+	if (client->state == SSH_READY) {
+		tmate_encoder_set_ready_callback(&client->tmate_session->encoder, NULL, NULL);
+		tmate_decoder_destroy(&client->tmate_session->decoder);
+
+		client->tmate_session->min_sx = -1;
+		client->tmate_session->min_sy = -1;
+		recalculate_sizes();
+	}
+
 	if (client->session) {
 		/* ssh_free() also frees the associated channels. */
 		ssh_free(client->session);
@@ -406,19 +428,8 @@ static void __kill_ssh_client(struct tmate_ssh_client *client,
 		client->channel = NULL;
 	}
 
-	client->state = SSH_NONE;
+	if (last_client)
-}
+		tmate_reconnect_session(client->tmate_session);
-
-static void kill_ssh_client(struct tmate_ssh_client *client,
-			    const char *fmt, ...)
-{
-	va_list ap;
-
-	TAILQ_REMOVE(&client->tmate_session->clients, client, node);
-
-	va_start(ap, fmt);
-	__kill_ssh_client(client, fmt, ap);
-	va_end(ap);
 
 	free(client->server_ip);
 	free(client);
@@ -432,33 +443,6 @@ static void connect_ssh_client(struct tmate_ssh_client *client)
 	}
 }
 
-static void on_reconnect_timer(__unused evutil_socket_t fd, __unused short what, void *arg)
-{
-	connect_ssh_client(arg);
-}
-
-static void reconnect_ssh_client(struct tmate_ssh_client *client,
-				 const char *fmt, ...)
-{
-	/* struct timeval tv; */
-	va_list ap;
-
-#if 1
-	TAILQ_REMOVE(&client->tmate_session->clients, client, node);
-#endif
-
-	va_start(ap, fmt);
-	__kill_ssh_client(client, fmt, ap);
-	va_end(ap);
-
-	/* Not yet implemented... */
-#if 0
-	tv.tv_sec = 1;
-	tv.tv_usec = 0;
-	evtimer_add(&client->ev_ssh_reconnect, &tv);
-#endif
-}
-
 static void ssh_log_function(int priority, const char *function,
 			     const char *buffer, __unused void *userdata)
 {
@@ -489,9 +473,6 @@ struct tmate_ssh_client *tmate_ssh_client_alloc(struct tmate_session *session,
 
 	client->has_init_conn_fd = false;
 
-	evtimer_assign(&client->ev_ssh_reconnect, session->ev_base,
-		       on_reconnect_timer, client);
-
 	connect_ssh_client(client);
 
 	return client;

tmate.h

@@ -30,6 +30,7 @@ struct tmate_encoder {
 extern void tmate_encoder_init(struct tmate_encoder *encoder,
 			       tmate_encoder_write_cb *callback,
 			       void *userdata);
+extern void tmate_encoder_destroy(struct tmate_encoder *encoder);
 extern void tmate_encoder_set_ready_callback(struct tmate_encoder *encoder,
 					     tmate_encoder_write_cb *callback,
 					     void *userdata);
@@ -50,6 +51,7 @@ struct tmate_decoder {
 };
 
 extern void tmate_decoder_init(struct tmate_decoder *decoder, tmate_decoder_reader *reader, void *userdata);
+extern void tmate_decoder_destroy(struct tmate_decoder *decoder);
 extern void tmate_decoder_get_buffer(struct tmate_decoder *decoder, char **buf, size_t *len);
 extern void tmate_decoder_commit(struct tmate_decoder *decoder, size_t len);
 
@@ -75,6 +77,8 @@ extern void unpack_array(struct tmate_unpacker *uk, struct tmate_unpacker *neste
 
 #define TMATE_PROTOCOL_VERSION 6
 
+struct tmate_session;
+
 extern void tmate_write_header(void);
 extern void tmate_write_ready(void);
 extern void tmate_sync_layout(void);
@@ -86,6 +90,7 @@ extern void tmate_status(const char *left, const char *right);
 extern void tmate_sync_copy_mode(struct window_pane *wp);
 extern void tmate_write_copy_mode(struct window_pane *wp, const char *str);
 extern void tmate_write_fin(void);
+extern void tmate_send_reconnection_state(struct tmate_session *session);
 
 /* tmate-decoder.c */
 
@@ -135,7 +140,6 @@ struct tmate_ssh_client {
 
 	bool has_init_conn_fd;
 	struct event ev_ssh;
-	struct event ev_ssh_reconnect;
 };
 TAILQ_HEAD(tmate_ssh_clients, tmate_ssh_client);
 
@@ -166,11 +170,17 @@ struct tmate_session {
 	struct tmate_ssh_clients clients;
 	int need_passphrase;
 	char *passphrase;
+
+	bool reconnected;
+	struct event ev_connection_retry;
+	char *last_server_ip;
+	char *reconnection_data;
 };
 
 extern struct tmate_session tmate_session;
 extern void tmate_session_init(struct event_base *base);
 extern void tmate_session_start(void);
+extern void tmate_reconnect_session(struct tmate_session *session);
 
 /* tmate-debug.c */
 extern void tmate_print_stack_trace(void);