2023-10-17 12:31:13 +08:00
// Copyright 2023, Command Line Inc.
// SPDX-License-Identifier: Apache-2.0
2022-06-28 15:04:08 -07:00
package server
import (
2022-07-06 18:59:46 -07:00
"context"
2023-08-31 22:03:38 -07:00
"errors"
2022-06-28 15:04:08 -07:00
"fmt"
2023-08-31 22:03:38 -07:00
"io"
2023-09-05 21:21:34 -07:00
"io/fs"
2022-06-28 15:04:08 -07:00
"os"
2022-08-09 14:23:59 -07:00
"os/exec"
2023-08-31 22:03:38 -07:00
"path/filepath"
2022-08-11 10:21:11 -07:00
"sort"
2022-08-09 14:23:59 -07:00
"strings"
2022-06-28 15:04:08 -07:00
"sync"
2022-11-27 13:47:18 -08:00
"time"
2022-06-28 15:04:08 -07:00
2022-08-09 14:23:59 -07:00
"github.com/alessio/shellescape"
2023-10-16 13:25:53 -07:00
"github.com/wavetermdev/waveterm/waveshell/pkg/base"
"github.com/wavetermdev/waveterm/waveshell/pkg/packet"
2024-01-16 16:11:04 -08:00
"github.com/wavetermdev/waveterm/waveshell/pkg/shellapi"
2023-10-16 13:25:53 -07:00
"github.com/wavetermdev/waveterm/waveshell/pkg/shexec"
2024-01-16 16:11:04 -08:00
"github.com/wavetermdev/waveterm/waveshell/pkg/utilfn"
2022-06-28 15:04:08 -07:00
)
2023-08-31 22:03:38 -07:00
const MaxFileDataPacketSize = 16 * 1024
const WriteFileContextTimeout = 30 * time . Second
const cleanLoopTime = 5 * time . Second
const MaxWriteFileContextData = 100
2024-01-16 16:11:04 -08:00
type shellStateMapKey struct {
ShellType string
Hash string
}
type ShellStateMap struct {
Lock * sync . Mutex
StateMap map [ shellStateMapKey ] * packet . ShellState // shelltype+hash -> state
CurrentStateMap map [ string ] string // shelltype -> hash
}
2022-07-06 17:16:45 -07:00
// TODO create unblockable packet-sender (backed by an array) for clientproc
2022-06-28 15:04:08 -07:00
type MServer struct {
2023-08-31 22:03:38 -07:00
Lock * sync . Mutex
MainInput * packet . PacketParser
Sender * packet . PacketSender
ClientMap map [ base . CommandKey ] * shexec . ClientProc
Debug bool
2024-01-16 16:11:04 -08:00
StateMap * ShellStateMap
WriteErrorCh chan bool // closed if there is a I/O write error
2023-08-31 22:03:38 -07:00
WriteErrorChOnce * sync . Once
WriteFileContextMap map [ string ] * WriteFileContext
Done bool
}
type WriteFileContext struct {
CVar * sync . Cond
Data [] * packet . FileDataPacketType
LastActive time . Time
Err error
Done bool
2022-06-28 15:04:08 -07:00
}
func ( m * MServer ) Close () {
m . Sender . Close ()
m . Sender . WaitForDone ()
2023-08-31 22:03:38 -07:00
m . Lock . Lock ()
defer m . Lock . Unlock ()
m . Done = true
}
func ( m * MServer ) checkDone () bool {
m . Lock . Lock ()
defer m . Lock . Unlock ()
return m . Done
}
func ( m * MServer ) getWriteFileContext ( reqId string ) * WriteFileContext {
m . Lock . Lock ()
defer m . Lock . Unlock ()
wfc := m . WriteFileContextMap [ reqId ]
if wfc == nil {
wfc = & WriteFileContext {
CVar : sync . NewCond ( & sync . Mutex {}),
LastActive : time . Now (),
}
m . WriteFileContextMap [ reqId ] = wfc
}
return wfc
}
func ( m * MServer ) addFileDataPacket ( pk * packet . FileDataPacketType ) {
m . Lock . Lock ()
wfc := m . WriteFileContextMap [ pk . RespId ]
m . Lock . Unlock ()
if wfc == nil {
return
}
wfc . CVar . L . Lock ()
defer wfc . CVar . L . Unlock ()
if wfc . Done || wfc . Err != nil {
return
}
if len ( wfc . Data ) > MaxWriteFileContextData {
wfc . Err = errors . New ( "write-file buffer length exceeded" )
wfc . Data = nil
wfc . CVar . Broadcast ()
return
}
wfc . LastActive = time . Now ()
wfc . Data = append ( wfc . Data , pk )
wfc . CVar . Signal ()
}
func ( wfc * WriteFileContext ) setDone () {
wfc . CVar . L . Lock ()
defer wfc . CVar . L . Unlock ()
wfc . Done = true
wfc . Data = nil
wfc . CVar . Broadcast ()
}
func ( m * MServer ) cleanWriteFileContexts () {
now := time . Now ()
var staleWfcs [] * WriteFileContext
m . Lock . Lock ()
for reqId , wfc := range m . WriteFileContextMap {
if now . Sub ( wfc . LastActive ) > WriteFileContextTimeout {
staleWfcs = append ( staleWfcs , wfc )
delete ( m . WriteFileContextMap , reqId )
}
}
m . Lock . Unlock ()
// we do this outside of m.Lock just in case there is some lock contention (end of WriteFile could theoretically be slow)
for _ , wfc := range staleWfcs {
wfc . setDone ()
}
2022-06-28 15:04:08 -07:00
}
2022-06-28 19:01:33 -07:00
func ( m * MServer ) ProcessCommandPacket ( pk packet . CommandPacketType ) {
ck := pk . GetCK ()
if ck == "" {
2022-11-27 13:47:18 -08:00
m . Sender . SendMessageFmt ( "received '%s' packet without ck" , pk . GetType ())
2022-06-28 19:01:33 -07:00
return
}
m . Lock . Lock ()
2022-07-06 17:16:45 -07:00
cproc := m . ClientMap [ ck ]
2022-06-28 19:01:33 -07:00
m . Lock . Unlock ()
2022-07-06 17:16:45 -07:00
if cproc == nil {
2022-07-07 13:25:42 -07:00
m . Sender . SendCmdError ( ck , fmt . Errorf ( "no client proc for ck '%s', pk=%s" , ck , packet . AsString ( pk )))
2022-06-28 19:01:33 -07:00
return
}
2022-07-06 17:16:45 -07:00
cproc . Input . SendPacket ( pk )
return
2022-06-29 14:29:38 -07:00
}
2022-08-11 10:21:11 -07:00
func runSingleCompGen ( cwd string , compType string , prefix string ) ([] string , bool , error ) {
2024-01-16 16:11:04 -08:00
sapi , err := shellapi . MakeShellApi ( packet . ShellType_bash )
if err != nil {
return nil , false , err
}
2022-08-11 10:21:11 -07:00
if ! packet . IsValidCompGenType ( compType ) {
return nil , false , fmt . Errorf ( "invalid compgen type '%s'" , compType )
2022-08-09 14:23:59 -07:00
}
2022-08-11 10:21:11 -07:00
compGenCmdStr := fmt . Sprintf ( "cd %s; compgen -A %s -- %s | sort | uniq | head -n %d" , shellescape . Quote ( cwd ), shellescape . Quote ( compType ), shellescape . Quote ( prefix ), packet . MaxCompGenValues + 1 )
2024-01-16 16:11:04 -08:00
ecmd := exec . Command ( sapi . GetLocalShellPath (), "-c" , compGenCmdStr )
2022-08-09 14:23:59 -07:00
outputBytes , err := ecmd . Output ()
if err != nil {
2022-08-11 10:21:11 -07:00
return nil , false , fmt . Errorf ( "compgen error: %w" , err )
2022-08-09 14:23:59 -07:00
}
outputStr := string ( outputBytes )
parts := strings . Split ( outputStr , "\n" )
if len ( parts ) > 0 && parts [ len ( parts ) - 1 ] == "" {
parts = parts [ 0 : len ( parts ) - 1 ]
}
2022-08-10 16:07:41 -07:00
hasMore := false
if len ( parts ) > packet . MaxCompGenValues {
hasMore = true
parts = parts [ 0 : packet . MaxCompGenValues ]
}
2022-08-11 10:21:11 -07:00
return parts , hasMore , nil
}
func appendSlashes ( comps [] string ) {
for idx , comp := range comps {
comps [ idx ] = comp + "/"
}
}
func strArrToMap ( strs [] string ) map [ string ] bool {
rtn := make ( map [ string ] bool )
for _ , s := range strs {
rtn [ s ] = true
}
return rtn
}
2022-11-23 10:51:51 -08:00
func ( m * MServer ) runMixedCompGen ( compPk * packet . CompGenPacketType ) {
2022-08-11 10:21:11 -07:00
// get directories and files, unique them and put slashes on directories for completion
reqId := compPk . GetReqId ()
compDirs , hasMoreDirs , err := runSingleCompGen ( compPk . Cwd , "directory" , compPk . Prefix )
if err != nil {
m . Sender . SendErrorResponse ( reqId , err )
return
}
2022-11-23 10:51:51 -08:00
compFiles , hasMoreFiles , err := runSingleCompGen ( compPk . Cwd , compPk . CompType , compPk . Prefix )
2022-08-11 10:21:11 -07:00
if err != nil {
m . Sender . SendErrorResponse ( reqId , err )
return
}
dirMap := strArrToMap ( compDirs )
// seed comps with dirs (but append slashes)
comps := compDirs
appendSlashes ( comps )
// add files that are not directories (look up in dirMap)
for _ , file := range compFiles {
if dirMap [ file ] {
continue
}
comps = append ( comps , file )
}
sort . Strings ( comps ) // resort
m . Sender . SendResponse ( reqId , map [ string ] interface {}{ "comps" : comps , "hasmore" : ( hasMoreFiles || hasMoreDirs )})
return
}
func ( m * MServer ) runCompGen ( compPk * packet . CompGenPacketType ) {
reqId := compPk . GetReqId ()
2022-11-23 10:51:51 -08:00
if compPk . CompType == "file" || compPk . CompType == "command" {
m . runMixedCompGen ( compPk )
2022-08-11 10:21:11 -07:00
return
}
comps , hasMore , err := runSingleCompGen ( compPk . Cwd , compPk . CompType , compPk . Prefix )
if err != nil {
m . Sender . SendErrorResponse ( reqId , err )
return
}
if compPk . CompType == "directory" {
appendSlashes ( comps )
}
m . Sender . SendResponse ( reqId , map [ string ] interface {}{ "comps" : comps , "hasmore" : hasMore })
2022-08-09 14:23:59 -07:00
return
}
2024-01-16 16:11:04 -08:00
func ( m * MServer ) reinit ( reqId string , shellType string ) {
ssPk , err := shexec . MakeShellStatePacket ( shellType )
2022-10-27 21:59:17 -07:00
if err != nil {
m . Sender . SendErrorResponse ( reqId , fmt . Errorf ( "error creating init packet: %w" , err ))
return
}
2024-01-16 16:11:04 -08:00
err = m . StateMap . SetCurrentState ( ssPk . State . GetShellType (), ssPk . State )
if err != nil {
m . Sender . SendErrorResponse ( reqId , fmt . Errorf ( "error setting current state: %w" , err ))
return
}
ssPk . RespId = reqId
m . Sender . SendPacket ( ssPk )
2022-10-27 21:59:17 -07:00
}
2023-09-05 21:21:34 -07:00
func makeTemp ( path string , mode fs . FileMode ) ( * os . File , error ) {
dirName := filepath . Dir ( path )
baseName := filepath . Base ( path )
baseTempName := baseName + ".tmp."
writeFd , err := os . CreateTemp ( dirName , baseTempName )
if err != nil {
return nil , err
}
err = writeFd . Chmod ( mode )
if err != nil {
writeFd . Close ()
os . Remove ( writeFd . Name ())
return nil , fmt . Errorf ( "error setting tempfile permissions: %w" , err )
}
return writeFd , nil
}
2023-09-06 21:45:15 -07:00
func checkFileWritable ( path string ) error {
finfo , err := os . Stat ( path ) // ok to follow symlinks
if errors . Is ( err , fs . ErrNotExist ) {
dirName := filepath . Dir ( path )
dirInfo , err := os . Stat ( dirName )
if err != nil {
return fmt . Errorf ( "file does not exist, error trying to stat parent directory: %w" , err )
}
if ! dirInfo . IsDir () {
return fmt . Errorf ( "file does not exist, parent path [%s] is not a directory" , dirName )
}
return nil
} else {
if err != nil {
return fmt . Errorf ( "cannot stat: %w" , err )
}
if finfo . IsDir () {
return fmt . Errorf ( "invalid path, cannot write a directory" )
}
if ( finfo . Mode () & fs . ModeSymlink ) != 0 {
return fmt . Errorf ( "writefile does not support symlinks" ) // note this shouldn't happen because we're using Stat (not Lstat)
}
if ( finfo . Mode () & ( fs . ModeNamedPipe | fs . ModeSocket | fs . ModeDevice )) != 0 {
return fmt . Errorf ( "writefile does not support special files (named pipes, sockets, devices): mode=%v" , finfo . Mode ())
}
writePerm := ( finfo . Mode (). Perm () & 0 o222 )
if writePerm == 0 {
return fmt . Errorf ( "file is not writable, perms: %v" , finfo . Mode (). Perm ())
}
return nil
}
}
func copyFile ( dstName string , srcName string ) error {
srcFd , err := os . Open ( srcName )
if err != nil {
return err
}
defer srcFd . Close ()
2023-09-06 21:57:41 -07:00
dstFd , err := os . OpenFile ( dstName , os . O_CREATE | os . O_WRONLY | os . O_TRUNC , 0 o666 ) // use 666 because OpenFile respects umask
2023-09-06 21:45:15 -07:00
if err != nil {
return err
}
// we don't defer dstFd.Close() so we can return an error if dstFd.Close() returns an error
_ , err = io . Copy ( dstFd , srcFd )
if err != nil {
dstFd . Close ()
return err
}
return dstFd . Close ()
}
2023-08-31 22:03:38 -07:00
func ( m * MServer ) writeFile ( pk * packet . WriteFilePacketType , wfc * WriteFileContext ) {
defer wfc . setDone ()
if pk . Path == "" {
resp := packet . MakeWriteFileReadyPacket ( pk . ReqId )
resp . Error = "invalid write-file request, no path specified"
m . Sender . SendPacket ( resp )
return
}
2023-09-06 21:45:15 -07:00
err := checkFileWritable ( pk . Path )
2023-08-31 22:03:38 -07:00
if err != nil {
resp := packet . MakeWriteFileReadyPacket ( pk . ReqId )
resp . Error = err . Error ()
m . Sender . SendPacket ( resp )
return
}
var writeFd * os . File
if pk . UseTemp {
2023-09-06 21:45:15 -07:00
writeFd , err = os . CreateTemp ( "" , "mshell.writefile.*" ) // "" means make this file in standard TempDir
2023-08-31 22:03:38 -07:00
if err != nil {
resp := packet . MakeWriteFileReadyPacket ( pk . ReqId )
2023-09-06 21:45:15 -07:00
resp . Error = fmt . Sprintf ( "cannot create temp file: %v" , err )
2023-08-31 22:03:38 -07:00
m . Sender . SendPacket ( resp )
return
}
} else {
2023-09-06 21:57:41 -07:00
writeFd , err = os . OpenFile ( pk . Path , os . O_CREATE | os . O_WRONLY | os . O_TRUNC , 0 o666 ) // use 666 because OpenFile respects umask
2023-08-31 22:03:38 -07:00
if err != nil {
resp := packet . MakeWriteFileReadyPacket ( pk . ReqId )
resp . Error = fmt . Sprintf ( "write-file could not open file: %v" , err )
m . Sender . SendPacket ( resp )
return
}
}
// ok, so now writeFd is valid, send the "ready" response
resp := packet . MakeWriteFileReadyPacket ( pk . ReqId )
m . Sender . SendPacket ( resp )
// now we wait for data (cond var)
// this Unlock() runs first (because it is a later defer) so we can still run wfc.setDone() safely
wfc . CVar . L . Lock ()
defer wfc . CVar . L . Unlock ()
var doneErr error
for {
if wfc . Done {
break
}
if wfc . Err != nil {
doneErr = wfc . Err
break
}
if len ( wfc . Data ) == 0 {
wfc . CVar . Wait ()
continue
}
dataPk := wfc . Data [ 0 ]
wfc . Data = wfc . Data [ 1 :]
if dataPk . Error != "" {
doneErr = fmt . Errorf ( "error received from client: %v" , errors . New ( dataPk . Error ))
break
}
if len ( dataPk . Data ) > 0 {
_ , err := writeFd . Write ( dataPk . Data )
if err != nil {
doneErr = fmt . Errorf ( "error writing data to file: %v" , err )
break
}
}
if dataPk . Eof {
break
}
}
closeErr := writeFd . Close ()
if doneErr == nil && closeErr != nil {
doneErr = fmt . Errorf ( "error closing file: %v" , closeErr )
}
if pk . UseTemp {
if doneErr != nil {
os . Remove ( writeFd . Name ())
} else {
2023-09-06 21:45:15 -07:00
// copy file between writeFd.Name() and pk.Path
copyErr := copyFile ( pk . Path , writeFd . Name ())
if err != nil {
doneErr = fmt . Errorf ( "error writing file: %v" , copyErr )
2023-08-31 22:03:38 -07:00
}
2023-09-06 21:45:15 -07:00
os . Remove ( writeFd . Name ())
2023-08-31 22:03:38 -07:00
}
}
donePk := packet . MakeWriteFileDonePacket ( pk . ReqId )
if doneErr != nil {
donePk . Error = doneErr . Error ()
}
m . Sender . SendPacket ( donePk )
}
2023-09-06 21:45:15 -07:00
func ( m * MServer ) returnStreamFileNewFileResponse ( pk * packet . StreamFilePacketType ) {
// ok, file doesn't exist, so try to check the directory at least to see if we can write a file here
resp := packet . MakeStreamFileResponse ( pk . ReqId )
defer func () {
if resp . Error == "" {
resp . Done = true
}
m . Sender . SendPacket ( resp )
}()
dirName := filepath . Dir ( pk . Path )
dirInfo , err := os . Stat ( dirName )
if err != nil {
resp . Error = fmt . Sprintf ( "file does not exist, error trying to stat parent directory: %v" , err )
return
}
if ! dirInfo . IsDir () {
resp . Error = fmt . Sprintf ( "file does not exist, parent path [%s] is not a directory" , dirName )
return
}
resp . Info = & packet . FileInfo {
Name : pk . Path ,
Size : 0 ,
ModTs : 0 ,
IsDir : false ,
Perm : int ( dirInfo . Mode (). Perm ()),
NotFound : true ,
}
return
}
2023-08-31 22:03:38 -07:00
func ( m * MServer ) streamFile ( pk * packet . StreamFilePacketType ) {
resp := packet . MakeStreamFileResponse ( pk . ReqId )
finfo , err := os . Stat ( pk . Path )
2023-09-06 21:45:15 -07:00
if errors . Is ( err , fs . ErrNotExist ) {
// special return
m . returnStreamFileNewFileResponse ( pk )
return
}
2023-08-31 22:03:38 -07:00
if err != nil {
resp . Error = fmt . Sprintf ( "cannot stat file %q: %v" , pk . Path , err )
m . Sender . SendPacket ( resp )
return
}
resp . Info = & packet . FileInfo {
Name : pk . Path ,
Size : finfo . Size (),
ModTs : finfo . ModTime (). UnixMilli (),
IsDir : finfo . IsDir (),
Perm : int ( finfo . Mode (). Perm ()),
}
if pk . StatOnly {
resp . Done = true
m . Sender . SendPacket ( resp )
return
}
// like the http Range header. range header is end inclusive. for us, endByte is non-inclusive (so we add 1)
var startByte , endByte int64
if len ( pk . ByteRange ) == 0 {
endByte = finfo . Size ()
} else if len ( pk . ByteRange ) == 1 && pk . ByteRange [ 0 ] >= 0 {
startByte = pk . ByteRange [ 0 ]
endByte = finfo . Size ()
} else if len ( pk . ByteRange ) == 1 && pk . ByteRange [ 0 ] < 0 {
startByte = finfo . Size () + pk . ByteRange [ 0 ] // "+" since ByteRange[0] is less than 0
endByte = finfo . Size ()
} else if len ( pk . ByteRange ) == 2 {
startByte = pk . ByteRange [ 0 ]
endByte = pk . ByteRange [ 1 ] + 1
} else {
resp . Error = fmt . Sprintf ( "invalid byte range (%d entries)" , len ( pk . ByteRange ))
m . Sender . SendPacket ( resp )
return
}
if startByte < 0 {
startByte = 0
}
if endByte > finfo . Size () {
endByte = finfo . Size ()
}
if startByte >= endByte {
resp . Done = true
m . Sender . SendPacket ( resp )
return
}
fd , err := os . Open ( pk . Path )
if err != nil {
resp . Error = fmt . Sprintf ( "opening file: %v" , err )
m . Sender . SendPacket ( resp )
return
}
defer fd . Close ()
m . Sender . SendPacket ( resp )
var buffer [ MaxFileDataPacketSize ] byte
var sentDone bool
first := true
for ; startByte < endByte ; startByte += MaxFileDataPacketSize {
if ! first {
// throttle packet sending @ 1000 packets/s, or 16M/s
time . Sleep ( 1 * time . Millisecond )
}
first = false
readLen := int64Min ( MaxFileDataPacketSize , endByte - startByte )
bufSlice := buffer [ 0 : readLen ]
nr , err := fd . ReadAt ( bufSlice , startByte )
dataPk := packet . MakeFileDataPacket ( pk . ReqId )
dataPk . Data = make ([] byte , nr )
copy ( dataPk . Data , bufSlice )
if err == io . EOF {
dataPk . Eof = true
} else if err != nil {
dataPk . Error = err . Error ()
}
m . Sender . SendPacket ( dataPk )
if dataPk . GetResponseDone () {
sentDone = true
break
}
}
if ! sentDone {
dataPk := packet . MakeFileDataPacket ( pk . ReqId )
dataPk . Eof = true
m . Sender . SendPacket ( dataPk )
}
return
}
func int64Min ( v1 int64 , v2 int64 ) int64 {
if v1 < v2 {
return v1
}
return v2
}
2022-08-09 14:23:59 -07:00
func ( m * MServer ) ProcessRpcPacket ( pk packet . RpcPacketType ) {
reqId := pk . GetReqId ()
if cdPk , ok := pk .( * packet . CdPacketType ); ok {
err := os . Chdir ( cdPk . Dir )
if err != nil {
m . Sender . SendErrorResponse ( reqId , fmt . Errorf ( "cannot change directory: %w" , err ))
return
}
m . Sender . SendResponse ( reqId , true )
return
}
if compPk , ok := pk .( * packet . CompGenPacketType ); ok {
go m . runCompGen ( compPk )
return
}
2024-01-16 16:11:04 -08:00
if reinitPk , ok := pk .( * packet . ReInitPacketType ); ok {
go m . reinit ( reqId , reinitPk . ShellType )
2022-10-27 21:59:17 -07:00
return
}
2023-08-31 22:03:38 -07:00
if streamPk , ok := pk .( * packet . StreamFilePacketType ); ok {
go m . streamFile ( streamPk )
return
}
if writePk , ok := pk .( * packet . WriteFilePacketType ); ok {
wfc := m . getWriteFileContext ( writePk . ReqId )
go m . writeFile ( writePk , wfc )
return
}
2022-08-09 14:23:59 -07:00
m . Sender . SendErrorResponse ( reqId , fmt . Errorf ( "invalid rpc type '%s'" , pk . GetType ()))
return
}
2024-01-16 16:11:04 -08:00
func ( m * MServer ) clientPacketCallback ( shellType string , pk packet . PacketType ) {
2022-11-27 13:47:18 -08:00
if pk . GetType () != packet . CmdDonePacketStr {
return
}
donePk := pk .( * packet . CmdDonePacketType )
if donePk . FinalState == nil {
return
}
2024-01-16 16:11:04 -08:00
stateHash , curState := m . StateMap . GetCurrentState ( shellType )
2022-11-27 13:47:18 -08:00
if curState == nil {
return
}
2024-01-16 16:11:04 -08:00
sapi , err := shellapi . MakeShellApi ( curState . GetShellType ())
if err != nil {
return
}
diff , err := sapi . MakeShellStateDiff ( curState , stateHash , donePk . FinalState )
2022-11-27 13:47:18 -08:00
if err != nil {
return
}
donePk . FinalState = nil
2024-01-16 16:11:04 -08:00
donePk . FinalStateDiff = diff
}
func ( m * MServer ) isShellInitialized ( shellType string ) bool {
_ , curState := m . StateMap . GetCurrentState ( shellType )
return curState != nil
2022-11-27 13:47:18 -08:00
}
2022-06-28 17:20:01 -07:00
func ( m * MServer ) runCommand ( runPacket * packet . RunPacketType ) {
2022-06-28 19:01:33 -07:00
if err := runPacket . CK . Validate ( "packet" ); err != nil {
2022-07-06 00:21:44 -07:00
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "server run packets require valid ck: %s" , err ))
2022-06-28 19:01:33 -07:00
return
}
2024-01-16 16:11:04 -08:00
if runPacket . ShellType == "" {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "server run packets require shell type" ))
return
}
_ , curInitState := m . StateMap . GetCurrentState ( runPacket . ShellType )
if curInitState == nil {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "shell type %q is not initialized" , runPacket . ShellType ))
return
}
if runPacket . State == nil {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "server run packets require state" ))
return
}
_ , _ , err := packet . ParseShellStateVersion ( runPacket . State . Version )
if err != nil {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "invalid shellstate version: %w" , err ))
return
}
if ! packet . StateVersionsCompatible ( runPacket . State . Version , curInitState . Version ) {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "shellstate version %q is not compatible with current shell version %q" , runPacket . State . Version , curInitState . Version ))
return
}
ecmd , err := shexec . MakeMShellSingleCmd ()
2022-07-06 18:59:46 -07:00
if err != nil {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "server run packets require valid ck: %s" , err ))
return
}
2024-01-25 10:18:11 -08:00
cproc , _ , err := shexec . MakeClientProc ( context . Background (), shexec . CmdWrap { Cmd : ecmd })
2022-07-06 17:16:45 -07:00
if err != nil {
m . Sender . SendErrorResponse ( runPacket . ReqId , fmt . Errorf ( "starting mshell client: %s" , err ))
return
}
m . Lock . Lock ()
m . ClientMap [ runPacket . CK ] = cproc
m . Lock . Unlock ()
2022-06-28 17:20:01 -07:00
go func () {
2022-07-06 17:16:45 -07:00
defer func () {
2022-11-27 13:47:18 -08:00
r := recover ()
finalPk := packet . MakeCmdFinalPacket ( runPacket . CK )
finalPk . Ts = time . Now (). UnixMilli ()
if r != nil {
finalPk . Error = fmt . Sprintf ( "%s" , r )
}
m . Sender . SendPacket ( finalPk )
2022-07-06 17:16:45 -07:00
m . Lock . Lock ()
delete ( m . ClientMap , runPacket . CK )
m . Lock . Unlock ()
cproc . Close ()
}()
2022-07-06 18:59:46 -07:00
shexec . SendRunPacketAndRunData ( context . Background (), cproc . Input , runPacket )
2024-01-16 16:11:04 -08:00
cproc . ProxySingleOutput ( runPacket . CK , m . Sender , func ( pk packet . PacketType ) {
m . clientPacketCallback ( runPacket . ShellType , pk )
})
2022-06-28 17:20:01 -07:00
}()
}
2022-11-27 13:47:18 -08:00
func ( m * MServer ) packetSenderErrorHandler ( sender * packet . PacketSender , pk packet . PacketType , err error ) {
if serr , ok := err .( * packet . SendError ); ok && serr . IsMarshalError {
msg := packet . MakeMessagePacket ( err . Error ())
if cpk , ok := pk .( packet . CommandPacketType ); ok {
msg . CK = cpk . GetCK ()
}
sender . SendPacket ( msg )
2022-12-05 15:38:44 -08:00
return
} else {
// I/O error: close the WriteErrorCh to signal that we are dead (cannot continue if we can't write output)
m . WriteErrorChOnce . Do ( func () {
close ( m . WriteErrorCh )
})
2022-11-27 13:47:18 -08:00
}
}
2022-12-05 15:38:44 -08:00
func ( server * MServer ) runReadLoop () {
2022-06-28 22:05:47 -07:00
builder := packet . MakeRunPacketBuilder ()
2022-06-28 15:04:08 -07:00
for pk := range server . MainInput . MainCh {
2022-06-28 19:01:33 -07:00
if server . Debug {
fmt . Printf ( "PK> %s\n" , packet . AsString ( pk ))
}
2022-06-28 22:05:47 -07:00
ok , runPacket := builder . ProcessPacket ( pk )
if ok {
if runPacket != nil {
server . runCommand ( runPacket )
continue
}
2022-06-28 17:20:01 -07:00
continue
}
2022-06-28 19:01:33 -07:00
if cmdPk , ok := pk .( packet . CommandPacketType ); ok {
server . ProcessCommandPacket ( cmdPk )
2022-06-28 15:04:08 -07:00
continue
}
2022-08-09 14:23:59 -07:00
if rpcPk , ok := pk .( packet . RpcPacketType ); ok {
server . ProcessRpcPacket ( rpcPk )
continue
}
2023-08-31 22:03:38 -07:00
if fileDataPk , ok := pk .( * packet . FileDataPacketType ); ok {
server . addFileDataPacket ( fileDataPk )
continue
}
2022-11-27 13:47:18 -08:00
server . Sender . SendMessageFmt ( "invalid packet '%s' sent to mshell server" , packet . AsString ( pk ))
2022-06-28 15:04:08 -07:00
continue
}
2022-12-05 15:38:44 -08:00
}
func RunServer () ( int , error ) {
debug := false
if len ( os . Args ) >= 3 && os . Args [ 2 ] == "--debug" {
debug = true
}
server := & MServer {
2023-08-31 22:03:38 -07:00
Lock : & sync . Mutex {},
ClientMap : make ( map [ base . CommandKey ] * shexec . ClientProc ),
2024-01-16 16:11:04 -08:00
StateMap : MakeShellStateMap (),
2023-08-31 22:03:38 -07:00
Debug : debug ,
WriteErrorCh : make ( chan bool ),
WriteErrorChOnce : & sync . Once {},
WriteFileContextMap : make ( map [ string ] * WriteFileContext ),
2022-12-05 15:38:44 -08:00
}
2023-08-31 22:03:38 -07:00
go func () {
for {
if server . checkDone () {
return
}
time . Sleep ( cleanLoopTime )
server . cleanWriteFileContexts ()
}
}()
2022-12-05 15:38:44 -08:00
if debug {
packet . GlobalDebug = true
}
2023-12-18 12:42:40 -08:00
server . MainInput = packet . MakePacketParser ( os . Stdin , nil )
2022-12-05 15:38:44 -08:00
server . Sender = packet . MakePacketSender ( os . Stdout , server . packetSenderErrorHandler )
defer server . Close ()
var err error
initPacket , err := shexec . MakeServerInitPacket ()
if err != nil {
return 1 , err
}
server . Sender . SendPacket ( initPacket )
ticker := time . NewTicker ( 1 * time . Minute )
go func () {
for range ticker . C {
server . Sender . SendPacket ( packet . MakePingPacket ())
}
}()
defer ticker . Stop ()
readLoopDoneCh := make ( chan bool )
go func () {
defer close ( readLoopDoneCh )
server . runReadLoop ()
}()
select {
case <- readLoopDoneCh :
break
case <- server . WriteErrorCh :
break
}
2022-06-28 15:04:08 -07:00
return 0 , nil
}
2024-01-16 16:11:04 -08:00
func MakeShellStateMap () * ShellStateMap {
return & ShellStateMap {
Lock : & sync . Mutex {},
StateMap : make ( map [ shellStateMapKey ] * packet . ShellState ),
CurrentStateMap : make ( map [ string ] string ),
}
}
func ( sm * ShellStateMap ) GetCurrentState ( shellType string ) ( string , * packet . ShellState ) {
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
hval := sm . CurrentStateMap [ shellType ]
return hval , sm . StateMap [ shellStateMapKey { ShellType : shellType , Hash : hval }]
}
func ( sm * ShellStateMap ) SetCurrentState ( shellType string , state * packet . ShellState ) error {
if state == nil {
return fmt . Errorf ( "cannot set nil state" )
}
if shellType != state . GetShellType () {
return fmt . Errorf ( "shell type mismatch: %s != %s" , shellType , state . GetShellType ())
}
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
hval , _ := state . EncodeAndHash ()
key := shellStateMapKey { ShellType : shellType , Hash : hval }
sm . StateMap [ key ] = state
sm . CurrentStateMap [ shellType ] = hval
return nil
}
func ( sm * ShellStateMap ) GetStateByHash ( shellType string , hash string ) * packet . ShellState {
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
return sm . StateMap [ shellStateMapKey { ShellType : shellType , Hash : hash }]
}
func ( sm * ShellStateMap ) Clear () {
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
sm . StateMap = make ( map [ shellStateMapKey ] * packet . ShellState )
sm . CurrentStateMap = make ( map [ string ] string )
}
func ( sm * ShellStateMap ) GetShells () [] string {
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
return utilfn . GetMapKeys ( sm . CurrentStateMap )
}
func ( sm * ShellStateMap ) HasShell ( shellType string ) bool {
sm . Lock . Lock ()
defer sm . Lock . Unlock ()
_ , found := sm . CurrentStateMap [ shellType ]
return found
}