1// Package sshstdio adapts an ssh subprocess' stdin/stdout into vev's framed
2// Transport interface. It intentionally builds argv slices for os/exec instead
3// of shell command strings so remote targets and session names are never shell-
4// interpolated locally.
5package sshstdio
6
7import (
8 "bytes"
9 "context"
10 "encoding/binary"
11 "errors"
12 "fmt"
13 "io"
14 "log/slog"
15 "os"
16 "os/exec"
17 "strings"
18 "sync"
19 "time"
20
21 "github.com/bnema/vev/internal/ports"
22)
23
24const (
25 maxFrameLen = ports.MaxFrameLen
26 frameHeaderLen = 4
27 sshCloseTimeout = 3 * time.Second
28)
29
30var (
31 ErrZeroLengthFrame = errors.New("sshstdio: zero-length frame")
32 ErrFrameTooLarge = errors.New("sshstdio: frame exceeds maximum length")
33)
34
35type closeFunc func() error
36type eofErrFunc func() error
37type operationKind uint8
38
39const (
40 operationSend operationKind = iota
41 operationReceive
42 operationKindCount
43)
44
45type operationWait struct {
46 send <-chan struct{}
47 receive <-chan struct{}
48}
49
50type Option func(*transport)
51
52// WithRuntimeObserver enables process-local transport marks; it deliberately
53// takes only an observer so a carriage adapter never owns trace time.
54func WithRuntimeObserver(observer ports.SerializedRuntimeObserver) Option {
55 return func(t *transport) { t.observer = observer }
56}
57
58// NewTransport wraps separate reader/writer streams as a framed Transport.
59func NewTransport(r io.Reader, w io.Writer, closeFn closeFunc, opts ...Option) ports.Transport {
60 t := newTransport(r, w, closeFn, nil)
61 for _, opt := range opts {
62 if opt != nil {
63 opt(t)
64 }
65 }
66 return t
67}
68
69func newTransport(r io.Reader, w io.Writer, closeFn closeFunc, eofErr eofErrFunc) *transport {
70 if closeFn == nil {
71 closeFn = func() error { return nil }
72 }
73 t := &transport{
74 r: r,
75 w: w,
76 close: closeFn,
77 eofErr: eofErr,
78 done: make(chan struct{}),
79 }
80 if file, ok := r.(*os.File); ok && file == os.Stdin {
81 t.r = newProcessStdinReader(file, t.done)
82 } else if closer, ok := r.(io.Closer); ok {
83 t.readerCloser = closer
84 } else {
85 t.r = newUnownedReader(r, t.done)
86 }
87 if file, ok := w.(*os.File); !ok || file != os.Stdout {
88 if closer, ok := w.(io.Closer); ok {
89 t.writerCloser = closer
90 }
91 }
92 return t
93}
94
95type transport struct {
96 r io.Reader
97 w io.Writer
98 close closeFunc
99 eofErr eofErrFunc
100 readerCloser io.Closer
101 writerCloser io.Closer
102 done chan struct{}
103 observer ports.SerializedRuntimeObserver
104 operationMu sync.Mutex
105 operationCount [operationKindCount]int
106 closing bool
107 operationsDone [operationKindCount]chan struct{}
108 closeOnce sync.Once
109 closeErr error
110
111 mu sync.Mutex
112 readBuf []byte
113}
114
115type readResult struct {
116 data []byte
117 err error
118}
119
120// unownedReader is the generic cancellation boundary for readers that cannot
121// be closed by the transport. The worker may remain blocked until the source
122// itself is released; process stdin uses the singleton pump below instead.
123// ponytail: arbitrary io.Reader has no cancellation contract, so one fallback
124// worker may remain until that source returns; add a cancellable reader port if
125// another production unowned source appears.
126type unownedReader struct {
127 r io.Reader
128 done <-chan struct{}
129}
130
131func newUnownedReader(r io.Reader, done <-chan struct{}) *unownedReader {
132 return &unownedReader{r: r, done: done}
133}
134
135func (r *unownedReader) Read(dst []byte) (int, error) {
136 select {
137 case <-r.done:
138 return 0, io.ErrClosedPipe
139 default:
140 }
141 result := make(chan readResult, 1)
142 go func() {
143 buf := make([]byte, len(dst))
144 n, err := r.r.Read(buf)
145 result <- readResult{data: buf[:n], err: err}
146 }()
147 select {
148 case <-r.done:
149 return 0, io.ErrClosedPipe
150 case result := <-result:
151 select {
152 case <-r.done:
153 return 0, io.ErrClosedPipe
154 default:
155 }
156 copy(dst, result.data)
157 return len(result.data), result.err
158 }
159}
160
161type processStdinPump struct {
162 chunks chan readResult
163 mu sync.Mutex
164 buf []byte
165 err error
166}
167
168var processStdinPumps struct {
169 sync.Mutex
170 file *os.File
171 pump *processStdinPump
172}
173
174// processStdinPumpFor keeps one read worker for the process stdin lifetime.
175// Reconnects therefore do not strand one blocked worker per transport.
176func processStdinPumpFor(file *os.File) *processStdinPump {
177 processStdinPumps.Lock()
178 defer processStdinPumps.Unlock()
179 if processStdinPumps.pump != nil && processStdinPumps.file == file {
180 return processStdinPumps.pump
181 }
182 pump := &processStdinPump{chunks: make(chan readResult, 1)}
183 processStdinPumps.file = file
184 processStdinPumps.pump = pump
185 go pump.run(file)
186 return pump
187}
188
189func (p *processStdinPump) run(file *os.File) {
190 for {
191 buf := make([]byte, 32*1024)
192 n, err := file.Read(buf)
193 if n > 0 {
194 p.chunks <- readResult{data: append([]byte(nil), buf[:n]...)}
195 }
196 if err != nil {
197 p.chunks <- readResult{err: err}
198 close(p.chunks)
199 return
200 }
201 }
202}
203
204func (p *processStdinPump) read(done <-chan struct{}, dst []byte) (int, error) {
205 select {
206 case <-done:
207 return 0, io.ErrClosedPipe
208 default:
209 }
210 p.mu.Lock()
211 if len(p.buf) != 0 {
212 n := copy(dst, p.buf)
213 p.buf = p.buf[n:]
214 p.mu.Unlock()
215 return n, nil
216 }
217 if p.err != nil {
218 err := p.err
219 p.mu.Unlock()
220 return 0, err
221 }
222 p.mu.Unlock()
223
224 select {
225 case <-done:
226 return 0, io.ErrClosedPipe
227 case result, ok := <-p.chunks:
228 if !ok {
229 return 0, io.EOF
230 }
231 if len(result.data) != 0 {
232 select {
233 case <-done:
234 return 0, io.ErrClosedPipe
235 default:
236 }
237 n := copy(dst, result.data)
238 p.mu.Lock()
239 p.buf = append(p.buf[:0], result.data[n:]...)
240 p.mu.Unlock()
241 return n, nil
242 }
243 p.mu.Lock()
244 p.err = result.err
245 p.mu.Unlock()
246 return 0, result.err
247 }
248}
249
250type processStdinReader struct {
251 pump *processStdinPump
252 done <-chan struct{}
253}
254
255func newProcessStdinReader(file *os.File, done <-chan struct{}) io.Reader {
256 return processStdinReader{pump: processStdinPumpFor(file), done: done}
257}
258
259func (r processStdinReader) Read(dst []byte) (int, error) {
260 return r.pump.read(r.done, dst)
261}
262
263func (t *transport) Send(f ports.Frame) error {
264 end := t.beginOperation(ports.RuntimeAdapterSendStart, uint64(len(f.Payload)))
265 n := 1 + len(f.Payload)
266 if n > maxFrameLen {
267 end(false)
268 return ErrFrameTooLarge
269 }
270
271 buf := make([]byte, frameHeaderLen+n)
272 binary.BigEndian.PutUint32(buf[:frameHeaderLen], uint32(n))
273 buf[frameHeaderLen] = byte(f.Type)
274 copy(buf[frameHeaderLen+1:], f.Payload)
275
276 t.mu.Lock()
277 _, err := t.w.Write(buf)
278 t.mu.Unlock()
279 end(err == nil)
280 return err
281}
282
283func (t *transport) Recv() (ports.Frame, error) {
284 end := t.beginOperation(ports.RuntimeAdapterReceiveStart, 0)
285 var hdr [frameHeaderLen]byte
286 if _, err := io.ReadFull(t.r, hdr[:]); err != nil {
287 err = t.mapEOFError(err)
288 end(false)
289 return ports.Frame{}, err
290 }
291
292 n := binary.BigEndian.Uint32(hdr[:])
293 if n == 0 {
294 end(false)
295 return ports.Frame{}, ErrZeroLengthFrame
296 }
297 if n > maxFrameLen {
298 end(false)
299 return ports.Frame{}, ErrFrameTooLarge
300 }
301
302 if cap(t.readBuf) < int(n) {
303 t.readBuf = make([]byte, n)
304 } else {
305 t.readBuf = t.readBuf[:n]
306 }
307 if _, err := io.ReadFull(t.r, t.readBuf); err != nil {
308 err = t.mapEOFError(err)
309 end(false)
310 return ports.Frame{}, err
311 }
312
313 payload := append([]byte(nil), t.readBuf[1:]...)
314 end(true)
315 return ports.Frame{Type: ports.MsgType(t.readBuf[0]), Payload: payload}, nil
316}
317
318func (t *transport) beginOperation(start ports.RuntimeMarkKind, bytes uint64) func(bool) {
319 if t.observer == nil {
320 return func(bool) {}
321 }
322 kind := operationSend
323 if start == ports.RuntimeAdapterReceiveStart {
324 kind = operationReceive
325 }
326 if !t.beginObservedOperation(kind) {
327 return func(bool) {}
328 }
329 correlation := ports.NewRuntimeCorrelation()
330 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, start, bytes, true))
331 end := ports.RuntimeAdapterSendEnd
332 if start == ports.RuntimeAdapterReceiveStart {
333 end = ports.RuntimeAdapterReceiveEnd
334 }
335 return func(valid bool) {
336 defer t.finishObservedOperation(kind)
337 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, end, bytes, valid))
338 }
339}
340
341func (t *transport) mapEOFError(err error) error {
342 if t.eofErr == nil || (!errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF)) {
343 return err
344 }
345 if sshErr := t.eofErr(); sshErr != nil {
346 return sshErr
347 }
348 return err
349}
350
351func (t *transport) Close() error {
352 t.closeOnce.Do(func() {
353 wait := t.beginShutdown()
354 // Close owned streams before waiting for in-flight operations. The
355 // writer is what releases a Send blocked on a child stdin pipe; the
356 // reader releases a Recv blocked in io.ReadFull. Unowned streams cannot
357 // be interrupted safely, so the cancellation boundary releases Recv
358 // without closing the source.
359 var writeErr, readErr error
360 if t.writerCloser != nil {
361 writeErr = t.writerCloser.Close()
362 }
363 if t.readerCloser != nil {
364 readErr = t.readerCloser.Close()
365 }
366 closeErr := t.close()
367 if t.writerCloser != nil && wait.send != nil {
368 <-wait.send
369 }
370 if t.readerCloser != nil && wait.receive != nil {
371 <-wait.receive
372 }
373 switch {
374 case closeErr != nil:
375 t.closeErr = closeErr
376 case readErr != nil:
377 t.closeErr = readErr
378 default:
379 t.closeErr = writeErr
380 }
381 })
382 return t.closeErr
383}
384
385func (t *transport) beginObservedOperation(kind operationKind) bool {
386 t.operationMu.Lock()
387 defer t.operationMu.Unlock()
388 if t.closing {
389 return false
390 }
391 if t.operationCount[kind] == 0 {
392 t.operationsDone[kind] = make(chan struct{})
393 }
394 t.operationCount[kind]++
395 return true
396}
397
398func (t *transport) finishObservedOperation(kind operationKind) {
399 t.operationMu.Lock()
400 defer t.operationMu.Unlock()
401 t.operationCount[kind]--
402 if t.operationCount[kind] == 0 {
403 close(t.operationsDone[kind])
404 }
405}
406
407func (t *transport) beginShutdown() operationWait {
408 t.operationMu.Lock()
409 defer t.operationMu.Unlock()
410 t.closing = true
411 close(t.done)
412 var wait operationWait
413 if t.operationCount[operationSend] != 0 {
414 wait.send = t.operationsDone[operationSend]
415 }
416 if t.operationCount[operationReceive] != 0 {
417 wait.receive = t.operationsDone[operationReceive]
418 }
419 return wait
420}
421
422type processWaiter struct {
423 cmd *exec.Cmd
424 stdin io.Closer
425 stderr *bytes.Buffer
426 timeout time.Duration
427 log *slog.Logger
428 target string
429 session string
430
431 waitOnce sync.Once
432 waitErr error
433}
434
435func newProcessWaiter(cmd *exec.Cmd, stdin io.Closer, stderr *bytes.Buffer, timeout time.Duration, log *slog.Logger, target, session string) *processWaiter {
436 w := &processWaiter{
437 cmd: cmd,
438 stdin: stdin,
439 stderr: stderr,
440 timeout: timeout,
441 log: log,
442 target: target,
443 session: session,
444 }
445 return w
446}
447
448func (w *processWaiter) close() error {
449 _ = w.stdin.Close()
450 return w.wait(w.timeout)
451}
452
453func (w *processWaiter) eofErr() error {
454 return w.wait(sshCloseTimeout)
455}
456
457func (w *processWaiter) wait(timeout time.Duration) error {
458 w.waitOnce.Do(func() {
459 // cmd.Wait must start only after stdout is no longer being read. os/exec
460 // requires callers to finish reading StdoutPipe before Wait, so DialContext
461 // calls this from transport Close or after Recv has already observed EOF.
462 waitCh := make(chan error, 1)
463 go func() { waitCh <- w.cmd.Wait() }()
464 select {
465 case w.waitErr = <-waitCh:
466 case <-time.After(timeout):
467 _ = w.cmd.Process.Kill()
468 w.waitErr = <-waitCh
469 }
470 w.waitErr = formatProcessWaitError(w.waitErr, w.stderr, w.log, w.target, w.session)
471 })
472 return w.waitErr
473}
474
475func formatProcessWaitError(err error, stderr *bytes.Buffer, log *slog.Logger, target, session string) error {
476 if err == nil {
477 return nil
478 }
479 stderrText := strings.TrimSpace(stderr.String())
480 if log != nil {
481 attrs := []any{"target", target, "session", session, "err", err}
482 if stderrText != "" {
483 attrs = append(attrs, "stderr", stderrText)
484 }
485 log.Warn("ssh exited non-cleanly", attrs...)
486 }
487 if stderrText != "" {
488 return fmt.Errorf("sshstdio: ssh exited: %w: %s", err, stderrText)
489 }
490 return fmt.Errorf("sshstdio: ssh exited: %w", err)
491}
492
493func newProcessCloser(cmd *exec.Cmd, stdin io.Closer, stderr *bytes.Buffer, timeout time.Duration, log *slog.Logger, target, session string) closeFunc {
494 return newProcessWaiter(cmd, stdin, stderr, timeout, log, target, session).close
495}
496
497// CommandSpec is the exact ssh argv vev will execute locally.
498type CommandSpec struct {
499 Path string
500 Args []string
501}
502
503// BuildCommand constructs the local ssh subprocess argv without invoking a
504// local shell. OpenSSH sends the remote command as one string for the remote
505// user's shell to interpret, so every remote argv word is POSIX single-quoted.
506func BuildCommand(target, _ string) CommandSpec {
507 return BuildCommandForMode(target, "_stdio", "")
508}
509
510// BuildCommandForMode constructs the local ssh subprocess argv for a hidden vev
511// remote mode such as _stdio or _udp-bootstrap.
512func BuildCommandForMode(target, mode, _ string) CommandSpec {
513 return BuildCommandForRemoteCommand(target, "vev", mode)
514}
515
516// BuildCommandForRemoteCommand constructs ssh argv for an arbitrary remote
517// command. Every remote word is POSIX single-quoted; the target remains one
518// local argv word after the option terminator.
519func BuildCommandForRemoteCommand(target string, command ...string) CommandSpec {
520 remote := make([]string, 0, len(command))
521 for _, word := range command {
522 remote = append(remote, shellQuote(word))
523 }
524 args := []string{"--", target, strings.Join(remote, " ")}
525 return CommandSpec{Path: "ssh", Args: args}
526}
527
528func shellQuote(s string) string {
529 return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
530}
531
532// Dial starts ssh target vev _stdio and returns a Transport over the child
533// process' stdio. Session selection is carried in Hello. The subprocess is
534// started with exec.Command argv, never through a shell.
535func Dial(target, session string) (ports.Transport, error) {
536 return DialContext(context.Background(), target, session)
537}
538
539// DialContextWithRuntimeObserver is the option-bearing production constructor.
540func DialContextWithRuntimeObserver(ctx context.Context, target, session string, logger *slog.Logger, opts ...Option) (ports.Transport, error) {
541 return dialContext(ctx, target, session, logger, opts...)
542}
543
544// DialContext is like Dial, but the context gates ssh startup. Once the
545// transport is returned, its Close method owns the subprocess lifetime; the
546// handshake context must not kill an already-established carriage. Callers may
547// pass a logger to record ssh start failures and non-clean exits without logging
548// the generated command line.
549func DialContext(ctx context.Context, target, session string, logger ...*slog.Logger) (ports.Transport, error) {
550 var log *slog.Logger
551 if len(logger) > 0 {
552 log = logger[0]
553 }
554 return dialContext(ctx, target, session, log)
555}
556
557func dialContext(ctx context.Context, target, session string, log *slog.Logger, opts ...Option) (ports.Transport, error) {
558 if err := ctx.Err(); err != nil {
559 return nil, err
560 }
561 spec := BuildCommand(target, session)
562 // The caller's context is also the bounded protocol-handshake context and is
563 // canceled after the first committed publication. Binding it to the command
564 // would kill a healthy long-lived carriage at that boundary. Transport.Close
565 // owns cancellation after Start; the preflight check above keeps canceled
566 // attempts from starting a process.
567 cmd := exec.Command(spec.Path, spec.Args...)
568 stdin, err := cmd.StdinPipe()
569 if err != nil {
570 return nil, fmt.Errorf("sshstdio: stdin pipe: %w", err)
571 }
572 stdout, err := cmd.StdoutPipe()
573 if err != nil {
574 return nil, fmt.Errorf("sshstdio: stdout pipe: %w", err)
575 }
576 var stderr bytes.Buffer
577 cmd.Stderr = &stderr
578 if err := cmd.Start(); err != nil {
579 if log != nil {
580 log.Error("ssh start failed", "target", target, "session", session, "err", err)
581 }
582 return nil, fmt.Errorf("sshstdio: start ssh: %w", err)
583 }
584
585 waiter := newProcessWaiter(cmd, stdin, &stderr, sshCloseTimeout, log, target, session)
586 transport := newTransport(stdout, stdin, waiter.close, waiter.eofErr)
587 if err := ctx.Err(); err != nil {
588 _ = transport.Close()
589 return nil, err
590 }
591 for _, opt := range opts {
592 if opt != nil {
593 opt(transport)
594 }
595 }
596 return transport, nil
597}