byre/internal/server/terminal.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"syscall"

	"github.com/coder/websocket"
	"github.com/creack/pty"

	"g.fc.run/misc/byre/internal/api"
)

const (
	terminalBacklog   = 256 << 10
	terminalClientBuf = 256
	terminalReadLimit = 1 << 20
	terminalMaxSize   = 1 << 15
	terminalExited    = websocket.StatusCode(4000)
)

type terminals struct {
	mu     sync.Mutex
	byTask map[string]*terminal
}

type terminal struct {
	ptmx *os.File
	cmd  *exec.Cmd

	mu      sync.Mutex
	backlog []byte
	clients map[chan []byte]bool
	exit    string
	exited  bool
}

func newTerminals() *terminals {
	return &terminals{byTask: map[string]*terminal{}}
}

func (ts *terminals) attach(t api.Task, size api.TerminalSize) (*terminal, error) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	term := ts.byTask[t.Id]
	if term != nil {
		return term, nil
	}

	shell := os.Getenv("SHELL")
	if shell == "" {
		return nil, errors.New("SHELL is not set in byre's environment")
	}
	cmd := exec.Command(shell, "-l")
	cmd.Dir = t.WorktreePath
	cmd.Env = append(os.Environ(), "TERM=xterm-256color", "COLORTERM=truecolor", fmt.Sprintf("PORT=%d", t.PortBase))
	ptmx, err := pty.StartWithSize(cmd, winsize(size))
	if err != nil {
		return nil, fmt.Errorf("start %s: %w", shell, err)
	}

	term = &terminal{ptmx: ptmx, cmd: cmd, clients: map[chan []byte]bool{}}
	ts.byTask[t.Id] = term
	slog.Info("terminal started", "task", t.Id, "shell", shell, "pid", cmd.Process.Pid)
	go func() {
		term.pump()
		ts.mu.Lock()
		if ts.byTask[t.Id] == term {
			delete(ts.byTask, t.Id)
		}
		ts.mu.Unlock()
		slog.Info("terminal ended", "task", t.Id, "exit", term.exit)
	}()
	return term, nil
}

func (ts *terminals) close(taskID string) {
	ts.mu.Lock()
	term := ts.byTask[taskID]
	delete(ts.byTask, taskID)
	ts.mu.Unlock()
	if term != nil {
		term.kill()
	}
}

func (ts *terminals) closeAll() {
	ts.mu.Lock()
	all := ts.byTask
	ts.byTask = map[string]*terminal{}
	ts.mu.Unlock()
	for _, term := range all {
		term.kill()
	}
}

func winsize(size api.TerminalSize) *pty.Winsize {
	return &pty.Winsize{Cols: uint16(size.Cols), Rows: uint16(size.Rows)}
}

func validSize(size api.TerminalSize) bool {
	return size.Cols >= 1 && size.Rows >= 1 && size.Cols <= terminalMaxSize && size.Rows <= terminalMaxSize
}

func (term *terminal) pump() {
	buf := make([]byte, 32<<10)
	for {
		n, err := term.ptmx.Read(buf)
		if n > 0 {
			term.output(bytes.Clone(buf[:n]))
		}
		if err != nil {
			break
		}
	}

	err := term.cmd.Wait()
	term.ptmx.Close()
	exit := fmt.Sprint(err)
	if term.cmd.ProcessState != nil {
		exit = term.cmd.ProcessState.String()
	}

	term.mu.Lock()
	defer term.mu.Unlock()
	term.exit = exit
	term.exited = true
	for c := range term.clients {
		close(c)
	}
	term.clients = nil
}

func (term *terminal) output(b []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	term.backlog = append(term.backlog, b...)
	if len(term.backlog) > 2*terminalBacklog {
		term.backlog = bytes.Clone(term.backlog[len(term.backlog)-terminalBacklog:])
	}
	for c := range term.clients {
		select {
		case c <- b:
		default:
			delete(term.clients, c)
			close(c)
		}
	}
}

func (term *terminal) subscribe() ([]byte, chan []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	if term.exited {
		return nil, nil
	}
	c := make(chan []byte, terminalClientBuf)
	term.clients[c] = true
	return bytes.Clone(term.backlog[max(0, len(term.backlog)-terminalBacklog):]), c
}

func (term *terminal) unsubscribe(c chan []byte) {
	term.mu.Lock()
	defer term.mu.Unlock()

	if term.clients[c] {
		delete(term.clients, c)
		close(c)
	}
}

func (term *terminal) exitStatus() (string, bool) {
	term.mu.Lock()
	defer term.mu.Unlock()
	return term.exit, term.exited
}

func (term *terminal) kill() {
	err := syscall.Kill(-term.cmd.Process.Pid, syscall.SIGHUP)
	if err != nil && !errors.Is(err, syscall.ESRCH) {
		slog.Warn("hang up terminal", "pid", term.cmd.Process.Pid, "err", err)
	}
	term.ptmx.Close()
}

func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID string, params api.GetTaskTerminalParams) {
	t, _, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}
	size := api.TerminalSize{Cols: params.Cols, Rows: params.Rows}
	if !validSize(size) {
		writeError(w, http.StatusBadRequest, fmt.Errorf("cols and rows must be between 1 and %d", terminalMaxSize))
		return
	}

	conn, err := websocket.Accept(w, r, nil)
	if err != nil {
		slog.Warn("accept terminal websocket", "task", taskID, "err", err)
		return
	}
	defer conn.CloseNow()
	conn.SetReadLimit(terminalReadLimit)

	term, err := s.terminals.attach(t, size)
	if err != nil {
		slog.Error("start terminal", "task", taskID, "err", err)
		conn.Close(websocket.StatusInternalError, "couldn't start the shell; see byre's log")
		return
	}

	backlog, out := term.subscribe()
	if out == nil {
		exit, _ := term.exitStatus()
		conn.Close(terminalExited, exit)
		return
	}
	defer term.unsubscribe(out)

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	go func() {
		defer cancel()
		terminalInput(ctx, conn, term)
	}()

	if len(backlog) > 0 {
		err = conn.Write(ctx, websocket.MessageBinary, backlog)
		if err != nil {
			return
		}
	}
	for {
		select {
		case <-ctx.Done():
			return
		case b, open := <-out:
			if open {
				err = conn.Write(ctx, websocket.MessageBinary, b)
				if err != nil {
					return
				}
				continue
			}
			exit, exited := term.exitStatus()
			if exited {
				conn.Close(terminalExited, exit)
				return
			}
			conn.Close(websocket.StatusTryAgainLater, "fell behind the terminal's output")
			return
		}
	}
}

func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal) {
	for {
		typ, b, err := conn.Read(ctx)
		if err != nil {
			return
		}
		if typ == websocket.MessageBinary {
			_, err = term.ptmx.Write(b)
			if err != nil {
				slog.Debug("terminal input", "err", err)
			}
			continue
		}
		size := api.TerminalSize{}
		err = json.Unmarshal(b, &size)
		if err != nil || !validSize(size) {
			conn.Close(websocket.StatusUnsupportedData, "resize must be a TerminalSize")
			return
		}
		err = pty.Setsize(term.ptmx, winsize(size))
		if err != nil {
			slog.Debug("terminal resize", "err", err)
		}
	}
}