From 4555d8d6fb2441ee3ca0ed5e979c3f77a897712d Mon Sep 17 00:00:00 2001
From: Nicolas Viennot <nicolas@viennot.biz>
Date: Sat, 24 Aug 2019 17:46:43 -0400
Subject: [PATCH] Add PROXY protocol support

---
 docker-entrypoint.sh |  6 ++++-
 tmate-main.c         | 10 +++++---
 tmate-ssh-server.c   | 60 ++++++++++++++++++++++++++++++++++++++++++--
 tmate.h              |  1 +
 4 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
index 7b877779..72c744a2 100755
--- a/docker-entrypoint.sh
+++ b/docker-entrypoint.sh
@@ -1,8 +1,12 @@
 #!/bin/sh
 set -e
 
+if [ "${USE_PROXY_PROTOCOL}" == "1" ]; then
+  set -- -x
+fi
+
 if [ "${HAS_WEBSOCKET}" == "1" ]; then
-  set -- -x localhost "$@"
+  set -- -w localhost "$@"
 fi
 
 if [ ! -z "${SSH_HOSTNAME}" ]; then
diff --git a/tmate-main.c b/tmate-main.c
index ed3c4456..0cbd0c75 100644
--- a/tmate-main.c
+++ b/tmate-main.c
@@ -41,6 +41,7 @@ struct tmate_settings _tmate_settings = {
 	.websocket_port      	= TMATE_DEFAULT_WEBSOCKET_PORT,
 	.tmate_host      	= NULL,
 	.log_level      	= LOG_NOTICE,
+	.use_proxy_protocol	= false,
 	.use_syslog      	= false,
 };
 
@@ -103,7 +104,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] [-x websocket_hostname] [-q websocket_port] [-s] [-v]\n");
+	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");
 }
 
 static char* get_full_hostname(void)
@@ -156,7 +157,7 @@ int main(int argc, char **argv, char **envp)
 {
 	int opt;
 
-	while ((opt = getopt(argc, argv, "b:h:k:a:p:x:q:sv")) != -1) {
+	while ((opt = getopt(argc, argv, "b:h:k:a:p:w:q:xsv")) != -1) {
 		switch (opt) {
 		case 'b':
 			tmate_settings->bind_addr = xstrdup(optarg);
@@ -173,12 +174,15 @@ int main(int argc, char **argv, char **envp)
 		case 'p':
 			tmate_settings->ssh_port = atoi(optarg);
 			break;
-		case 'x':
+		case 'w':
 			tmate_settings->websocket_hostname = xstrdup(optarg);
 			break;
 		case 'q':
 			tmate_settings->websocket_port = atoi(optarg);
 			break;
+		case 'x':
+			tmate_settings->use_proxy_protocol = true;
+			break;
 		case 's':
 			tmate_settings->use_syslog = true;
 			break;
diff --git a/tmate-ssh-server.c b/tmate-ssh-server.c
index 4952335b..11b158ab 100644
--- a/tmate-ssh-server.c
+++ b/tmate-ssh-server.c
@@ -278,7 +278,7 @@ static void client_bootstrap(struct tmate_session *_session)
 	/* never reached */
 }
 
-static int get_client_ip(int fd, char *dst, size_t len)
+static int get_client_ip_socket(int fd, char *dst, size_t len)
 {
 	struct sockaddr sa;
 	socklen_t sa_len = sizeof(sa);
@@ -305,6 +305,62 @@ static int get_client_ip(int fd, char *dst, size_t len)
 	return 0;
 }
 
+static void read_single_line(int fd, char *dst, size_t len)
+{
+	/*
+	 * This reads exactly one line from fd.
+	 * We cannot read bytes after the new line.
+	 * We could use recv() with MSG_PEEK to do this more efficiently.
+	 */
+	for (size_t i = 0; i < len; i++) {
+		if (read(fd, &dst[i], 1) <= 0)
+			break;
+
+		if (dst[i] == '\r')
+			i--;
+
+		if (dst[i] == '\n') {
+			dst[i] = '\0';
+			return;
+		}
+	}
+
+	tmate_fatal("Cannot read proxy header. Load balancer may be misconfigured");
+}
+
+static int get_client_ip_proxy_protocol(int fd, char *dst, size_t len)
+{
+	char header[110];
+	const char *signature = "PROXY ";
+
+	if (read(fd, header, strlen(signature)) != (ssize_t)strlen(signature))
+		tmate_fatal("Cannot read proxy header");
+
+	if (memcmp(header, signature, strlen(signature)))
+		tmate_fatal("No proxy header found. Load balancer may be misconfigured");
+
+	read_single_line(fd, header, sizeof(header));
+
+	int tok_num = 0;
+	for (char *tok = strtok(header, " "); tok; tok = strtok(NULL, " "), tok_num++) {
+		if (tok_num == 1)
+			strncpy(dst, tok, len);
+	}
+
+	if (tok_num != 5)
+		tmate_fatal("Proxy header is invalid");
+
+	return 0;
+}
+
+static int get_client_ip(int fd, char *dst, size_t len)
+{
+	if (tmate_settings->use_proxy_protocol)
+		return get_client_ip_proxy_protocol(fd, dst, len);
+	else
+		return get_client_ip_socket(fd, dst, len);
+}
+
 static void ssh_log_function(int priority, const char *function,
 			     const char *buffer, __unused void *userdata)
 {
@@ -426,7 +482,7 @@ void tmate_ssh_server_main(struct tmate_session *session, const char *keys_dir,
 		alarm(TMATE_SSH_GRACE_PERIOD);
 
 		if (get_client_ip(fd, client->ip_address, sizeof(client->ip_address)) < 0)
-			tmate_fatal("Error getting Client IP from connection");
+			tmate_fatal("Error getting client IP from connection");
 
 		tmate_info("Connection accepted ip=%s", client->ip_address);
 
diff --git a/tmate.h b/tmate.h
index d9fc5b70..042303fe 100644
--- a/tmate.h
+++ b/tmate.h
@@ -216,6 +216,7 @@ struct tmate_settings {
 	const char *tmate_host;
 	const char *bind_addr;
 	int log_level;
+	bool use_proxy_protocol;
 	bool use_syslog;
 };
 extern struct tmate_settings *tmate_settings;