From 00b690ce8c9aa020b816bc47dae3be9fdef89294 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <bradley.rydzewski@harness.io>
Date: Fri, 20 Nov 2020 16:47:52 -0500
Subject: [PATCH] support for remote tmate servers

---
 CHANGELOG.md                   |  4 ++++
 command/compile.go             | 12 ++++++++++++
 command/daemon/config.go       |  4 ++++
 command/daemon/daemon.go       |  4 ++++
 command/exec.go                | 12 ++++++++++++
 engine/compiler/compiler.go    | 12 ++++++++++++
 engine/compiler/shell/shell.go | 20 +++++++-------------
 7 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 739034d..92a9862 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+### Added
+- support for self-hosted tmate instances
+
 ## 1.6.1
 ### Changed
 - restrict temporary volumes used with docker plugins
diff --git a/command/compile.go b/command/compile.go
index 01c5ec7..749e828 100644
--- a/command/compile.go
+++ b/command/compile.go
@@ -202,6 +202,18 @@ func registerCompile(app *kingpin.Application) {
 	cmd.Flag("tmate-enabled", "tmate enabled").
 		BoolVar(&c.Tmate.Enabled)
 
+	cmd.Flag("tmate-server-host", "tmate server host").
+		StringVar(&c.Tmate.Server)
+
+	cmd.Flag("tmate-server-port", "tmate server port").
+		StringVar(&c.Tmate.Port)
+
+	cmd.Flag("tmate-server-rsa-fingerprint", "tmate server rsa fingerprint").
+		StringVar(&c.Tmate.RSA)
+
+	cmd.Flag("tmate-server-ed25519-fingerprint", "tmate server rsa fingerprint").
+		StringVar(&c.Tmate.ED25519)
+
 	// shared pipeline flags
 	c.Flags = internal.ParseFlags(cmd)
 }
diff --git a/command/daemon/config.go b/command/daemon/config.go
index a6d7400..8accf3b 100644
--- a/command/daemon/config.go
+++ b/command/daemon/config.go
@@ -110,6 +110,10 @@ type Config struct {
 	Tmate struct {
 		Enabled bool   `envconfig:"DRONE_TMATE_ENABLED" default:"false"`
 		Image   string `envconfig:"DRONE_TMATE_IMAGE"   default:"drone/drone-runner-docker:1"`
+		Server  string `envconfig:"DRONE_TMATE_HOST"`
+		Port    string `envconfig:"DRONE_TMATE_PORT"`
+		RSA     string `envconfig:"DRONE_TMATE_FINGERPRINT_RSA"`
+		ED25519 string `envconfig:"DRONE_TMATE_FINGERPRINT_ED25519"`
 	}
 }
 
diff --git a/command/daemon/daemon.go b/command/daemon/daemon.go
index 2570228..7d4cade 100644
--- a/command/daemon/daemon.go
+++ b/command/daemon/daemon.go
@@ -144,6 +144,10 @@ func (c *daemonCommand) run(*kingpin.ParseContext) error {
 			Tmate: compiler.Tmate{
 				Image:   config.Tmate.Image,
 				Enabled: config.Tmate.Enabled,
+				Server:  config.Tmate.Server,
+				Port:    config.Tmate.Port,
+				RSA:     config.Tmate.RSA,
+				ED25519: config.Tmate.ED25519,
 			},
 			Environ: provider.Combine(
 				provider.Static(config.Runner.Environ),
diff --git a/command/exec.go b/command/exec.go
index f7655b5..5927704 100644
--- a/command/exec.go
+++ b/command/exec.go
@@ -336,6 +336,18 @@ func registerExec(app *kingpin.Application) {
 	cmd.Flag("tmate-enabled", "tmate enabled").
 		BoolVar(&c.Tmate.Enabled)
 
+	cmd.Flag("tmate-server-host", "tmate server host").
+		StringVar(&c.Tmate.Server)
+
+	cmd.Flag("tmate-server-port", "tmate server port").
+		StringVar(&c.Tmate.Port)
+
+	cmd.Flag("tmate-server-rsa-fingerprint", "tmate server rsa fingerprint").
+		StringVar(&c.Tmate.RSA)
+
+	cmd.Flag("tmate-server-ed25519-fingerprint", "tmate server rsa fingerprint").
+		StringVar(&c.Tmate.ED25519)
+
 	cmd.Flag("debug", "enable debug logging").
 		BoolVar(&c.Debug)
 
diff --git a/engine/compiler/compiler.go b/engine/compiler/compiler.go
index 6d9362c..ff57e22 100644
--- a/engine/compiler/compiler.go
+++ b/engine/compiler/compiler.go
@@ -57,6 +57,10 @@ type Resources struct {
 type Tmate struct {
 	Image   string
 	Enabled bool
+	Server  string
+	Port    string
+	RSA     string
+	ED25519 string
 }
 
 // Compiler compiles the Yaml configuration file to an
@@ -236,6 +240,14 @@ func (c *Compiler) Compile(ctx context.Context, args runtime.CompilerArgs) runti
 		envs["DRONE_DOCKER_VOLUME_PATH"] = volume.HostPath.Path
 	}
 
+	// create tmate variables
+	if c.Tmate.Server != "" {
+		envs["DRONE_TMATE_HOST"] = c.Tmate.Server
+		envs["DRONE_TMATE_PORT"] = c.Tmate.Port
+		envs["DRONE_TMATE_FINGERPRINT_RSA"] = c.Tmate.RSA
+		envs["DRONE_TMATE_FINGERPRINT_ED25519"] = c.Tmate.ED25519
+	}
+
 	// create the .netrc environment variables if not
 	// explicitly disabled
 	if c.NetrcCloneOnly == false {
diff --git a/engine/compiler/shell/shell.go b/engine/compiler/shell/shell.go
index c37a600..a064a12 100644
--- a/engine/compiler/shell/shell.go
+++ b/engine/compiler/shell/shell.go
@@ -63,20 +63,14 @@ remote_debug() {
 	fi
 }
 
+if [ ! -z "${DRONE_TMATE_HOST}" ]; then
+	echo "set -g tmate-server-host $DRONE_TMATE_HOST" >> $HOME/.tmate.conf
+	echo "set -g tmate-server-port $DRONE_TMATE_PORT" >> $HOME/.tmate.conf
+	echo "set -g tmate-server-rsa-fingerprint $DRONE_TMATE_FINGERPRINT_RSA" >> $HOME/.tmate.conf
+	echo "set -g tmate-server-ed25519-fingerprint $DRONE_TMATE_FINGERPRINT_ED25519" >> $HOME/.tmate.conf
+fi
+
 if [ "${DRONE_BUILD_DEBUG}" = "true" ]; then
 	trap remote_debug EXIT
 fi
 `
-
-// TODO(bradrydzewski) add timeout 30m to tmate
-
-//
-// TODO(bradrydzewski) support custom tmate config options:
-//
-//   if [ ! -z "${DRONE_TMATE_SERVER_HOST}" ]; then
-// 	  echo "set -g tmate-server-host \"$DRONE_TMATE_SERVER_HOST\"" >> $HOME/.tmate.conf
-// 	  echo "set -g tmate-server-port $DRONE_TMATE_SERVER_PORT" >> $HOME/.tmate.conf
-// 	  echo "set -g tmate-server-rsa-fingerprint \"$DRONE_TMATE_SERVER_RSA_FINGERPRINT\"" >> $HOME/.tmate.conf
-// 	  echo "set -g tmate-server-ed25519-fingerprint \"$DRONE_TMATE_SERVER_ED25519_FINGERPRINT\"" >> $HOME/.tmate.conf
-//   fi
-//