From 22c40a6f72a2a8f736c59a8e415c4db7a2208885 Mon Sep 17 00:00:00 2001
From: yunlzheng <yunl.zheng@gmail.com>
Date: Sun, 8 Mar 2020 13:33:17 +0800
Subject: [PATCH] fix: fixed GetRandomSSHPort fail and add test case.

---
 pkg/kt/util/network.go      | 14 +++++++++++-
 pkg/kt/util/network_test.go | 43 +++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 pkg/kt/util/network_test.go

diff --git a/pkg/kt/util/network.go b/pkg/kt/util/network.go
index e8fc9ea..5a98a2a 100644
--- a/pkg/kt/util/network.go
+++ b/pkg/kt/util/network.go
@@ -3,13 +3,25 @@ package util
 import (
 	"fmt"
 	"net"
+	"strings"
 
 	"github.com/rs/zerolog/log"
 )
 
 // GetRandomSSHPort get pod random ssh port
 func GetRandomSSHPort(podIP string) string {
-	return fmt.Sprintf("22%s", podIP[len(podIP)-2:])
+	parts := strings.Split(podIP, ".")
+	rdm := parts[len(parts)-1]
+
+	if len(rdm) == 1 {
+		rdm = fmt.Sprintf("0%s", rdm)
+	}
+
+	if len(rdm) > 2 {
+		rdm = rdm[len(rdm)-2:]
+	}
+
+	return fmt.Sprintf("22%s", rdm)
 }
 
 // GetOutboundIP Get preferred outbound ip of this machine
diff --git a/pkg/kt/util/network_test.go b/pkg/kt/util/network_test.go
new file mode 100644
index 0000000..59a32b1
--- /dev/null
+++ b/pkg/kt/util/network_test.go
@@ -0,0 +1,43 @@
+package util
+
+import "testing"
+
+func TestGetRandomSSHPort(t *testing.T) {
+	type args struct {
+		podIP string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "GetRandomSSHPortPodIP",
+			args: args{
+				podIP: "172.16.0.9",
+			},
+			want: "2209",
+		},
+		{
+			name: "GetRandomSSHPortPodIP",
+			args: args{
+				podIP: "172.16.0.91",
+			},
+			want: "2291",
+		},
+		{
+			name: "GetRandomSSHPortPodIP",
+			args: args{
+				podIP: "172.16.0.191",
+			},
+			want: "2291",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := GetRandomSSHPort(tt.args.podIP); got != tt.want {
+				t.Errorf("GetRandomSSHPort() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}