muxplex/muxterm/pool.go

package main

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

	"github.com/creack/pty"
)

// Session represents a PTY-backed tmux attach process.
type Session struct {
	PTY  *os.File
	Cmd  *exec.Cmd
	Cols uint16
	Rows uint16
}

// Pool manages named tmux PTY sessions.
type Pool struct {
	mu       sync.Mutex
	sessions map[string]*Session
}

// NewPool returns an empty Pool.
func NewPool() *Pool {
	return &Pool{
		sessions: make(map[string]*Session),
	}
}

// Attach reuses an existing live session or spawns a new tmux attach process.
// Dead sessions are cleaned up before spawning a replacement.
func (p *Pool) Attach(name string, cols, rows uint16) (*Session, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if sess, ok := p.sessions[name]; ok {
		// Check if process is still alive via Signal(0).
		if sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil {
			return sess, nil
		}
		// Dead session — clean up.
		sess.PTY.Close()
		delete(p.sessions, name)
	}

	cmd := exec.Command("tmux", "attach", "-t", name)
	// Always set TERM=xterm-256color so tmux can render to the PTY.
	// When muxterm runs as a systemd service TERM is typically unset, and
	// even when set to a basic type (e.g. "dumb") tmux fails immediately
	// with "open terminal failed: terminal does not support clear screen".
	// For a PTY-backed web terminal xterm-256color is always the right value.
	env := make([]string, 0, len(os.Environ())+1)
	for _, e := range os.Environ() {
		if !strings.HasPrefix(e, "TERM=") {
			env = append(env, e)
		}
	}
	env = append(env, "TERM=xterm-256color")
	cmd.Env = env
	winSize := &pty.Winsize{Cols: cols, Rows: rows}

	ptmx, err := pty.StartWithSize(cmd, winSize)
	if err != nil {
		return nil, fmt.Errorf("pty start: %w", err)
	}

	sess := &Session{
		PTY:  ptmx,
		Cmd:  cmd,
		Cols: cols,
		Rows: rows,
	}
	p.sessions[name] = sess

	// Background goroutine: wait for process exit and clean up.
	go func() {
		cmd.Wait()
		p.mu.Lock()
		defer p.mu.Unlock()
		// Only delete if the map still holds this exact session
		// (Attach may have replaced it already).
		if cur, ok := p.sessions[name]; ok && cur == sess {
			sess.PTY.Close()
			delete(p.sessions, name)
		}
	}()

	return sess, nil
}

// Resize changes the PTY window size for a named session.
func (p *Pool) Resize(name string, cols, rows uint16) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	sess, ok := p.sessions[name]
	if !ok {
		return fmt.Errorf("session %q not found", name)
	}

	if err := pty.Setsize(sess.PTY, &pty.Winsize{Cols: cols, Rows: rows}); err != nil {
		return fmt.Errorf("setsize: %w", err)
	}
	sess.Cols = cols
	sess.Rows = rows
	return nil
}

// Get returns the session for name, or nil if not present.
func (p *Pool) Get(name string) *Session {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.sessions[name]
}

// IsAlive returns true if the named session exists and its process is alive.
func (p *Pool) IsAlive(name string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	sess, ok := p.sessions[name]
	if !ok {
		return false
	}
	return sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil
}

// CloseAll sends SIGTERM to all sessions and closes their PTYs.
func (p *Pool) CloseAll() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for name, sess := range p.sessions {
		if sess.Cmd.Process != nil {
			sess.Cmd.Process.Signal(syscall.SIGTERM)
		}
		sess.PTY.Close()
		delete(p.sessions, name)
	}
}