feat: initial webssh tooling

2026-08-08 13:23:13 +00:00 · 2025-05-30 19:35:47 -04:00
parent f6d8ad25e1
commit 7dfcbef774
16 changed files with 792 additions and 83 deletions
@@ -0,0 +1,161 @@
+package sshutil
+
+import (
+	"context"
+	"errors"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/tale/headplane/agent/internal/tsnet"
+	"github.com/tale/headplane/agent/internal/util"
+	"golang.org/x/crypto/ssh"
+)
+
+type SshConnectParams struct {
+	Hostname string
+	Port     int
+	Username string
+	Id 	 string
+}
+
+func dialAndValidateTailscaleSSH(agent *tsnet.TSAgent, params SshConnectParams) (*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)
+	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)
+	config := &ssh.ClientConfig{
+		User: params.Username,
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	clientConn, chans, reqs, err := ssh.NewClientConn(conn, 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())
+
+	if !strings.Contains(sVer, "Tailscale") {
+		log.Error("Connected to non-Tailscale SSH server: %s", sVer)
+		conn.Close()
+		return nil, errors.New("not a Tailscale SSH server")
+	}
+
+	log.Info("Connected to SSH server running %s at %s", client.ServerVersion(), addr)
+	return client, 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) {
+	log := util.GetLogger()
+
+	if agent == nil {
+		log.Error("Tailscale agent is nil")
+		return nil, errors.New("tailscale agent is nil")
+	}
+
+	if params.Hostname == "" || params.Port <= 0 || params.Username == "" {
+		log.Error("Invalid SSH connection parameters: %+v", params)
+		return nil, errors.New("invalid SSH connection parameters")
+	}
+
+	client, err := dialAndValidateTailscaleSSH(agent, params)
+	if err != nil {
+		log.Error("Failed to open SSH pty: %s", err)
+		return nil, err
+	}
+
+	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
+			}
+		}
+	}();
+
+	if err := sess.Shell(); err != nil {
+		log.Error("Failed to start shell: %s", err)
+		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)
+	}()
+
+	return client, nil
+}
@@ -0,0 +1,41 @@
+package sshutil
+
+import (
+	"encoding/binary"
+	"fmt"
+)
+
+func encodeFrame(id string, data []byte) ([]byte, error) {
+	sid := []byte(id)
+	if len(sid) > 255 {
+		return nil, fmt.Errorf("session ID too long")
+	}
+
+	payloadLen := len(data)
+	buf := make([]byte, 1+len(sid)+4+payloadLen)
+
+	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
+
+	return buf, nil
+}
+
+func decodeFrame(buf []byte) (id string, payload []byte, ok bool) {
+	if len(buf) < 5 {
+		return "", nil, false
+	}
+	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
+	}
+	payload = buf[1+sidLen+4 : 1+sidLen+4+payloadLen]
+	return id, payload, true
+}
@@ -0,0 +1,56 @@
+package sshutil
+
+import (
+	"io"
+
+	"github.com/tale/headplane/agent/internal/util"
+)
+
+func StartInputReader(fd3 io.Reader) {
+	log := util.GetLogger();
+
+	log.Info("Starting SSH fd3 input reader")
+	go func() {
+		buffer := make([]byte, 8192)
+		for {
+			n, err := fd3.Read(buffer)
+			if err != nil {
+				log.Error("fd3 read error: %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)
+				if err != nil {
+					log.Error("failed to write to session stdin: %v", err)
+					continue
+				}
+			}
+		}
+	}()
+}
+
+func streamSSHOutput(id string, r io.Reader, fd4 io.Writer) {
+	buf := make([]byte, 1024)
+	for {
+		n, err := r.Read(buf)
+		if err != nil {
+			break
+		}
+		frame, _ := encodeFrame(id, buf[:n])
+		fd4.Write(frame)
+	}
+}
@@ -0,0 +1,67 @@
+package sshutil
+
+import (
+	"io"
+	"sync"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type SessionContext struct {
+	ID      string
+	Session *ssh.Session
+	Stdin   io.WriteCloser
+	Stdout  io.Reader
+	InputCh chan []byte
+}
+
+var sessions = make(map[string]*SessionContext)
+var sessionsMu sync.RWMutex
+
+func addSession(id string, session *ssh.Session) *SessionContext {
+	sessionsMu.Lock()
+	defer sessionsMu.Unlock()
+
+	if _, exists := sessions[id]; exists {
+		return nil // Session with this ID already exists
+	}
+
+	stdin, err := session.StdinPipe()
+	if err != nil {
+		return nil // Handle error appropriately in production code
+	}
+
+	stdout, err := session.StdoutPipe()
+	if err != nil {
+		stdin.Close() // Close stdin if stdout pipe creation fails
+		return nil // Handle error appropriately in production code
+	}
+
+	sessionContext := &SessionContext{
+		ID:      id,
+		Session: session,
+		Stdin:   stdin,
+		Stdout:  stdout,
+	}
+
+	sessions[id] = sessionContext
+	return sessionContext
+}
+
+func GetSession(id string) *SessionContext {
+	sessionsMu.RLock()
+	defer sessionsMu.RUnlock()
+
+	return sessions[id] // Returns nil if session does not exist
+}
+
+func RemoveSession(id string) {
+	sessionsMu.Lock()
+	defer sessionsMu.Unlock()
+
+	if sessionContext, exists := sessions[id]; exists {
+		sessionContext.Stdin.Close() // Close the stdin pipe
+		sessionContext.Session.Close() // Close the SSH session
+		delete(sessions, id) // Remove from the map
+	}
+}