From e059b33708a73dd1e7d021dff4ce1d2236a274c3 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Tue, 17 Nov 2020 10:14:11 +0800
Subject: [PATCH] chore: support multiple port (#168)

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
---
 plugin.go      | 14 +++++++++++-
 plugin_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/plugin.go b/plugin.go
index 18e26e6..3317966 100644
--- a/plugin.go
+++ b/plugin.go
@@ -54,13 +54,25 @@ func escapeArg(arg string) string {
 	return "'" + strings.Replace(arg, "'", `'\''`, -1) + "'"
 }
 
+func (p Plugin) hostPort(host string) (string, string) {
+	hosts := strings.Split(host, ":")
+	port := strconv.Itoa(p.Config.Port)
+	if len(hosts) > 1 {
+		host = hosts[0]
+		port = hosts[1]
+	}
+
+	return host, port
+}
+
 func (p Plugin) exec(host string, wg *sync.WaitGroup, errChannel chan error) {
+	host, port := p.hostPort(host)
 	// Create MakeConfig instance with remote username, server address and path to private key.
 	ssh := &easyssh.MakeConfig{
 		Server:            host,
 		User:              p.Config.Username,
 		Password:          p.Config.Password,
-		Port:              strconv.Itoa(p.Config.Port),
+		Port:              port,
 		Key:               p.Config.Key,
 		KeyPath:           p.Config.KeyPath,
 		Passphrase:        p.Config.Passphrase,
diff --git a/plugin_test.go b/plugin_test.go
index be1cf54..65956ee 100644
--- a/plugin_test.go
+++ b/plugin_test.go
@@ -715,3 +715,62 @@ func TestUseInsecureCipher(t *testing.T) {
 
 	assert.Equal(t, unindent(expected), unindent(buffer.String()))
 }
+
+func TestPlugin_hostPort(t *testing.T) {
+	type fields struct {
+		Config Config
+		Writer io.Writer
+	}
+	type args struct {
+		h string
+	}
+	tests := []struct {
+		name     string
+		fields   fields
+		args     args
+		wantHost string
+		wantPort string
+	}{
+		{
+			name: "default host and port",
+			fields: fields{
+				Config: Config{
+					Port: 22,
+				},
+			},
+			args: args{
+				h: "localhost",
+			},
+			wantHost: "localhost",
+			wantPort: "22",
+		},
+		{
+			name: "different port",
+			fields: fields{
+				Config: Config{
+					Port: 22,
+				},
+			},
+			args: args{
+				h: "localhost:443",
+			},
+			wantHost: "localhost",
+			wantPort: "443",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := Plugin{
+				Config: tt.fields.Config,
+				Writer: tt.fields.Writer,
+			}
+			gotHost, gotPort := p.hostPort(tt.args.h)
+			if gotHost != tt.wantHost {
+				t.Errorf("Plugin.hostPort() gotHost = %v, want %v", gotHost, tt.wantHost)
+			}
+			if gotPort != tt.wantPort {
+				t.Errorf("Plugin.hostPort() gotPort = %v, want %v", gotPort, tt.wantPort)
+			}
+		})
+	}
+}