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 "errors"
9 "fmt"
10 "io"
11 "log/slog"
12 "os"
13 "os/exec"
14 "strings"
15 "sync"
16 "time"
17
18 "github.com/bnema/vev/internal/adapters/streamframe"
19 "github.com/bnema/vev/internal/ports"
20 "github.com/bnema/vev/internal/protocol/wire"
21)
22
23const sshCloseTimeout = 3 * time.Second
24
25var (
26 ErrZeroLengthFrame = streamframe.ErrZeroLength
27 ErrFrameTooLarge = streamframe.ErrTooLarge
28)
29
30type closeFunc func() error
31type eofErrFunc func() error
32type operationKind uint8
33
34const (
35 operationSend operationKind = iota
36 operationReceive
37 operationKindCount
38)
39
40type operationWait struct {
41 send <-chan struct{}
42 receive <-chan struct{}
43}
44
45type Option func(*transport)
46
47// WithRuntimeObserver enables process-local transport marks; it deliberately
48// takes only an observer so a carriage adapter never owns trace time.
49func WithRuntimeObserver(observer ports.SerializedRuntimeObserver) Option {
50 return func(t *transport) { t.observer = observer }
51}
52
53// NewTransport wraps separate reader/writer streams as a framed Transport.
54func NewTransport(r io.Reader, w io.Writer, closeFn closeFunc, opts ...Option) wire.Transport {
55 t := newTransport(r, w, closeFn, nil)
56 for _, opt := range opts {
57 if opt != nil {
58 opt(t)
59 }
60 }
61 return t
62}
63
64func newTransport(r io.Reader, w io.Writer, closeFn closeFunc, eofErr eofErrFunc) *transport {
65 if closeFn == nil {
66 closeFn = func() error { return nil }
67 }
68 t := &transport{
69 r: r,
70 w: w,
71 close: closeFn,
72 eofErr: eofErr,
73 done: make(chan struct{}),
74 framer: nil,
75 }
76 if file, ok := r.(*os.File); ok && file == os.Stdin {
77 t.r = newProcessStdinReader(file, t.done)
78 } else if closer, ok := r.(io.Closer); ok {
79 t.readerCloser = closer
80 } else {
81 t.r = newUnownedReader(r, t.done)
82 }
83 if file, ok := w.(*os.File); !ok || file != os.Stdout {
84 if closer, ok := w.(io.Closer); ok {
85 t.writerCloser = closer
86 }
87 }
88 // Framing, queueing, and the writer live in streamframe; this adapter
89 // keeps reader cancellation, process waiting, and EOF mapping. The
90 // framer's closer releases owned streams so Close unblocks Recv.
91 t.framer = streamframe.NewFramer(t.r, &serializedWriter{w: t.w}, t.closeOwnedStreams)
92 return t
93}
94
95var _ wire.BoundedTransport = (*transport)(nil)
96
97// serializedWriter serializes concurrent streamframe writes over one child
98// stdin pipe; streamframe already serializes queue admission, this guards
99// the raw Write against interleaving.
100type serializedWriter struct {
101 mu sync.Mutex
102 w io.Writer
103}
104
105func (w *serializedWriter) Write(p []byte) (int, error) {
106 w.mu.Lock()
107 defer w.mu.Unlock()
108 return w.w.Write(p)
109}
110
111type transport struct {
112 r io.Reader
113 w io.Writer
114 framer *streamframe.Framer
115 close closeFunc
116 eofErr eofErrFunc
117 readerCloser io.Closer
118 writerCloser io.Closer
119 done chan struct{}
120 observer ports.SerializedRuntimeObserver
121 operationMu sync.Mutex
122 operationCount [operationKindCount]int
123 closing bool
124 operationsDone [operationKindCount]chan struct{}
125 closeOnce sync.Once
126 closeErr error
127}
128
129type readResult struct {
130 data []byte
131 err error
132}
133
134// unownedReader is the generic cancellation boundary for readers that cannot
135// be closed by the transport. The worker may remain blocked until the source
136// itself is released; process stdin uses the singleton pump below instead.
137// ponytail: arbitrary io.Reader has no cancellation contract, so one fallback
138// worker may remain until that source returns; add a cancellable reader port if
139// another production unowned source appears.
140type unownedReader struct {
141 r io.Reader
142 done <-chan struct{}
143}
144
145func newUnownedReader(r io.Reader, done <-chan struct{}) *unownedReader {
146 return &unownedReader{r: r, done: done}
147}
148
149func (r *unownedReader) Read(dst []byte) (int, error) {
150 select {
151 case <-r.done:
152 return 0, io.ErrClosedPipe
153 default:
154 }
155 result := make(chan readResult, 1)
156 go func() {
157 buf := make([]byte, len(dst))
158 n, err := r.r.Read(buf)
159 result <- readResult{data: buf[:n], err: err}
160 }()
161 select {
162 case <-r.done:
163 return 0, io.ErrClosedPipe
164 case result := <-result:
165 select {
166 case <-r.done:
167 return 0, io.ErrClosedPipe
168 default:
169 }
170 copy(dst, result.data)
171 return len(result.data), result.err
172 }
173}
174
175type processStdinPump struct {
176 chunks chan readResult
177 mu sync.Mutex
178 buf []byte
179 err error
180}
181
182var processStdinPumps struct {
183 sync.Mutex
184 file *os.File
185 pump *processStdinPump
186}
187
188// processStdinPumpFor keeps one read worker for the process stdin lifetime.
189// Reconnects therefore do not strand one blocked worker per transport.
190func processStdinPumpFor(file *os.File) *processStdinPump {
191 processStdinPumps.Lock()
192 defer processStdinPumps.Unlock()
193 if processStdinPumps.pump != nil && processStdinPumps.file == file {
194 return processStdinPumps.pump
195 }
196 pump := &processStdinPump{chunks: make(chan readResult, 1)}
197 processStdinPumps.file = file
198 processStdinPumps.pump = pump
199 go pump.run(file)
200 return pump
201}
202
203func (p *processStdinPump) run(file *os.File) {
204 for {
205 buf := make([]byte, 32*1024)
206 n, err := file.Read(buf)
207 if n > 0 {
208 p.chunks <- readResult{data: append([]byte(nil), buf[:n]...)}
209 }
210 if err != nil {
211 p.chunks <- readResult{err: err}
212 close(p.chunks)
213 return
214 }
215 }
216}
217
218func (p *processStdinPump) read(done <-chan struct{}, dst []byte) (int, error) {
219 select {
220 case <-done:
221 return 0, io.ErrClosedPipe
222 default:
223 }
224 p.mu.Lock()
225 if len(p.buf) != 0 {
226 n := copy(dst, p.buf)
227 p.buf = p.buf[n:]
228 p.mu.Unlock()
229 return n, nil
230 }
231 if p.err != nil {
232 err := p.err
233 p.mu.Unlock()
234 return 0, err
235 }
236 p.mu.Unlock()
237
238 select {
239 case <-done:
240 return 0, io.ErrClosedPipe
241 case result, ok := <-p.chunks:
242 if !ok {
243 return 0, io.EOF
244 }
245 if len(result.data) != 0 {
246 select {
247 case <-done:
248 return 0, io.ErrClosedPipe
249 default:
250 }
251 n := copy(dst, result.data)
252 p.mu.Lock()
253 p.buf = append(p.buf[:0], result.data[n:]...)
254 p.mu.Unlock()
255 return n, nil
256 }
257 p.mu.Lock()
258 p.err = result.err
259 p.mu.Unlock()
260 return 0, result.err
261 }
262}
263
264type processStdinReader struct {
265 pump *processStdinPump
266 done <-chan struct{}
267}
268
269func newProcessStdinReader(file *os.File, done <-chan struct{}) io.Reader {
270 return processStdinReader{pump: processStdinPumpFor(file), done: done}
271}
272
273func (r processStdinReader) Read(dst []byte) (int, error) {
274 return r.pump.read(r.done, dst)
275}
276
277func (t *transport) Send(envelope wire.Envelope) error {
278 end := t.beginOperation(ports.RuntimeAdapterSendStart, uint64(len(envelope.Payload)))
279 err := t.framer.Send(envelope.Payload)
280 end(err == nil)
281 return err
282}
283
284func (t *transport) Recv() (wire.Envelope, error) {
285 end := t.beginOperation(ports.RuntimeAdapterReceiveStart, 0)
286 payload, err := t.framer.Recv()
287 if err != nil {
288 err = t.mapEOFError(err)
289 end(false)
290 return wire.Envelope{}, err
291 }
292 end(true)
293 return wire.Envelope{Payload: payload}, nil
294}
295
296// RecvBounded reads one complete envelope bounded to limit before allocation,
297// sharing the framer's validation with Recv: an over-limit length prefix is
298// refused without reading its body. It is the bounded raw-carriage capability
299// a preamble-negotiated multiplexer consumes.
300func (t *transport) RecvBounded(limit uint64) (wire.Envelope, error) {
301 end := t.beginOperation(ports.RuntimeAdapterReceiveStart, 0)
302 payload, err := t.framer.RecvBounded(limit)
303 if err != nil {
304 err = t.mapEOFError(err)
305 end(false)
306 return wire.Envelope{}, err
307 }
308 end(true)
309 return wire.Envelope{Payload: payload}, nil
310}
311
312func (t *transport) beginOperation(start ports.RuntimeMarkKind, bytes uint64) func(bool) {
313 if t.observer == nil {
314 return func(bool) {}
315 }
316 kind := operationSend
317 if start == ports.RuntimeAdapterReceiveStart {
318 kind = operationReceive
319 }
320 if !t.beginObservedOperation(kind) {
321 return func(bool) {}
322 }
323 correlation := ports.NewRuntimeCorrelation()
324 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, start, bytes, true))
325 end := ports.RuntimeAdapterSendEnd
326 if start == ports.RuntimeAdapterReceiveStart {
327 end = ports.RuntimeAdapterReceiveEnd
328 }
329 return func(valid bool) {
330 defer t.finishObservedOperation(kind)
331 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, end, bytes, valid))
332 }
333}
334
335func (t *transport) mapEOFError(err error) error {
336 if t.eofErr == nil || (!errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF)) {
337 return err
338 }
339 if sshErr := t.eofErr(); sshErr != nil {
340 return sshErr
341 }
342 return err
343}
344
345func (t *transport) Close() error {
346 t.closeOnce.Do(func() {
347 wait := t.beginShutdown()
348 // The framer's closer (closeOwnedStreams below) releases owned
349 // streams before waiting for in-flight operations. Unowned
350 // streams cannot be interrupted safely, so the cancellation
351 // boundary releases Recv without closing the source.
352 t.closeErr = t.framer.Close()
353 if t.writerCloser != nil && wait.send != nil {
354 <-wait.send
355 }
356 if t.readerCloser != nil && wait.receive != nil {
357 <-wait.receive
358 }
359 })
360 return t.closeErr
361}
362
363// closeOwnedStreams closes owned streams before the framer waits for its
364// writer, then waits for the child process. It runs exactly once via the
365// framer's Close.
366func (t *transport) closeOwnedStreams() error {
367 var writeErr, readErr error
368 if t.writerCloser != nil {
369 writeErr = t.writerCloser.Close()
370 }
371 if t.readerCloser != nil {
372 readErr = t.readerCloser.Close()
373 }
374 closeErr := t.close()
375 switch {
376 case closeErr != nil:
377 return closeErr
378 case readErr != nil:
379 return readErr
380 default:
381 return writeErr
382 }
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
422// stderrSink is the bounded diagnostic capture a process waiter reads for a
423// non-clean exit. A plain *bytes.Buffer keeps the session dial's existing raw
424// behavior; the mux dial supplies a bounded, sanitizing capture and sets
425// sanitize so the public error never carries remote stderr.
426type stderrSink interface {
427 io.Writer
428 String() string
429}
430
431type processWaiter struct {
432 cmd *exec.Cmd
433 stdin io.Closer
434 stderr stderrSink
435 timeout time.Duration
436 log *slog.Logger
437 target string
438 session string
439 sanitize bool
440
441 waitOnce sync.Once
442 waitErr error
443}
444
445func newProcessWaiter(cmd *exec.Cmd, stdin io.Closer, stderr stderrSink, timeout time.Duration, log *slog.Logger, target, session string) *processWaiter {
446 w := &processWaiter{
447 cmd: cmd,
448 stdin: stdin,
449 stderr: stderr,
450 timeout: timeout,
451 log: log,
452 target: target,
453 session: session,
454 }
455 return w
456}
457
458// newMuxProcessWaiter builds the waiter for one mux child. It reuses the
459// session wait/kill contract unchanged but sanitizes the public error: the
460// bounded, sanitized stderr is logged, never returned.
461func newMuxProcessWaiter(cmd *exec.Cmd, stdin io.Closer, stderr stderrSink, timeout time.Duration, log *slog.Logger) *processWaiter {
462 w := newProcessWaiter(cmd, stdin, stderr, timeout, log, "", "")
463 w.sanitize = true
464 return w
465}
466
467func (w *processWaiter) close() error {
468 _ = w.stdin.Close()
469 return w.wait(w.timeout)
470}
471
472func (w *processWaiter) eofErr() error {
473 return w.wait(sshCloseTimeout)
474}
475
476func (w *processWaiter) wait(timeout time.Duration) error {
477 w.waitOnce.Do(func() {
478 // cmd.Wait must start only after stdout is no longer being read. os/exec
479 // requires callers to finish reading StdoutPipe before Wait, so DialContext
480 // calls this from transport Close or after Recv has already observed EOF.
481 waitCh := make(chan error, 1)
482 go func() { waitCh <- w.cmd.Wait() }()
483 select {
484 case w.waitErr = <-waitCh:
485 case <-time.After(timeout):
486 _ = w.cmd.Process.Kill()
487 w.waitErr = <-waitCh
488 }
489 w.waitErr = formatProcessWaitError(w.waitErr, w.stderr, w.log, w.target, w.session, w.sanitize)
490 })
491 return w.waitErr
492}
493
494func formatProcessWaitError(err error, stderr stderrSink, log *slog.Logger, target, session string, sanitize bool) error {
495 if err == nil {
496 return nil
497 }
498 var stderrText string
499 if stderr != nil {
500 stderrText = strings.TrimSpace(stderr.String())
501 }
502 if log != nil {
503 attrs := []any{"target", target, "session", session, "err", err}
504 if stderrText != "" {
505 attrs = append(attrs, "stderr", stderrText)
506 }
507 log.Warn("ssh exited non-cleanly", attrs...)
508 }
509 if sanitize {
510 // The mux carriage's public error carries only the typed outcome; the
511 // bounded, sanitized diagnostic above is local-only.
512 return fmt.Errorf("%w: %w", ErrMuxSSHExit, err)
513 }
514 if stderrText != "" {
515 return fmt.Errorf("sshstdio: ssh exited: %w: %s", err, stderrText)
516 }
517 return fmt.Errorf("sshstdio: ssh exited: %w", err)
518}
519
520func newProcessCloser(cmd *exec.Cmd, stdin io.Closer, stderr stderrSink, timeout time.Duration, log *slog.Logger, target, session string) closeFunc {
521 return newProcessWaiter(cmd, stdin, stderr, timeout, log, target, session).close
522}
523
524// CommandSpec is the exact ssh argv vev will execute locally.
525type CommandSpec struct {
526 Path string
527 Args []string
528}
529
530// BuildCommandForRemoteCommand constructs ssh argv for an arbitrary remote
531// command. Every remote word is POSIX single-quoted; the target remains one
532// local argv word after the option terminator.
533func BuildCommandForRemoteCommand(target string, command ...string) CommandSpec {
534 remote := make([]string, 0, len(command))
535 for _, word := range command {
536 remote = append(remote, shellQuote(word))
537 }
538 args := []string{"--", target, strings.Join(remote, " ")}
539 return CommandSpec{Path: "ssh", Args: args}
540}
541
542func shellQuote(s string) string {
543 return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
544}