waveterm/backup
0
0
Fork 0
mirror of https://github.com/wavetermdev/backup.git synced 2026-08-05 13:57:07 -07:00
Code Issues Packages Projects Releases Wiki Activity

queue entire functions when cmdline has not been committed to DB yet. new inputpacket handling (for winsize)

Browse Source
This commit is contained in:
sawka
2022-09-05 16:31:22 -07:00
parent a20ee78e6d
commit 54d2f5d761
5 changed files with 77 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/main-server.go
+13
View File
@@ -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)
pkg/remote/remote.go
+21 -16
View File
@@ -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 {
pkg/remote/updatequeue.go
+22 -23
View File
@@ -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)
}
pkg/scpacket/scpacket.go
+1 -13
View File
@@ -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
}
pkg/scws/scws.go
+20 -14
View File
@@ -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
}