hamlab/internal/term/manager.go

// Package term owns PTY-backed shell sessions that outlive any one browser tab.
package term

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"os"
	"os/exec"
	"sort"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/creack/pty"
)

const (
	ringSize   = 256 * 1024
	clientBuf  = 512
	killGrace  = 3 * time.Second
	readChunk  = 32 * 1024
	defaultCol = 120
	defaultRow = 30
)

var ErrNotFound = errors.New("terminal not found")

// Msg is one unit of output to a client: raw PTY bytes, or a control event.
type Msg struct {
	Data  []byte // non-nil for PTY output
	Event string // "exit" when the shell has ended
	Code  int
}

type Client struct {
	Ch      chan Msg
	session *Session
	dropped bool
}

type Info struct {
	ID       string    `json:"id"`
	Name     string    `json:"name"`
	Running  bool      `json:"running"`
	ExitCode int       `json:"exit_code"`
	Created  time.Time `json:"created"`
	Clients  int       `json:"clients"`
}

type Session struct {
	ID      string
	Name    string
	Created time.Time

	mu       sync.Mutex
	ptmx     *os.File
	cmd      *exec.Cmd
	ring     *Ring
	clients  map[*Client]struct{}
	running  bool
	exitCode int
}

type Manager struct {
	mu       sync.RWMutex
	sessions map[string]*Session
	shell    string
	home     string
	OnChange func([]Info) // called (without locks held) when the list changes
	closing  atomic.Bool
}

func NewManager(shell, home string) *Manager {
	return &Manager{sessions: map[string]*Session{}, shell: shell, home: home}
}

func newID() string {
	b := make([]byte, 4)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

// Create starts a new login shell. id may be empty.
func (m *Manager) Create(id, name string) (*Session, error) {
	if id == "" {
		id = newID()
	}
	cmd := exec.Command(m.shell, "-l")
	cmd.Dir = m.home
	cmd.Env = append(os.Environ(), "TERM=xterm-256color", "COLORTERM=truecolor", "HAMLAB=1")
	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: defaultCol, Rows: defaultRow})
	if err != nil {
		return nil, err
	}
	s := &Session{
		ID: id, Name: name, Created: time.Now(),
		ptmx: ptmx, cmd: cmd, ring: NewRing(ringSize),
		clients: map[*Client]struct{}{}, running: true,
	}
	m.mu.Lock()
	m.sessions[id] = s
	m.mu.Unlock()
	go m.pump(s)
	m.changed()
	return s, nil
}

func (m *Manager) pump(s *Session) {
	buf := make([]byte, readChunk)
	for {
		n, err := s.ptmx.Read(buf)
		if n > 0 {
			chunk := make([]byte, n)
			copy(chunk, buf[:n])
			s.mu.Lock()
			s.ring.Write(chunk)
			s.broadcast(Msg{Data: chunk})
			s.mu.Unlock()
		}
		if err != nil {
			break // EIO when the slave side closes
		}
	}
	_ = s.cmd.Wait()
	code := s.cmd.ProcessState.ExitCode()
	s.mu.Lock()
	s.running = false
	s.exitCode = code
	s.broadcast(Msg{Event: "exit", Code: code})
	s.mu.Unlock()
	_ = s.ptmx.Close()
	m.changed()
	// An exited shell stays listed for a while so the exit is visible, then reaps itself.
	time.AfterFunc(10*time.Minute, func() {
		m.mu.Lock()
		cur, ok := m.sessions[s.ID]
		m.mu.Unlock()
		if ok && cur == s && !s.running {
			_ = m.Close(s.ID)
		}
	})
}

// broadcast requires s.mu. Slow clients are dropped rather than blocking the PTY.
func (s *Session) broadcast(msg Msg) {
	for c := range s.clients {
		select {
		case c.Ch <- msg:
		default:
			if !c.dropped {
				c.dropped = true
				close(c.Ch)
			}
			delete(s.clients, c)
		}
	}
}

func (m *Manager) get(id string) (*Session, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	s, ok := m.sessions[id]
	if !ok {
		return nil, ErrNotFound
	}
	return s, nil
}

func (m *Manager) List() []Info {
	m.mu.RLock()
	out := make([]Info, 0, len(m.sessions))
	for _, s := range m.sessions {
		s.mu.Lock()
		out = append(out, Info{ID: s.ID, Name: s.Name, Running: s.running, ExitCode: s.exitCode,
			Created: s.Created, Clients: len(s.clients)})
		s.mu.Unlock()
	}
	m.mu.RUnlock()
	sort.Slice(out, func(i, j int) bool { return out[i].Created.Before(out[j].Created) })
	return out
}

func (m *Manager) Rename(id, name string) error {
	s, err := m.get(id)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.Name = name
	s.mu.Unlock()
	m.changed()
	return nil
}

// Attach subscribes a client and returns the scrollback snapshot to replay first.
func (m *Manager) Attach(id string) (*Client, []byte, bool, error) {
	s, err := m.get(id)
	if err != nil {
		return nil, nil, false, err
	}
	c := &Client{Ch: make(chan Msg, clientBuf), session: s}
	s.mu.Lock()
	snap := s.ring.Bytes()
	running := s.running
	s.clients[c] = struct{}{}
	s.mu.Unlock()
	m.changed()
	return c, snap, running, nil
}

func (m *Manager) Detach(c *Client) {
	s := c.session
	s.mu.Lock()
	if _, ok := s.clients[c]; ok {
		delete(s.clients, c)
		if !c.dropped {
			c.dropped = true
			close(c.Ch)
		}
	}
	s.mu.Unlock()
	m.changed()
}

func (m *Manager) Write(id string, p []byte) error {
	s, err := m.get(id)
	if err != nil {
		return err
	}
	s.mu.Lock()
	running := s.running
	s.mu.Unlock()
	if !running {
		return nil
	}
	_, err = s.ptmx.Write(p)
	return err
}

func (m *Manager) Resize(id string, cols, rows uint16) error {
	s, err := m.get(id)
	if err != nil {
		return err
	}
	if cols == 0 || rows == 0 {
		return nil
	}
	return pty.Setsize(s.ptmx, &pty.Winsize{Cols: cols, Rows: rows})
}

// Close ends the shell (SIGHUP to its process group, SIGKILL after a grace
// period) and forgets the session.
func (m *Manager) Close(id string) error {
	m.mu.Lock()
	s, ok := m.sessions[id]
	if ok {
		delete(m.sessions, id)
	}
	m.mu.Unlock()
	if !ok {
		return ErrNotFound
	}
	s.mu.Lock()
	running := s.running
	for c := range s.clients {
		if !c.dropped {
			c.dropped = true
			close(c.Ch)
		}
		delete(s.clients, c)
	}
	s.mu.Unlock()
	if running && s.cmd.Process != nil {
		pid := s.cmd.Process.Pid
		_ = syscall.Kill(-pid, syscall.SIGHUP)
		go func() {
			time.Sleep(killGrace)
			_ = syscall.Kill(-pid, syscall.SIGKILL)
		}()
	}
	m.changed()
	return nil
}

func (m *Manager) changed() {
	if m.OnChange != nil && !m.closing.Load() {
		m.OnChange(m.List())
	}
}

// Shutdown hangs up every shell. Called on daemon exit. The saved list is left as it
// was, because that list is what the next start re-opens.
func (m *Manager) Shutdown() {
	m.closing.Store(true)
	for _, s := range m.List() {
		_ = m.Close(s.ID)
	}
}