packer-cn/packer/communicator.go

package packer

import (
	"context"
	"io"
	"os"
	"sync"

	"golang.org/x/sync/errgroup"

	"github.com/mitchellh/iochan"
)

// CmdDisconnect is a sentinel value to indicate a RemoteCmd
// exited because the remote side disconnected us.
const CmdDisconnect int = 2300218

// RemoteCmd represents a remote command being prepared or run.
type RemoteCmd struct {
	// Command is the command to run remotely. This is executed as if
	// it were a shell command, so you are expected to do any shell escaping
	// necessary.
	Command string

	// Stdin specifies the process's standard input. If Stdin is
	// nil, the process reads from an empty bytes.Buffer.
	Stdin io.Reader

	// Stdout and Stderr represent the process's standard output and
	// error.
	//
	// If either is nil, it will be set to ioutil.Discard.
	Stdout io.Writer
	Stderr io.Writer

	// Once Exited is true, this will contain the exit code of the process.
	exitStatus int

	// This thing is a mutex, lock when making modifications concurrently
	sync.Mutex

	exitChInit sync.Once
	exitCh     chan interface{}
}

// A Communicator is the interface used to communicate with the machine
// that exists that will eventually be packaged into an image. Communicators
// allow you to execute remote commands, upload files, etc.
//
// Communicators must be safe for concurrency, meaning multiple calls to
// Start or any other method may be called at the same time.
type Communicator interface {
	// Start takes a RemoteCmd and starts it. The RemoteCmd must not be
	// modified after being used with Start, and it must not be used with
	// Start again. The Start method returns immediately once the command
	// is started. It does not wait for the command to complete. The
	// RemoteCmd.Exited field should be used for this.
	Start(context.Context, *RemoteCmd) error

	// Upload uploads a file to the machine to the given path with the
	// contents coming from the given reader. This method will block until
	// it completes.
	Upload(string, io.Reader, *os.FileInfo) error

	// UploadDir uploads the contents of a directory recursively to
	// the remote path. It also takes an optional slice of paths to
	// ignore when uploading.
	//
	// The folder name of the source folder should be created unless there
	// is a trailing slash on the source "/". For example: "/tmp/src" as
	// the source will create a "src" directory in the destination unless
	// a trailing slash is added. This is identical behavior to rsync(1).
	UploadDir(dst string, src string, exclude []string) error

	// Download downloads a file from the machine from the given remote path
	// with the contents writing to the given writer. This method will
	// block until it completes.
	Download(string, io.Writer) error

	DownloadDir(src string, dst string, exclude []string) error
}

// RunWithUi runs the remote command and streams the output to any configured
// Writers for stdout/stderr, while also writing each line as it comes to a Ui.
// RunWithUi will not return until the command finishes or is cancelled.
func (r *RemoteCmd) RunWithUi(ctx context.Context, c Communicator, ui Ui) error {
	r.initchan()

	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()
	defer stdout_w.Close()
	defer stderr_w.Close()

	// Retain the original stdout/stderr that we can replace back in.
	originalStdout := r.Stdout
	originalStderr := r.Stderr
	defer func() {
		r.Lock()
		defer r.Unlock()

		r.Stdout = originalStdout
		r.Stderr = originalStderr
	}()

	// Set the writers for the output so that we get it streamed to us
	if r.Stdout == nil {
		r.Stdout = stdout_w
	} else {
		r.Stdout = io.MultiWriter(r.Stdout, stdout_w)
	}

	if r.Stderr == nil {
		r.Stderr = stderr_w
	} else {
		r.Stderr = io.MultiWriter(r.Stderr, stderr_w)
	}

	// Loop and get all our output until done.
	printFn := func(in io.Reader, out func(string)) error {
		for output := range iochan.LineReader(in) {
			if output != "" {
				out(output)
			}
		}
		return nil
	}

	wg, ctx := errgroup.WithContext(ctx)

	wg.Go(func() error { return printFn(stdout_r, ui.Message) })
	wg.Go(func() error { return printFn(stderr_r, ui.Error) })

	if err := c.Start(ctx, r); err != nil {
		return err
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-r.exitCh:
		return nil
	}
}

// SetExited is a helper for setting that this process is exited. This
// should be called by communicators who are running a remote command in
// order to set that the command is done.
func (r *RemoteCmd) SetExited(status int) {
	r.initchan()

	r.Lock()
	r.exitStatus = status
	r.Unlock()

	close(r.exitCh)
}

// Wait for command exit and return exit status
func (r *RemoteCmd) Wait() int {
	r.initchan()
	<-r.exitCh
	r.Lock()
	defer r.Unlock()
	return r.exitStatus
}

func (r *RemoteCmd) ExitStatus() int {
	return r.Wait()
}

func (r *RemoteCmd) initchan() {
	r.exitChInit.Do(func() {
		if r.exitCh == nil {
			r.exitCh = make(chan interface{})
		}
	})
}