fix bugs:退出资产无法获取第一个输入

pkg/handler/banner.go

@@ -6,8 +6,6 @@ import (
 	"io"
 	"text/template"
 
-	"github.com/gliderlabs/ssh"
-
 	"github.com/jumpserver/koko/pkg/config"
 	"github.com/jumpserver/koko/pkg/i18n"
 	"github.com/jumpserver/koko/pkg/logger"
@@ -62,7 +60,7 @@ type ColorMeta struct {
 	ColorEnd       string
 }
 
-func displayBanner(sess ssh.Session, user string) {
+func displayBanner(sess io.ReadWriter, user string) {
 	title := defaultTitle
 	cf := config.GetConf()
 	if cf.HeaderTitle != "" {

pkg/handler/pagination.go

@@ -145,14 +145,14 @@ func (p *AssetPagination) displayPageAssets() {
 			names[i] = systemUser[i].Name
 		}
 		row["systemUsers"] = strings.Join(names, ",")
-		comments := make([]string,0)
-		for _, item := range strings.Split(strings.TrimSpace(j.Comment), "\r\n"){
-			if strings.TrimSpace(item) == ""{
+		comments := make([]string, 0)
+		for _, item := range strings.Split(strings.TrimSpace(j.Comment), "\r\n") {
+			if strings.TrimSpace(item) == "" {
 				continue
 			}
-			comments = append(comments,strings.ReplaceAll(strings.TrimSpace(item)," ",","))
+			comments = append(comments, strings.ReplaceAll(strings.TrimSpace(item), " ", ","))
 		}
-		row["comment"] = strings.Join(comments,"|")
+		row["comment"] = strings.Join(comments, "|")
 		data[i] = row
 	}
 	w, _ := p.term.GetSize()

pkg/handler/session.go

@@ -60,9 +60,10 @@ func SessionHandler(sess ssh.Session) {
 }
 
 func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandler {
-	term := utils.NewTerminal(sess, "Opt> ")
+	wrapperSess := NewWrapperSession(sess)
+	term := utils.NewTerminal(wrapperSess, "Opt> ")
 	handler := &interactiveHandler{
-		sess:           sess,
+		sess:           wrapperSess,
 		user:           user,
 		term:           term,
 		mu:             new(sync.RWMutex),
@@ -73,7 +74,7 @@ func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandl
 }
 
 type interactiveHandler struct {
-	sess         ssh.Session
+	sess         *WrapperSession
 	user         *model.User
 	term         *utils.Terminal
 	winWatchChan chan bool
@@ -116,11 +117,11 @@ func (h *interactiveHandler) displayBanner() {
 }
 
 func (h *interactiveHandler) watchWinSizeChange() {
-	_, sessChan, _ := h.sess.Pty()
+	sessChan := h.sess.WinCh()
 	winChan := sessChan
 	for {
 		select {
-		case <-h.sess.Context().Done():
+		case <-h.sess.Sess.Context().Done():
 			return
 		case sig, ok := <-h.winWatchChan:
 			if !ok {
@@ -366,9 +367,8 @@ func (h *interactiveHandler) searchNodeAssets(num int) (assets []model.Asset) {
 }
 
 func (h *interactiveHandler) Proxy(ctx context.Context) {
-	userConn := &proxy.UserSSHConnection{Session: h.sess}
 	p := proxy.ProxyServer{
-		UserConn:   userConn,
+		UserConn:   h.sess,
 		User:       h.user,
 		Asset:      h.assetSelect,
 		SystemUser: h.systemUserSelect,

pkg/handler/test.go

-package handler
-
-import (
-	"fmt"
-	"io"
-	"os"
-	"os/exec"
-	"syscall"
-	"unsafe"
-
-	"github.com/gliderlabs/ssh"
-	"github.com/kr/pty"
-)
-
-func setWinsize(f *os.File, w, h int) {
-	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
-		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
-}
-
-func TestHandler(s ssh.Session) {
-	cmd := exec.Command("top")
-	ptyReq, winCh, isPty := s.Pty()
-	if isPty {
-		cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
-		f, err := pty.Start(cmd)
-		if err != nil {
-			panic(err)
-		}
-		go func() {
-			for win := range winCh {
-				setWinsize(f, win.Width, win.Height)
-			}
-		}()
-		go func() {
-			io.Copy(f, s) // stdin
-		}()
-		io.Copy(s, f) // stdout
-	} else {
-		io.WriteString(s, "No PTY requested.\n")
-		s.Exit(1)
-	}
-}

pkg/handler/wrappersession.go

+package handler
+
+import (
+	"io"
+	"strings"
+	"sync"
+
+	"github.com/gliderlabs/ssh"
+
+	"github.com/jumpserver/koko/pkg/logger"
+)
+
+type WrapperSession struct {
+	Sess      ssh.Session
+	inWriter  io.WriteCloser
+	outReader io.ReadCloser
+	mux       *sync.RWMutex
+}
+
+func (w *WrapperSession) initial() {
+	w.initReadPip()
+	go w.readLoop()
+}
+
+func (w *WrapperSession) readLoop() {
+	defer logger.Debug("session loop break")
+	buf := make([]byte, 1024*8)
+	for {
+		nr, err := w.Sess.Read(buf)
+
+		if nr > 0 {
+			w.mux.RLock()
+			_, _ = w.inWriter.Write(buf[:nr])
+			w.mux.RUnlock()
+		}
+		if err != nil {
+			break
+		}
+	}
+	_ = w.inWriter.Close()
+
+}
+
+func (w *WrapperSession) Read(p []byte) (int, error) {
+	w.mux.RLock()
+	defer w.mux.RUnlock()
+	return w.outReader.Read(p)
+}
+
+func (w *WrapperSession) Close() error {
+	var err error
+	err = w.inWriter.Close()
+	w.initReadPip()
+	return err
+}
+
+func (w *WrapperSession) Write(p []byte) (int, error) {
+	return w.Sess.Write(p)
+}
+
+func (w *WrapperSession) initReadPip() {
+	w.mux.Lock()
+	defer w.mux.Unlock()
+	w.outReader, w.inWriter = io.Pipe()
+}
+
+func (w *WrapperSession) Protocol() string {
+	return "ssh"
+}
+
+func (w *WrapperSession) User() string {
+	return w.Sess.User()
+}
+
+func (w *WrapperSession) WinCh() (winch <-chan ssh.Window) {
+	_, winch, ok := w.Sess.Pty()
+	if ok {
+		return
+	}
+	return nil
+}
+
+func (w *WrapperSession) LoginFrom() string {
+	return "ST"
+}
+
+func (w *WrapperSession) RemoteAddr() string {
+	return strings.Split(w.Sess.RemoteAddr().String(), ":")[0]
+}
+
+func (w *WrapperSession) Pty() ssh.Pty {
+	pty, _, _ := w.Sess.Pty()
+	return pty
+}
+
+func NewWrapperSession(sess ssh.Session) *WrapperSession {
+	w := &WrapperSession{
+		Sess: sess,
+		mux:  new(sync.RWMutex),
+	}
+	w.initial()
+	return w
+}

pkg/httpd/client.go

@@ -22,20 +22,14 @@ type Client struct {
 	pty       ssh.Pty
 }
 
-func (c *Client) Protocol() string {
-	return "ws"
-}
-
 func (c *Client) WinCh() <-chan ssh.Window {
 	return c.WinChan
 }
 
-func (c *Client) User() string {
-	return c.user.Username
-}
 func (c *Client) LoginFrom() string {
 	return "WT"
 }
+
 func (c *Client) RemoteAddr() string {
 	return c.addr
 }

pkg/proxy/parser.go

@@ -152,6 +152,8 @@ func (p *Parser) parseCmdOutput() {
 
 // ParseUserInput 解析用户的输入
 func (p *Parser) ParseUserInput(b []byte) []byte {
+	p.lock.Lock()
+	defer p.lock.Unlock()
 	p.once.Do(func() {
 		p.inputInitial = true
 	})
@@ -165,8 +167,6 @@ func (p *Parser) parseZmodemState(b []byte) {
 	if len(b) < 20 {
 		return
 	}
-	p.lock.Lock()
-	defer p.lock.Unlock()
 	if p.zmodemState == "" {
 		if len(b) > 25 && bytes.Contains(b[:50], zmodemRecvStartMark) {
 			p.zmodemState = zmodemStateRecv
@@ -202,8 +202,6 @@ func (p *Parser) parseVimState(b []byte) {
 func (p *Parser) splitCmdStream(b []byte) {
 	p.parseVimState(b)
 	p.parseZmodemState(b)
-	p.lock.RLock()
-	defer p.lock.RUnlock()
 	if p.zmodemState != "" || p.inVimState || !p.inputInitial {
 		return
 	}
@@ -220,6 +218,8 @@ func (p *Parser) splitCmdStream(b []byte) {
 
 // ParseServerOutput 解析服务器输出
 func (p *Parser) ParseServerOutput(b []byte) []byte {
+	p.lock.Lock()
+	defer p.lock.Unlock()
 	p.splitCmdStream(b)
 	return b
 }
@@ -246,8 +246,8 @@ func (p *Parser) IsCommandForbidden() (string, bool) {
 }
 
 func (p *Parser) IsInZmodemRecvState() bool {
-	p.lock.RLock()
-	defer p.lock.RUnlock()
+	p.lock.Lock()
+	defer p.lock.Unlock()
 	return p.zmodemState != ""
 }
 

pkg/proxy/transport.go

@@ -14,7 +14,7 @@ type Transport interface {
 
 type DirectTransport struct {
 	name       string
-	readWriter io.ReadWriter
+	readWriter io.ReadWriteCloser
 	ch         chan []byte
 	closed     bool
 }
@@ -28,13 +28,12 @@ func (dt *DirectTransport) Write(p []byte) (n int, err error) {
 }
 
 func (dt *DirectTransport) Close() error {
-	logger.Debug("Close transport")
 	if dt.closed {
 		return nil
 	}
+	logger.Debug("Close transport")
 	dt.closed = true
-	close(dt.ch)
-	return nil
+	return dt.readWriter.Close()
 }
 
 func (dt *DirectTransport) Chan() <-chan []byte {
@@ -49,18 +48,13 @@ func (dt *DirectTransport) Keep() {
 			_ = dt.Close()
 			break
 		}
-		if !dt.closed {
-			dt.ch <- buf[:n]
-		} else {
-			logger.Debug("Transport closed")
-			break
-		}
+		dt.ch <- buf[:n]
 	}
-	return
+	close(dt.ch)
 }
 
-func NewDirectTransport(name string, readWriter io.ReadWriter) Transport {
-	ch := make(chan []byte, 1024*32)
+func NewDirectTransport(name string, readWriter io.ReadWriteCloser) Transport {
+	ch := make(chan []byte, 1024)
 	tr := DirectTransport{readWriter: readWriter, ch: ch}
 	go tr.Keep()
 	return &tr

pkg/proxy/userconn.go

@@ -2,50 +2,14 @@ package proxy
 
 import (
 	"io"
-	"strings"
 
 	"github.com/gliderlabs/ssh"
 )
 
 type UserConnection interface {
 	io.ReadWriteCloser
-	Protocol() string
 	WinCh() <-chan ssh.Window
-	User() string
 	LoginFrom() string
 	RemoteAddr() string
 	Pty() ssh.Pty
 }
-
-type UserSSHConnection struct {
-	ssh.Session
-}
-
-func (uc *UserSSHConnection) Protocol() string {
-	return "ssh"
-}
-
-func (uc *UserSSHConnection) User() string {
-	return uc.Session.User()
-}
-
-func (uc *UserSSHConnection) WinCh() (winch <-chan ssh.Window) {
-	_, winch, ok := uc.Session.Pty()
-	if ok {
-		return
-	}
-	return nil
-}
-
-func (uc *UserSSHConnection) LoginFrom() string {
-	return "ST"
-}
-
-func (uc *UserSSHConnection) RemoteAddr() string {
-	return strings.Split(uc.Session.RemoteAddr().String(), ":")[0]
-}
-
-func (uc *UserSSHConnection) Pty() ssh.Pty {
-	pty, _, _ := uc.Session.Pty()
-	return pty
-}