package session
import (
"bytes"
"cmp"
"crypto/rand"
"encoding/base32"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/url"
"os"
"os/exec"
"os/signal"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"syscall"
"time"
"github.com/creack/pty"
xterm "github.com/gitpod-io/xterm-go"
"github.com/samber/lo"
"golang.org/x/term"
"tobi/internal/protocol"
"tobi/internal/target"
"tobi/internal/termcolor"
)
var (
ErrDetached = errors.New("session detached")
b32 = base32.NewEncoding("abcdefghijklmnopqrstuvwxyz234567").WithPadding(base32.NoPadding)
)
// how long an owner outlives its last viewer; atomic because a test shortens
// it while owners from earlier tests are still running
var detachedTTL atomic.Int64
func init() { detachedTTL.Store(int64(24 * time.Hour)) }
type Session struct {
ID string `json:"id"`
Host string `json:"host"`
Command string `json:"command"`
Cwd string `json:"cwd"`
StartedAt time.Time `json:"started_at"`
Attached bool `json:"attached,omitempty"`
}
func (s Session) Canonical() string { return s.Host + "." + s.ID }
type HopRequest struct {
Target string
}
func (h HopRequest) Error() string { return "hop to " + h.Target }
func ValidateID(id string) bool {
return len(id) >= 1 && len(id) <= 64 && strings.TrimLeft(id, "abcdefghijklmnopqrstuvwxyz0123456789-_") == ""
}
func GenerateID() string {
b := make([]byte, 10)
rand.Read(b)
return b32.EncodeToString(b)
}
func RuntimeDir() string {
base := filepath.Join(os.TempDir(), fmt.Sprintf("tobi-%d", os.Getuid()))
if xdg := cmp.Or(os.Getenv("XDG_RUNTIME_DIR"), loginRuntimeDir()); xdg != "" {
base = filepath.Join(xdg, "tobi")
}
if fi, err := os.Lstat(base); err == nil && (fi.Mode()&os.ModeSymlink != 0 || int(fi.Sys().(*syscall.Stat_t).Uid) != os.Getuid()) {
_ = os.RemoveAll(base)
}
_ = os.MkdirAll(base, 0700)
_ = os.Chmod(base, 0700)
return base
}
func loginRuntimeDir() string {
d := fmt.Sprintf("/run/user/%d", os.Getuid())
if fi, err := os.Stat(d); err == nil && fi.IsDir() && int(fi.Sys().(*syscall.Stat_t).Uid) == os.Getuid() {
return d
}
return ""
}
func SocketPath(id string) string {
if !ValidateID(id) {
return ""
}
return filepath.Join(RuntimeDir(), id+".sock")
}
func IsLocal(host string) bool {
return strings.EqualFold(host, ShortHostname()) || strings.EqualFold(host, "localhost")
}
func ShortHostname() string {
h, _ := os.Hostname()
h, _, _ = strings.Cut(h, ".")
return lo.Ternary(h != "", h, "localhost")
}
func EnsureOwner(id string, argv []string) error {
if !ValidateID(id) {
return fmt.Errorf("invalid session id: %q", id)
}
if len(argv) == 0 {
return errors.New("empty command")
}
sockPath := SocketPath(id)
if answers(sockPath, 30*time.Millisecond) {
return nil
}
if _, err := exec.LookPath(argv[0]); err != nil {
return err
}
cmd := exec.Command(Executable(), argv...)
cmd.Env = append(os.Environ(), "TOBI_OWNER=1", "TOBI_ID="+id)
cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
// the owner is detached, so its stderr is the only way it can explain itself
var stderr syncBuffer
cmd.Stderr = &stderr
if err := cmd.Start(); err != nil {
return fmt.Errorf("start owner: %w", err)
}
exited := make(chan struct{})
go func() { _ = cmd.Wait(); close(exited) }()
fail := func(reason string) error {
if msg := stderr.lastLine(); msg != "" {
return fmt.Errorf("owner for session %q %s: %s", id, reason, msg)
}
return fmt.Errorf("owner for session %q %s", id, reason)
}
for deadline := time.Now().Add(2 * time.Second); time.Now().Before(deadline); time.Sleep(10 * time.Millisecond) {
if answers(sockPath, 20*time.Millisecond) {
return nil
}
select {
case <-exited:
return fail("exited before it was ready")
default:
}
}
return fail("did not start")
}
// Executable is this binary, to re-run as an owner or ship to a remote
func Executable() string {
if self, err := os.Executable(); err == nil {
return self
}
return os.Args[0]
}
// read while the child may still be writing, so the copier must not be raced
type syncBuffer struct {
mu sync.Mutex
buf bytes.Buffer
}
func (b *syncBuffer) Write(p []byte) (int, error) {
b.mu.Lock()
defer b.mu.Unlock()
if b.buf.Len() > 8<<10 {
return len(p), nil
}
return b.buf.Write(p)
}
func (b *syncBuffer) lastLine() string {
b.mu.Lock()
defer b.mu.Unlock()
lines := strings.Split(strings.TrimSpace(b.buf.String()), "\n")
return strings.TrimSpace(lines[len(lines)-1])
}
func ListLocal() []Session {
dir := RuntimeDir()
entries, _ := os.ReadDir(dir)
var list []Session
for _, e := range entries {
id, ok := strings.CutSuffix(e.Name(), ".sock")
if !ok || e.IsDir() || !ValidateID(id) {
continue
}
conn, err := net.DialTimeout("unix", filepath.Join(dir, e.Name()), 250*time.Millisecond)
if err != nil {
_ = os.Remove(filepath.Join(dir, e.Name()))
continue
}
_ = conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
_ = protocol.WriteFrame(conn, protocol.TypeInfo, nil)
t, p, err := protocol.ReadFrame(conn)
_ = conn.Close()
var s Session
if err == nil && t == protocol.TypeInfoResp && json.Unmarshal(p, &s) == nil {
list = append(list, s)
}
}
return list
}
func Attach(stream io.ReadWriteCloser, inputCh <-chan []byte, colors map[int]termcolor.Color) error {
defer stream.Close()
size := func(t protocol.FrameType) error {
c, r, _ := term.GetSize(int(os.Stdout.Fd()))
return protocol.WriteFrame(stream, t, append(protocol.EncodeSize(c, r), termcolor.EncodeColors(colors)...))
}
if err := size(protocol.TypeAttach); err != nil {
return err
}
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGWINCH)
defer signal.Stop(sigCh)
done := make(chan struct{})
defer close(done)
go func() {
for {
select {
case input, ok := <-inputCh:
if !ok || protocol.WriteFrame(stream, protocol.TypeInput, input) != nil {
return
}
case <-sigCh:
if size(protocol.TypeResize) != nil {
return
}
case <-done:
return
}
}
}()
for {
t, payload, err := protocol.ReadFrame(stream)
if err != nil {
return ErrDetached
}
switch t {
case protocol.TypeSnapshot, protocol.TypeOutput:
_, _ = os.Stdout.Write(payload)
case protocol.TypeHop:
return HopRequest{Target: string(payload)}
}
}
}
func SendHop(controlSock, target string) error {
conn, err := net.Dial("unix", controlSock)
if err != nil {
return err
}
defer conn.Close()
return protocol.WriteFrame(conn, protocol.TypeHop, []byte(target))
}
// an owner that predates TypeKill just hangs up, so one that still answers gets a signal instead
func Kill(id string) error {
sock := SocketPath(id)
if sock == "" {
return fmt.Errorf("invalid session id: %q", id)
}
conn, err := net.DialTimeout("unix", sock, 250*time.Millisecond)
if err != nil {
return fmt.Errorf("session %q is not running", id)
}
pid := peerPID(conn)
_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
err = protocol.WriteFrame(conn, protocol.TypeKill, nil)
_, _ = io.Copy(io.Discard, conn)
_ = conn.Close()
switch {
case err != nil:
return err
case !answers(sock, 250*time.Millisecond):
return nil
case pid <= 0:
return fmt.Errorf("session %q outlived the kill and its owner pid is unknown", id)
}
if err := syscall.Kill(pid, syscall.SIGTERM); err != nil {
return fmt.Errorf("signal the owner of %q: %w", id, err)
}
for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); time.Sleep(20 * time.Millisecond) {
if !answers(sock, 250*time.Millisecond) {
return nil
}
}
return fmt.Errorf("session %q outlived SIGTERM", id)
}
func answers(sock string, timeout time.Duration) bool {
c, err := net.DialTimeout("unix", sock, timeout)
if err == nil {
_ = c.Close()
}
return err == nil
}
func peerPID(c net.Conn) int {
uc, ok := c.(*net.UnixConn)
if !ok {
return 0
}
raw, err := uc.SyscallConn()
if err != nil {
return 0
}
pid := 0
_ = raw.Control(func(fd uintptr) { pid = sockPeerPID(int(fd)) })
return pid
}
type outMsg struct {
t protocol.FrameType
p []byte
}
type viewerClient struct {
conn net.Conn
ch chan outMsg
done chan struct{}
once sync.Once
cols int
rows int
colors map[int]termcolor.Color
}
// send queues msg, dropping a viewer too slow to keep up
func (vc *viewerClient) send(msg outMsg) {
select {
case vc.ch <- msg:
default:
vc.Close()
}
}
func (vc *viewerClient) Close() {
vc.once.Do(func() {
close(vc.done)
_ = vc.conn.Close()
})
}
func pump(vc *viewerClient, snap []byte) {
_ = protocol.WriteFrame(vc.conn, protocol.TypeSnapshot, snap)
for {
select {
case <-vc.done:
return
case msg := <-vc.ch:
if protocol.WriteFrame(vc.conn, msg.t, msg.p) != nil {
vc.Close()
return
}
}
}
}
type Owner struct {
sess Session
ptyFile *os.File
vt *xterm.Terminal
addon *xterm.SerializeAddon
cmd *exec.Cmd
listener net.Listener
mu sync.Mutex
viewers map[*viewerClient]struct{}
// the viewer that sent input last, so a paste lands on the terminal it came from
active atomic.Pointer[viewerClient]
// the osc color handler reads colorSets while mu is held for vt writes,
// so it gets its own lock
colorMu sync.Mutex
colorSets map[*viewerClient]map[int]termcolor.Color
idleTimer *time.Timer
// set once a hop has moved viewers elsewhere: the shell that asked for it
// only ran tobi, so the owner goes with its last viewer
retiring bool
closed bool
}
func (o *Owner) applySizeLocked() {
cols, rows := 0, 0
for vc := range o.viewers {
if cols == 0 || vc.cols < cols {
cols = vc.cols
}
if rows == 0 || vc.rows < rows {
rows = vc.rows
}
}
if cols < 1 || rows < 1 {
return
}
_ = pty.Setsize(o.ptyFile, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows)})
o.vt.Resize(cols, rows)
}
func (o *Owner) outputLocked(p []byte) {
for vc := range o.viewers {
vc.send(outMsg{t: protocol.TypeOutput, p: p})
}
}
// activeOutputLocked sends p to the viewer that typed last, or to everyone when
// that viewer is gone or nobody has typed yet
func (o *Owner) activeOutputLocked(p []byte) {
if vc := o.active.Load(); vc != nil {
if _, ok := o.viewers[vc]; ok {
vc.send(outMsg{t: protocol.TypeOutput, p: p})
return
}
}
o.outputLocked(p)
}
func (o *Owner) hop(p []byte) {
o.mu.Lock()
viewers := lo.Keys(o.viewers)
if t := target.Parse(string(p)); len(viewers) > 0 && !(t.ID == o.sess.ID && IsLocal(t.Host)) {
o.retiring = true
}
o.mu.Unlock()
for _, vc := range viewers {
vc.send(outMsg{t: protocol.TypeHop, p: p})
}
}
func (o *Owner) addViewer(conn net.Conn, cols, rows int, colors map[int]termcolor.Color) (*viewerClient, []byte) {
vc := &viewerClient{conn: conn, ch: make(chan outMsg, 128), done: make(chan struct{}), cols: cols, rows: rows, colors: colors}
o.mu.Lock()
defer o.mu.Unlock()
if o.closed {
return nil, nil
}
if o.idleTimer != nil {
o.idleTimer.Stop()
o.idleTimer = nil
}
o.viewers[vc] = struct{}{}
o.colorMu.Lock()
o.colorSets[vc] = colors
o.colorMu.Unlock()
o.applySizeLocked()
return vc, append(o.addon.Serialize(nil), mouseEncodings[o.vt.DecPrivateModes().MouseEncoding]...)
}
// the serializer restores mouse tracking but not its encoding, so a new viewer would report in X10
var mouseEncodings = map[string]string{"SGR": "\x1b[?1006h", "SGR_PIXELS": "\x1b[?1016h"}
func (o *Owner) removeViewer(vc *viewerClient) {
o.mu.Lock()
defer o.mu.Unlock()
delete(o.viewers, vc)
o.colorMu.Lock()
delete(o.colorSets, vc)
o.colorMu.Unlock()
switch {
case len(o.viewers) > 0 || o.closed:
case o.retiring:
go o.Close()
default:
o.idleTimer = time.AfterFunc(time.Duration(detachedTTL.Load()), o.Close)
}
o.applySizeLocked()
}
var csiQueries = map[string]struct{}{
"\x1b[?6n": {}, "\x1b[6n": {}, "\x1b[=0c": {}, "\x1b[=c": {},
"\x1b[>0c": {}, "\x1b[>c": {}, "\x1b[0c": {}, "\x1b[c": {},
"\x1b[?2026$p": {}, "\x1b[>0q": {}, "\x1b[>q": {}, "\x1b[?u": {},
}
func terminalQuery(seq []byte) bool {
if len(seq) < 3 {
return false
}
if seq[1] == '[' {
_, ok := csiQueries[string(seq)]
return ok
}
if seq[1] != ']' {
return false
}
end := len(seq) - 1
if seq[end] == '\a' {
seq = seq[2:end]
} else {
seq = seq[2 : end-1] // querySpan only calls us for a complete ST
}
s := string(seq)
if s == "10;?" || s == "11;?" || s == "12;?" {
return true
}
if !strings.HasPrefix(s, "4;") {
return false
}
parts := strings.Split(s[2:], ";")
if len(parts) < 2 || len(parts)%2 != 0 {
return false
}
for i := 1; i < len(parts); i += 2 {
if parts[i] != "?" {
return false
}
}
return true
}
func querySpan(b []byte) (int, bool) {
if len(b) < 2 {
return 0, false
}
switch b[1] {
case '[':
i := 2
for i < len(b) && b[i] >= 0x20 && b[i] <= 0x3f {
i++
}
if i == len(b) {
return 0, false
}
if b[i] < 0x40 || b[i] > 0x7e {
return 1, false
}
return i + 1, terminalQuery(b[:i+1])
case ']':
for i := 2; i < len(b); i++ {
if b[i] == '\a' {
return i + 1, terminalQuery(b[:i+1])
}
if b[i] != '\x1b' {
continue
}
if i+1 == len(b) {
return 0, false
}
if b[i+1] == '\\' {
return i + 2, terminalQuery(b[:i+2])
}
return i, true
}
default:
return 2, false
}
return 0, false
}
type queryFilter struct{ held []byte }
// what one chunk of pty output becomes
type filtered struct {
emulator []byte // the owner's own terminal state, which answers the queries
shared []byte // every viewer
routed []byte // only the active viewer
}
const clipboardOSC = "\x1b]5522;"
// clipboardProbe asks whether the terminal speaks the kitty clipboard protocol.
const clipboardProbe = "\x1b[?5522$p"
// routedToActive reports whether seq belongs to the terminal the user is typing
// into. a clipboard (osc 5522) request prompts there and the reply goes into
// the pty, which two replying terminals would corrupt. the probe is that
// terminal's to answer, not the emulator's, which would report the mode unknown.
func routedToActive(seq []byte) bool {
return bytes.HasPrefix(seq, []byte(clipboardOSC)) || string(seq) == clipboardProbe
}
// Feed splits pty output. terminal queries stay out of shared because the
// emulator answers them; routed sequences stay out of the emulator.
func (f *queryFilter) Feed(chunk []byte) (out filtered) {
f.held = append(f.held, chunk...)
for {
i := bytes.IndexByte(f.held, 0x1b)
if i < 0 {
out.emulator = append(out.emulator, f.held...)
out.shared = append(out.shared, f.held...)
f.held = f.held[:0]
return out
}
out.emulator = append(out.emulator, f.held[:i]...)
out.shared = append(out.shared, f.held[:i]...)
f.held = f.held[i:]
n, query := querySpan(f.held)
if n == 0 && len(f.held) <= 64<<10 {
return out
}
if n == 0 {
n, query = len(f.held), true
}
seq := f.held[:n]
switch {
case query:
out.emulator = append(out.emulator, seq...)
case routedToActive(seq):
out.routed = append(out.routed, seq...)
default:
out.emulator = append(out.emulator, seq...)
out.shared = append(out.shared, seq...)
}
f.held = f.held[n:]
}
}
func (o *Owner) readPTYLoop() {
var filter queryFilter
buf := make([]byte, 4096)
for {
n, err := o.ptyFile.Read(buf)
if n > 0 {
o.mu.Lock()
out := filter.Feed(buf[:n])
_, _ = o.vt.Write(out.emulator)
if len(out.shared) > 0 {
o.outputLocked(out.shared)
}
if len(out.routed) > 0 {
o.activeOutputLocked(out.routed)
}
o.mu.Unlock()
}
if err != nil {
o.Close()
return
}
}
}
func (o *Owner) handleConn(conn net.Conn) {
t, payload, err := protocol.ReadFrame(conn)
if err != nil {
_ = conn.Close()
return
}
if t != protocol.TypeAttach {
switch t {
case protocol.TypeInfo:
o.mu.Lock()
s := o.sess
s.Attached = len(o.viewers) > 0
data, _ := json.Marshal(s)
o.mu.Unlock()
_ = protocol.WriteFrame(conn, protocol.TypeInfoResp, data)
case protocol.TypeHop:
o.hop(payload)
case protocol.TypeKill:
o.Close()
}
_ = conn.Close()
return
}
cols, rows, _ := protocol.DecodeSize(payload)
colors, _ := termcolor.DecodeColors(payload[min(4, len(payload)):])
vc, snap := o.addViewer(conn, cmp.Or(cols, 80), cmp.Or(rows, 24), colors)
if vc == nil {
_ = conn.Close()
return
}
go pump(vc, snap)
for {
ft, fp, err := protocol.ReadFrame(conn)
if err != nil {
break
}
switch ft {
case protocol.TypeInput:
o.active.Store(vc)
_, _ = o.ptyFile.Write(fp)
case protocol.TypeResize:
if c, r, err := protocol.DecodeSize(fp); err == nil && c >= 1 && r >= 1 {
o.mu.Lock()
vc.cols, vc.rows = c, r
o.applySizeLocked()
o.mu.Unlock()
}
}
}
o.removeViewer(vc)
vc.Close()
}
func (o *Owner) Close() {
o.mu.Lock()
defer o.mu.Unlock()
if o.closed {
return
}
o.closed = true
if o.idleTimer != nil {
o.idleTimer.Stop()
}
for vc := range o.viewers {
vc.Close()
}
_ = o.listener.Close()
_ = o.ptyFile.Close()
_ = o.cmd.Process.Kill()
o.vt.Dispose()
_ = os.Remove(o.listener.Addr().String())
}
func RunOwner(id string, command []string) error {
if !ValidateID(id) || len(command) == 0 {
return errors.New("invalid id or command")
}
sockPath := SocketPath(id)
_ = os.Remove(sockPath)
listener, err := net.Listen("unix", sockPath)
if err != nil {
return err
}
_ = os.Chmod(sockPath, 0700)
vt := xterm.New(xterm.WithCols(80), xterm.WithRows(24), xterm.WithScrollback(5000))
cmd := exec.Command(command[0], command[1:]...)
env := lo.Reject(os.Environ(), func(kv string, _ int) bool {
return strings.HasPrefix(kv, "CI=")
})
if t := os.Getenv("TERM"); t == "" || t == "dumb" {
env = append(env, "TERM=xterm-256color", "COLORTERM=truecolor")
}
cmd.Env = append(env, "TOBI_SESSION="+id, "TOBI_CONTROL="+sockPath)
ptyFile, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: 80, Rows: 24})
if err != nil {
_ = listener.Close()
_ = os.Remove(sockPath)
return err
}
cwd, _ := os.Getwd()
owner := &Owner{
sess: Session{
ID: id,
Host: ShortHostname(),
Command: strings.Join(command, " "),
Cwd: cwd,
StartedAt: time.Now(),
},
ptyFile: ptyFile,
vt: vt,
addon: xterm.NewSerializeAddon(vt),
cmd: cmd,
listener: listener,
viewers: map[*viewerClient]struct{}{},
colorSets: map[*viewerClient]map[int]termcolor.Color{},
}
vt.OnTitleChange(func(title string) {
if title != "" {
owner.sess.Command = title
}
})
vt.RegisterOscHandler(7, xterm.NewOscStringHandler(func(payload string) bool {
if u, err := url.Parse(payload); err == nil && u.Scheme == "file" && u.Path != "" {
owner.sess.Cwd = u.Path
}
return true
}))
vt.OnData(func(s string) {
_, _ = owner.ptyFile.Write([]byte(s))
})
vt.OnColor(func(events []xterm.ColorEvent) {
for _, ev := range events {
if ev.Type != xterm.ColorRequestReport || ev.Index < 0 || ev.Index > int(xterm.SpecialColorCursor) {
continue
}
owner.colorMu.Lock()
sets := lo.Values(owner.colorSets)
owner.colorMu.Unlock()
if c, ok := termcolor.Agreed(ev.Index, sets); ok {
fmt.Fprint(owner.ptyFile, termcolor.Reply(ev.Index, c))
}
}
})
owner.idleTimer = time.AfterFunc(time.Duration(detachedTTL.Load()), owner.Close)
go owner.readPTYLoop()
go func() { _ = cmd.Wait() }()
for {
conn, err := listener.Accept()
if err != nil {
return nil
}
go owner.handleConn(conn)
}
}