diff --git a/cmd/kubevpn/cmds/ssh.go b/cmd/kubevpn/cmds/ssh.go index 4a2b6be7..c7a4e3cc 100644 --- a/cmd/kubevpn/cmds/ssh.go +++ b/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(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) + config.Header.Set("ssh", string(bytes)) 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 } diff --git a/pkg/daemon/handler/ssh.go b/pkg/daemon/handler/ssh.go index 63dc9c76..7e98055d 100644 --- a/pkg/daemon/handler/ssh.go +++ b/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 - if v := conn.Request().Header.Get("extra-cidr"); v != "" { - 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) + defer delete(SessionMap, conf.SessionID) + defer delete(CondReady, conf.SessionID) ctx, cancelFunc := context.WithCancel(conn.Request().Context()) h := &wsHandler{ - sshConfig: &sshConfig, + sshConfig: &conf.Config, conn: conn, - cidr: extraCIDR, - width: width, - height: height, - sessionId: sessionID, + cidr: conf.ExtraCIDR, + width: conf.Width, + height: conf.Height, + sessionId: conf.SessionID, condReady: cancelFunc, } - CondReady[sessionID] = ctx + CondReady[conf.SessionID] = ctx defer conn.Close() h.handle(conn.Request().Context()) }))