Merge pull request #62 from balena-io/proxy-protocol

add Proxy Protocol support
balena-io · Jan 14, 2020 · d03476a · d03476a
2 parents a4b9f07 + 2917a14
commit d03476a
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 36 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -4,7 +4,7 @@ jobs:
   build:
     working_directory: /go/src/github.com/balena-io/sshproxy
     docker:
-      - image: golang:1.13.5
+      - image: golang:1.13.6
     steps:
       - checkout
       - run:

diff --git a/README.md b/README.md
@@ -17,37 +17,40 @@ they can all be set via commandline, environment or config file.
 
 | Name               | Commandline                 | Environment                   | Config               |
 |--------------------|-----------------------------|-------------------------------|----------------------|
-| API Host           | `--apihost` `-H`            | `BALENA_API_HOST`             | `apihost`            |
-| API Port           | `--apiport` `-P`            | `BALENA_API_PORT`             | `apiport`            |
-| API Key            | `--apikey` `-K`             | `SSHPROXY_API_KEY`            | `apikey`             |
-| Dir                | `--dir` `-d`                | `SSHPROXY_DIR`                |                      |
-| Port               | `--port` `-p`               | `SSHPROXY_PORT`               | `port`               |
-| Shell              | `--shell` `-s`              | `SSHPROXY_SHELL`              | `shell`              |
-| Shell UID          | `--shell-uid` `-u`          | `SSHPROXY_SHELL_UID`          | `shell-uid`          |
-| Shell GID          | `--shell-gid` `-g`          | `SSHPROXY_SHELL_GID`          | `shell-gid`          |
-| Idle Timeout       | `--idle-timeout`, `-i`      | `SSHPROXY_IDLE_TIMEOUT`       | `idle-timeout`       |
+| Allow Env          | `--allow-env` `-E`          | `SSHPROXY_ALLOW_ENV`          | `allow-env`          |
+| API Host           | `--apihost`, `-H`           | `BALENA_API_HOST`             | `apihost`            |
+| API Key            | `--apikey`, `-K`            | `SSHPROXY_API_KEY`            | `apikey`             |
+| API Port           | `--apiport`, `-P`           | `BALENA_API_PORT`             | `apiport`            |
 | Auth Failed Banner | `--auth-failed-banner` `-b` | `SSHPROXY_AUTH_FAILED_BANNER` | `auth-failed-banner` |
+| Bind               | `--bind`, `-b`              | `SSHPROXY_BIND`               | `bind`               |
+| Dir                | `--dir`, `-d`               | `SSHPROXY_DIR`                |                      |
+| Idle Timeout       | `--idle-timeout`, `-i`      | `SSHPROXY_IDLE_TIMEOUT`       | `idle-timeout`       |
 | Max Auth Tries     | `--max-auth-tries` `-m`     | `SSHPROXY_MAX_AUTH_TRIES`     | `max-auth-tries`     |
-| Allow Env          | `--allow-env` `-E`          | `SSHPROXY_ALLOW_ENV`          | `allow-env`          |
+| Metrics Bind       | `--metrics-bind`, `-M`      | `SSHPROXY_METRICS_BIND`       | `metrics-bind`       |
 | Sentry DSN         | `--sentry-dsn` `-S`         | `SSHPROXY_SENTRY_DSN`         | `sentry-dsn`         |
+| Shell              | `--shell`, `-s`             | `SSHPROXY_SHELL`              | `shell`              |
+| Shell GID          | `--shell-gid`, `-g`         | `SSHPROXY_SHELL_GID`          | `shell-gid`          |
+| Shell UID          | `--shell-uid`, `-u`         | `SSHPROXY_SHELL_UID`          | `shell-uid`          |
+| Use Proxy Protocol | `--use-proxyprotocol`, `-p` | `SSHPROXY_USE_PROXYPROTOCOL`  | `use-proxyprotocol`  |
 | Verbosity          | `--verbosity`, `-v`         | `SSHPROXY_VERBOSITY`          | `verbosity`	 	  |
-| Metrics Port       | `--metrics-port`, `-M`      | `SSHPROXY_METRICS_PORT`       | `metrics-port`       |
 
 ```
 Usage of sshproxy:
   -E, --allow-env string            List of environment variables to pass from client to shell (default: None)
   -H, --apihost string              Balena API Host (default "api.balena-cloud.com")
   -K, --apikey string               Balena API Key (required)
   -P, --apiport string              Balena API Port (default "443")
-  -b, --auth-failed-banner string   Path to template displayed after failed authentication
+  -B, --auth-failed-banner string   Path to template displayed after failed authentication
+  -b, --bind string                 Address the ssh service will bind to (default ":22")
   -d, --dir string                  Work dir, holds ssh keys and sshproxy config (default "/etc/sshproxy")
   -i, --idle-timeout int            Idle timeout (seconds, 0 = none)
   -m, --max-auth-tries int          Maximum number of authentication attempts per connection (default 0; unlimited)
-  -p, --port int                    Port the ssh service will listen on (default 22)
+  -M, --metrics-bind string         Address the prometheus metrics server should bind to (default: disabled)
   -S, --sentry-dsn string           Sentry DSN for error reporting
   -s, --shell string                Path to shell to execute post-authentication (default "shell.sh")
   -g, --shell-gid int               Group to run shell as (default: current gid) (default -1)
   -u, --shell-uid int               User to run shell as (default: current uid) (default -1)
+  -p, --use-proxyprotocol           Enable Proxy Protocol support
   -v, --verbosity int               Set verbosity level (0 = quiet, 1 = normal, 2 = verbose, 3 = debug, default: 1) (default 1)
       --version                     Display version and exit
 ```

diff --git a/main.go b/main.go
@@ -31,6 +31,7 @@ import (
 
 	raven "github.com/getsentry/raven-go"
 	"github.com/gliderlabs/ssh"
+	proxyproto "github.com/wrboyce/go-proxyproto"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
@@ -40,20 +41,21 @@ import (
 var version string
 
 func init() {
+	pflag.CommandLine.StringP("allow-env", "E", "", "List of environment variables to pass from client to shell (default: None)")
 	pflag.CommandLine.StringP("apihost", "H", "api.balena-cloud.com", "Balena API Host")
-	pflag.CommandLine.StringP("apiport", "P", "443", "Balena API Port")
 	pflag.CommandLine.StringP("apikey", "K", "", "Balena API Key (required)")
-	pflag.CommandLine.StringP("dir", "d", "/etc/sshproxy", "Work dir, holds ssh keys and sshproxy config")
+	pflag.CommandLine.StringP("apiport", "P", "443", "Balena API Port")
+	pflag.CommandLine.StringP("auth-failed-banner", "B", "", "Path to template displayed after failed authentication")
 	pflag.CommandLine.StringP("bind", "b", ":22", "Address the ssh service will bind to")
-	pflag.CommandLine.StringP("shell", "s", "shell.sh", "Path to shell to execute post-authentication")
-	pflag.CommandLine.Int64P("shell-uid", "u", -1, "User to run shell as (default: current uid)")
-	pflag.CommandLine.Int64P("shell-gid", "g", -1, "Group to run shell as (default: current gid)")
+	pflag.CommandLine.StringP("dir", "d", "/etc/sshproxy", "Work dir, holds ssh keys and sshproxy config")
 	pflag.CommandLine.IntP("idle-timeout", "i", 0, "Idle timeout (seconds, 0 = none)")
-	pflag.CommandLine.StringP("auth-failed-banner", "B", "", "Path to template displayed after failed authentication")
 	pflag.CommandLine.IntP("max-auth-tries", "m", 0, "Maximum number of authentication attempts per connection (default 0; unlimited)")
-	pflag.CommandLine.StringP("allow-env", "E", "", "List of environment variables to pass from client to shell (default: None)")
 	pflag.CommandLine.StringP("metrics-bind", "M", "", "Address the prometheus metrics server should bind to (default: disabled)")
 	pflag.CommandLine.StringP("sentry-dsn", "S", "", "Sentry DSN for error reporting")
+	pflag.CommandLine.StringP("shell", "s", "shell.sh", "Path to shell to execute post-authentication")
+	pflag.CommandLine.Int64P("shell-gid", "g", -1, "Group to run shell as (default: current gid)")
+	pflag.CommandLine.Int64P("shell-uid", "u", -1, "User to run shell as (default: current uid)")
+	pflag.CommandLine.BoolP("use-proxyprotocol", "p", false, "Enable Proxy Protocol support")
 	pflag.CommandLine.IntP("verbosity", "v", 1, "Set verbosity level (0 = quiet, 1 = normal, 2 = verbose, 3 = debug, default: 1)")
 	pflag.CommandLine.BoolP("version", "", false, "Display version and exit")
 
@@ -63,52 +65,59 @@ func init() {
 		if err := viper.BindPFlags(pflag.CommandLine); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("apihost", "BALENA_API_HOST"); err != nil {
+		if err := viper.BindEnv("allow-env", "SSHPROXY_ALLOW_ENV"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("apiport", "BALENA_API_PORT"); err != nil {
+		if err := viper.BindEnv("apihost", "BALENA_API_HOST"); err != nil {
 			return err
 		}
 		if err := viper.BindEnv("apikey", "SSHPROXY_API_KEY"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("dir"); err != nil {
+		if err := viper.BindEnv("apiport", "BALENA_API_PORT"); err != nil {
+			return err
+		}
+		if err := viper.BindEnv("auth-failed-banner", "SSHPROXY_AUTH_FAILED_BANNER"); err != nil {
 			return err
 		}
 		if err := viper.BindEnv("bind"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("shell"); err != nil {
+		if err := viper.BindEnv("dir"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("verbosity"); err != nil {
+		if err := viper.BindEnv("idle-timeout", "SSHPROXY_IDLE_TIMEOUT"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("shell-uid", "SSHPROXY_SHELL_UID"); err != nil {
+		if err := viper.BindEnv("max-auth-tries", "SSHPROXY_MAX_AUTH_TRIES"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("shell-gid", "SSHPROXY_SHELL_GID"); err != nil {
+		if err := viper.BindEnv("metrics-bind", "SSHPROXY_METRICS_BIND"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("auth-failed-banner", "SSHPROXY_AUTH_FAILED_BANNER"); err != nil {
+		if err := viper.BindEnv("sentry-dsn", "SSHPROXY_SENTRY_DSN"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("max-auth-tries", "SSHPROXY_MAX_AUTH_TRIES"); err != nil {
+		if err := viper.BindEnv("shell"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("allow-env", "SSHPROXY_ALLOW_ENV"); err != nil {
+		if err := viper.BindEnv("shell-gid", "SSHPROXY_SHELL_GID"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("metrics-bind", "SSHPROXY_METRICS_BIND"); err != nil {
+		if err := viper.BindEnv("shell-uid", "SSHPROXY_SHELL_UID"); err != nil {
 			return err
 		}
-		if err := viper.BindEnv("sentry-dsn", "SSHPROXY_SENTRY_DSN"); err != nil {
+		if err := viper.BindEnv("use-proxyprotocol", "SSHPROXY_USE_PROXYPROTOCOL"); err != nil {
+			return err
+		}
+		if err := viper.BindEnv("verbosity"); err != nil {
 			return err
 		}
 		return nil
 	}()
 	if err != nil {
-		log.Fatal("Initialisation failed", err)
+		fmt.Fprintf(os.Stderr, "%s", err)
+		os.Exit(1)
 	}
 }
 
@@ -174,7 +183,6 @@ func main() {
 		MaxAuthTries: viper.GetInt("max-auth-tries"),
 	}
 	server := ssh.Server{
-		Addr:                 viper.GetString("bind"),
 		PublicKeyHandler:     auth.publicKeyHandler,
 		ServerConfigCallback: func(session ssh.Context) *gossh.ServerConfig { return sshConfig },
 	}
@@ -221,10 +229,19 @@ func main() {
 		go serveMetrics(metricsBind)
 	}
 
+	listener, err := net.Listen("tcp", viper.GetString("bind"))
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "%s", err)
+		os.Exit(1)
+	}
+	if viper.GetBool("use-proxyprotocol") {
+		listener = &proxyproto.Listener{Listener: listener}
+	}
+
 	if verbosity >= 1 {
 		log.Printf("starting ssh server on %s", viper.GetString("bind"))
 	}
-	if err := server.ListenAndServe(); err != nil {
+	if err := server.Serve(listener); err != nil {
 		fmt.Fprintf(os.Stderr, "%s", err)
 		os.Exit(1)
 	}