zrepl/util/iocommand.go

package util

import (
	"bytes"
	"fmt"
	"golang.org/x/sys/unix"
	"io"
	"os/exec"
	"syscall"
)

// An IOCommand exposes a forked process's std(in|out|err) through the io.ReadWriteCloser interface.
type IOCommand struct {
	Cmd        *exec.Cmd
	Stdin      io.Writer
	Stdout     io.Reader
	StderrBuf  *bytes.Buffer
	ExitResult IOCommandExitResult
}

const IOCommandStderrBufSize = 1024

type IOCommandError struct {
	WaitErr error
	Stderr  []byte
}

type IOCommandExitResult struct {
	Error      error
	WaitStatus syscall.WaitStatus
}

func (e IOCommandError) Error() string {
	return fmt.Sprintf("underlying process exited with error: %s\nstderr: %s\n", e.WaitErr, e.Stderr)
}

func RunIOCommand(command string, args ...string) (c *IOCommand, err error) {
	c, err = NewIOCommand(command, args, IOCommandStderrBufSize)
	if err != nil {
		return
	}
	err = c.Start()
	return
}

func NewIOCommand(command string, args []string, stderrBufSize int) (c *IOCommand, err error) {

	if stderrBufSize == 0 {
		stderrBufSize = IOCommandStderrBufSize
	}

	c = &IOCommand{}

	c.Cmd = exec.Command(command, args...)

	if c.Stdout, err = c.Cmd.StdoutPipe(); err != nil {
		return
	}

	if c.Stdin, err = c.Cmd.StdinPipe(); err != nil {
		return
	}

	c.StderrBuf = bytes.NewBuffer(make([]byte, 0, stderrBufSize))
	c.Cmd.Stderr = c.StderrBuf

	return

}

func (c *IOCommand) Start() (err error) {
	if err = c.Cmd.Start(); err != nil {
		return
	}
	return
}

// Read from process's stdout.
// The behavior after Close()ing is undefined
func (c *IOCommand) Read(buf []byte) (n int, err error) {
	n, err = c.Stdout.Read(buf)
	if err == io.EOF {
		if waitErr := c.doWait(); waitErr != nil {
			err = waitErr
		}
	}
	return
}

func (c *IOCommand) doWait() (err error) {
	waitErr := c.Cmd.Wait()
	waitStatus := c.Cmd.ProcessState.Sys().(syscall.WaitStatus) // Fail hard if we're not on UNIX
	if waitErr != nil {

		// https://support.ssh.com/manuals/client-user/44/ssh2_Return_Values.html
		// If ssh is terminated via signal, its exit status is 128 + signal number
		if waitStatus.ExitStatus() == 128+int(syscall.SIGTERM) {
			// discard wait err, we assume this is due to earlier c.Close()
			goto out
		}

		err = IOCommandError{
			WaitErr: waitErr,
			Stderr:  c.StderrBuf.Bytes(),
		}
	}

out:
	c.ExitResult = IOCommandExitResult{
		Error:      err,
		WaitStatus: waitStatus,
	}
	return
}

// Write to process's stdin.
// The behavior after Close()ing is undefined
func (c *IOCommand) Write(buf []byte) (n int, err error) {
	return c.Stdin.Write(buf)
}

// Terminate the child process and collect its exit status
// It is safe to call Close() multiple times.
func (c *IOCommand) Close() (err error) {
	if c.Cmd.ProcessState == nil {
		// racy...
		err = unix.Kill(c.Cmd.Process.Pid, syscall.SIGTERM)
		if err != nil {
			return
		}
		return c.doWait()
	} else {
		return c.ExitResult.Error
	}
}