From babe6841418a64e6e3fea90a2e09df4558097496 Mon Sep 17 00:00:00 2001
From: Michael Borkenstein <mborkenstein@cloudflare.com>
Date: Wed, 16 Oct 2019 12:08:33 -0500
Subject: [PATCH] AUTH-2135: Adds support for IPv6 and tests

---
 sshserver/preamble_test.go  | 29 +++++++++++++++++++++++++++++
 sshserver/sshserver_unix.go | 30 +++++++++++++++++++++++++-----
 2 files changed, 54 insertions(+), 5 deletions(-)
 create mode 100644 sshserver/preamble_test.go

diff --git a/sshserver/preamble_test.go b/sshserver/preamble_test.go
new file mode 100644
index 00000000..04fd447e
--- /dev/null
+++ b/sshserver/preamble_test.go
@@ -0,0 +1,29 @@
+package sshserver
+
+import (
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestHasPort(t *testing.T) {
+	type testCase struct {
+		input          string
+		expectedOutput string
+	}
+
+	tests := []testCase{
+		{"localhost", "localhost:22"},
+		{"other.addr:22", "other.addr:22"},
+		{"[2001:db8::1]:8080", "[2001:db8::1]:8080"},
+		{"[::1]", "[::1]:22"},
+		{"2001:0db8:3c4d:0015:0000:0000:1a2f:1234", "[2001:0db8:3c4d:0015:0000:0000:1a2f:1234]:22"},
+		{"::1", "[::1]:22"},
+	}
+
+	for _, test := range tests {
+		out, err := canonicalizeDest(test.input)
+		require.Nil(t, err)
+		assert.Equal(t, test.expectedOutput, out)
+	}
+}
diff --git a/sshserver/sshserver_unix.go b/sshserver/sshserver_unix.go
index e445527c..1ea26f8a 100644
--- a/sshserver/sshserver_unix.go
+++ b/sshserver/sshserver_unix.go
@@ -11,7 +11,6 @@ import (
 	"fmt"
 	"io"
 	"net"
-	"regexp"
 	"runtime"
 	"strings"
 	"time"
@@ -37,6 +36,7 @@ const (
 	sshContextPreamble    = "sshPreamble"
 	sshContextSSHClient   = "sshClient"
 	SSHPreambleLength     = 2
+	defaultSSHPort        = "22"
 )
 
 type auditEvent struct {
@@ -283,16 +283,36 @@ func (s *SSHProxy) readPreamble(conn net.Conn) (*SSHPreamble, error) {
 		return nil, err
 	}
 
-	ok, err := regexp.Match(`^[^:]*:\d+$`, []byte(preamble.Destination))
+	preamble.Destination, err = canonicalizeDest(preamble.Destination)
 	if err != nil {
 		return nil, err
 	}
+	return &preamble, nil
+}
 
-	if !ok {
-		preamble.Destination += ":22"
+// canonicalizeDest adds a default port if one doesnt exist
+func canonicalizeDest(dest string) (string, error) {
+	_, _, err := net.SplitHostPort(dest)
+	// if host and port are split without error, a port exists.
+	if err != nil {
+		addrErr, ok := err.(*net.AddrError)
+		if !ok {
+			return "", err
+		}
+		// If the port is missing, append it.
+		if addrErr.Err == "missing port in address" {
+			return fmt.Sprintf("%s:%s", dest, defaultSSHPort), nil
+		}
+
+		// If there are too many colons and address is IPv6, wrap in brackets and append port. Otherwise invalid address
+		ip := net.ParseIP(dest)
+		if addrErr.Err == "too many colons in address" && ip != nil && ip.To4() == nil {
+			return fmt.Sprintf("[%s]:%s", dest, defaultSSHPort), nil
+		}
+		return "", addrErr
 	}
 
-	return &preamble, nil
+	return dest, nil
 }
 
 // dialDestination creates a new SSH client and dials the destination server