From 79fdc99c8170b86ca51c8deac9fc291a9716a9e2 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Sat, 20 Jan 2024 22:02:59 +0800
Subject: [PATCH] update

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
---
 plugin.go      |  4 +++-
 plugin_test.go | 19 +++++++++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/plugin.go b/plugin.go
index 56d5b76..d99639d 100644
--- a/plugin.go
+++ b/plugin.go
@@ -61,7 +61,9 @@ func escapeArg(arg string) string {
 func (p Plugin) hostPort(host string) (string, string) {
 	hosts := strings.Split(host, ":")
 	port := strconv.Itoa(p.Config.Port)
-	if len(hosts) > 1 {
+	if len(hosts) > 1 &&
+		(p.Config.Protocol == easyssh.PROTOCOL_TCP ||
+			p.Config.Protocol == easyssh.PROTOCOL_TCP4) {
 		host = hosts[0]
 		port = hosts[1]
 	}
diff --git a/plugin_test.go b/plugin_test.go
index d3b6af1..db087f1 100644
--- a/plugin_test.go
+++ b/plugin_test.go
@@ -797,7 +797,8 @@ func TestPlugin_hostPort(t *testing.T) {
 			name: "different port",
 			fields: fields{
 				Config: Config{
-					Port: 22,
+					Port:     22,
+					Protocol: easyssh.PROTOCOL_TCP4,
 				},
 			},
 			args: args{
@@ -806,6 +807,20 @@ func TestPlugin_hostPort(t *testing.T) {
 			wantHost: "localhost",
 			wantPort: "443",
 		},
+		{
+			name: "ipv6",
+			fields: fields{
+				Config: Config{
+					Port:     22,
+					Protocol: easyssh.PROTOCOL_TCP6,
+				},
+			},
+			args: args{
+				h: "::1",
+			},
+			wantHost: "::1",
+			wantPort: "22",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -963,7 +978,7 @@ func TestCommandWithIPv6(t *testing.T) {
 
 	plugin := Plugin{
 		Config: Config{
-			Host:     []string{"127.0.0.1"},
+			Host:     []string{"::1"},
 			Username: "drone-scp",
 			Port:     22,
 			KeyPath:  "./tests/.ssh/id_rsa",