add labels parameter for connect command

cmd/ktctl/main.go

@@ -20,6 +20,7 @@ var (
 	namespace  string
 	debug      bool
 	image      string
+	labels     string
 
 	// connect
 	disableDNS  bool
@@ -31,7 +32,7 @@ var (
 	// exchange
 	expose string
 
-	//context
+	// context
 	pidFile  string
 	userHome string
 )
@@ -79,6 +80,11 @@ func main() {
 			Usage:       "debug mode",
 			Destination: &debug,
 		},
+		cli.StringFlag{
+			Name:        "label,l",
+			Usage:       "Extra labels on proxy pod e.g. 'label1=val1,label2=val2'",
+			Destination: &labels,
+		},
 	}
 
 	app.Commands = []cli.Command{
@@ -123,6 +129,7 @@ func main() {
 					Image:      image,
 					PidFile:    pidFile,
 					UserHome:   userHome,
+					Labels:     labels,
 				}
 				if debug {
 					zerolog.SetGlobalLevel(zerolog.DebugLevel)
@@ -149,6 +156,7 @@ func main() {
 					Image:      image,
 					PidFile:    pidFile,
 					UserHome:   userHome,
+					Labels:     labels,
 				}
 				if debug {
 					zerolog.SetGlobalLevel(zerolog.DebugLevel)
@@ -175,6 +183,7 @@ func main() {
 					Image:      image,
 					PidFile:    pidFile,
 					UserHome:   userHome,
+					Labels:     labels,
 				}
 				if debug {
 					zerolog.SetGlobalLevel(zerolog.DebugLevel)

pkg/kt/action/action.go

@@ -8,4 +8,5 @@ type Action struct {
 	Image      string
 	PidFile    string
 	UserHome   string
+	Labels     string
 }

pkg/kt/action/conect_action.go

@@ -48,15 +48,19 @@ func (action *Action) Connect(sshPort int, method string, socke5Proxy int, disab
 	}
 
 	workload := fmt.Sprintf("kt-connect-daemon-%s", strings.ToLower(util.RandomString(5)))
+	labels := map[string]string{
+		"kt":           workload,
+		"kt-component": "connect",
+		"control-by":   "kt",
+	}
+	for k, v := range util.String2Map(action.Labels) {
+		labels[k] = v
+	}
 
 	endPointIP, podName, err := factory.CreateEndpoint(
 		clientSet,
 		workload,
-		map[string]string{
-			"kt":           workload,
-			"kt-component": "connect",
-			"control-by":   "kt",
-		},
+		labels,
 		action.Image,
 		action.Namespace,
 	)

pkg/kt/connect/connect.go

@@ -88,7 +88,8 @@ func (c *Connect) PrepareSSHPrivateKey() (err error) {
 }
 
 // CreateEndpoint create a endpoint to connect from local
-func (c *Connect) CreateEndpoint(clientset *kubernetes.Clientset, name string, labels map[string]string, image string, namespace string) (podIP string, podName string, err error) {
+func (c *Connect) CreateEndpoint(clientset *kubernetes.Clientset, name string, labels map[string]string, image string,
+	namespace string) (podIP string, podName string, err error) {
 	return createAndWait(clientset, namespace, name, labels, image)
 }
 

pkg/kt/util/util.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"strings"
 )
 
 // IsDaemonRunning check daemon is running or not
@@ -92,3 +93,16 @@ func SSHUttle(remoteHost string, remotePort int, DNSServer string, disableDNS bo
 	args = append(args, cidrs...)
 	return exec.Command("sshuttle", args...)
 }
+
+// Convert parameter string to real map "k1=v1,k2=v2" -> {"k1":"v1","k2","v2"}
+func String2Map(str string) map[string]string {
+	res := make(map[string]string)
+	splitStr := strings.Split(str, ",")
+	for _, item := range splitStr {
+		index := strings.Index(item, "=")
+		if index > 0 {
+			res[item[0:index]] = item[index+1:]
+		}
+	}
+	return res
+}