github.com/StackPointCloud/packer@v0.10.2-0.20180716202532-b28098e0f79b/packer/communicator.go (about)

     1  package packer
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"strings"
     7  	"sync"
     8  	"unicode"
     9  
    10  	"github.com/mitchellh/iochan"
    11  )
    12  
    13  // CmdDisconnect is a sentinel value to indicate a RemoteCmd
    14  // exited because the remote side disconnected us.
    15  const CmdDisconnect int = 2300218
    16  
    17  // RemoteCmd represents a remote command being prepared or run.
    18  type RemoteCmd struct {
    19  	// Command is the command to run remotely. This is executed as if
    20  	// it were a shell command, so you are expected to do any shell escaping
    21  	// necessary.
    22  	Command string
    23  
    24  	// Stdin specifies the process's standard input. If Stdin is
    25  	// nil, the process reads from an empty bytes.Buffer.
    26  	Stdin io.Reader
    27  
    28  	// Stdout and Stderr represent the process's standard output and
    29  	// error.
    30  	//
    31  	// If either is nil, it will be set to ioutil.Discard.
    32  	Stdout io.Writer
    33  	Stderr io.Writer
    34  
    35  	// This will be set to true when the remote command has exited. It
    36  	// shouldn't be set manually by the user, but there is no harm in
    37  	// doing so.
    38  	Exited bool
    39  
    40  	// Once Exited is true, this will contain the exit code of the process.
    41  	ExitStatus int
    42  
    43  	// Internal fields
    44  	exitCh chan struct{}
    45  
    46  	// This thing is a mutex, lock when making modifications concurrently
    47  	sync.Mutex
    48  }
    49  
    50  // A Communicator is the interface used to communicate with the machine
    51  // that exists that will eventually be packaged into an image. Communicators
    52  // allow you to execute remote commands, upload files, etc.
    53  //
    54  // Communicators must be safe for concurrency, meaning multiple calls to
    55  // Start or any other method may be called at the same time.
    56  type Communicator interface {
    57  	// Start takes a RemoteCmd and starts it. The RemoteCmd must not be
    58  	// modified after being used with Start, and it must not be used with
    59  	// Start again. The Start method returns immediately once the command
    60  	// is started. It does not wait for the command to complete. The
    61  	// RemoteCmd.Exited field should be used for this.
    62  	Start(*RemoteCmd) error
    63  
    64  	// Upload uploads a file to the machine to the given path with the
    65  	// contents coming from the given reader. This method will block until
    66  	// it completes.
    67  	Upload(string, io.Reader, *os.FileInfo) error
    68  
    69  	// UploadDir uploads the contents of a directory recursively to
    70  	// the remote path. It also takes an optional slice of paths to
    71  	// ignore when uploading.
    72  	//
    73  	// The folder name of the source folder should be created unless there
    74  	// is a trailing slash on the source "/". For example: "/tmp/src" as
    75  	// the source will create a "src" directory in the destination unless
    76  	// a trailing slash is added. This is identical behavior to rsync(1).
    77  	UploadDir(dst string, src string, exclude []string) error
    78  
    79  	// Download downloads a file from the machine from the given remote path
    80  	// with the contents writing to the given writer. This method will
    81  	// block until it completes.
    82  	Download(string, io.Writer) error
    83  
    84  	DownloadDir(src string, dst string, exclude []string) error
    85  }
    86  
    87  // StartWithUi runs the remote command and streams the output to any
    88  // configured Writers for stdout/stderr, while also writing each line
    89  // as it comes to a Ui.
    90  func (r *RemoteCmd) StartWithUi(c Communicator, ui Ui) error {
    91  	stdout_r, stdout_w := io.Pipe()
    92  	stderr_r, stderr_w := io.Pipe()
    93  	defer stdout_w.Close()
    94  	defer stderr_w.Close()
    95  
    96  	// Retain the original stdout/stderr that we can replace back in.
    97  	originalStdout := r.Stdout
    98  	originalStderr := r.Stderr
    99  	defer func() {
   100  		r.Lock()
   101  		defer r.Unlock()
   102  
   103  		r.Stdout = originalStdout
   104  		r.Stderr = originalStderr
   105  	}()
   106  
   107  	// Set the writers for the output so that we get it streamed to us
   108  	if r.Stdout == nil {
   109  		r.Stdout = stdout_w
   110  	} else {
   111  		r.Stdout = io.MultiWriter(r.Stdout, stdout_w)
   112  	}
   113  
   114  	if r.Stderr == nil {
   115  		r.Stderr = stderr_w
   116  	} else {
   117  		r.Stderr = io.MultiWriter(r.Stderr, stderr_w)
   118  	}
   119  
   120  	// Start the command
   121  	if err := c.Start(r); err != nil {
   122  		return err
   123  	}
   124  
   125  	// Create the channels we'll use for data
   126  	exitCh := make(chan struct{})
   127  	stdoutCh := iochan.DelimReader(stdout_r, '\n')
   128  	stderrCh := iochan.DelimReader(stderr_r, '\n')
   129  
   130  	// Start the goroutine to watch for the exit
   131  	go func() {
   132  		defer close(exitCh)
   133  		defer stdout_w.Close()
   134  		defer stderr_w.Close()
   135  		r.Wait()
   136  	}()
   137  
   138  	// Loop and get all our output
   139  OutputLoop:
   140  	for {
   141  		select {
   142  		case output := <-stderrCh:
   143  			if output != "" {
   144  				ui.Message(r.cleanOutputLine(output))
   145  			}
   146  		case output := <-stdoutCh:
   147  			if output != "" {
   148  				ui.Message(r.cleanOutputLine(output))
   149  			}
   150  		case <-exitCh:
   151  			break OutputLoop
   152  		}
   153  	}
   154  
   155  	// Make sure we finish off stdout/stderr because we may have gotten
   156  	// a message from the exit channel before finishing these first.
   157  	for output := range stdoutCh {
   158  		ui.Message(r.cleanOutputLine(output))
   159  	}
   160  
   161  	for output := range stderrCh {
   162  		ui.Message(r.cleanOutputLine(output))
   163  	}
   164  
   165  	return nil
   166  }
   167  
   168  // SetExited is a helper for setting that this process is exited. This
   169  // should be called by communicators who are running a remote command in
   170  // order to set that the command is done.
   171  func (r *RemoteCmd) SetExited(status int) {
   172  	r.Lock()
   173  	defer r.Unlock()
   174  
   175  	if r.exitCh == nil {
   176  		r.exitCh = make(chan struct{})
   177  	}
   178  
   179  	r.Exited = true
   180  	r.ExitStatus = status
   181  	close(r.exitCh)
   182  }
   183  
   184  // Wait waits for the remote command to complete.
   185  func (r *RemoteCmd) Wait() {
   186  	// Make sure our condition variable is initialized.
   187  	r.Lock()
   188  	if r.exitCh == nil {
   189  		r.exitCh = make(chan struct{})
   190  	}
   191  	r.Unlock()
   192  
   193  	<-r.exitCh
   194  }
   195  
   196  // cleanOutputLine cleans up a line so that '\r' don't muck up the
   197  // UI output when we're reading from a remote command.
   198  func (r *RemoteCmd) cleanOutputLine(line string) string {
   199  	// Trim surrounding whitespace
   200  	line = strings.TrimRightFunc(line, unicode.IsSpace)
   201  
   202  	// Trim up to the first carriage return, since that text would be
   203  	// lost anyways.
   204  	idx := strings.LastIndex(line, "\r")
   205  	if idx > -1 {
   206  		line = line[idx+1:]
   207  	}
   208  
   209  	return line
   210  }