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/exec"
16 "strings"
17 "sync"
18 "time"
19
20 "github.com/bnema/vev/internal/ports"
21)
22
23const (
24 maxFrameLen = 16 << 20
25 frameHeaderLen = 4
26 sshCloseTimeout = 3 * time.Second
27)
28
29var (
30 ErrZeroLengthFrame = errors.New("sshstdio: zero-length frame")
31 ErrFrameTooLarge = errors.New("sshstdio: frame exceeds maximum length")
32)
33
34type closeFunc func() error
35type eofErrFunc func() error
36
37type Option func(*transport)
38
39// WithRuntimeObserver enables process-local transport marks; it deliberately
40// takes only an observer so a carriage adapter never owns trace time.
41func WithRuntimeObserver(observer ports.SerializedRuntimeObserver) Option {
42 return func(t *transport) { t.observer = observer }
43}
44
45// NewTransport wraps separate reader/writer streams as a framed Transport.
46func NewTransport(r io.Reader, w io.Writer, closeFn closeFunc, opts ...Option) ports.Transport {
47 t := newTransport(r, w, closeFn, nil)
48 for _, opt := range opts {
49 if opt != nil {
50 opt(t)
51 }
52 }
53 return t
54}
55
56func newTransport(r io.Reader, w io.Writer, closeFn closeFunc, eofErr eofErrFunc) *transport {
57 if closeFn == nil {
58 closeFn = func() error { return nil }
59 }
60 return &transport{r: r, w: w, close: closeFn, eofErr: eofErr}
61}
62
63type transport struct {
64 r io.Reader
65 w io.Writer
66 close closeFunc
67 eofErr eofErrFunc
68 observer ports.SerializedRuntimeObserver
69 operationMu sync.Mutex
70 operationCount int
71 closing bool
72 operationsDone chan struct{}
73
74 mu sync.Mutex
75 readBuf []byte
76}
77
78func (t *transport) Send(f ports.Frame) error {
79 end := t.beginOperation(ports.RuntimeAdapterSendStart, uint64(len(f.Payload)))
80 n := 1 + len(f.Payload)
81 if n > maxFrameLen {
82 end(false)
83 return ErrFrameTooLarge
84 }
85
86 buf := make([]byte, frameHeaderLen+n)
87 binary.BigEndian.PutUint32(buf[:frameHeaderLen], uint32(n))
88 buf[frameHeaderLen] = byte(f.Type)
89 copy(buf[frameHeaderLen+1:], f.Payload)
90
91 t.mu.Lock()
92 _, err := t.w.Write(buf)
93 t.mu.Unlock()
94 end(err == nil)
95 return err
96}
97
98func (t *transport) Recv() (ports.Frame, error) {
99 end := t.beginOperation(ports.RuntimeAdapterReceiveStart, 0)
100 var hdr [frameHeaderLen]byte
101 if _, err := io.ReadFull(t.r, hdr[:]); err != nil {
102 err = t.mapEOFError(err)
103 end(false)
104 return ports.Frame{}, err
105 }
106
107 n := binary.BigEndian.Uint32(hdr[:])
108 if n == 0 {
109 end(false)
110 return ports.Frame{}, ErrZeroLengthFrame
111 }
112 if n > maxFrameLen {
113 end(false)
114 return ports.Frame{}, ErrFrameTooLarge
115 }
116
117 if cap(t.readBuf) < int(n) {
118 t.readBuf = make([]byte, n)
119 } else {
120 t.readBuf = t.readBuf[:n]
121 }
122 if _, err := io.ReadFull(t.r, t.readBuf); err != nil {
123 err = t.mapEOFError(err)
124 end(false)
125 return ports.Frame{}, err
126 }
127
128 payload := append([]byte(nil), t.readBuf[1:]...)
129 end(true)
130 return ports.Frame{Type: ports.MsgType(t.readBuf[0]), Payload: payload}, nil
131}
132
133func (t *transport) beginOperation(start ports.RuntimeMarkKind, bytes uint64) func(bool) {
134 if t.observer == nil {
135 return func(bool) {}
136 }
137 if !t.beginObservedOperation() {
138 return func(bool) {}
139 }
140 correlation := ports.NewRuntimeCorrelation()
141 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, start, bytes, true))
142 end := ports.RuntimeAdapterSendEnd
143 if start == ports.RuntimeAdapterReceiveStart {
144 end = ports.RuntimeAdapterReceiveEnd
145 }
146 return func(valid bool) {
147 defer t.finishObservedOperation()
148 t.observer.ObserveRuntime(ports.NewRuntimeMarkWithCorrelation("sshstdio", correlation, end, bytes, valid))
149 }
150}
151
152func (t *transport) mapEOFError(err error) error {
153 if t.eofErr == nil || (!errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF)) {
154 return err
155 }
156 if sshErr := t.eofErr(); sshErr != nil {
157 return sshErr
158 }
159 return err
160}
161
162func (t *transport) Close() error {
163 done := t.beginShutdown()
164 // Closing the reader first releases a Recv blocked in io.ReadFull even if
165 // the subprocess shutdown callback reports an error. This lets Recv emit
166 // its failed end mark before the process observer is closed.
167 var readErr error
168 if r, ok := t.r.(io.Closer); ok {
169 readErr = r.Close()
170 }
171 closeErr := t.close()
172 if done != nil {
173 <-done
174 }
175 if closeErr != nil {
176 return closeErr
177 }
178 return readErr
179}
180
181func (t *transport) beginObservedOperation() bool {
182 t.operationMu.Lock()
183 defer t.operationMu.Unlock()
184 if t.closing {
185 return false
186 }
187 if t.operationCount == 0 {
188 t.operationsDone = make(chan struct{})
189 }
190 t.operationCount++
191 return true
192}
193
194func (t *transport) finishObservedOperation() {
195 t.operationMu.Lock()
196 defer t.operationMu.Unlock()
197 t.operationCount--
198 if t.operationCount == 0 {
199 close(t.operationsDone)
200 }
201}
202
203func (t *transport) beginShutdown() <-chan struct{} {
204 t.operationMu.Lock()
205 defer t.operationMu.Unlock()
206 t.closing = true
207 if t.operationCount == 0 {
208 return nil
209 }
210 return t.operationsDone
211}
212
213type processWaiter struct {
214 cmd *exec.Cmd
215 stdin io.Closer
216 stderr *bytes.Buffer
217 timeout time.Duration
218 log *slog.Logger
219 target string
220 session string
221
222 waitOnce sync.Once
223 waitErr error
224}
225
226func newProcessWaiter(cmd *exec.Cmd, stdin io.Closer, stderr *bytes.Buffer, timeout time.Duration, log *slog.Logger, target, session string) *processWaiter {
227 w := &processWaiter{
228 cmd: cmd,
229 stdin: stdin,
230 stderr: stderr,
231 timeout: timeout,
232 log: log,
233 target: target,
234 session: session,
235 }
236 return w
237}
238
239func (w *processWaiter) close() error {
240 _ = w.stdin.Close()
241 return w.wait(w.timeout)
242}
243
244func (w *processWaiter) eofErr() error {
245 return w.wait(sshCloseTimeout)
246}
247
248func (w *processWaiter) wait(timeout time.Duration) error {
249 w.waitOnce.Do(func() {
250 // cmd.Wait must start only after stdout is no longer being read. os/exec
251 // requires callers to finish reading StdoutPipe before Wait, so DialContext
252 // calls this from transport Close or after Recv has already observed EOF.
253 waitCh := make(chan error, 1)
254 go func() { waitCh <- w.cmd.Wait() }()
255 select {
256 case w.waitErr = <-waitCh:
257 case <-time.After(timeout):
258 _ = w.cmd.Process.Kill()
259 w.waitErr = <-waitCh
260 }
261 w.waitErr = formatProcessWaitError(w.waitErr, w.stderr, w.log, w.target, w.session)
262 })
263 return w.waitErr
264}
265
266func formatProcessWaitError(err error, stderr *bytes.Buffer, log *slog.Logger, target, session string) error {
267 if err == nil {
268 return nil
269 }
270 stderrText := strings.TrimSpace(stderr.String())
271 if log != nil {
272 attrs := []any{"target", target, "session", session, "err", err}
273 if stderrText != "" {
274 attrs = append(attrs, "stderr", stderrText)
275 }
276 log.Warn("ssh exited non-cleanly", attrs...)
277 }
278 if stderrText != "" {
279 return fmt.Errorf("sshstdio: ssh exited: %w: %s", err, stderrText)
280 }
281 return fmt.Errorf("sshstdio: ssh exited: %w", err)
282}
283
284func newProcessCloser(cmd *exec.Cmd, stdin io.Closer, stderr *bytes.Buffer, timeout time.Duration, log *slog.Logger, target, session string) closeFunc {
285 return newProcessWaiter(cmd, stdin, stderr, timeout, log, target, session).close
286}
287
288// CommandSpec is the exact ssh argv vev will execute locally.
289type CommandSpec struct {
290 Path string
291 Args []string
292}
293
294// BuildCommand constructs the local ssh subprocess argv without invoking a
295// local shell. OpenSSH sends the remote command as one string for the remote
296// user's shell to interpret, so every remote argv word is POSIX single-quoted.
297func BuildCommand(target, session string) CommandSpec {
298 return BuildCommandForMode(target, "_stdio", session)
299}
300
301// BuildCommandForMode constructs the local ssh subprocess argv for a hidden vev
302// remote mode such as _stdio or _udp-bootstrap.
303func BuildCommandForMode(target, mode, session string) CommandSpec {
304 remote := []string{shellQuote("vev"), shellQuote(mode)}
305 if session != "" {
306 remote = append(remote, shellQuote(session))
307 }
308 args := []string{"--", target, strings.Join(remote, " ")}
309 return CommandSpec{Path: "ssh", Args: args}
310}
311
312func shellQuote(s string) string {
313 return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
314}
315
316// Dial starts ssh target vev _stdio [session] and returns a Transport over the
317// child process' stdio. The subprocess is started with exec.Command argv, never
318// through a shell.
319func Dial(target, session string) (ports.Transport, error) {
320 return DialContext(context.Background(), target, session)
321}
322
323// DialContextWithRuntimeObserver is the option-bearing production constructor.
324func DialContextWithRuntimeObserver(ctx context.Context, target, session string, logger *slog.Logger, opts ...Option) (ports.Transport, error) {
325 return dialContext(ctx, target, session, logger, opts...)
326}
327
328// DialContext is like Dial, but the context is propagated to ssh startup so a
329// canceled attach attempt interrupts the local ssh process. Callers may pass a
330// logger to record ssh start failures and non-clean exits without logging the
331// generated command line.
332func DialContext(ctx context.Context, target, session string, logger ...*slog.Logger) (ports.Transport, error) {
333 var log *slog.Logger
334 if len(logger) > 0 {
335 log = logger[0]
336 }
337 return dialContext(ctx, target, session, log)
338}
339
340func dialContext(ctx context.Context, target, session string, log *slog.Logger, opts ...Option) (ports.Transport, error) {
341 if err := ctx.Err(); err != nil {
342 return nil, err
343 }
344 spec := BuildCommand(target, session)
345 cmd := exec.CommandContext(ctx, spec.Path, spec.Args...)
346 stdin, err := cmd.StdinPipe()
347 if err != nil {
348 return nil, fmt.Errorf("sshstdio: stdin pipe: %w", err)
349 }
350 stdout, err := cmd.StdoutPipe()
351 if err != nil {
352 return nil, fmt.Errorf("sshstdio: stdout pipe: %w", err)
353 }
354 var stderr bytes.Buffer
355 cmd.Stderr = &stderr
356 if err := cmd.Start(); err != nil {
357 if log != nil {
358 log.Error("ssh start failed", "target", target, "session", session, "err", err)
359 }
360 return nil, fmt.Errorf("sshstdio: start ssh: %w", err)
361 }
362
363 waiter := newProcessWaiter(cmd, stdin, &stderr, sshCloseTimeout, log, target, session)
364 transport := newTransport(stdout, stdin, waiter.close, waiter.eofErr)
365 for _, opt := range opts {
366 if opt != nil {
367 opt(transport)
368 }
369 }
370 return transport, nil
371}