From 29dbf5dde9b94fb0725dc95823e8b4bd7bb408d2 Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Thu, 27 Oct 2022 22:52:38 -0700
Subject: [PATCH] runcommand now sets the state at the time of execution

---
 pkg/cmdrunner/cmdrunner.go |  3 +--
 pkg/remote/remote.go       | 54 +++++++++++++++++++++++++++++++++++---
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/pkg/cmdrunner/cmdrunner.go b/pkg/cmdrunner/cmdrunner.go
index acd7a8ffe..de84d097b 100644
--- a/pkg/cmdrunner/cmdrunner.go
+++ b/pkg/cmdrunner/cmdrunner.go
@@ -245,8 +245,7 @@ func RunCommand(ctx context.Context, pk *scpacket.FeCommandPacketType) (sstore.U
 	runPacket := packet.MakeRunPacket()
 	runPacket.ReqId = uuid.New().String()
 	runPacket.CK = base.MakeCommandKey(ids.SessionId, scbase.GenSCUUID())
-	runPacket.State = ids.Remote.RemoteState
-	runPacket.StateComplete = true
+	// runPacket.State is set in remote.RunCommand()
 	runPacket.UsePty = true
 	runPacket.TermOpts = getUITermOpts(pk.UIContext)
 	runPacket.Command = strings.TrimSpace(cmdStr)
diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go
index cf89e9bdc..e197076ab 100644
--- a/pkg/remote/remote.go
+++ b/pkg/remote/remote.go
@@ -901,6 +901,9 @@ func (msh *MShellProc) ReInit(ctx context.Context) (*packet.InitPacketType, erro
 	if !ok {
 		return nil, fmt.Errorf("invalid reinit response (not an initpacket): %T", resp)
 	}
+	if initPk.State == nil {
+		return nil, fmt.Errorf("invalid reinit response initpk does not contain remote state")
+	}
 	msh.WithLock(func() {
 		msh.Remote.InitPk = initPk
 	})
@@ -960,6 +963,7 @@ func (msh *MShellProc) Launch() {
 		go msh.NotifyRemoteUpdate()
 	})
 	cproc, initPk, err := shexec.MakeClientProc(makeClientCtx, ecmd)
+	// TODO check if initPk.State is not nil
 	var mshellVersion string
 	msh.WithLock(func() {
 		msh.MakeClientCancelFn = nil
@@ -1135,8 +1139,8 @@ func RunCommand(ctx context.Context, sessionId string, windowId string, remotePt
 	if !msh.IsConnected() {
 		return nil, nil, fmt.Errorf("remote '%s' is not connected", remotePtr.RemoteId)
 	}
-	if runPacket.State == nil {
-		return nil, nil, fmt.Errorf("no remote state passed to RunCommand")
+	if runPacket.State != nil {
+		return nil, nil, fmt.Errorf("runPacket.State should not be set, it is set in RunCommand")
 	}
 	var newPSC *base.CommandKey
 	if runPacket.ReturnState {
@@ -1162,8 +1166,21 @@ func RunCommand(ctx context.Context, sessionId string, windowId string, remotePt
 			}
 		}
 	}()
+	// get current remote-instance state
+	currentState, err := sstore.GetRemoteState(ctx, sessionId, windowId, remotePtr)
+	if err != nil {
+		return nil, nil, fmt.Errorf("cannot get current remote state: %w", err)
+	}
+	if currentState == nil {
+		currentState = msh.ServerProc.InitPk.State
+	}
+	if currentState == nil {
+		return nil, nil, fmt.Errorf("cannot run command, no valid remote state")
+	}
+	runPacket.State = currentState
+	runPacket.StateComplete = true
 	msh.ServerProc.Output.RegisterRpc(runPacket.ReqId)
-	err := shexec.SendRunPacketAndRunData(ctx, msh.ServerProc.Input, runPacket)
+	err = shexec.SendRunPacketAndRunData(ctx, msh.ServerProc.Input, runPacket)
 	if err != nil {
 		return nil, nil, fmt.Errorf("sending run packet to remote: %w", err)
 	}
@@ -1208,6 +1225,37 @@ func RunCommand(ctx context.Context, sessionId string, windowId string, remotePt
 	return cmd, func() { removeCmdWait(runPacket.CK) }, nil
 }
 
+func (msh *MShellProc) AddWaitingCmd(rct RunCmdType) {
+	msh.Lock.Lock()
+	defer msh.Lock.Unlock()
+	msh.WaitingCmds = append(msh.WaitingCmds, rct)
+}
+
+func (msh *MShellProc) reExecSingle(rct RunCmdType) {
+	// TODO fixme
+	ctx, cancelFn := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancelFn()
+	_, callback, _ := RunCommand(ctx, rct.SessionId, rct.WindowId, rct.RemotePtr, rct.RunPacket)
+	if callback != nil {
+		defer callback()
+	}
+}
+
+func (msh *MShellProc) ReExecWaitingCmds() {
+	msh.Lock.Lock()
+	defer msh.Lock.Unlock()
+	for len(msh.WaitingCmds) > 0 {
+		rct := msh.WaitingCmds[0]
+		go msh.reExecSingle(rct)
+		if rct.RunPacket.ReturnState {
+			break
+		}
+	}
+	if len(msh.WaitingCmds) == 0 {
+		msh.WaitingCmds = nil
+	}
+}
+
 func (msh *MShellProc) AddRunningCmd(rct RunCmdType) {
 	msh.Lock.Lock()
 	defer msh.Lock.Unlock()