patterm/internal/pty/pty.go

// Package pty wraps creack/pty with the small surface the spike needs.
package pty

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
	"syscall"

	cpty "github.com/creack/pty"
)

// PTY holds a child process attached to a pseudo-terminal master fd.
//
// mu guards the master field only. Read/Write/Resize capture the *os.File
// under the lock and then do the (potentially blocking) I/O without holding
// it, so Close can swap master to nil and close the fd concurrently — closing
// the captured *os.File unblocks an in-flight Read. This avoids a data race
// between pumpChild's Read and Session.Shutdown's Close, which the daemon now
// hits routinely (daemon stop, not just process exit).
type PTY struct {
	mu     sync.Mutex
	master *os.File
	cmd    *exec.Cmd
}

// Start spawns argv with stdin/stdout/stderr attached to a new PTY sized
// (cols, rows). The returned PTY exposes the master fd for the parent to
// read from and write to.
func Start(argv []string, env []string, workDir string, cols, rows uint16) (*PTY, error) {
	if len(argv) == 0 {
		return nil, fmt.Errorf("pty: empty argv")
	}
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = workDir
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true, Setctty: true}
	if env != nil {
		cmd.Env = ensureTerm(env)
	} else {
		// Default to the parent environment but force TERM to xterm-256color
		// so child programs assume something modern and we observe SGR + alt
		// screen sequences.
		cmd.Env = ensureTerm(os.Environ())
	}

	ws := &cpty.Winsize{Cols: cols, Rows: rows}
	master, err := cpty.StartWithSize(cmd, ws)
	if err != nil {
		return nil, fmt.Errorf("pty: start %v: %w", argv, err)
	}
	return &PTY{master: master, cmd: cmd}, nil
}

func (p *PTY) Read(b []byte) (int, error) {
	p.mu.Lock()
	m := p.master
	p.mu.Unlock()
	if m == nil {
		return 0, io.ErrClosedPipe
	}
	return m.Read(b)
}

func (p *PTY) Write(b []byte) (int, error) {
	p.mu.Lock()
	m := p.master
	p.mu.Unlock()
	if m == nil {
		return 0, io.ErrClosedPipe
	}
	return m.Write(b)
}

func (p *PTY) Resize(cols, rows uint16) error {
	p.mu.Lock()
	m := p.master
	p.mu.Unlock()
	if m == nil {
		return io.ErrClosedPipe
	}
	return cpty.Setsize(m, &cpty.Winsize{Cols: cols, Rows: rows})
}

// Wait blocks until the child exits and returns its exit error if any.
func (p *PTY) Wait() error {
	if p.cmd == nil {
		return nil
	}
	return p.cmd.Wait()
}

// Pid returns the child's PID, or -1 if the process is not running.
func (p *PTY) Pid() int {
	if p.cmd == nil || p.cmd.Process == nil {
		return -1
	}
	return p.cmd.Process.Pid
}

// Close terminates the child (best effort) and releases the master fd.
func (p *PTY) Close() error {
	p.mu.Lock()
	m := p.master
	p.master = nil
	p.mu.Unlock()
	var firstErr error
	if m != nil {
		if err := m.Close(); err != nil {
			firstErr = err
		}
	}
	if p.cmd != nil && p.cmd.Process != nil {
		pid := p.cmd.Process.Pid
		if pid > 0 {
			_ = syscall.Kill(-pid, syscall.SIGKILL)
		}
		_ = p.cmd.Process.Kill()
	}
	return firstErr
}

// envDefaults are added to the child's environment unless the parent already
// set them. Modern agent TUIs check these and silently downgrade rendering
// when they're missing (no truecolor, ASCII-only banners, etc.).
var envDefaults = map[string]string{
	"TERM":      "xterm-256color",
	"COLORTERM": "truecolor",
}

// envStrip names variables we DROP before launching a child. COLUMNS /
// LINES inherited from the parent shell describe the *host* terminal,
// not the PTY we created — when they leak through, TUIs that prefer
// env over TIOCGWINSZ render past the PTY's actual cell grid and
// overwrite our chrome.
var envStrip = map[string]bool{
	"COLUMNS": true,
	"LINES":   true,
}

func ensureTerm(env []string) []string {
	have := make(map[string]bool, len(envDefaults))
	out := make([]string, 0, len(env)+len(envDefaults))
	for _, kv := range env {
		key := envKey(kv)
		if envStrip[key] {
			continue
		}
		if _, isDefault := envDefaults[key]; isDefault {
			have[key] = true
		}
		out = append(out, kv)
	}
	for k, v := range envDefaults {
		if !have[k] {
			out = append(out, k+"="+v)
		}
	}
	return out
}

func envKey(kv string) string {
	for i := 0; i < len(kv); i++ {
		if kv[i] == '=' {
			return kv[:i]
		}
	}
	return kv
}