[Update] 修改config

pkg/common/client.go

@@ -59,6 +59,10 @@ func (c *Client) SetAuth(auth ClientAuth) {
 	c.Auth = auth
 }
 
+func (c *Client) SetHeader(k, v string) {
+	c.Headers[k] = v
+}
+
 func (c *Client) marshalData(data interface{}) (reader io.Reader, error error) {
 	dataRaw, err := json.Marshal(data)
 	if err != nil {

pkg/handler/session.go

@@ -27,7 +27,7 @@ func SessionHandler(sess ssh.Session) {
 		ctx, cancel := cctx.NewContext(sess)
 		defer cancel()
 		handler := newInteractiveHandler(sess, ctx.User())
-		logger.Debugf("User Request pty: %s %s", sess.User(), pty.Term)
+		logger.Debugf("User Request pty %s: %s", pty.Term, sess.User())
 		handler.Dispatch(ctx)
 	} else {
 		utils.IgnoreErrWriteString(sess, "No PTY requested.\n")

pkg/proxy/parser.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"bytes"
+	"cocogo/pkg/i18n"
 	"fmt"
 	"sync"
 
@@ -25,6 +26,12 @@ var (
 	charEnter = []byte("\r")
 )
 
+func newParser() *Parser {
+	parser := &Parser{}
+	parser.initial()
+	return parser
+}
+
 // Parse 解析用户输入输出, 拦截过滤用户输入输出
 type Parser struct {
 	inputBuf  *bytes.Buffer
@@ -35,8 +42,7 @@ type Parser struct {
 	userOutputChan chan []byte
 	srvInputChan   chan []byte
 	srvOutputChan  chan []byte
-
-	cmdChan chan [2]string
+	cmdChan        chan [2]string
 
 	inputInitial  bool
 	inputPreState bool
@@ -54,7 +60,7 @@ type Parser struct {
 	closed         bool
 }
 
-func (p *Parser) Initial() {
+func (p *Parser) initial() {
 	p.inputBuf = new(bytes.Buffer)
 	p.cmdBuf = new(bytes.Buffer)
 	p.outputBuf = new(bytes.Buffer)
@@ -65,6 +71,12 @@ func (p *Parser) Initial() {
 	p.cmdOutputParser = &CmdParser{}
 	p.cmdInputParser.Initial()
 	p.cmdOutputParser.Initial()
+
+	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.cmdChan = make(chan [2]string, 1024)
 }
 
 func (p *Parser) Parse() {
@@ -103,7 +115,8 @@ func (p *Parser) parseInputState(b []byte) []byte {
 		// 用户输入了Enter，开始结算命令
 		p.parseCmdInput()
 		if p.IsCommandForbidden() {
-			p.srvOutputChan <- []byte("\r\nCommand ls is forbidden")
+			fbdMsg := utils.WrapperWarn(fmt.Sprintf(i18n.T("Command `%s` is forbidden"), p.command))
+			p.srvOutputChan <- []byte("\r\n" + fbdMsg)
 			return []byte{utils.CharCleanLine, '\r'}
 		}
 	} else {

pkg/proxy/recorder.go

 package proxy
 
 import (
-	"cocogo/pkg/common"
-	"cocogo/pkg/logger"
+	"encoding/json"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 	"time"
 
+	"cocogo/pkg/common"
 	"cocogo/pkg/config"
+	"cocogo/pkg/logger"
 	"cocogo/pkg/model"
 )
 
@@ -25,13 +27,16 @@ func NewCommandRecorder(sess *SwitchSession) (recorder *CommandRecorder) {
 	storage := NewCommandStorage()
 	recorder = &CommandRecorder{Session: sess, queue: make(chan *model.Command, 10), storage: storage}
 	go recorder.record()
+	recorder.Start()
 	return recorder
 }
 
-func NewReplyRecord(sess *SwitchSession) *ReplyRecorder {
+func NewReplyRecord(sess *SwitchSession) (recorder *ReplyRecorder) {
 	storage := NewReplayStorage()
 	srvStorage := &ServerReplayStorage{}
-	return &ReplyRecorder{SessionID: sess.Id, storage: storage, backOffStorage: srvStorage}
+	recorder = &ReplyRecorder{SessionID: sess.Id, storage: storage, backOffStorage: srvStorage}
+	recorder.Start()
+	return recorder
 }
 
 func (c *CommandRecorder) Record(command [2]string) {
@@ -62,7 +67,11 @@ func (c *CommandRecorder) record() {
 	cmdList := make([]*model.Command, 0)
 	for {
 		select {
-		case p := <-c.queue:
+		case p, ok := <-c.queue:
+			if !ok {
+				logger.Debug("Session command recorder close: ", c.Session.Id)
+				return
+			}
 			cmdList = append(cmdList, p)
 			if len(cmdList) < 5 {
 				continue
@@ -91,13 +100,18 @@ type ReplyRecorder struct {
 	absGzFilePath string
 	target        string
 	file          *os.File
-	StartTime     time.Time
+	timeStartNano int64
 
 	storage        ReplayStorage
 	backOffStorage ReplayStorage
 }
 
 func (r *ReplyRecorder) Record(b []byte) {
+	if len(b) > 0 {
+		delta := float64(time.Now().UnixNano()-r.timeStartNano) / 1000 / 1000 / 1000
+		data, _ := json.Marshal(string(b))
+		_, _ = r.file.WriteString(fmt.Sprintf(`"%.3f":%s`, delta, data))
+	}
 }
 
 func (r *ReplyRecorder) Start() {
@@ -109,6 +123,7 @@ func (r *ReplyRecorder) Start() {
 	r.absFilePath = filepath.Join(replayDir, r.SessionID)
 	r.absGzFilePath = filepath.Join(replayDir, today, gzFileName)
 	r.target = strings.Join([]string{today, gzFileName}, "/")
+	r.timeStartNano = time.Now().UnixNano()
 
 	err := common.EnsureDirExist(replayDir)
 	if err != nil {
@@ -116,6 +131,7 @@ func (r *ReplyRecorder) Start() {
 		return
 	}
 
+	logger.Debug("Replay file path: ", r.absFilePath)
 	r.file, err = os.Create(r.absFilePath)
 	if err != nil {
 		logger.Errorf("Create file %s error: %s\n", r.absFilePath, err)
@@ -125,22 +141,25 @@ func (r *ReplyRecorder) Start() {
 
 func (r *ReplyRecorder) End() {
 	_ = r.file.Close()
+	go r.uploadReplay()
+}
+
+func (r *ReplyRecorder) uploadReplay() {
+	maxRetry := 3
 	if !common.FileExists(r.absFilePath) {
+		logger.Debug("Replay file not found, passed: ", r.absFilePath)
 		return
 	}
-	if stat, err := os.Stat(r.absGzFilePath); err == nil && stat.Size() == 0 {
+	if stat, err := os.Stat(r.absFilePath); err == nil && stat.Size() == 0 {
+		logger.Debug("Replay file is empty, removed: ", r.absFilePath)
 		_ = os.Remove(r.absFilePath)
 		return
 	}
-	go r.uploadReplay()
 	if !common.FileExists(r.absGzFilePath) {
+		logger.Debug("Compress replay file: ", r.absFilePath)
 		_ = common.GzipCompressFile(r.absFilePath, r.absGzFilePath)
 		_ = os.Remove(r.absFilePath)
 	}
-}
-
-func (r *ReplyRecorder) uploadReplay() {
-	maxRetry := 3
 
 	for i := 0; i <= maxRetry; i++ {
 		logger.Debug("Upload replay file: ", r.absGzFilePath)

pkg/proxy/switch.go

@@ -34,12 +34,11 @@ type SwitchSession struct {
 	replayRecorder *ReplyRecorder
 	parser         *Parser
 
-	cmdRecordChan chan [2]string
-	userConn      UserConnection
-	srvConn       ServerConnection
-	userChan      Transport
-	srvChan       Transport
-	cancelFunc    context.CancelFunc
+	userConn   UserConnection
+	srvConn    ServerConnection
+	userChan   Transport
+	srvChan    Transport
+	cancelFunc context.CancelFunc
 }
 
 func (s *SwitchSession) Initial() {
@@ -51,19 +50,10 @@ func (s *SwitchSession) Initial() {
 	s.RemoteAddr = s.userConn.RemoteAddr()
 	s.DateStart = time.Now()
 
-	s.cmdRecordChan = make(chan [2]string, 1024)
 	s.cmdRecorder = NewCommandRecorder(s)
 	s.replayRecorder = NewReplyRecord(s)
 
-	parser := &Parser{
-		userInputChan:  make(chan []byte, 1024),
-		userOutputChan: make(chan []byte, 1024),
-		srvInputChan:   make(chan []byte, 1024),
-		srvOutputChan:  make(chan []byte, 1024),
-		cmdChan:        s.cmdRecordChan,
-	}
-	parser.Initial()
-	s.parser = parser
+	s.parser = newParser()
 }
 
 func (s *SwitchSession) watchWindowChange(ctx context.Context, winCh <-chan ssh.Window) {

pkg/service/init.go

 package service
 
 import (
+	"encoding/json"
 	"os"
 	"path"
 	"path/filepath"
@@ -18,6 +19,8 @@ func Initial() {
 	keyPath := config.Conf.AccessKeyFile
 	client.BaseHost = config.Conf.CoreHost
 	authClient.BaseHost = config.Conf.CoreHost
+	client.SetHeader("X-JMS-ORG", "ROOT")
+	authClient.SetHeader("X-JMS-ORG", "ROOT")
 
 	if !path.IsAbs(config.Conf.AccessKeyFile) {
 		keyPath = filepath.Join(config.Conf.RootPath, keyPath)
@@ -26,6 +29,7 @@ func Initial() {
 	_ = ak.Load()
 	authClient.Auth = ak
 	validateAccessAuth()
+	MustLoadServerConfigOnce()
 	go KeepSyncConfigWithServer()
 }
 
@@ -52,22 +56,37 @@ func validateAccessAuth() {
 }
 
 func MustLoadServerConfigOnce() {
-
+	var data map[string]interface{}
+	err := authClient.Get(TerminalConfigURL, &data)
+	if err != nil {
+		logger.Error("Load config from server error: ", err)
+		return
+	}
+	data["TERMINAL_HOST_KEY"] = "Hidden"
+	msg, err := json.Marshal(data)
+	if err != nil {
+		logger.Error("Marsha server config error: %s", err)
+		return
+	}
+	logger.Debug("Load config from server: " + string(msg))
+	err = LoadConfigFromServer()
+	if err != nil {
+		logger.Error("Load config from server error: ", err)
+	}
 }
 
-func LoadConfigFromServer(conf *config.Config) (err error) {
+func LoadConfigFromServer() (err error) {
+	conf := config.Conf
 	conf.Mux.Lock()
 	defer conf.Mux.Unlock()
 	err = authClient.Get(TerminalConfigURL, conf)
-	if err != nil {
-		logger.Warn("Sync config with server error: ", err)
-	}
 	return err
 }
 
 func KeepSyncConfigWithServer() {
 	for {
-		err := LoadConfigFromServer(config.Conf)
+		logger.Debug("Sync config from server")
+		err := LoadConfigFromServer()
 		if err != nil {
 			logger.Warn("Sync config with server error: ", err)
 		}

pkg/service/terminal.go

@@ -84,6 +84,5 @@ func PushSessionReplay(sessionID, gZipFile string) {
 }
 
 func PushSessionCommand(commands []*model.Command) (err error) {
-	fmt.Println("Commands: ", commands)
 	return
 }

pkg/utils/util.go

@@ -41,6 +41,6 @@ func WrapperTitle(text string) string {
 }
 
 func WrapperWarn(text string) string {
-	text += "\r\n\r\n"
+	text += "\r\n"
 	return WrapperString(text, Red)
 }