feat: expand frame type to support stdout/stderr chan

agent/internal/sshutil/connect.go

@@ -3,8 +3,6 @@ package sshutil
 import (
 	"context"
 	"errors"
-	"io"
-	"os"
 	"strconv"
 	"strings"
 
@@ -13,149 +11,128 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-type SshConnectParams struct {
-	Hostname string
-	Port     int
-	Username string
-	Id 	 string
+type SSHConnectPayload struct {
+	SessionId string `cbor:"sessionId"`
+	Username  string `cbor:"username"`
+	Hostname  string `cbor:"hostname"`
+	Port      int    `cbor:"port"`
 }
 
-func dialAndValidateTailscaleSSH(agent *tsnet.TSAgent, params SshConnectParams) (*ssh.Client, error) {
+func connectToTailscaleSSH(agent *tsnet.TSAgent, params SSHConnectPayload) (*ssh.Client, error) {
 	log := util.GetLogger()
-
 	addr := strings.Join([]string{params.Hostname, ":", strconv.Itoa(params.Port)}, "")
 
-	log.Debug("Attempting to dial %s via Tailscale SSH", addr)
-	conn, err := agent.Dial(context.Background(), "tcp", addr)
+	log.Debug("Initiating Tailscale SSH connection to %s@%s", params.Username, addr)
+	tailnetConn, err := agent.Dial(context.Background(), "tcp", addr)
 	if err != nil {
-		log.Error("Failed to connect to Tailscale SSH: %s", err)
 		return nil, err
 	}
 
-	log.Debug("Connected to Tailscale SSH at %s", addr)
+	log.Debug("Routed connection via tsnet to %s", addr)
 	config := &ssh.ClientConfig{
 		User: params.Username,
+		// This isn't a concern because we are only dialing within the Tailnet
+		// and every device is trusted and *should* be ACL accessible.
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
 
-	clientConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
+	conn, chans, reqs, err := ssh.NewClientConn(tailnetConn, addr, config)
 	if err != nil {
-		log.Error("Failed to create SSH client connection: %s", err)
 		conn.Close()
 		return nil, err
 	}
 
-	client := ssh.NewClient(clientConn, chans, reqs)
-	sVer := string(client.ServerVersion())
+	// At this point we have successfully connected to the node
+	sshClient := ssh.NewClient(conn, chans, reqs)
+	version := string(sshClient.ServerVersion())
 
-	if !strings.Contains(sVer, "Tailscale") {
-		log.Error("Connected to non-Tailscale SSH server: %s", sVer)
+	if !strings.Contains(version, "Tailscale") {
 		conn.Close()
-		return nil, errors.New("not a Tailscale SSH server")
+		return nil, errors.New("server is not running Tailscale SSH")
 	}
 
-	log.Info("Connected to SSH server running %s at %s", client.ServerVersion(), addr)
-	return client, nil
+	log.Info("Connected to %s@%s:%d via Tailscale SSH (%s)", params.Username, params.Hostname, params.Port, version)
+	return sshClient, nil
 }
 
-func bindStdinToFd(sess *ssh.Session, fd int) error {
-	log := util.GetLogger()
-
-	sshIn := os.NewFile(uintptr(fd), "sshInput")
-	if sshIn == nil {
-		log.Error("Failed to create file from stdin fd %d", fd)
-		return errors.New("failed to create file from stdin fd")
-	}
-
-	stdin, err := sess.StdinPipe()
-	if err != nil {
-		log.Error("Failed to get stdin pipe: %s", err)
-		return err
-	}
-
-	go io.Copy(stdin, sshIn) // From Node → SSH session
-	return nil
-}
-
-func bindStdoutToFd(sess *ssh.Session, fd int) error {
-	log := util.GetLogger()
-
-	sshOut := os.NewFile(uintptr(fd), "sshOutput")
-	if sshOut == nil {
-		log.Error("Failed to create file from stdout fd %d", fd)
-		return errors.New("failed to create file from stdout fd")
-	}
-
-	stdout, err := sess.StdoutPipe()
-	if err != nil {
-		log.Error("Failed to get stdout pipe: %s", err)
-		return err
-	}
-
-	go io.Copy(sshOut, stdout) // From SSH → Node
-	return nil
-}
-
-func OpenSshPty(agent *tsnet.TSAgent, params SshConnectParams) (*ssh.Client, error) {
+func StartWebSSH(agent *tsnet.TSAgent, params SSHConnectPayload) {
 	log := util.GetLogger()
 
 	if agent == nil {
-		log.Error("Tailscale agent is nil")
-		return nil, errors.New("tailscale agent is nil")
+		log.Error("tsnet.TSAgent is not initialized correctly")
+		return
 	}
 
-	if params.Hostname == "" || params.Port <= 0 || params.Username == "" {
-		log.Error("Invalid SSH connection parameters: %+v", params)
-		return nil, errors.New("invalid SSH connection parameters")
+	if params.Hostname == "" || params.Port <= 0 || params.Username == "" || params.SessionId == "" {
+		log.Error("Invalid SSH connection parameters: %v", params)
+		return
 	}
 
-	client, err := dialAndValidateTailscaleSSH(agent, params)
+	client, err := connectToTailscaleSSH(agent, params)
 	if err != nil {
-		log.Error("Failed to open SSH pty: %s", err)
-		return nil, err
+		log.Error("Failed to connect to Tailscale SSH for (%s): %s", params.SessionId, err)
+		return
 	}
 
+	// Everything in the func is related to the SSH session.
+	// Each session runs in its own goroutine, allowing concurrency.
 	go func() {
-	sess, err := client.NewSession()
-	if err != nil {
-		log.Error("Failed to create new SSH session: %s", err)
-		client.Close()
-	}
-
-	modes := ssh.TerminalModes{
-		ssh.ECHO:          1, // enable echoing
-		ssh.TTY_OP_ISPEED: 14400,
-		ssh.TTY_OP_OSPEED: 14400,
-	}
-
-	if err := sess.RequestPty("xterm-256color", 80, 40, modes); err != nil {
-		log.Error("Failed to request PTY: %s", err)
-		client.Close()
-	}
-
-	ctx := addSession(params.Id, sess)
-	go func() {
-		for data := range ctx.InputCh {
-			if _, err := ctx.Stdin.Write(data); err != nil {
-				log.Error("Failed to write to SSH stdin: %s", err)
-				return
-			}
+		log.Debug("Creating SSH session for session ID: %s", params.SessionId)
+		sess, err := client.NewSession()
+		if err != nil {
+			log.Error("Failed to create new SSH session: %s", err)
+			client.Close()
+			return
 		}
-	}();
 
-	if err := sess.Shell(); err != nil {
-		log.Error("Failed to start shell: %s", err)
+		modes := ssh.TerminalModes{
+			ssh.ECHO:          1,
+			ssh.TTY_OP_ISPEED: 14400,
+			ssh.TTY_OP_OSPEED: 14400,
+		}
+
+		// Resize event is possible via the control channel later
+		err = sess.RequestPty("xterm-256color", 80, 40, modes)
+		if err != nil {
+			log.Error("Failed to request PTY for (%s): %s", params.SessionId, err)
+			return
+		}
+
+		ctx, err := registerSessionChans(params.SessionId, sess)
+		if err != nil {
+			log.Error("Failed to register session channels for (%s): %s", params.SessionId, err)
+			client.Close()
+			return
+		}
+
+		// Input buffer handler
+		go func() {
+			for data := range ctx.InputCh {
+				_, err := ctx.Stdin.Write(data)
+				if err != nil {
+					log.Error("Failed to write to SSH stdin: %s", err)
+					return
+				}
+			}
+		}()
+
+		// This spawns 2 goroutins for stdout and stderr
+		dispatchSSHStdout(params.SessionId, ctx.Stdout, ctx.Stderr)
+
+		// Spin up a shell and wait for the pty to terminate
+		err = sess.Shell()
+		if err != nil {
+			log.Error("Failed to start shell for (%s): %s", params.SessionId, err)
+			client.Close()
+			return
+		}
+
+		log.Info("Opened an SSH PTY for %s", params.SessionId)
+		sess.Wait()
+		sess.Close()
 		client.Close()
-	}
 
-
-	log.Info("Successfully opened SSH pty for %s@%s:%d", params.Username, params.Hostname, params.Port)
-	go streamSSHOutput(params.Id, ctx.Stdout, os.NewFile(4, "sshOutput"))
-		sess.Wait();
-		sess.Close();
-		log.Info("SSH session %s closed (goSide)", params.Id)
-		RemoveSession(params.Id)
+		log.Info("SSH session for %s closed", params.SessionId)
+		RemoveSession(params.SessionId)
 	}()
-
-	return client, nil
 }

agent/internal/sshutil/encoder.go

@@ -5,37 +5,121 @@ import (
 	"fmt"
 )
 
-func encodeFrame(id string, data []byte) ([]byte, error) {
-	sid := []byte(id)
-	if len(sid) > 255 {
-		return nil, fmt.Errorf("session ID too long")
+// An SSH frame is used to wrap raw binary data to and from an SSH session
+// in order to allow multiplexing connections over a single file descriptor.
+//
+// In practice, this is how we can easily support multiple SSH connections
+// through the 2 file descriptors created by the parent node process.
+//
+// This is the format of an SSH frame:
+// - Magic: The first 4 bytes are HPLS (0x48504C53) to identify the frame.
+// - Version Byte: The first byte is the version of the frame format.
+// - Channel Type: The second byte indicates the type of channel.
+// - Session ID: The next bytes are the length and actual session ID.
+// - Payload: The remaining bytes are the payload length and actual data.
+//
+// +---------+----------+--------------+-------------+----------+
+// | Magic   | Version  | Channel Type | SID Length  | SID      |
+// | 4 bytes | 1 byte   | 1 byte       | 1 byte (S)  | S bytes  |
+// +---------+----------+--------------+------------------------+
+// | Payload Length     | Payload                               |
+// | 4 bytes (u32, P)   | P bytes                               |
+// +--------------------+---------------------------------------+
+
+const (
+	MagicString = "HPLS"
+	VersionByte = 1
+)
+
+type ChannelType int
+
+const (
+	ChannelTypeStdin ChannelType = iota
+	ChannelTypeStdout
+	ChannelTypeStderr
+)
+
+type SSHFrame struct {
+	ChannelType ChannelType
+	SessionID   string
+	Payload     []byte
+
+	Length func() int
+}
+
+type HPLSFrame1 struct{}
+
+func (t HPLSFrame1) Encode(frame SSHFrame) ([]byte, error) {
+	frameChan := frame.ChannelType
+	switch frameChan {
+	case ChannelTypeStdin, ChannelTypeStdout, ChannelTypeStderr:
+	default:
+		return nil, fmt.Errorf("invalid channel type: %d", frameChan)
 	}
 
-	payloadLen := len(data)
-	buf := make([]byte, 1+len(sid)+4+payloadLen)
+	if len(frame.SessionID) == 0 {
+		return nil, fmt.Errorf("session ID cannot be empty")
+	}
 
-	buf[0] = byte(len(sid))                                          // SID length
-	copy(buf[1:], sid)                                               // SID
-	binary.BigEndian.PutUint32(buf[1+len(sid):], uint32(payloadLen)) // Payload length
-	copy(buf[1+len(sid)+4:], data)                                   // Payload
+	if len(frame.Payload) == 0 {
+		return nil, fmt.Errorf("payload cannot be empty")
+	}
 
+	sid := []byte(frame.SessionID)
+	if len(sid) > 255 {
+		return nil, fmt.Errorf("session ID exceeds 255 byte limit")
+	}
+
+	if len(frame.Payload) > 0xFFFFFFFF {
+		return nil, fmt.Errorf("payload exceeds 4GB limit")
+	}
+
+	frameLen := 4                      // Magic
+	frameLen += 1                      // Version byte
+	frameLen += 1                      // Channel type
+	frameLen += 1 + len(sid)           // Session ID length + SID
+	frameLen += 4 + len(frame.Payload) // Payload length + Payload
+
+	buf := make([]byte, frameLen)
+	copy(buf[0:4], []byte(MagicString))
+	buf[4] = VersionByte
+	buf[5] = byte(frameChan)
+	buf[6] = byte(len(sid))
+
+	offset := 7 + len(sid)
+	copy(buf[7:offset], sid)
+
+	binary.BigEndian.PutUint32(buf[offset:offset+4], uint32(len(frame.Payload)))
+	copy(buf[offset+4:], frame.Payload)
 	return buf, nil
 }
 
-func decodeFrame(buf []byte) (id string, payload []byte, ok bool) {
-	if len(buf) < 5 {
-		return "", nil, false
+func (t HPLSFrame1) Decode(buf []byte) (SSHFrame, error) {
+	frame := SSHFrame{}
+	if len(buf) < 5 || string(buf[0:4]) != MagicString || buf[4] != VersionByte {
+		return frame, fmt.Errorf("illegal HPLS1 frame format")
 	}
-	sidLen := int(buf[0])
-	if len(buf) < 1+sidLen+4 {
-		return "", nil, false
-	}
-	id = string(buf[1 : 1+sidLen])
-	payloadLen := int(binary.BigEndian.Uint32(buf[1+sidLen:]))
 
-	if len(buf) < 1+sidLen+4+payloadLen {
-		return "", nil, false
+	frame.ChannelType = ChannelType(buf[5])
+	if frame.ChannelType < ChannelTypeStdin || frame.ChannelType > ChannelTypeStderr {
+		return frame, fmt.Errorf("invalid channel type: %d", frame.ChannelType)
 	}
-	payload = buf[1+sidLen+4 : 1+sidLen+4+payloadLen]
-	return id, payload, true
+
+	sidLen := int(buf[6])
+	if len(buf) < 7+sidLen+4 {
+		return frame, fmt.Errorf("buffer too short for session ID and payload length")
+	}
+
+	frame.SessionID = string(buf[7 : 7+sidLen])
+	payloadLen := int(binary.BigEndian.Uint32(buf[7+sidLen:]))
+	if len(buf) < 7+sidLen+4+payloadLen {
+		return frame, fmt.Errorf("buffer too short for payload")
+	}
+
+	frame.Payload = buf[7+sidLen+4 : 7+sidLen+4+payloadLen]
+	frame.Length = func() int {
+		return 7 + sidLen + 4 + payloadLen
+	}
+
+	return frame, nil
 }

agent/internal/sshutil/handler.go

@@ -2,55 +2,145 @@ package sshutil
 
 import (
 	"io"
+	"os"
 
 	"github.com/tale/headplane/agent/internal/util"
 )
 
-func StartInputReader(fd3 io.Reader) {
-	log := util.GetLogger();
+// The file descriptors attached by the parent node process
+const (
+	InputFd  = 3
+	OutputFd = 4
+)
 
-	log.Info("Starting SSH fd3 input reader")
+// 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 {
-			n, err := fd3.Read(buffer)
+			// 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("fd3 read error: %v", err)
+				log.Error("Failed to read from SSH stdin: %v", err)
 				return
 			}
-			offset := 0
-			for offset < n {
-				id, payload, ok := decodeFrame(buffer[offset:n])
-				if !ok {
-					break // Wait for more data
-				}
-				offset += 1 + len(id) + 4 + len(payload)
 
-				sessionsMu.RLock()
-				sess, ok := sessions[id]
-				sessionsMu.RUnlock()
-				if !ok {
-					log.Error("invalid session id: %s", id)
-					continue
-				}
-				_, err := sess.Stdin.Write(payload)
+			// 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 {
-					log.Error("failed to write to session stdin: %v", err)
+					// 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 streamSSHOutput(id string, r io.Reader, fd4 io.Writer) {
+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
+	}
+
+	go readerStreamRoutine(StreamRoutine{
+		SessionID:   id,
+		Reader:      stdout,
+		Writer:      fd,
+		ChannelType: ChannelTypeStdout,
+	})
+
+	go readerStreamRoutine(StreamRoutine{
+		SessionID:   id,
+		Reader:      stderr,
+		Writer:      fd,
+		ChannelType: ChannelTypeStderr,
+	})
+}
+
+type StreamRoutine struct {
+	SessionID   string
+	Reader      io.Reader
+	Writer      io.Writer
+	ChannelType ChannelType
+}
+
+func readerStreamRoutine(routine StreamRoutine) {
+	hpls1 := HPLSFrame1{}
 	buf := make([]byte, 1024)
 	for {
-		n, err := r.Read(buf)
+		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
 		}
-		frame, _ := encodeFrame(id, buf[:n])
-		fd4.Write(frame)
 	}
 }

agent/internal/sshutil/sessions.go

@@ -1,9 +1,11 @@
 package sshutil
 
 import (
+	"errors"
 	"io"
 	"sync"
 
+	"github.com/tale/headplane/agent/internal/util"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -12,56 +14,70 @@ type SessionContext struct {
 	Session *ssh.Session
 	Stdin   io.WriteCloser
 	Stdout  io.Reader
+	Stderr  io.Reader
 	InputCh chan []byte
 }
 
 var sessions = make(map[string]*SessionContext)
-var sessionsMu sync.RWMutex
+var sessionsLock sync.RWMutex
 
-func addSession(id string, session *ssh.Session) *SessionContext {
-	sessionsMu.Lock()
-	defer sessionsMu.Unlock()
+func registerSessionChans(id string, session *ssh.Session) (*SessionContext, error) {
+	log := util.GetLogger()
+
+	sessionsLock.Lock()
+	defer sessionsLock.Unlock()
 
 	if _, exists := sessions[id]; exists {
-		return nil // Session with this ID already exists
+		return sessions[id], nil
 	}
 
 	stdin, err := session.StdinPipe()
 	if err != nil {
-		return nil // Handle error appropriately in production code
+		return nil, errors.New("failed to create stdin pipe: " + err.Error())
 	}
 
 	stdout, err := session.StdoutPipe()
 	if err != nil {
-		stdin.Close() // Close stdin if stdout pipe creation fails
-		return nil // Handle error appropriately in production code
+		stdin.Close()
+		return nil, errors.New("failed to create stdout pipe: " + err.Error())
 	}
 
-	sessionContext := &SessionContext{
+	stderr, err := session.StderrPipe()
+	if err != nil {
+		stdin.Close()
+		return nil, errors.New("failed to create stderr pipe: " + err.Error())
+	}
+
+	ctx := &SessionContext{
 		ID:      id,
 		Session: session,
 		Stdin:   stdin,
 		Stdout:  stdout,
+		Stderr:  stderr,
+		// Buffered channel to queue input data
+		InputCh: make(chan []byte, 256),
 	}
 
-	sessions[id] = sessionContext
-	return sessionContext
+	sessions[id] = ctx
+	log.Debug("Registered session %s with stdin, stdout, and stderr pipes", id)
+	return ctx, nil
 }
 
-func GetSession(id string) *SessionContext {
-	sessionsMu.RLock()
-	defer sessionsMu.RUnlock()
+func lookupSession(id string) (*SessionContext, bool) {
+	sessionsLock.RLock()
+	defer sessionsLock.RUnlock()
 
-	return sessions[id] // Returns nil if session does not exist
+	sessionContext, exists := sessions[id]
+	return sessionContext, exists
 }
 
 func RemoveSession(id string) {
-	sessionsMu.Lock()
-	defer sessionsMu.Unlock()
+	sessionsLock.Lock()
+	defer sessionsLock.Unlock()
 
 	if sessionContext, exists := sessions[id]; exists {
-		sessionContext.Stdin.Close() // Close the stdin pipe
+		sessionContext.Stdin.Close()   // Close the stdin pipe
 		sessionContext.Session.Close() // Close the SSH session
-		delete(sessions, id) // Remove from the map
+		delete(sessions, id)           // Remove from the map
 	}
 }