Commit b9394148 authored by Eric's avatar Eric Committed by Eric_Lee

fix bug: cpu problem

parent 110b75c6
......@@ -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)
}
......@@ -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)
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment