refactor: cmd ssh client

cmd/kubevpn/cmds/ssh.go

@@ -6,8 +6,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"strconv"
-	"strings"
 
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
@@ -20,6 +18,7 @@ import (
 	"k8s.io/kubectl/pkg/util/term"
 
 	"github.com/wencaiwulue/kubevpn/v2/pkg/daemon"
+	"github.com/wencaiwulue/kubevpn/v2/pkg/daemon/handler"
 	pkgssh "github.com/wencaiwulue/kubevpn/v2/pkg/ssh"
 	"github.com/wencaiwulue/kubevpn/v2/pkg/util"
 )
@@ -28,7 +27,7 @@ import (
 // Remember to use network mask 32, because ssh using unique network CIDR 223.255.0.0/16
 func CmdSSH(_ cmdutil.Factory) *cobra.Command {
 	var sshConf = &pkgssh.SshConfig{}
-	var ExtraCIDR []string
+	var extraCIDR []string
 	cmd := &cobra.Command{
 		Use:   "ssh",
 		Short: "Ssh to jump server",
@@ -72,16 +71,19 @@ func CmdSSH(_ cmdutil.Factory) *cobra.Command {
 			if err != nil {
 				return fmt.Errorf("terminal get size: %s", err)
 			}
-			marshal, err := json.Marshal(sshConf)
+			sessionID := uuid.NewString()
+			ssh := handler.Ssh{
+				Config:    *sshConf,
+				ExtraCIDR: extraCIDR,
+				Width:     width,
+				Height:    height,
+				SessionID: sessionID,
+			}
+			bytes, err := json.Marshal(ssh)
 			if err != nil {
 				return err
 			}
-			sessionID := uuid.NewString()
+			config.Header.Set("ssh", string(bytes))
-			config.Header.Set("ssh", string(marshal))
-			config.Header.Set("extra-cidr", strings.Join(ExtraCIDR, ","))
-			config.Header.Set("terminal-width", strconv.Itoa(width))
-			config.Header.Set("terminal-height", strconv.Itoa(height))
-			config.Header.Set("session-id", sessionID)
 			client := daemon.GetTCPClient(true)
 			if client == nil {
 				return fmt.Errorf("client is nil")
@@ -114,7 +116,7 @@ func CmdSSH(_ cmdutil.Factory) *cobra.Command {
 		},
 	}
 	pkgssh.AddSshFlags(cmd.Flags(), sshConf)
-	cmd.Flags().StringArrayVar(&ExtraCIDR, "extra-cidr", []string{}, "Extra network CIDR string, eg: --extra-cidr 192.168.0.159/24 --extra-cidr 192.168.1.160/32")
+	cmd.Flags().StringArrayVar(&extraCIDR, "extra-cidr", []string{}, "Extra network CIDR string, eg: --extra-cidr 192.168.0.159/24 --extra-cidr 192.168.1.160/32")
 	return cmd
 }
 

pkg/daemon/handler/ssh.go

@@ -341,36 +341,38 @@ func (w *wsHandler) PrintLine(msg string) {
 var SessionMap = make(map[string]*ssh.Session)
 var CondReady = make(map[string]context.Context)
 
+type Ssh struct {
+	Config    pkgssh.SshConfig
+	ExtraCIDR []string
+	Width     int
+	Height    int
+	SessionID string
+}
+
 func init() {
 	http.Handle("/ws", websocket.Handler(func(conn *websocket.Conn) {
-		var sshConfig pkgssh.SshConfig
 		b := conn.Request().Header.Get("ssh")
-		if err := json.Unmarshal([]byte(b), &sshConfig); err != nil {
+		var conf Ssh
+		err := json.Unmarshal([]byte(b), &conf)
+		if err != nil {
 			_, _ = conn.Write([]byte(err.Error()))
 			_ = conn.Close()
 			return
 		}
-		var extraCIDR []string
+		defer delete(SessionMap, conf.SessionID)
-		if v := conn.Request().Header.Get("extra-cidr"); v != "" {
+		defer delete(CondReady, conf.SessionID)
-			extraCIDR = strings.Split(v, ",")
-		}
-		width, _ := strconv.Atoi(conn.Request().Header.Get("width"))
-		height, _ := strconv.Atoi(conn.Request().Header.Get("height"))
-		sessionID := conn.Request().Header.Get("session-id")
-		defer delete(SessionMap, sessionID)
-		defer delete(CondReady, sessionID)
 
 		ctx, cancelFunc := context.WithCancel(conn.Request().Context())
 		h := &wsHandler{
-			sshConfig: &sshConfig,
+			sshConfig: &conf.Config,
 			conn:      conn,
-			cidr:      extraCIDR,
+			cidr:      conf.ExtraCIDR,
-			width:     width,
+			width:     conf.Width,
-			height:    height,
+			height:    conf.Height,
-			sessionId: sessionID,
+			sessionId: conf.SessionID,
 			condReady: cancelFunc,
 		}
-		CondReady[sessionID] = ctx
+		CondReady[conf.SessionID] = ctx
 		defer conn.Close()
 		h.handle(conn.Request().Context())
 	}))