Allow to advertize a specific SSH port (for load balancer)

docker-entrypoint.sh

@@ -13,4 +13,7 @@ if [ ! -z "${SSH_HOSTNAME}" ]; then
   set -- -h "${SSH_HOSTNAME}" "$@"
 fi
 
-exec tmate-ssh-server -p ${SSH_PORT:-2200} -k ${SSH_KEYS_PATH} "$@"
+SSH_PORT_LISTEN=${SSH_PORT_LISTEN:-2200}
+SSH_PORT_ADVERTIZE=${SSH_PORT_ADVERTIZE:-${SSH_PORT_LISTEN}}
+
+exec tmate-ssh-server -p ${SSH_PORT_LISTEN} -q ${SSH_PORT_ADVERTIZE} -k ${SSH_KEYS_PATH} "$@"

tmate-daemon-decoder.c

@@ -9,8 +9,7 @@ char *tmate_left_status, *tmate_right_status;
 static void tmate_header(struct tmate_session *session,
 			 struct tmate_unpacker *uk)
 {
-	char port_arg[16] = {0};
+	char *ssh_conn_str;
-	char tmp[512];
 
 	session->client_protocol_version = unpack_int(uk);
 	if (session->client_protocol_version <= 4) {
@@ -32,19 +31,16 @@ static void tmate_header(struct tmate_session *session,
 		return;
 	}
 
-	if (tmate_settings->ssh_port != 22)
+	ssh_conn_str = get_ssh_conn_string(session->session_token_ro);
-		sprintf(port_arg, " -p%d", tmate_settings->ssh_port);
-
-	sprintf(tmp, "ssh%s %s@%s", port_arg, session->session_token_ro, tmate_settings->tmate_host);
-
 	tmate_notify("Note: clear your terminal before sharing readonly access");
-	tmate_notify("ssh session read only: %s", tmp);
+	tmate_notify("ssh session read only: %s", ssh_conn_str);
+	tmate_set_env("tmate_ssh_ro", ssh_conn_str);
+	free(ssh_conn_str);
 
-	sprintf(tmp, "ssh%s %s@%s", port_arg, session->session_token, tmate_settings->tmate_host);
+	ssh_conn_str = get_ssh_conn_string(session->session_token);
-	tmate_notify("ssh session: %s", tmp);
+	tmate_notify("ssh session: %s", ssh_conn_str);
-
+	tmate_set_env("tmate_ssh", ssh_conn_str);
-	tmate_set_env("tmate_ssh_ro", tmp);
+	free(ssh_conn_str);
-	tmate_set_env("tmate_ssh", tmp);
 
 	tmate_send_client_ready();
 }

tmate-debug.c

@@ -47,7 +47,7 @@ static int print_resolved_stack_frame(const char *frame)
 		return -1;
 
 	line[strlen(line)-1] = 0; /* remove \n */
-	tmate_info("%s(%s) [%s]", file, line, address);
+	tmate_crit("%s(%s) [%s]", file, line, address);
 	return 0;
 }
 #endif
@@ -62,13 +62,13 @@ void tmate_print_stack_trace(void)
 	size = backtrace(array, 20);
 	strings = backtrace_symbols(array, size);
 
-	tmate_info("============ %zd stack frames ============", size);
+	tmate_crit("============ %zd stack frames ============", size);
 
 	for (i = 1; i < size; i++) {
 #if DEBUG
 		if (print_resolved_stack_frame(strings[i]) < 0)
 #endif
-			tmate_info("%s", strings[i]);
+			tmate_crit("%s", strings[i]);
 	}
 
 	free(strings);

tmate-main.c

@@ -36,6 +36,7 @@ struct tmate_settings _tmate_settings = {
 	.keys_dir        	= TMATE_SSH_DEFAULT_KEYS_DIR,
 	.authorized_keys_path 	= NULL,
 	.ssh_port        	= TMATE_SSH_DEFAULT_PORT,
+	.ssh_port_advertized    = -1,
 	.websocket_hostname  	= NULL,
 	.bind_addr	 	= NULL,
 	.websocket_port      	= TMATE_DEFAULT_WEBSOCKET_PORT,
@@ -104,7 +105,7 @@ void request_server_termination(void)
 
 static void usage(void)
 {
-	fprintf(stderr, "usage: tmate-ssh-server [-b ip] [-h hostname] [-k keys_dir] [-a authorized_keys_path] [-p port] [-w websocket_hostname] [-q websocket_port] [-x] [-s] [-v]\n");
+	fprintf(stderr, "usage: tmate-ssh-server [-b ip] [-h hostname] [-k keys_dir] [-a authorized_keys_path] [-p listen_port] [-q ssh_port_advertized] [-w websocket_hostname] [-z websocket_port] [-x] [-s] [-v]\n");
 }
 
 static char* get_full_hostname(void)
@@ -157,7 +158,7 @@ int main(int argc, char **argv, char **envp)
 {
 	int opt;
 
-	while ((opt = getopt(argc, argv, "b:h:k:a:p:w:q:xsv")) != -1) {
+	while ((opt = getopt(argc, argv, "b:h:k:a:p:q:w:z:xsv")) != -1) {
 		switch (opt) {
 		case 'b':
 			tmate_settings->bind_addr = xstrdup(optarg);
@@ -174,10 +175,13 @@ int main(int argc, char **argv, char **envp)
 		case 'p':
 			tmate_settings->ssh_port = atoi(optarg);
 			break;
+		case 'q':
+			tmate_settings->ssh_port_advertized = atoi(optarg);
+			break;
 		case 'w':
 			tmate_settings->websocket_hostname = xstrdup(optarg);
 			break;
-		case 'q':
+		case 'z':
 			tmate_settings->websocket_port = atoi(optarg);
 			break;
 		case 'x':

tmate-ssh-server.c

@@ -14,6 +14,21 @@
 
 #include "tmate.h"
 
+char *get_ssh_conn_string(const char *session_token)
+{
+	char port_arg[16] = {0};
+	char *ret;
+
+	int ssh_port_advertized = tmate_settings->ssh_port_advertized == -1 ?
+		tmate_settings->ssh_port :
+		tmate_settings->ssh_port_advertized;
+
+	if (ssh_port_advertized != 22)
+		sprintf(port_arg, " -p%d", ssh_port_advertized);
+	xasprintf(&ret, "ssh%s %s@%s", port_arg, session_token, tmate_settings->tmate_host);
+	return ret;
+}
+
 static int pty_request(__unused ssh_session session,
 		       __unused ssh_channel channel,
 		       __unused const char *term,
@@ -330,6 +345,7 @@ static void read_single_line(int fd, char *dst, size_t len)
 static int get_client_ip_proxy_protocol(int fd, char *dst, size_t len)
 {
 	char header[110];
+	int tok_num;
 	const char *signature = "PROXY ";
 
 	if (read(fd, header, strlen(signature)) != (ssize_t)strlen(signature))
@@ -342,10 +358,12 @@ static int get_client_ip_proxy_protocol(int fd, char *dst, size_t len)
 
 	tmate_debug("proxy header: %s", header);
 
-	int tok_num = 0;
+	tok_num = 0;
 	for (char *tok = strtok(header, " "); tok; tok = strtok(NULL, " "), tok_num++) {
-		if (tok_num == 1)
+		if (tok_num == 1) {
 			strncpy(dst, tok, len);
+			dst[len-1] = '\0';
+		}
 	}
 
 	if (tok_num != 5)
@@ -436,7 +454,7 @@ static void handle_sigchld(__unused int sig)
 
 static void handle_sigsegv(__unused int sig)
 {
-	tmate_info("CRASH, printing stack trace");
+	tmate_crit("CRASH, printing stack trace");
 	tmate_print_stack_trace();
 	tmate_fatal("CRASHED");
 }

tmate-websocket.c

@@ -280,9 +280,6 @@ void tmate_send_websocket_daemon_msg(__unused struct tmate_session *session,
 
 void tmate_send_websocket_header(struct tmate_session *session)
 {
-	char port_arg[16] = {0};
-	char ssh_cmd_fmt[512];
-
 	if (!tmate_has_websocket())
 		return;
 
@@ -294,10 +291,9 @@ void tmate_send_websocket_header(struct tmate_session *session)
 	pack(string, session->session_token);
 	pack(string, session->session_token_ro);
 
-	if (tmate_settings->ssh_port != 22)
+	char *ssh_cmd_fmt = get_ssh_conn_string("%s");
-		sprintf(port_arg, " -p%d", tmate_settings->ssh_port);
-	sprintf(ssh_cmd_fmt, "ssh%s %%s@%s", port_arg, tmate_settings->tmate_host);
 	pack(string, ssh_cmd_fmt);
+	free(ssh_cmd_fmt);
 
 	pack(string, session->client_version);
 	pack(int, session->client_protocol_version);

tmate.h

@@ -21,14 +21,15 @@ extern void printflike(2, 3) tmate_log(int level, const char *msg, ...);
 #define tmate_info(str, ...)	tmate_log(LOG_INFO, str, ##__VA_ARGS__)
 #define tmate_notice(str, ...)	tmate_log(LOG_NOTICE, str, ##__VA_ARGS__)
 #define tmate_warn(str, ...)	tmate_log(LOG_WARNING, str, ##__VA_ARGS__)
+#define tmate_crit(str, ...)	tmate_log(LOG_CRIT, str, ##__VA_ARGS__)
 #define tmate_fatal(str, ...)					\
 ({								\
-	tmate_log(LOG_CRIT, "fatal: " str, ##__VA_ARGS__);	\
+	tmate_crit("fatal: " str, ##__VA_ARGS__);		\
  	exit(1);						\
 })
 #define tmate_fatal_info(str, ...)				\
 ({								\
-	tmate_log(LOG_INFO, "fatal: " str, ##__VA_ARGS__);	\
+	tmate_info("fatal: " str, ##__VA_ARGS__);		\
  	exit(1);						\
 })
 
@@ -154,6 +155,7 @@ extern void tmate_flush_pty(struct tmate_session *session);
 
 struct tmate_ssh_client;
 typedef void ssh_client_latency_cb(void *userdata, int latency_ms);
+extern char *get_ssh_conn_string(const char *session_token);
 extern void tmate_start_ssh_latency_probes(struct tmate_ssh_client *client,
 					   struct ssh_server_callbacks_struct *server_callbacks,
 					   int keepalive_interval_ms);
@@ -216,6 +218,7 @@ struct tmate_settings {
 	const char *keys_dir;
 	const char *authorized_keys_path;
 	int ssh_port;
+	int ssh_port_advertized;
 	const char *websocket_hostname;
 	int websocket_port;
 	const char *tmate_host;