1package sshstdio
2
3import (
4 "bytes"
5 "context"
6 "encoding/binary"
7 "errors"
8 "io"
9 "log/slog"
10 "os/exec"
11 "strings"
12 "sync"
13 "testing"
14 "time"
15
16 "github.com/bnema/vev/internal/domain"
17 "github.com/bnema/vev/internal/ports"
18)
19
20func TestBuildCommandUsesExecArgs(t *testing.T) {
21 // BuildCommand deliberately starts `ssh -- target 'vev' '_stdio'`.
22 // Session selection and terminal color capability are carried in vev's Hello
23 // message, not in ssh env.
24 tests := []struct {
25 name string
26 target string
27 session string
28 want []string
29 }{
30 {name: "no session", target: "user@example.com", want: []string{"--", "user@example.com", "'vev' '_stdio'"}},
31 {name: "session is carried by Hello rather than ssh command", target: "user@example.com", session: "work; rm -rf /", want: []string{"--", "user@example.com", "'vev' '_stdio'"}},
32 {name: "target kept as single local arg after option terminator", target: "user@host; touch /tmp/pwn", session: "work", want: []string{"--", "user@host; touch /tmp/pwn", "'vev' '_stdio'"}},
33 }
34
35 for _, tt := range tests {
36 t.Run(tt.name, func(t *testing.T) {
37 got := BuildCommand(tt.target, tt.session)
38 if got.Path != "ssh" {
39 t.Fatalf("Path = %q, want ssh", got.Path)
40 }
41 if len(got.Args) != len(tt.want) {
42 t.Fatalf("Args len = %d, want %d (%q)", len(got.Args), len(tt.want), got.Args)
43 }
44 for i := range tt.want {
45 if got.Args[i] != tt.want[i] {
46 t.Fatalf("Args[%d] = %q, want %q (all args %q)", i, got.Args[i], tt.want[i], got.Args)
47 }
48 }
49 })
50 }
51}
52
53func TestBuildCommandForModeUsesCanonicalSSHArgs(t *testing.T) {
54 got := BuildCommandForMode("user@example.com", "_udp-bootstrap", "work")
55 want := []string{"--", "user@example.com", "'vev' '_udp-bootstrap'"}
56 if got.Path != "ssh" {
57 t.Fatalf("Path = %q, want ssh", got.Path)
58 }
59 for i := range want {
60 if got.Args[i] != want[i] {
61 t.Fatalf("Args[%d] = %q, want %q (all args %q)", i, got.Args[i], want[i], got.Args)
62 }
63 }
64}
65
66func TestBuildCommandForRemoteCommandQuotesEveryWord(t *testing.T) {
67 tests := []struct {
68 name string
69 target string
70 command []string
71 want []string
72 }{
73 {
74 name: "catalog command",
75 target: "arch",
76 command: []string{"vev", "cmd", "remote-catalog", "--json"},
77 want: []string{"--", "arch", "'vev' 'cmd' 'remote-catalog' '--json'"},
78 },
79 {
80 name: "metacharacters in remote words are quoted",
81 target: "user@host; touch /tmp/pwn",
82 command: []string{"vev", "cmd", "remote-catalog; rm -rf /", "--json"},
83 want: []string{"--", "user@host; touch /tmp/pwn", "'vev' 'cmd' 'remote-catalog; rm -rf /' '--json'"},
84 },
85 {
86 name: "single quotes are posix escaped",
87 target: "arch",
88 command: []string{"it's", "fine"},
89 want: []string{"--", "arch", "'it'\\''s' 'fine'"},
90 },
91 }
92
93 for _, tt := range tests {
94 t.Run(tt.name, func(t *testing.T) {
95 got := BuildCommandForRemoteCommand(tt.target, tt.command...)
96 if got.Path != "ssh" {
97 t.Fatalf("Path = %q, want ssh", got.Path)
98 }
99 if len(got.Args) != len(tt.want) {
100 t.Fatalf("Args len = %d, want %d (%q)", len(got.Args), len(tt.want), got.Args)
101 }
102 for i := range tt.want {
103 if got.Args[i] != tt.want[i] {
104 t.Fatalf("Args[%d] = %q, want %q (all args %q)", i, got.Args[i], tt.want[i], got.Args)
105 }
106 }
107 })
108 }
109}
110
111func TestCloseInterruptsBlockedSend(t *testing.T) {
112 writer := &blockedWriter{started: make(chan struct{}), released: make(chan struct{})}
113 transport := NewTransport(nil, writer, nil)
114
115 sendDone := make(chan error, 1)
116 go func() { sendDone <- transport.Send(ports.Frame{Type: ports.MsgOutput, Payload: []byte("blocked")}) }()
117 <-writer.started
118 if err := transport.Close(); err != nil {
119 t.Fatalf("Close() error = %v", err)
120 }
121 if err := <-sendDone; err == nil {
122 t.Fatal("Send() error = nil after Close()")
123 }
124}
125
126type blockedWriter struct {
127 started chan struct{}
128 released chan struct{}
129 once sync.Once
130}
131
132func (w *blockedWriter) Write([]byte) (int, error) {
133 w.once.Do(func() { close(w.started) })
134 <-w.released
135 return 0, io.ErrClosedPipe
136}
137
138func (w *blockedWriter) Close() error {
139 select {
140 case <-w.released:
141 default:
142 close(w.released)
143 }
144 return nil
145}
146
147func TestTransportRoundTripAndVersionMismatchFrame(t *testing.T) {
148 clientRead, serverWrite := io.Pipe()
149 serverRead, clientWrite := io.Pipe()
150 client := NewTransport(clientRead, clientWrite, func() error {
151 _ = clientRead.Close()
152 return clientWrite.Close()
153 })
154 server := NewTransport(serverRead, serverWrite, func() error {
155 _ = serverRead.Close()
156 return serverWrite.Close()
157 })
158 defer func() { _ = client.Close() }()
159 defer func() { _ = server.Close() }()
160
161 hello := ports.Hello{Version: ports.ProtocolVersion + 1, Intent: ports.IntentAttach, Name: "work", Size: domain.Size{Cols: 80, Rows: 24}}
162 go func() {
163 f, err := server.Recv()
164 if err != nil {
165 t.Errorf("server Recv: %v", err)
166 return
167 }
168 got, err := ports.UnmarshalHello(f.Payload)
169 if err != nil {
170 t.Errorf("UnmarshalHello: %v", err)
171 return
172 }
173 if got.Version == ports.ProtocolVersion {
174 t.Errorf("test did not send a mismatched version")
175 }
176 _ = server.Send(ports.Frame{Type: ports.MsgError, Payload: ports.MarshalErrorMsg(ports.ErrorMsg{Code: ports.ErrVersionMismatch, Text: "protocol version mismatch"})})
177 }()
178
179 if err := client.Send(ports.Frame{Type: ports.MsgHello, Payload: ports.MarshalHello(hello)}); err != nil {
180 t.Fatalf("client Send: %v", err)
181 }
182 reply, err := client.Recv()
183 if err != nil {
184 t.Fatalf("client Recv: %v", err)
185 }
186 if reply.Type != ports.MsgError {
187 t.Fatalf("reply type = %d, want MsgError", reply.Type)
188 }
189 em, err := ports.UnmarshalErrorMsg(reply.Payload)
190 if err != nil {
191 t.Fatalf("UnmarshalErrorMsg: %v", err)
192 }
193 if em.Code != ports.ErrVersionMismatch {
194 t.Fatalf("error code = %d, want ErrVersionMismatch", em.Code)
195 }
196}
197
198func TestTransportRejectsZeroLengthFrame(t *testing.T) {
199 tr := NewTransport(strings.NewReader("\x00\x00\x00\x00"), io.Discard, nil)
200 _, err := tr.Recv()
201 if !errors.Is(err, ErrZeroLengthFrame) {
202 t.Fatalf("Recv error = %v, want ErrZeroLengthFrame", err)
203 }
204}
205
206func TestTransportUsesCanonicalFrameMaximum(t *testing.T) {
207 boundaryPayload := make([]byte, ports.MaxFrameLen-1)
208 boundaryWire := &bytes.Buffer{}
209 boundarySend := NewTransport(nil, boundaryWire, nil)
210 if err := boundarySend.Send(ports.Frame{Type: ports.MsgOutput, Payload: boundaryPayload}); err != nil {
211 t.Fatalf("boundary Send error = %v", err)
212 }
213 if got := binary.BigEndian.Uint32(boundaryWire.Bytes()[:frameHeaderLen]); got != ports.MaxFrameLen {
214 t.Fatalf("boundary frame length = %d, want %d", got, ports.MaxFrameLen)
215 }
216 boundaryRecv := NewTransport(bytes.NewReader(boundaryWire.Bytes()), io.Discard, nil)
217 boundaryFrame, err := boundaryRecv.Recv()
218 if err != nil {
219 t.Fatalf("boundary Recv error = %v", err)
220 }
221 if boundaryFrame.Type != ports.MsgOutput || len(boundaryFrame.Payload) != len(boundaryPayload) {
222 t.Fatalf("boundary frame = type %d, payload %d bytes; want type %d, payload %d bytes", boundaryFrame.Type, len(boundaryFrame.Payload), ports.MsgOutput, len(boundaryPayload))
223 }
224 if !bytes.Equal(boundaryFrame.Payload, boundaryPayload) {
225 t.Fatal("boundary payload was corrupted")
226 }
227
228 send := NewTransport(nil, io.Discard, nil)
229 err = send.Send(ports.Frame{Type: ports.MsgOutput, Payload: make([]byte, ports.MaxFrameLen)})
230 if !errors.Is(err, ErrFrameTooLarge) {
231 t.Fatalf("Send error = %v, want ErrFrameTooLarge", err)
232 }
233
234 var header [frameHeaderLen]byte
235 binary.BigEndian.PutUint32(header[:], ports.MaxFrameLen+1)
236 recv := NewTransport(bytes.NewReader(header[:]), io.Discard, nil)
237 _, err = recv.Recv()
238 if !errors.Is(err, ErrFrameTooLarge) {
239 t.Fatalf("Recv error = %v, want ErrFrameTooLarge", err)
240 }
241}
242
243func TestDialContextCanceledBeforeStart(t *testing.T) {
244 ctx, cancel := context.WithCancel(context.Background())
245 cancel()
246 _, err := DialContext(ctx, "example.com", "work")
247 if !errors.Is(err, context.Canceled) {
248 t.Fatalf("DialContext error = %v, want context.Canceled", err)
249 }
250}
251
252func TestRecvMapsUnexpectedEOFDuringHeaderToSSHExit(t *testing.T) {
253 sshErr := errors.New("ssh exited with status 255")
254 tr := newTransport(strings.NewReader("\x00"), io.Discard, nil, func() error { return sshErr })
255
256 _, err := tr.Recv()
257 if !errors.Is(err, sshErr) {
258 t.Fatalf("Recv error = %v, want %v", err, sshErr)
259 }
260}
261
262func TestRecvReportsSSHExitWhenProcessClosesBeforeFrame(t *testing.T) {
263 cmd := exec.CommandContext(context.Background(), "sh", "-c", "echo connection refused >&2; exit 255")
264 stdin, err := cmd.StdinPipe()
265 if err != nil {
266 t.Fatalf("StdinPipe: %v", err)
267 }
268 stdout, err := cmd.StdoutPipe()
269 if err != nil {
270 t.Fatalf("StdoutPipe: %v", err)
271 }
272 var stderr bytes.Buffer
273 cmd.Stderr = &stderr
274 if err := cmd.Start(); err != nil {
275 t.Fatalf("Start: %v", err)
276 }
277
278 waiter := newProcessWaiter(cmd, stdin, &stderr, time.Second, nil, "user@example.com", "work")
279 tr := newTransport(stdout, stdin, waiter.close, waiter.eofErr)
280
281 _, err = tr.Recv()
282 if err == nil {
283 t.Fatal("Recv error = nil, want ssh exit error")
284 }
285 for _, want := range []string{"sshstdio: ssh exited:", "connection refused"} {
286 if !strings.Contains(err.Error(), want) {
287 t.Fatalf("Recv error = %q, want substring %q", err, want)
288 }
289 }
290}
291
292func TestProcessCloserLogsNonCleanExitStderr(t *testing.T) {
293 cmd := exec.Command("sh", "-c", "echo remote failure >&2; exit 7")
294 stdin, err := cmd.StdinPipe()
295 if err != nil {
296 t.Fatalf("StdinPipe: %v", err)
297 }
298 var stderr bytes.Buffer
299 cmd.Stderr = &stderr
300 if err := cmd.Start(); err != nil {
301 t.Fatalf("Start: %v", err)
302 }
303
304 var logBuf bytes.Buffer
305 log := slog.New(slog.NewJSONHandler(&logBuf, nil))
306 err = newProcessCloser(cmd, stdin, &stderr, time.Second, log, "user@example.com", "work")()
307 if err == nil {
308 t.Fatal("Close error = nil, want non-clean ssh exit")
309 }
310 entry := logBuf.String()
311 for _, want := range []string{"ssh exited non-cleanly", "remote failure", "user@example.com", "work"} {
312 if !strings.Contains(entry, want) {
313 t.Fatalf("log entry = %q, want substring %q", entry, want)
314 }
315 }
316 if strings.Contains(entry, "'vev' '_stdio'") || strings.Contains(entry, "-- user@example.com") {
317 t.Fatalf("log entry includes generated command line: %q", entry)
318 }
319}
320
321func TestProcessCloser(t *testing.T) {
322 tests := []struct {
323 name string
324 cmd *exec.Cmd
325 timeout time.Duration
326 wantErrs []string
327 wantBound bool
328 }{
329 {
330 name: "wedged process is killed after timeout and reaped",
331 cmd: exec.Command("sleep", "30"),
332 timeout: 50 * time.Millisecond,
333 wantErrs: []string{"sshstdio: ssh exited:"},
334 wantBound: true,
335 },
336 {
337 name: "stderr from failing shell command is included in error",
338 cmd: exec.Command("sh", "-c", "echo kaboom >&2; exit 7"),
339 timeout: time.Second,
340 wantErrs: []string{"sshstdio: ssh exited:", "kaboom"},
341 },
342 }
343
344 for _, tt := range tests {
345 t.Run(tt.name, func(t *testing.T) {
346 stdin, err := tt.cmd.StdinPipe()
347 if err != nil {
348 t.Fatalf("StdinPipe: %v", err)
349 }
350 var stderr bytes.Buffer
351 tt.cmd.Stderr = &stderr
352 if err := tt.cmd.Start(); err != nil {
353 t.Fatalf("Start: %v", err)
354 }
355
356 started := time.Now()
357 err = newProcessCloser(tt.cmd, stdin, &stderr, tt.timeout, nil, "", "")()
358 elapsed := time.Since(started)
359
360 if err == nil {
361 t.Fatalf("Close error = nil, want error containing %q", tt.wantErrs)
362 }
363 for _, want := range tt.wantErrs {
364 if !strings.Contains(err.Error(), want) {
365 t.Fatalf("Close error = %q, want substring %q", err.Error(), want)
366 }
367 }
368 if tt.wantBound && elapsed > time.Second {
369 t.Fatalf("Close took %s, want bounded below 1s", elapsed)
370 }
371 if tt.cmd.ProcessState == nil {
372 t.Fatalf("ProcessState is nil, want child reaped")
373 }
374 })
375 }
376}