diff --git a/main-mshell.go b/main-mshell.go index 65d7a8c9..014b5a76 100644 --- a/main-mshell.go +++ b/main-mshell.go @@ -153,7 +153,7 @@ func readFullRunPacket(packetParser *packet.PacketParser) (*packet.RunPacketType return nil, fmt.Errorf("no run packet received") } -func handleSingle() { +func handleSingle(fromServer bool) { packetParser := packet.MakePacketParser(os.Stdin) sender := packet.MakePacketSender(os.Stdout) defer func() { @@ -175,6 +175,12 @@ func handleSingle() { sender.SendErrorResponse(runPacket.ReqId, err) return } + if fromServer { + err = runPacket.CK.Validate("run packet") + if err != nil { + sender.SendErrorResponse(runPacket.ReqId, fmt.Errorf("run packets from server must have a CK: %v", err)) + } + } if runPacket.Detached { cmd, startPk, err := shexec.RunCommandDetached(runPacket, sender) if err != nil { @@ -187,7 +193,7 @@ func handleSingle() { cmd.DetachedWait(startPk) return } else { - cmd, err := shexec.RunCommandSimple(runPacket, sender) + cmd, err := shexec.RunCommandSimple(runPacket, sender, true) if err != nil { sender.SendErrorResponse(runPacket.ReqId, fmt.Errorf("error running command: %w", err)) return @@ -504,7 +510,10 @@ func main() { os.Exit(rtnCode) } } else if firstArg == "--single" { - handleSingle() + handleSingle(false) + return + } else if firstArg == "--single-from-server" { + handleSingle(true) return } else if firstArg == "--server" { rtnCode, err := server.RunServer() diff --git a/pkg/packet/packet.go b/pkg/packet/packet.go index 3584d5e8..3ec863b8 100644 --- a/pkg/packet/packet.go +++ b/pkg/packet/packet.go @@ -412,8 +412,9 @@ func MakeRawPacket(val string) *RawPacketType { } type MessagePacketType struct { - Type string `json:"type"` - Message string `json:"message"` + Type string `json:"type"` + CK base.CommandKey `json:"ck,omitempty"` + Message string `json:"message"` } func (*MessagePacketType) GetType() string { @@ -813,7 +814,17 @@ func (DefaultUPR) UnknownPacket(pk PacketType) { } else { fmt.Fprintf(os.Stderr, "[error] invalid packet received '%s'", AsExtType(pk)) } +} +type MessageUPR struct { + CK base.CommandKey + Sender *PacketSender +} + +func (upr MessageUPR) UnknownPacket(pk PacketType) { + msg := FmtMessagePacket("[error] invalid packet received %s", AsString(pk)) + msg.CK = upr.CK + upr.Sender.SendPacket(msg) } // todo: clean hanging entries in RunMap when in server mode diff --git a/pkg/server/server.go b/pkg/server/server.go index c1cd6bb4..ad4b3367 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -161,7 +161,7 @@ func (m *MServer) runCommand(runPacket *packet.RunPacketType) { m.Sender.SendErrorResponse(runPacket.ReqId, fmt.Errorf("server run packets require valid ck: %s", err)) return } - ecmd, err := shexec.SSHOpts{}.MakeMShellSingleCmd() + ecmd, err := shexec.SSHOpts{}.MakeMShellSingleCmd(true) if err != nil { m.Sender.SendErrorResponse(runPacket.ReqId, fmt.Errorf("server run packets require valid ck: %s", err)) return diff --git a/pkg/shexec/shexec.go b/pkg/shexec/shexec.go index 5a841a88..a7080007 100644 --- a/pkg/shexec/shexec.go +++ b/pkg/shexec/shexec.go @@ -375,13 +375,18 @@ func (opts SSHOpts) MakeMShellServerCmd() (*exec.Cmd, error) { return ecmd, nil } -func (opts SSHOpts) MakeMShellSingleCmd() (*exec.Cmd, error) { +func (opts SSHOpts) MakeMShellSingleCmd(fromServer bool) (*exec.Cmd, error) { if opts.SSHHost == "" { execFile, err := os.Executable() if err != nil { return nil, fmt.Errorf("cannot find local mshell executable: %w", err) } - ecmd := exec.Command(execFile, "--single") + var ecmd *exec.Cmd + if fromServer { + ecmd = exec.Command(execFile, "--single-from-server") + } else { + ecmd = exec.Command(execFile, "--single") + } return ecmd, nil } return opts.MakeSSHExecCmd(ClientCommand), nil @@ -657,7 +662,7 @@ func HasDupStdin(fds []packet.RemoteFd) bool { func RunClientSSHCommandAndWait(runPacket *packet.RunPacketType, fdContext FdContext, sshOpts SSHOpts, upr packet.UnknownPacketReporter, debug bool) (*packet.CmdDonePacketType, error) { cmd := MakeShExec(runPacket.CK, upr) - ecmd, err := sshOpts.MakeMShellSingleCmd() + ecmd, err := sshOpts.MakeMShellSingleCmd(false) if err != nil { return nil, err } @@ -829,8 +834,12 @@ func getTermType(pk *packet.RunPacketType) string { return termType } -func RunCommandSimple(pk *packet.RunPacketType, sender *packet.PacketSender) (*ShExecType, error) { - cmd := MakeShExec(pk.CK, nil) +func RunCommandSimple(pk *packet.RunPacketType, sender *packet.PacketSender, fromServer bool) (*ShExecType, error) { + var upr packet.UnknownPacketReporter + if fromServer { + upr = packet.MessageUPR{CK: pk.CK, Sender: sender} + } + cmd := MakeShExec(pk.CK, upr) if pk.UsePty { cmd.Cmd = exec.Command("bash", "-i", "-c", pk.Command) } else {