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

fix bugs

parent 9de03f5e
...@@ -63,11 +63,12 @@ func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandl ...@@ -63,11 +63,12 @@ func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandl
wrapperSess := NewWrapperSession(sess) wrapperSess := NewWrapperSession(sess)
term := utils.NewTerminal(wrapperSess, "Opt> ") term := utils.NewTerminal(wrapperSess, "Opt> ")
handler := &interactiveHandler{ handler := &interactiveHandler{
sess: wrapperSess, sess: wrapperSess,
user: user, user: user,
term: term, term: term,
mu: new(sync.RWMutex), mu: new(sync.RWMutex),
finishedLoaded: make(chan struct{}), finishedLoaded: make(chan struct{}),
assetDataLoaded: make(chan struct{}),
} }
handler.Initial() handler.Initial()
return handler return handler
...@@ -85,7 +86,8 @@ type interactiveHandler struct { ...@@ -85,7 +86,8 @@ type interactiveHandler struct {
searchResult model.AssetList searchResult model.AssetList
nodes model.NodeList nodes model.NodeList
mu *sync.RWMutex mu *sync.RWMutex
finishedLoaded chan struct{} nodeDataLoaded chan struct{}
assetDataLoaded chan struct{}
} }
func (h *interactiveHandler) Initial() { func (h *interactiveHandler) Initial() {
...@@ -98,18 +100,24 @@ func (h *interactiveHandler) Initial() { ...@@ -98,18 +100,24 @@ func (h *interactiveHandler) Initial() {
func (h *interactiveHandler) loadAssetsFromCache() { func (h *interactiveHandler) loadAssetsFromCache() {
if assets, ok := userAssetsCached.Get(h.user.ID); ok { if assets, ok := userAssetsCached.Get(h.user.ID); ok {
h.assets = assets h.assets = assets
go h.firstLoadAssetAndNodes() close(h.assetDataLoaded)
} else { } else {
h.assets = make([]model.Asset, 0) h.assets = make([]model.Asset, 0)
h.firstLoadAssetAndNodes()
} }
go h.firstLoadAssetAndNodes()
} }
func (h *interactiveHandler) firstLoadAssetAndNodes() { func (h *interactiveHandler) firstLoadAssetAndNodes() {
h.loadUserAssets("1") h.loadUserAssets("1")
h.loadUserAssetNodes("1") h.loadUserAssetNodes("1")
close(h.finishedLoaded)
logger.Debug("first Load Asset And Nodes done") logger.Debug("first Load Asset And Nodes done")
close(h.nodeDataLoaded)
select {
case <-h.assetDataLoaded:
return
default:
close(h.assetDataLoaded)
}
} }
func (h *interactiveHandler) displayBanner() { func (h *interactiveHandler) displayBanner() {
...@@ -166,6 +174,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) { ...@@ -166,6 +174,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
break break
} }
line = strings.TrimSpace(line) line = strings.TrimSpace(line)
<-h.assetDataLoaded
switch len(line) { switch len(line) {
case 0, 1: case 0, 1:
switch strings.ToLower(line) { switch strings.ToLower(line) {
...@@ -174,7 +183,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) { ...@@ -174,7 +183,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
h.displayAssets(h.assets) h.displayAssets(h.assets)
h.mu.RUnlock() h.mu.RUnlock()
case "g": case "g":
<-h.finishedLoaded <-h.nodeDataLoaded
h.displayNodes(h.nodes) h.displayNodes(h.nodes)
case "h": case "h":
h.displayBanner() h.displayBanner()
...@@ -197,6 +206,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) { ...@@ -197,6 +206,7 @@ func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
searchWord := strings.TrimSpace(strings.TrimPrefix(line, "g")) searchWord := strings.TrimSpace(strings.TrimPrefix(line, "g"))
if num, err := strconv.Atoi(searchWord); err == nil { if num, err := strconv.Atoi(searchWord); err == nil {
if num >= 0 { if num >= 0 {
<-h.nodeDataLoaded
assets := h.searchNodeAssets(num) assets := h.searchNodeAssets(num)
h.displayAssets(assets) h.displayAssets(assets)
continue continue
......
...@@ -12,6 +12,7 @@ import ( ...@@ -12,6 +12,7 @@ import (
"github.com/jumpserver/koko/pkg/config" "github.com/jumpserver/koko/pkg/config"
"github.com/jumpserver/koko/pkg/logger" "github.com/jumpserver/koko/pkg/logger"
"github.com/jumpserver/koko/pkg/model" "github.com/jumpserver/koko/pkg/model"
"github.com/jumpserver/koko/pkg/service"
) )
func NewCommandRecorder(sid string) (recorder *CommandRecorder) { func NewCommandRecorder(sid string) (recorder *CommandRecorder) {
...@@ -80,6 +81,7 @@ func (c *CommandRecorder) record() { ...@@ -80,6 +81,7 @@ func (c *CommandRecorder) record() {
err := c.storage.BulkSave(cmdList) err := c.storage.BulkSave(cmdList)
if err == nil { if err == nil {
cmdList = cmdList[:0] cmdList = cmdList[:0]
maxRetry = 0
continue continue
} }
...@@ -168,7 +170,7 @@ func (r *ReplyRecorder) uploadReplay() { ...@@ -168,7 +170,7 @@ func (r *ReplyRecorder) uploadReplay() {
} }
func (r *ReplyRecorder) UploadGzipFile(maxRetry int) { func (r *ReplyRecorder) UploadGzipFile(maxRetry int) {
if r.storage == nil{ if r.storage == nil {
r.backOffStorage = defaultReplayStorage r.backOffStorage = defaultReplayStorage
r.storage = NewReplayStorage() r.storage = NewReplayStorage()
} }
...@@ -177,6 +179,7 @@ func (r *ReplyRecorder) UploadGzipFile(maxRetry int) { ...@@ -177,6 +179,7 @@ func (r *ReplyRecorder) UploadGzipFile(maxRetry int) {
err := r.storage.Upload(r.AbsGzFilePath, r.Target) err := r.storage.Upload(r.AbsGzFilePath, r.Target)
if err == nil { if err == nil {
_ = os.Remove(r.AbsGzFilePath) _ = os.Remove(r.AbsGzFilePath)
service.FinishReply(r.sessionID)
break break
} }
// 如果还是失败,使用备用storage再传一次 // 如果还是失败,使用备用storage再传一次
......
...@@ -93,6 +93,5 @@ func postSession(s *SwitchSession) bool { ...@@ -93,6 +93,5 @@ func postSession(s *SwitchSession) bool {
func finishSession(s *SwitchSession) { func finishSession(s *SwitchSession) {
data := s.MapData() data := s.MapData()
service.FinishSession(data) service.FinishSession(data)
service.FinishReply(s.ID)
logger.Debugf("Finish session: %s", s.ID) logger.Debugf("Finish session: %s", s.ID)
} }
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