fix bugs

pkg/handler/session.go

@@ -63,11 +63,12 @@ func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandl
 	wrapperSess := NewWrapperSession(sess)
 	term := utils.NewTerminal(wrapperSess, "Opt> ")
 	handler := &interactiveHandler{
-		sess:           wrapperSess,
-		user:           user,
-		term:           term,
-		mu:             new(sync.RWMutex),
-		finishedLoaded: make(chan struct{}),
+		sess:            wrapperSess,
+		user:            user,
+		term:            term,
+		mu:              new(sync.RWMutex),
+		nodeDataLoaded:  make(chan struct{}),
+		assetDataLoaded: make(chan struct{}),
 	}
 	handler.Initial()
 	return handler
@@ -85,7 +86,8 @@ type interactiveHandler struct {
 	searchResult     model.AssetList
 	nodes            model.NodeList
 	mu               *sync.RWMutex
-	finishedLoaded   chan struct{}
+	nodeDataLoaded   chan struct{}
+	assetDataLoaded  chan struct{}
 }
 
 func (h *interactiveHandler) Initial() {
@@ -98,18 +100,24 @@ func (h *interactiveHandler) Initial() {
 func (h *interactiveHandler) loadAssetsFromCache() {
 	if assets, ok := userAssetsCached.Get(h.user.ID); ok {
 		h.assets = assets
-		go h.firstLoadAssetAndNodes()
+		close(h.assetDataLoaded)
 	} else {
 		h.assets = make([]model.Asset, 0)
-		go h.firstLoadAssetAndNodes()
 	}
+	go h.firstLoadAssetAndNodes()
 }
 
 func (h *interactiveHandler) firstLoadAssetAndNodes() {
 	h.loadUserAssets("1")
 	h.loadUserAssetNodes("1")
-	close(h.finishedLoaded)
 	logger.Debug("first Load Asset And Nodes done")
+	close(h.nodeDataLoaded)
+	select {
+	case <-h.assetDataLoaded:
+		return
+	default:
+		close(h.assetDataLoaded)
+	}
 }
 
 func (h *interactiveHandler) displayBanner() {
@@ -166,6 +174,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
 			break
 		}
 		line = strings.TrimSpace(line)
+		<-h.assetDataLoaded
 		switch len(line) {
 		case 0, 1:
 			switch strings.ToLower(line) {
@@ -174,7 +183,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
 				h.displayAssets(h.assets)
 				h.mu.RUnlock()
 			case "g":
-				<-h.finishedLoaded
+				<-h.nodeDataLoaded
 				h.displayNodes(h.nodes)
 			case "h":
 				h.displayBanner()
@@ -197,6 +206,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
 				searchWord := strings.TrimSpace(strings.TrimPrefix(line, "g"))
 				if num, err := strconv.Atoi(searchWord); err == nil {
 					if num >= 0 {
+						<-h.nodeDataLoaded
 						assets := h.searchNodeAssets(num)
 						h.displayAssets(assets)
 						continue

pkg/proxy/recorder.go

@@ -12,6 +12,7 @@ import (
 	"github.com/jumpserver/koko/pkg/config"
 	"github.com/jumpserver/koko/pkg/logger"
 	"github.com/jumpserver/koko/pkg/model"
+	"github.com/jumpserver/koko/pkg/service"
 )
 
 func NewCommandRecorder(sid string) (recorder *CommandRecorder) {
@@ -80,6 +81,7 @@ func (c *CommandRecorder) record() {
 		err := c.storage.BulkSave(cmdList)
 		if err == nil {
 			cmdList = cmdList[:0]
+			maxRetry = 0
 			continue
 		}
 
@@ -168,7 +170,7 @@ func (r *ReplyRecorder) uploadReplay() {
 }
 
 func (r *ReplyRecorder) UploadGzipFile(maxRetry int) {
-	if r.storage == nil{
+	if r.storage == nil {
 		r.backOffStorage = defaultReplayStorage
 		r.storage = NewReplayStorage()
 	}
@@ -177,6 +179,7 @@ func (r *ReplyRecorder) UploadGzipFile(maxRetry int) {
 		err := r.storage.Upload(r.AbsGzFilePath, r.Target)
 		if err == nil {
 			_ = os.Remove(r.AbsGzFilePath)
+			service.FinishReply(r.sessionID)
 			break
 		}
 		// 如果还是失败，使用备用storage再传一次

pkg/proxy/sessmanager.go

@@ -93,6 +93,5 @@ func postSession(s *SwitchSession) bool {
 func finishSession(s *SwitchSession) {
 	data := s.MapData()
 	service.FinishSession(data)
-	service.FinishReply(s.ID)
 	logger.Debugf("Finish session: %s", s.ID)
 }