fix bug: cpu problem

pkg/proxy/parser.go

@@ -34,9 +34,7 @@ func newParser() *Parser {
 
 // Parse 解析用户输入输出, 拦截过滤用户输入输出
 type Parser struct {
-	userInputChan  chan []byte
 	userOutputChan chan []byte
-	srvInputChan   chan []byte
 	srvOutputChan  chan []byte
 	cmdRecordChan  chan [2]string
 
@@ -65,41 +63,44 @@ func (p *Parser) initial() {
 	p.cmdOutputParser = NewCmdParser()
 
 	p.closed = make(chan struct{})
-	p.userInputChan = make(chan []byte, 1024)
-	p.userOutputChan = make(chan []byte, 1024)
-	p.srvInputChan = make(chan []byte, 1024)
-	p.srvOutputChan = make(chan []byte, 1024)
 	p.cmdRecordChan = make(chan [2]string, 1024)
 }
 
 // ParseStream 解析数据流
-func (p *Parser) ParseStream() {
-	defer func() {
-		close(p.userOutputChan)
-		close(p.srvOutputChan)
-		close(p.cmdRecordChan)
-		_ = p.cmdOutputParser.Close()
-		_ = p.cmdInputParser.Close()
-		logger.Debug("Parser parse stream routine done")
-	}()
-	for {
-		select {
-		case <-p.closed:
-			return
-		case b, ok := <-p.userInputChan:
-			if !ok {
-				return
-			}
-			b = p.ParseUserInput(b)
-			p.userOutputChan <- b
-		case b, ok := <-p.srvInputChan:
-			if !ok {
+func (p *Parser) ParseStream(userInChan, srvInChan <-chan []byte) (userOut, srvOut <-chan []byte) {
+
+	p.userOutputChan = make(chan []byte, 1)
+	p.srvOutputChan = make(chan []byte, 1)
+
+	go func() {
+		defer func() {
+			close(p.cmdRecordChan)
+			close(p.userOutputChan)
+			close(p.srvOutputChan)
+			_ = p.cmdOutputParser.Close()
+			_ = p.cmdInputParser.Close()
+			logger.Debug("Parser parse stream routine done")
+		}()
+		for {
+			select {
+			case <-p.closed:
 				return
+			case b, ok := <-userInChan:
+				if !ok {
+					return
+				}
+				b = p.ParseUserInput(b)
+				p.userOutputChan <- b
+			case b, ok := <-srvInChan:
+				if !ok {
+					return
+				}
+				b = p.ParseServerOutput(b)
+				p.srvOutputChan <- b
 			}
-			b = p.ParseServerOutput(b)
-			p.srvOutputChan <- b
 		}
-	}
+	}()
+	return p.userOutputChan, p.srvOutputChan
 }
 
 // Todo: parseMultipleInput 依然存在问题
@@ -170,7 +171,7 @@ func (p *Parser) parseZmodemState(b []byte) {
 		if bytes.Contains(b[:24], zmodemEndMark) {
 			logger.Debug("Zmodem end")
 			p.zmodemState = ""
-		} else if bytes.Contains(b[:24], zmodemCancelMark) {
+		} else if bytes.Contains(b, zmodemCancelMark) {
 			logger.Debug("Zmodem cancel")
 			p.zmodemState = ""
 		}
@@ -200,9 +201,7 @@ func (p *Parser) splitCmdStream(b []byte) {
 		p.cmdInputParser.WriteData(b)
 		return
 	}
-	// outputBuff 最大存储1024， 否则可能撑爆内存
-	// 如果最后一个字符不是ascii, 可以截断了某个中文字符的一部分，为了安全继续添加
-		p.cmdOutputParser.WriteData(b)
+	p.cmdOutputParser.WriteData(b)
 }
 
 // ParseServerOutput 解析服务器输出
@@ -249,6 +248,4 @@ func (p *Parser) Close() {
 		close(p.closed)
 
 	}
-	close(p.userInputChan)
-	close(p.srvInputChan)
 }

pkg/proxy/switch.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"fmt"
+	"io"
 	"strings"
 	"time"
 
@@ -38,9 +39,6 @@ type SwitchSession struct {
 	replayRecorder *ReplyRecorder
 	parser         *Parser
 
-	userTran Transport
-	srvTran  Transport
-
 	ctx    context.Context
 	cancel context.CancelFunc
 }
@@ -111,8 +109,6 @@ func (s *SwitchSession) postBridge() {
 	s.parser.Close()
 	s.replayRecorder.End()
 	s.cmdRecorder.End()
-	_ = s.userTran.Close()
-	_ = s.srvTran.Close()
 }
 
 // SetFilterRules 设置命令过滤规则
@@ -123,26 +119,30 @@ func (s *SwitchSession) SetFilterRules(cmdRules []model.SystemUserFilterRule) {
 // Bridge 桥接两个链接
 func (s *SwitchSession) Bridge(userConn UserConnection, srvConn srvconn.ServerConnection) (err error) {
 	winCh := userConn.WinCh()
-	// 将ReadWriter转换为Channel读写
-	s.srvTran = NewDirectTransport(s.ID, srvConn)
-	s.userTran = NewDirectTransport(s.ID, userConn)
-
 	defer func() {
 		logger.Info("Session bridge done: ", s.ID)
+		_ = userConn.Close()
+		_ = srvConn.Close()
 		s.postBridge()
 	}()
 
+	userInChan := make(chan []byte, 10)
+	srvInChan := make(chan []byte, 10)
+
 	// 处理数据流
-	go s.parser.ParseStream()
+	userOutChan, srvOutChan := s.parser.ParseStream(userInChan, srvInChan)
 	// 记录命令
 	go s.recordCommand()
+	go LoopRead(userConn, userInChan)
+	go LoopRead(srvConn, srvInChan)
+
 	for {
 		select {
 		// 检测是否超过最大空闲时间
 		case <-time.After(s.MaxIdleTime * time.Minute):
 			msg := fmt.Sprintf(i18n.T("Connect idle more than %d minutes, disconnect"), s.MaxIdleTime)
 			msg = utils.WrapperWarn(msg)
-			utils.IgnoreErrWriteString(s.userTran, "\n\r"+msg)
+			utils.IgnoreErrWriteString(userConn, "\n\r"+msg)
 			return
 		// 手动结束
 		case <-s.ctx.Done():
@@ -151,36 +151,27 @@ func (s *SwitchSession) Bridge(userConn UserConnection, srvConn srvconn.ServerCo
 			utils.IgnoreErrWriteString(userConn, "\n\r"+msg)
 			return
 		// 监控窗口大小变化
-		case win := <-winCh:
-			_ = srvConn.SetWinSize(win.Height, win.Width)
-			logger.Debugf("Window server change: %d*%d", win.Height, win.Width)
-		// Server发来数据流入parser中
-		case p, ok := <-s.srvTran.Chan():
+		case win, ok := <-winCh:
 			if !ok {
 				return
 			}
-			s.parser.srvInputChan <- p
-		// Server流入parser数据，经处理发给用户
-		case p, ok := <-s.parser.srvOutputChan:
+			_ = srvConn.SetWinSize(win.Height, win.Width)
+			logger.Debugf("Window server change: %d*%d", win.Height, win.Width)
+		// 经过parse处理的server数据，发给user
+		case p, ok := <-srvOutChan:
 			if !ok {
 				return
 			}
-			nw, _ := s.userTran.Write(p)
+			nw, _ := userConn.Write(p)
 			if !s.parser.IsInZmodemRecvState() {
 				s.replayRecorder.Record(p[:nw])
 			}
-		// User发来的数据流流入parser
-		case p, ok := <-s.userTran.Chan():
+		// 经过parse处理的user数据，发给server
+		case p, ok := <-userOutChan:
 			if !ok {
 				return
 			}
-			s.parser.userInputChan <- p
-		// User发来的数据经parser处理，发给Server
-		case p, ok := <-s.parser.userOutputChan:
-			if !ok {
-				return
-			}
-			_, err = s.srvTran.Write(p)
+			_, err = srvConn.Write(p)
 		}
 	}
 }
@@ -203,3 +194,18 @@ func (s *SwitchSession) MapData() map[string]interface{} {
 		"date_end":    dataEnd,
 	}
 }
+
+func LoopRead(read io.Reader, inChan chan<- []byte) {
+	defer logger.Debug("loop read end")
+	for {
+		buf := make([]byte, 1024)
+		nr, err := read.Read(buf)
+		if err != nil {
+			break
+		}
+		if nr > 0 {
+			inChan <- buf[:nr]
+		}
+	}
+	close(inChan)
+}