1package sshstdio
2
3import (
4 "bytes"
5 "context"
6 "errors"
7 "io"
8 "log/slog"
9 "os/exec"
10 "strings"
11 "testing"
12 "time"
13
14 "github.com/bnema/vev/internal/ports"
15)
16
17func TestBuildCommandUsesExecArgs(t *testing.T) {
18 // BuildCommand deliberately starts `ssh -- target 'vev' '_stdio' [session]`.
19 // Terminal color capability is carried in vev's Hello message, not in ssh env.
20 tests := []struct {
21 name string
22 target string
23 session string
24 want []string
25 }{
26 {name: "no session", target: "user@example.com", want: []string{"--", "user@example.com", "'vev' '_stdio'"}},
27 {name: "session shell metacharacters are quoted in remote command", target: "user@example.com", session: "work; rm -rf /", want: []string{"--", "user@example.com", "'vev' '_stdio' 'work; rm -rf /'"}},
28 {name: "single quote in session is posix escaped", target: "user@example.com", session: "it's fine", want: []string{"--", "user@example.com", "'vev' '_stdio' 'it'\\''s fine'"}},
29 {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' 'work'"}},
30 }
31
32 for _, tt := range tests {
33 t.Run(tt.name, func(t *testing.T) {
34 got := BuildCommand(tt.target, tt.session)
35 if got.Path != "ssh" {
36 t.Fatalf("Path = %q, want ssh", got.Path)
37 }
38 if len(got.Args) != len(tt.want) {
39 t.Fatalf("Args len = %d, want %d (%q)", len(got.Args), len(tt.want), got.Args)
40 }
41 for i := range tt.want {
42 if got.Args[i] != tt.want[i] {
43 t.Fatalf("Args[%d] = %q, want %q (all args %q)", i, got.Args[i], tt.want[i], got.Args)
44 }
45 }
46 })
47 }
48}
49
50func TestBuildCommandForModeUsesCanonicalSSHArgs(t *testing.T) {
51 got := BuildCommandForMode("user@example.com", "_udp-bootstrap", "work")
52 want := []string{"--", "user@example.com", "'vev' '_udp-bootstrap' 'work'"}
53 if got.Path != "ssh" {
54 t.Fatalf("Path = %q, want ssh", got.Path)
55 }
56 for i := range want {
57 if got.Args[i] != want[i] {
58 t.Fatalf("Args[%d] = %q, want %q (all args %q)", i, got.Args[i], want[i], got.Args)
59 }
60 }
61}
62
63func TestTransportRoundTripAndVersionMismatchFrame(t *testing.T) {
64 clientRead, serverWrite := io.Pipe()
65 serverRead, clientWrite := io.Pipe()
66 client := NewTransport(clientRead, clientWrite, func() error {
67 _ = clientRead.Close()
68 return clientWrite.Close()
69 })
70 server := NewTransport(serverRead, serverWrite, func() error {
71 _ = serverRead.Close()
72 return serverWrite.Close()
73 })
74 defer func() { _ = client.Close() }()
75 defer func() { _ = server.Close() }()
76
77 hello := ports.Hello{Version: ports.ProtocolVersion + 1, Intent: ports.IntentAttach, Name: "work"}
78 go func() {
79 f, err := server.Recv()
80 if err != nil {
81 t.Errorf("server Recv: %v", err)
82 return
83 }
84 got, err := ports.UnmarshalHello(f.Payload)
85 if err != nil {
86 t.Errorf("UnmarshalHello: %v", err)
87 return
88 }
89 if got.Version == ports.ProtocolVersion {
90 t.Errorf("test did not send a mismatched version")
91 }
92 _ = server.Send(ports.Frame{Type: ports.MsgError, Payload: ports.MarshalErrorMsg(ports.ErrorMsg{Code: ports.ErrVersionMismatch, Text: "protocol version mismatch"})})
93 }()
94
95 if err := client.Send(ports.Frame{Type: ports.MsgHello, Payload: ports.MarshalHello(hello)}); err != nil {
96 t.Fatalf("client Send: %v", err)
97 }
98 reply, err := client.Recv()
99 if err != nil {
100 t.Fatalf("client Recv: %v", err)
101 }
102 if reply.Type != ports.MsgError {
103 t.Fatalf("reply type = %d, want MsgError", reply.Type)
104 }
105 em, err := ports.UnmarshalErrorMsg(reply.Payload)
106 if err != nil {
107 t.Fatalf("UnmarshalErrorMsg: %v", err)
108 }
109 if em.Code != ports.ErrVersionMismatch {
110 t.Fatalf("error code = %d, want ErrVersionMismatch", em.Code)
111 }
112}
113
114func TestTransportRejectsZeroLengthFrame(t *testing.T) {
115 tr := NewTransport(strings.NewReader("\x00\x00\x00\x00"), io.Discard, nil)
116 _, err := tr.Recv()
117 if !errors.Is(err, ErrZeroLengthFrame) {
118 t.Fatalf("Recv error = %v, want ErrZeroLengthFrame", err)
119 }
120}
121
122func TestDialContextCanceledBeforeStart(t *testing.T) {
123 ctx, cancel := context.WithCancel(context.Background())
124 cancel()
125 _, err := DialContext(ctx, "example.com", "work")
126 if !errors.Is(err, context.Canceled) {
127 t.Fatalf("DialContext error = %v, want context.Canceled", err)
128 }
129}
130
131func TestRecvMapsUnexpectedEOFDuringHeaderToSSHExit(t *testing.T) {
132 sshErr := errors.New("ssh exited with status 255")
133 tr := newTransport(strings.NewReader("\x00"), io.Discard, nil, func() error { return sshErr })
134
135 _, err := tr.Recv()
136 if !errors.Is(err, sshErr) {
137 t.Fatalf("Recv error = %v, want %v", err, sshErr)
138 }
139}
140
141func TestRecvReportsSSHExitWhenProcessClosesBeforeFrame(t *testing.T) {
142 cmd := exec.CommandContext(context.Background(), "sh", "-c", "echo connection refused >&2; exit 255")
143 stdin, err := cmd.StdinPipe()
144 if err != nil {
145 t.Fatalf("StdinPipe: %v", err)
146 }
147 stdout, err := cmd.StdoutPipe()
148 if err != nil {
149 t.Fatalf("StdoutPipe: %v", err)
150 }
151 var stderr bytes.Buffer
152 cmd.Stderr = &stderr
153 if err := cmd.Start(); err != nil {
154 t.Fatalf("Start: %v", err)
155 }
156
157 waiter := newProcessWaiter(cmd, stdin, &stderr, time.Second, nil, "user@example.com", "work")
158 tr := newTransport(stdout, stdin, waiter.close, waiter.eofErr)
159
160 _, err = tr.Recv()
161 if err == nil {
162 t.Fatal("Recv error = nil, want ssh exit error")
163 }
164 for _, want := range []string{"sshstdio: ssh exited:", "connection refused"} {
165 if !strings.Contains(err.Error(), want) {
166 t.Fatalf("Recv error = %q, want substring %q", err, want)
167 }
168 }
169}
170
171func TestProcessCloserLogsNonCleanExitStderr(t *testing.T) {
172 cmd := exec.Command("sh", "-c", "echo remote failure >&2; exit 7")
173 stdin, err := cmd.StdinPipe()
174 if err != nil {
175 t.Fatalf("StdinPipe: %v", err)
176 }
177 var stderr bytes.Buffer
178 cmd.Stderr = &stderr
179 if err := cmd.Start(); err != nil {
180 t.Fatalf("Start: %v", err)
181 }
182
183 var logBuf bytes.Buffer
184 log := slog.New(slog.NewJSONHandler(&logBuf, nil))
185 err = newProcessCloser(cmd, stdin, &stderr, time.Second, log, "user@example.com", "work")()
186 if err == nil {
187 t.Fatal("Close error = nil, want non-clean ssh exit")
188 }
189 entry := logBuf.String()
190 for _, want := range []string{"ssh exited non-cleanly", "remote failure", "user@example.com", "work"} {
191 if !strings.Contains(entry, want) {
192 t.Fatalf("log entry = %q, want substring %q", entry, want)
193 }
194 }
195 if strings.Contains(entry, "'vev' '_stdio'") || strings.Contains(entry, "-- user@example.com") {
196 t.Fatalf("log entry includes generated command line: %q", entry)
197 }
198}
199
200func TestProcessCloser(t *testing.T) {
201 tests := []struct {
202 name string
203 cmd *exec.Cmd
204 timeout time.Duration
205 wantErrs []string
206 wantBound bool
207 }{
208 {
209 name: "wedged process is killed after timeout and reaped",
210 cmd: exec.Command("sleep", "30"),
211 timeout: 50 * time.Millisecond,
212 wantErrs: []string{"sshstdio: ssh exited:"},
213 wantBound: true,
214 },
215 {
216 name: "stderr from failing shell command is included in error",
217 cmd: exec.Command("sh", "-c", "echo kaboom >&2; exit 7"),
218 timeout: time.Second,
219 wantErrs: []string{"sshstdio: ssh exited:", "kaboom"},
220 },
221 }
222
223 for _, tt := range tests {
224 t.Run(tt.name, func(t *testing.T) {
225 stdin, err := tt.cmd.StdinPipe()
226 if err != nil {
227 t.Fatalf("StdinPipe: %v", err)
228 }
229 var stderr bytes.Buffer
230 tt.cmd.Stderr = &stderr
231 if err := tt.cmd.Start(); err != nil {
232 t.Fatalf("Start: %v", err)
233 }
234
235 started := time.Now()
236 err = newProcessCloser(tt.cmd, stdin, &stderr, tt.timeout, nil, "", "")()
237 elapsed := time.Since(started)
238
239 if err == nil {
240 t.Fatalf("Close error = nil, want error containing %q", tt.wantErrs)
241 }
242 for _, want := range tt.wantErrs {
243 if !strings.Contains(err.Error(), want) {
244 t.Fatalf("Close error = %q, want substring %q", err.Error(), want)
245 }
246 }
247 if tt.wantBound && elapsed > time.Second {
248 t.Fatalf("Close took %s, want bounded below 1s", elapsed)
249 }
250 if tt.cmd.ProcessState == nil {
251 t.Fatalf("ProcessState is nil, want child reaped")
252 }
253 })
254 }
255}