1package server
2
3import (
4 "encoding/base64"
5 "errors"
6 "fmt"
7 "io"
8 "net/http"
9 "os"
10 "os/exec"
11 "syscall"
12 "unsafe"
13
14 "github.com/coder/websocket"
15 "github.com/coder/websocket/wsjson"
16 "github.com/creack/pty"
17)
18
19// ExecMessage is the message format for terminal websocket communication
20type ExecMessage struct {
21 Type string `json:"type"`
22 Data string `json:"data,omitempty"`
23 Cols uint16 `json:"cols,omitempty"`
24 Rows uint16 `json:"rows,omitempty"`
25}
26
27// handleExecWS handles websocket connections for executing shell commands
28// Query params:
29// - cmd: the command to execute (required)
30// - cwd: working directory (optional, defaults to current dir)
31func (s *Server) handleExecWS(w http.ResponseWriter, r *http.Request) {
32 ctx := r.Context()
33
34 cmd := r.URL.Query().Get("cmd")
35 if cmd == "" {
36 http.Error(w, "cmd parameter required", http.StatusBadRequest)
37 return
38 }
39
40 cwd := r.URL.Query().Get("cwd")
41 if cwd == "" {
42 var err error
43 cwd, err = os.Getwd()
44 if err != nil {
45 cwd = "/"
46 }
47 }
48
49 // Upgrade to websocket
50 conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
51 CompressionMode: websocket.CompressionDisabled,
52 })
53 if err != nil {
54 s.logger.Error("Failed to upgrade websocket", "error", err)
55 return
56 }
57 defer conn.Close(websocket.StatusInternalError, "internal error")
58
59 // Wait for init message with terminal size
60 var initMsg ExecMessage
61 if err := wsjson.Read(ctx, conn, &initMsg); err != nil {
62 s.logger.Error("Failed to read init message", "error", err)
63 conn.Close(websocket.StatusPolicyViolation, "no init message")
64 return
65 }
66
67 if initMsg.Type != "init" {
68 conn.Close(websocket.StatusPolicyViolation, "expected init message")
69 return
70 }
71
72 cols := initMsg.Cols
73 rows := initMsg.Rows
74 if cols == 0 {
75 cols = 80
76 }
77 if rows == 0 {
78 rows = 24
79 }
80
81 // Create command
82 shellCmd := exec.CommandContext(ctx, "bash", "-c", cmd)
83 shellCmd.Dir = cwd
84 shellCmd.Env = append(os.Environ(), "TERM=xterm-256color")
85
86 // Start with pty
87 ptmx, err := pty.StartWithSize(shellCmd, &pty.Winsize{
88 Cols: cols,
89 Rows: rows,
90 })
91 if err != nil {
92 s.logger.Error("Failed to start command with pty", "error", err, "cmd", cmd)
93 errMsg := ExecMessage{
94 Type: "error",
95 Data: err.Error(),
96 }
97 wsjson.Write(ctx, conn, errMsg)
98 conn.Close(websocket.StatusInternalError, "failed to start command")
99 return
100 }
101 defer ptmx.Close()
102
103 // Channel to signal when process exits with all output sent
104 done := make(chan error, 1)
105
106 // Read from pty and send to websocket, then wait for process and signal done
107 go func() {
108 // First, read all output from pty
109 buf := make([]byte, 4096)
110 for {
111 n, err := ptmx.Read(buf)
112 if n > 0 {
113 msg := ExecMessage{
114 Type: "output",
115 Data: base64.StdEncoding.EncodeToString(buf[:n]),
116 }
117 if wsjson.Write(ctx, conn, msg) != nil {
118 return
119 }
120 }
121 if err != nil {
122 if !errors.Is(err, io.EOF) {
123 s.logger.Debug("PTY read error", "error", err)
124 }
125 break
126 }
127 }
128 // After pty is closed/EOF, wait for process exit status
129 err := shellCmd.Wait()
130 done <- err
131 }()
132
133 // Create a channel for incoming websocket messages
134 msgChan := make(chan ExecMessage)
135 errChan := make(chan error)
136
137 // Goroutine to read websocket messages
138 go func() {
139 for {
140 var msg ExecMessage
141 if err := wsjson.Read(ctx, conn, &msg); err != nil {
142 errChan <- err
143 return
144 }
145 msgChan <- msg
146 }
147 }()
148
149 // Process messages and handle exit
150 for {
151 select {
152 case <-ctx.Done():
153 if shellCmd.Process != nil {
154 shellCmd.Process.Kill()
155 }
156 return
157 case exitErr := <-done:
158 // Process exited
159 exitCode := 0
160 if exitErr != nil {
161 if exitError, ok := exitErr.(*exec.ExitError); ok {
162 exitCode = exitError.ExitCode()
163 }
164 }
165 exitMsg := ExecMessage{
166 Type: "exit",
167 Data: fmt.Sprintf("%d", exitCode),
168 }
169 s.logger.Info("Sending exit message", "exitCode", exitCode)
170 if err := wsjson.Write(ctx, conn, exitMsg); err != nil {
171 s.logger.Error("Failed to write exit message", "error", err)
172 }
173 s.logger.Info("Closing websocket normally")
174 conn.Close(websocket.StatusNormalClosure, "process exited")
175 return
176 case err := <-errChan:
177 if websocket.CloseStatus(err) != websocket.StatusNormalClosure {
178 s.logger.Debug("Websocket read error", "error", err)
179 }
180 if shellCmd.Process != nil {
181 shellCmd.Process.Kill()
182 }
183 return
184 case msg := <-msgChan:
185 switch msg.Type {
186 case "input":
187 if msg.Data != "" {
188 ptmx.Write([]byte(msg.Data))
189 }
190 case "resize":
191 if msg.Cols > 0 && msg.Rows > 0 {
192 setWinsize(ptmx, msg.Cols, msg.Rows)
193 }
194 }
195 }
196 }
197}
198
199// setWinsize sets the terminal window size
200func setWinsize(f *os.File, cols, rows uint16) error {
201 ws := struct {
202 Rows uint16
203 Cols uint16
204 X uint16
205 Y uint16
206 }{
207 Rows: rows,
208 Cols: cols,
209 }
210 _, _, errno := syscall.Syscall(
211 syscall.SYS_IOCTL,
212 f.Fd(),
213 uintptr(syscall.TIOCSWINSZ),
214 uintptr(unsafe.Pointer(&ws)),
215 )
216 if errno != 0 {
217 return errno
218 }
219 return nil
220}