153 lines
3.6 KiB
Go
153 lines
3.6 KiB
Go
package sshutil
|
|
|
|
import (
|
|
"io"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/tale/headplane/agent/internal/util"
|
|
)
|
|
|
|
// The file descriptors attached by the parent node process
|
|
const (
|
|
InputFd = 3
|
|
OutputFd = 4
|
|
)
|
|
|
|
// DispatchSSHStdin listens for SSH stdin frames on the InputFd file descriptor
|
|
// and writes the payload to the appropriate session's stdin.
|
|
//
|
|
// This function runs in a goroutine in main and is responsible for dispatching
|
|
// to ALL connections, not just its own like dispatchSSHStdout does.
|
|
func DispatchSSHStdin() {
|
|
log := util.GetLogger()
|
|
|
|
log.Debug("Opening file descriptor: %d for SSH stdin", InputFd)
|
|
fd := os.NewFile(InputFd, "ssh_stdin")
|
|
if fd == nil {
|
|
log.Error("Failed to open file descriptor %d for SSH stdin", InputFd)
|
|
return
|
|
}
|
|
|
|
log.Info("Listening for SSH stdin on fd %d", InputFd)
|
|
go func() {
|
|
buffer := make([]byte, 8192)
|
|
hpls1 := HPLSFrame1{}
|
|
|
|
for {
|
|
// This is the only check where we can detect if the descriptor
|
|
// was closed so we can return and exit the goroutine.
|
|
bufCount, err := fd.Read(buffer)
|
|
if err != nil {
|
|
log.Error("Failed to read from SSH stdin: %v", err)
|
|
return
|
|
}
|
|
|
|
// Check if we have 0 EOF, which means the descriptor was closed.
|
|
if bufCount == 0 {
|
|
log.Info("SSH stdin closed, stopping listener")
|
|
return
|
|
}
|
|
|
|
offset := 0
|
|
for offset < bufCount {
|
|
frame, err := hpls1.Decode(buffer[offset:bufCount])
|
|
if err != nil {
|
|
// We need to wait for more data to decode the frame
|
|
break
|
|
}
|
|
|
|
if frame.ChannelType != ChannelTypeStdin {
|
|
log.Error("Received invalid channel type: %d, expected %d", frame.ChannelType, ChannelTypeStdin)
|
|
continue
|
|
}
|
|
|
|
offset += frame.Length()
|
|
log.Debug("Received SSH stdin frame: %s", frame.SessionID)
|
|
sess, ok := lookupSession(frame.SessionID)
|
|
if !ok {
|
|
log.Error("Invalid session ID: %s", frame.SessionID)
|
|
continue
|
|
}
|
|
|
|
// Write the payload to the session's stdin
|
|
writeCount, err := sess.Stdin.Write(frame.Payload)
|
|
if err != nil {
|
|
log.Error("Failed to write to session stdin: %v", err)
|
|
continue
|
|
}
|
|
|
|
log.Debug("Wrote %d bytes to session %s stdin", writeCount, frame.SessionID)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func dispatchSSHStdout(id string, stdout io.Reader, stderr io.Reader) {
|
|
log := util.GetLogger()
|
|
|
|
log.Debug("Opening file descriptor: %d for SSH stdout", OutputFd)
|
|
fd := os.NewFile(OutputFd, "ssh_stdout")
|
|
if fd == nil {
|
|
log.Error("Failed to open file descriptor %d for SSH stdout", OutputFd)
|
|
return
|
|
}
|
|
|
|
batcher := NewFrameBatcher(fd, 10*time.Millisecond) // Roughly 60fps
|
|
|
|
go readerStreamRoutine(StreamRoutine{
|
|
SessionID: id,
|
|
Reader: stdout,
|
|
Writer: batcher,
|
|
ChannelType: ChannelTypeStdout,
|
|
})
|
|
|
|
go readerStreamRoutine(StreamRoutine{
|
|
SessionID: id,
|
|
Reader: stderr,
|
|
Writer: batcher,
|
|
ChannelType: ChannelTypeStderr,
|
|
})
|
|
}
|
|
|
|
type StreamRoutine struct {
|
|
SessionID string
|
|
Reader io.Reader
|
|
Writer *FrameBatcher
|
|
ChannelType ChannelType
|
|
}
|
|
|
|
func readerStreamRoutine(routine StreamRoutine) {
|
|
hpls1 := HPLSFrame1{}
|
|
buf := make([]byte, 16384) // 16 KiB buffer
|
|
for {
|
|
byteCount, err := routine.Reader.Read(buf)
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
util.GetLogger().Error("Failed to read from reader: %v", err)
|
|
}
|
|
|
|
break
|
|
}
|
|
|
|
frame, err := hpls1.Encode(SSHFrame{
|
|
ChannelType: routine.ChannelType,
|
|
SessionID: routine.SessionID,
|
|
Payload: buf[:byteCount],
|
|
})
|
|
|
|
if err != nil {
|
|
util.GetLogger().Error("Failed to encode frame: %v", err)
|
|
continue
|
|
}
|
|
|
|
// if _, err := routine.Writer.Write(frame); err != nil {
|
|
// util.GetLogger().Error("Failed to write frame to writer: %v", err)
|
|
// break
|
|
// }
|
|
//
|
|
|
|
routine.Writer.QueueMsg(frame)
|
|
}
|
|
}
|