From 54d2f5d76183b91184c1c3474b4a9a6263a45463 Mon Sep 17 00:00:00 2001 From: sawka Date: Mon, 5 Sep 2022 16:31:22 -0700 Subject: [PATCH] queue entire functions when cmdline has not been committed to DB yet. new inputpacket handling (for winsize) --- cmd/main-server.go | 13 +++++++++++ pkg/remote/remote.go | 37 ++++++++++++++++++-------------- pkg/remote/updatequeue.go | 45 +++++++++++++++++++-------------------- pkg/scpacket/scpacket.go | 14 +----------- pkg/scws/scws.go | 34 +++++++++++++++++------------ 5 files changed, 77 insertions(+), 66 deletions(-) diff --git a/cmd/main-server.go b/cmd/main-server.go index e7b91b3f..5add214a 100644 --- a/cmd/main-server.go +++ b/cmd/main-server.go @@ -261,7 +261,20 @@ func runWebSocketServer() { } } +func test() error { + return nil +} + func main() { + if len(os.Args) >= 2 && os.Args[1] == "--test" { + fmt.Printf("running test fn\n") + err := test() + if err != nil { + fmt.Printf("[error] %v\n", err) + } + return + } + scLock, err := scbase.AcquireSCLock() if err != nil || scLock == nil { fmt.Printf("[error] cannot acquire sh2 lock: %v\n", err) diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go index f6eeedab..ffebfa83 100644 --- a/pkg/remote/remote.go +++ b/pkg/remote/remote.go @@ -52,7 +52,7 @@ type Store struct { Lock *sync.Mutex Map map[string]*MShellProc // key=remoteid Log *CircleLog - CmdWaitMap map[base.CommandKey][]sstore.UpdatePacket + CmdWaitMap map[base.CommandKey][]func() } type MShellProc struct { @@ -127,7 +127,7 @@ func LoadRemotes(ctx context.Context) error { Lock: &sync.Mutex{}, Map: make(map[string]*MShellProc), Log: MakeCircleLog(100), - CmdWaitMap: make(map[base.CommandKey][]sstore.UpdatePacket), + CmdWaitMap: make(map[base.CommandKey][]func()), } allRemotes, err := sstore.GetAllRemotes(ctx) if err != nil { @@ -530,17 +530,13 @@ func (msh *MShellProc) IsCmdRunning(ck base.CommandKey) bool { return false } -func (msh *MShellProc) SendInput(pk *packet.InputPacketType) error { +func (msh *MShellProc) SendInput(dataPk *packet.DataPacketType) error { if !msh.IsConnected() { return fmt.Errorf("remote is not connected, cannot send input") } - if !msh.IsCmdRunning(pk.CK) { + if !msh.IsCmdRunning(dataPk.CK) { return fmt.Errorf("cannot send input, cmd is not running") } - dataPk := packet.MakeDataPacket() - dataPk.CK = pk.CK - dataPk.FdNum = 0 // stdin - dataPk.Data64 = pk.InputData64 return msh.ServerProc.Input.SendPacket(dataPk) } @@ -674,11 +670,7 @@ func (msh *MShellProc) handleCmdDonePacket(donePk *packet.CmdDonePacketType) { return } if update != nil { - // TODO fix timing issue (this update gets to the FE before run-command returns for short lived commands) - go func() { - time.Sleep(10 * time.Millisecond) - sendCmdUpdate(donePk.CK, update) - }() + sstore.MainBus.SendUpdate(donePk.CK.GetSessionId(), update) } return } @@ -711,7 +703,7 @@ func (msh *MShellProc) handleDataPacket(dataPk *packet.DataPacketType, dataPosMa } dataPosMap[dataPk.CK] += int64(len(realData)) if update != nil { - sendCmdUpdate(dataPk.CK, update) + sstore.MainBus.SendUpdate(dataPk.CK.GetSessionId(), update) } } if ack != nil { @@ -720,6 +712,18 @@ func (msh *MShellProc) handleDataPacket(dataPk *packet.DataPacketType, dataPosMa // fmt.Printf("data %s fd=%d len=%d eof=%v err=%v\n", dataPk.CK, dataPk.FdNum, len(realData), dataPk.Eof, dataPk.Error) } +func (msh *MShellProc) makeHandleDataPacketClosure(dataPk *packet.DataPacketType, dataPosMap map[base.CommandKey]int64) func() { + return func() { + msh.handleDataPacket(dataPk, dataPosMap) + } +} + +func (msh *MShellProc) makeHandleCmdDonePacketClosure(donePk *packet.CmdDonePacketType) func() { + return func() { + msh.handleCmdDonePacket(donePk) + } +} + func (msh *MShellProc) ProcessPackets() { defer msh.WithLock(func() { if msh.Status == StatusConnected { @@ -736,7 +740,7 @@ func (msh *MShellProc) ProcessPackets() { for pk := range msh.ServerProc.Output.MainCh { if pk.GetType() == packet.DataPacketStr { dataPk := pk.(*packet.DataPacketType) - msh.handleDataPacket(dataPk, dataPosMap) + runCmdUpdateFn(dataPk.CK, msh.makeHandleDataPacketClosure(dataPk, dataPosMap)) continue } if pk.GetType() == packet.DataAckPacketStr { @@ -750,7 +754,8 @@ func (msh *MShellProc) ProcessPackets() { continue } if pk.GetType() == packet.CmdDonePacketStr { - msh.handleCmdDonePacket(pk.(*packet.CmdDonePacketType)) + donePk := pk.(*packet.CmdDonePacketType) + runCmdUpdateFn(donePk.CK, msh.makeHandleCmdDonePacketClosure(donePk)) continue } if pk.GetType() == packet.CmdErrorPacketStr { diff --git a/pkg/remote/updatequeue.go b/pkg/remote/updatequeue.go index 105bb61f..bab359a1 100644 --- a/pkg/remote/updatequeue.go +++ b/pkg/remote/updatequeue.go @@ -2,65 +2,64 @@ package remote import ( "github.com/scripthaus-dev/mshell/pkg/base" - "github.com/scripthaus-dev/sh2-server/pkg/sstore" ) -func pushCmdWaitIfRequired(ck base.CommandKey, update sstore.UpdatePacket) bool { +func pushCmdWaitIfRequired(ck base.CommandKey, fn func()) bool { GlobalStore.Lock.Lock() defer GlobalStore.Lock.Unlock() - updates, ok := GlobalStore.CmdWaitMap[ck] + fns, ok := GlobalStore.CmdWaitMap[ck] if !ok { return false } - updates = append(updates, update) - GlobalStore.CmdWaitMap[ck] = updates + fns = append(fns, fn) + GlobalStore.CmdWaitMap[ck] = fns return true } -func sendCmdUpdate(ck base.CommandKey, update sstore.UpdatePacket) { - pushed := pushCmdWaitIfRequired(ck, update) +func runCmdUpdateFn(ck base.CommandKey, fn func()) { + pushed := pushCmdWaitIfRequired(ck, fn) if pushed { return } - sstore.MainBus.SendUpdate(ck.GetSessionId(), update) + fn() } -func runCmdWaitUpdates(ck base.CommandKey) { +func runCmdWaitFns(ck base.CommandKey) { for { - update := removeFirstCmdWaitUpdate(ck) - if update == nil { + fn := removeFirstCmdWaitFn(ck) + if fn == nil { break } - sstore.MainBus.SendUpdate(ck.GetSessionId(), update) + fn() } } -func removeFirstCmdWaitUpdate(ck base.CommandKey) sstore.UpdatePacket { +func removeFirstCmdWaitFn(ck base.CommandKey) func() { GlobalStore.Lock.Lock() defer GlobalStore.Lock.Unlock() - updates := GlobalStore.CmdWaitMap[ck] - if len(updates) == 0 { + fns := GlobalStore.CmdWaitMap[ck] + if len(fns) == 0 { delete(GlobalStore.CmdWaitMap, ck) return nil } - if len(updates) == 1 { + if len(fns) == 1 { delete(GlobalStore.CmdWaitMap, ck) - return updates[0] + return fns[0] } - update := updates[0] - GlobalStore.CmdWaitMap[ck] = updates[1:] - return update + fn := fns[0] + GlobalStore.CmdWaitMap[ck] = fns[1:] + return fn } func removeCmdWait(ck base.CommandKey) { GlobalStore.Lock.Lock() defer GlobalStore.Lock.Unlock() - updates := GlobalStore.CmdWaitMap[ck] - if len(updates) == 0 { + fns := GlobalStore.CmdWaitMap[ck] + if len(fns) == 0 { delete(GlobalStore.CmdWaitMap, ck) return } - go runCmdWaitUpdates(ck) + go runCmdWaitFns(ck) } diff --git a/pkg/scpacket/scpacket.go b/pkg/scpacket/scpacket.go index c496ed81..b6d2ff95 100644 --- a/pkg/scpacket/scpacket.go +++ b/pkg/scpacket/scpacket.go @@ -36,8 +36,7 @@ type FeInputPacketType struct { Remote sstore.RemotePtrType `json:"remote"` InputData64 string `json:"inputdata"` SigNum int `json:"signum,omitempty"` - WinSizeRows int `json:"winsizerows"` - WinSizeCols int `json:"winsizecols"` + WinSize *packet.WinSize `json:"winsize,omitempty"` } type WatchScreenPacketType struct { @@ -69,17 +68,6 @@ func MakeFeInputPacket() *FeInputPacketType { return &FeInputPacketType{Type: FeInputPacketStr} } -func (p *FeInputPacketType) ConvertToInputPacket() *packet.InputPacketType { - rtn := packet.MakeInputPacket() - rtn.CK = p.CK - rtn.RemoteId = p.Remote.RemoteId - rtn.InputData64 = p.InputData64 - rtn.SigNum = p.SigNum - rtn.WinSizeRows = p.WinSizeRows - rtn.WinSizeCols = p.WinSizeCols - return rtn -} - func (*WatchScreenPacketType) GetType() string { return WatchScreenPacketStr } diff --git a/pkg/scws/scws.go b/pkg/scws/scws.go index 917457b1..fb642696 100644 --- a/pkg/scws/scws.go +++ b/pkg/scws/scws.go @@ -188,9 +188,8 @@ func (ws *WSState) RunWSRead() { fmt.Printf("[error] invalid input packet, remoteid is not set\n") continue } - inputPk := feInputPk.ConvertToInputPacket() go func() { - err = sendCmdInput(inputPk) + err = sendCmdInput(feInputPk) if err != nil { fmt.Printf("[error] sending command input: %v\n", err) } @@ -210,24 +209,31 @@ func (ws *WSState) RunWSRead() { } } -func sendCmdInput(pk *packet.InputPacketType) error { +func sendCmdInput(pk *scpacket.FeInputPacketType) error { err := pk.CK.Validate("input packet") if err != nil { return err } - if pk.RemoteId == "" { + if pk.Remote.RemoteId == "" { return fmt.Errorf("input must set remoteid") } - if len(pk.InputData64) == 0 && pk.SigNum == 0 { - return fmt.Errorf("empty input packet") + if len(pk.InputData64) > 0 { + inputLen := packet.B64DecodedLen(pk.InputData64) + if inputLen > MaxInputDataSize { + return fmt.Errorf("input data size too large, len=%d (max=%d)", inputLen, MaxInputDataSize) + } + msh := remote.GetRemoteById(pk.Remote.RemoteId) + if msh == nil { + return fmt.Errorf("remote %d not found", pk.Remote.RemoteId) + } + dataPk := packet.MakeDataPacket() + dataPk.CK = pk.CK + dataPk.FdNum = 0 // stdin + dataPk.Data64 = pk.InputData64 + return msh.SendInput(dataPk) } - inputLen := packet.B64DecodedLen(pk.InputData64) - if inputLen > MaxInputDataSize { - return fmt.Errorf("input data size too large, len=%d (max=%d)", inputLen, MaxInputDataSize) + if pk.SigNum != 0 || pk.WinSize != nil { + return fmt.Errorf("signum / winsize not supported") } - msh := remote.GetRemoteById(pk.RemoteId) - if msh == nil { - return fmt.Errorf("cannot connect to remote") - } - return msh.SendInput(pk) + return nil }