Skip to content

common: add ssh.ExecWithOutput to allow streaming output #865

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danishprakash wants to merge 1 commit into
base: main
Choose a base branch
from
Open

common: add ssh.ExecWithOutput to allow streaming output #865

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 81 additions & 0 deletions common/pkg/ssh/connection_golang.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,87 @@ func golangConnectionScp(options ConnectionScpOptions) (*ConnectionScpReport, er
return &ConnectionScpReport{Response: remote.Name()}, nil
}

type golangExecOutputReader struct {
stdout io.Reader
sess *ssh.Session
client *ssh.Client
stderr *bytes.Buffer
waitCh <-chan error
}

func (r *golangExecOutputReader) Read(p []byte) (int, error) {
return r.stdout.Read(p)
}

func (r *golangExecOutputReader) Close() error {
r.sess.Close()
err := <-r.waitCh

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also probably safer to do this in Read when it hits EOF.

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See elsewhere, this could hang on an early Close().

// safe to close the connection now that the session is done
// and the exit status has been received.
Comment on lines +210 to +211

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I’m not sure what this is saying, is it somehow unusual / unexpected?

Perhaps, instead, the golangExecOutputReader.sess field could be documented to exist within client to make the relationship / lifetime explicit there.

r.client.Close()
if err != nil {
return fmt.Errorf("%v: %w", r.stderr.String(), err)
}
return nil
}

func golangConnectionExecWithOutput(options ConnectionExecOptions, input io.Reader) (io.ReadCloser, error) {
if !strings.HasPrefix(options.Host, "ssh://") {
options.Host = "ssh://" + options.Host
}
_, uri, err := Validate(options.User, options.Host, options.Port, options.Identity)
if err != nil {
return nil, err
}

cfg, err := ValidateAndConfigure(uri, options.Identity, false)
if err != nil {
return nil, err
}
dialAdd, err := ssh.Dial("tcp", uri.Host, cfg)
if err != nil {
return nil, fmt.Errorf("failed to connect: %w", err)
}
Comment on lines +220 to +235

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think sharing / not duplicating at least this part would be valuable. We already have 5 callers of ssh.Dial here in 3 different patterns — one of the differences looks justifiable, but it’s unclear.

I’m not at all asking for cleaning up the existing differences — but consolidating the existing code where it is common, and calling a shared helper, instead of adding another copy that can diverge over time would be valuable.


sess, err := dialAdd.NewSession()
if err != nil {
dialAdd.Close()
return nil, err
}

stderr := &bytes.Buffer{}
sess.Stderr = stderr
if input != nil {
sess.Stdin = input
}

stdout, err := sess.StdoutPipe()
if err != nil {
sess.Close()
dialAdd.Close()
return nil, err
}

if err := sess.Start(strings.Join(options.Args, " ")); err != nil {
sess.Close()
dialAdd.Close()
return nil, err
}

waitCh := make(chan error, 1)
go func() {
waitCh <- sess.Wait()
}()

return &golangExecOutputReader{
stdout: stdout,
sess: sess,
client: dialAdd,
stderr: stderr,
waitCh: waitCh,
}, nil
}

// ExecRemoteCommand takes a ssh client connection and a command to run and executes the
// command on the specified client. The function returns the Stdout from the client or the Stderr.
func ExecRemoteCommand(dial *ssh.Client, run string) ([]byte, error) {
Expand Down
68 changes: 57 additions & 11 deletions common/pkg/ssh/connection_native.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,27 +101,25 @@ func nativeConnectionCreate(options ConnectionCreateOptions) error {
})
}

func nativeConnectionExec(options ConnectionExecOptions, input io.Reader) (*ConnectionExecReport, error) {
func nativePrepareSSHCmd(options ConnectionExecOptions, input io.Reader) (*exec.Cmd, *bytes.Buffer, error) {
dst, uri, err := Validate(options.User, options.Host, options.Port, options.Identity)
if err != nil {
return nil, err
return nil, nil, err
}

ssh, err := exec.LookPath("ssh")
if err != nil {
return nil, err
return nil, nil, err
}

output := &bytes.Buffer{}
errors := &bytes.Buffer{}
if host, _, ok := strings.Cut(uri.Host, "/run"); ok {
uri.Host = host
}

options.Args = append([]string{uri.User.String() + "@" + uri.Hostname()}, options.Args...)
conf, err := config.Default()
if err != nil {
return nil, err
return nil, nil, err
}

args := []string{}
Expand All @@ -132,19 +130,67 @@ func nativeConnectionExec(options ConnectionExecOptions, input io.Reader) (*Conn
args = append(args, "-F", conf.Engine.SSHConfig)
}
args = append(args, options.Args...)
info := exec.Command(ssh, args...)
info.Stdout = output
info.Stderr = errors

stderr := &bytes.Buffer{}
cmd := exec.Command(ssh, args...)
cmd.Stderr = stderr
if input != nil {
info.Stdin = input
cmd.Stdin = input
}
err = info.Run()

return cmd, stderr, nil
}

func nativeConnectionExec(options ConnectionExecOptions, input io.Reader) (*ConnectionExecReport, error) {
cmd, stderr, err := nativePrepareSSHCmd(options, input)
if err != nil {
return nil, err
}

output := &bytes.Buffer{}
cmd.Stdout = output
if err := cmd.Run(); err != nil {

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AFAICS this could use cmd.Output() without having to implement the “collect stdout” and “include stderr on failure” parts; so nativePrepareSSHCmd can have a simpler interface.

return nil, fmt.Errorf("%v: %w", stderr.String(), err)
}
return &ConnectionExecReport{Response: output.String()}, nil
}

type nativeExecOutputReader struct {
stdout io.ReadCloser
cmd *exec.Cmd
stderr *bytes.Buffer
}

func (r *nativeExecOutputReader) Read(p []byte) (int, error) {
return r.stdout.Read(p)
}

func (r *nativeExecOutputReader) Close() error {
err := r.cmd.Wait()

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can work but it requires the caller to check errors on a ReadCloser.Close, which most ReadCloser consumers don’t bother to do.

I think doing this in Read at the point r.stdout.Read hits an EOF would be more reliable, admittedly at the cost of some complexity.

If this design remained, at the very least the API would need a fairly loud documentation that the caller must check errors on .Close().

@danishprakash danishprakash Jun 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right; I didn't think of usage semantics. I'm a bit unsure of waiting inside Read() that doesn't really scream idiomatic Go to me but I could be wrong. On the other hand, expecting users to check errors on Close() is also fragile as you mentioned.

I found that using io.Pipe might be suitable. With that, we use let the process write stdout to the write end of the pipe and we start the command in a goroutine while instantly returning the read end of the pipe.

func nativeConnectionExecWithOutput(...) (io.ReadCloser, error) {
  cmd, _ := nativePrepareSSHCmd(options, input)
  stderr := &bytes.Buffer{}
  cmd.Stderr = stderr

  r, w := io.Pipe()
  cmd.Stdout = w

  go func() {
      err := cmd.Run()
      if err != nil {
          err = fmt.Errorf("%v: %w", stderr.String(), err)
      }
      w.CloseWithError(err)
  }()

  return r, nil
}

CloseWithError ensures that upon an error, the next read on the pipe, the caller would receive the error so effectively moving error propogation from Close to Read.

@mtrmac mtrmac Jun 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit unsure of waiting inside Read() that doesn't really scream idiomatic Go to me but I could be wrong.

Intuitively, I think that if we decide what should happen on reader.Close without reading all of input (and thus blocking the producer child: are we fine with blocking forever? Do we return to the caller but leave the process running? Do we aim for a SIGPIPE/EPIPE? Do we explicitly kill the process?), that will more strictly drive the implementation of the other parts.

I found that using io.Pipe might be suitable.

If I’m reading the code correctly, closing the reader end should end up closing a pipe(2) reader end, with a SIGPIPE/EPIPE to the child process. That might be fine — but that goroutine would continue to exist until the command actually exits. So, see above about deciding the goal in that situation first.

@mtrmac mtrmac Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the caller fails due to an unrelated error and calls Close early, couldn’t this deadlock? We never consume the rest of r.stdout, the child process could block on writing to its stdout, Wait will never return.

@danishprakash danishprakash Jun 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, one option is for us to kill the process on Close, or to provide the command under a cancellable context and the caller then calls cancel()s on an early close.

if err != nil {
return fmt.Errorf("%v: %w", r.stderr.String(), err)
}
return nil
}

func nativeConnectionExecWithOutput(options ConnectionExecOptions, input io.Reader) (io.ReadCloser, error) {
cmd, stderr, err := nativePrepareSSHCmd(options, input)
if err != nil {
return nil, err
}

stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, err
}

if err := cmd.Start(); err != nil {
return nil, err
}

return &nativeExecOutputReader{stdout: stdout, cmd: cmd, stderr: stderr}, nil
}

func nativeConnectionScp(options ConnectionScpOptions) (*ConnectionScpReport, error) {
host, remotePath, localPath, swap, err := ParseScpArgs(options)
if err != nil {
Expand Down
7 changes: 7 additions & 0 deletions common/pkg/ssh/ssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,13 @@ func ExecWithInput(options *ConnectionExecOptions, kind EngineMode, input io.Rea
return rep.Response, nil
}

func ExecWithOutput(options *ConnectionExecOptions, kind EngineMode, input io.Reader) (io.ReadCloser, error) {
if kind == NativeMode {
return nativeConnectionExecWithOutput(*options, input)
}
return golangConnectionExecWithOutput(*options, input)
}

func Scp(options *ConnectionScpOptions, kind EngineMode) (string, error) {
var rep *ConnectionScpReport
var err error
Expand Down
21 changes: 21 additions & 0 deletions common/pkg/ssh/ssh_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,27 @@ func TestExecWithInput(t *testing.T) {
require.Error(t, err, "failed to connect: ssh: handshake failed: ssh: disconnect, reason 2: Too many authentication failures")
}

func TestExecWithOutput(t *testing.T) {
input, err := os.Open("/etc/fstab")
require.NoError(t, err)
defer input.Close()

options := ConnectionExecOptions{
Port: 22,
Host: "localhost",
Args: []string{"md5sum"},
}

// Native mode: Start() succeeds but the connection fails when the process runs.
// The error surfaces when the caller closes the reader.
rc, err := ExecWithOutput(&options, NativeMode, input)
require.NoError(t, err)
require.Error(t, rc.Close(), "exit status 255")

_, err = ExecWithOutput(&options, GolangMode, input)
require.Error(t, err, "failed to connect: ssh: handshake failed: ssh: disconnect, reason 2: Too many authentication failures")
}

func TestDial(t *testing.T) {
options := ConnectionDialOptions{
Port: 22,
Expand Down
Loading