1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

// Package terminal gives the browser an interactive shell: each WebSocket
// connection to its handler spawns one shell inside a pseudo-terminal and
// pipes bytes both ways, which is exactly what xterm.js expects to drive.
//
// Wire protocol: the server sends raw terminal output as binary frames; the
// client sends JSON text frames — {"type":"input","data":"…"} for keystrokes
// and {"type":"resize","cols":N,"rows":N} when the viewport changes.
package terminal

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"os/exec"

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

// Service spawns one shell per WebSocket connection.
type Service struct {
	// Shell is the command to run; when empty, $SHELL then bash then sh.
	Shell []string
	// Cwd is the directory each shell starts in.
	Cwd string

	logger *slog.Logger
}

// New creates a terminal service starting shells in cwd.
//
// Example:
//
//	mux.Handle("GET /ws/terminal", terminal.New(cfg.Cwd, logger).WebSocketHandler())
func New(cwd string, logger *slog.Logger) *Service {
	if logger == nil {
		logger = slog.Default()
	}
	return &Service{Cwd: cwd, logger: logger}
}

// controlMessage is what the client sends over text frames.
type controlMessage struct {
	Type string `json:"type"`
	Data string `json:"data,omitempty"`
	Cols uint16 `json:"cols,omitempty"`
	Rows uint16 `json:"rows,omitempty"`
}

// WebSocketHandler upgrades the request and runs one shell session until the
// socket closes or the shell exits.
func (s *Service) WebSocketHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
			OriginPatterns: []string{"localhost:*", "127.0.0.1:*"},
		})
		if err != nil {
			s.logger.Warn("terminal websocket accept failed", "error", err)
			return
		}
		defer func() { _ = conn.CloseNow() }()
		s.serveShell(r.Context(), conn)
	})
}

// serveShell runs the shell and the two pumps for one connection.
func (s *Service) serveShell(ctx context.Context, conn *websocket.Conn) {
	cmd := exec.Command(s.shellCommand()[0], s.shellCommand()[1:]...)
	cmd.Dir = s.Cwd
	cmd.Env = append(os.Environ(), "TERM=xterm-256color")

	ptmx, err := pty.Start(cmd)
	if err != nil {
		s.logger.Error("pty start failed", "error", err)
		_ = conn.Close(websocket.StatusInternalError, "pty start failed")
		return
	}
	defer func() {
		_ = ptmx.Close()
		_ = cmd.Process.Kill()
		_ = cmd.Wait()
	}()

	// Output pump: pty → socket, as binary frames.
	outputDone := make(chan struct{})
	go func() {
		defer close(outputDone)
		buffer := make([]byte, 8192)
		for {
			n, readErr := ptmx.Read(buffer)
			if n > 0 {
				if writeErr := conn.Write(ctx, websocket.MessageBinary, buffer[:n]); writeErr != nil {
					return
				}
			}
			if readErr != nil {
				// The shell exited (or the pty closed): tell the client.
				_ = conn.Close(websocket.StatusNormalClosure, "shell exited")
				return
			}
		}
	}()

	// Input pump: socket → pty, until the client disconnects.
	s.readControlMessages(ctx, conn, ptmx)
	<-outputDone
}

// readControlMessages feeds client frames into the pty.
func (s *Service) readControlMessages(ctx context.Context, conn *websocket.Conn, ptmx *os.File) {
	for {
		kind, data, err := conn.Read(ctx)
		if err != nil {
			return
		}
		if kind != websocket.MessageText {
			continue
		}
		var msg controlMessage
		if err := json.Unmarshal(data, &msg); err != nil {
			continue // a malformed frame must not kill the session
		}
		switch msg.Type {
		case "input":
			_, _ = ptmx.WriteString(msg.Data)
		case "resize":
			if msg.Cols > 0 && msg.Rows > 0 {
				_ = pty.Setsize(ptmx, &pty.Winsize{Cols: msg.Cols, Rows: msg.Rows})
			}
		}
	}
}

// shellCommand picks the shell to run: the configured one, else $SHELL, else
// bash, else sh.
func (s *Service) shellCommand() []string {
	if len(s.Shell) > 0 {
		return s.Shell
	}
	if shell := os.Getenv("SHELL"); shell != "" {
		return []string{shell, "-l"}
	}
	if path, err := exec.LookPath("bash"); err == nil {
		return []string{path, "-l"}
	}
	return []string{"sh"}
}