github.com/cycloidio/terraform@v1.1.10-0.20220513142504-76d5c768dc63/communicator/ssh/communicator.go (about)

     1  package ssh
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"context"
     7  	"errors"
     8  	"fmt"
     9  	"io"
    10  	"io/ioutil"
    11  	"log"
    12  	"math/rand"
    13  	"net"
    14  	"os"
    15  	"path/filepath"
    16  	"strconv"
    17  	"strings"
    18  	"sync"
    19  	"time"
    20  
    21  	"github.com/apparentlymart/go-shquot/shquot"
    22  	"github.com/cycloidio/terraform/communicator/remote"
    23  	"github.com/cycloidio/terraform/provisioners"
    24  	"github.com/zclconf/go-cty/cty"
    25  	"golang.org/x/crypto/ssh"
    26  	"golang.org/x/crypto/ssh/agent"
    27  
    28  	_ "github.com/cycloidio/terraform/logging"
    29  )
    30  
    31  const (
    32  	// DefaultShebang is added at the top of a SSH script file
    33  	DefaultShebang = "#!/bin/sh\n"
    34  )
    35  
    36  var (
    37  	// randShared is a global random generator object that is shared.  This must be
    38  	// shared since it is seeded by the current time and creating multiple can
    39  	// result in the same values. By using a shared RNG we assure different numbers
    40  	// per call.
    41  	randLock   sync.Mutex
    42  	randShared *rand.Rand
    43  
    44  	// enable ssh keeplive probes by default
    45  	keepAliveInterval = 2 * time.Second
    46  
    47  	// max time to wait for for a KeepAlive response before considering the
    48  	// connection to be dead.
    49  	maxKeepAliveDelay = 120 * time.Second
    50  )
    51  
    52  // Communicator represents the SSH communicator
    53  type Communicator struct {
    54  	connInfo        *connectionInfo
    55  	client          *ssh.Client
    56  	config          *sshConfig
    57  	conn            net.Conn
    58  	cancelKeepAlive context.CancelFunc
    59  
    60  	lock sync.Mutex
    61  }
    62  
    63  type sshConfig struct {
    64  	// The configuration of the Go SSH connection
    65  	config *ssh.ClientConfig
    66  
    67  	// connection returns a new connection. The current connection
    68  	// in use will be closed as part of the Close method, or in the
    69  	// case an error occurs.
    70  	connection func() (net.Conn, error)
    71  
    72  	// noPty, if true, will not request a pty from the remote end.
    73  	noPty bool
    74  
    75  	// sshAgent is a struct surrounding the agent.Agent client and the net.Conn
    76  	// to the SSH Agent. It is nil if no SSH agent is configured
    77  	sshAgent *sshAgent
    78  }
    79  
    80  type fatalError struct {
    81  	error
    82  }
    83  
    84  func (e fatalError) FatalError() error {
    85  	return e.error
    86  }
    87  
    88  // New creates a new communicator implementation over SSH.
    89  func New(v cty.Value) (*Communicator, error) {
    90  	connInfo, err := parseConnectionInfo(v)
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  
    95  	config, err := prepareSSHConfig(connInfo)
    96  	if err != nil {
    97  		return nil, err
    98  	}
    99  
   100  	// Set up the random number generator once. The seed value is the
   101  	// time multiplied by the PID. This can overflow the int64 but that
   102  	// is okay. We multiply by the PID in case we have multiple processes
   103  	// grabbing this at the same time. This is possible with Terraform and
   104  	// if we communicate to the same host at the same instance, we could
   105  	// overwrite the same files. Multiplying by the PID prevents this.
   106  	randLock.Lock()
   107  	defer randLock.Unlock()
   108  	if randShared == nil {
   109  		randShared = rand.New(rand.NewSource(
   110  			time.Now().UnixNano() * int64(os.Getpid())))
   111  	}
   112  
   113  	comm := &Communicator{
   114  		connInfo: connInfo,
   115  		config:   config,
   116  	}
   117  
   118  	return comm, nil
   119  }
   120  
   121  // Connect implementation of communicator.Communicator interface
   122  func (c *Communicator) Connect(o provisioners.UIOutput) (err error) {
   123  	// Grab a lock so we can modify our internal attributes
   124  	c.lock.Lock()
   125  	defer c.lock.Unlock()
   126  
   127  	if c.conn != nil {
   128  		c.conn.Close()
   129  	}
   130  
   131  	// Set the conn and client to nil since we'll recreate it
   132  	c.conn = nil
   133  	c.client = nil
   134  
   135  	if o != nil {
   136  		o.Output(fmt.Sprintf(
   137  			"Connecting to remote host via SSH...\n"+
   138  				"  Host: %s\n"+
   139  				"  User: %s\n"+
   140  				"  Password: %t\n"+
   141  				"  Private key: %t\n"+
   142  				"  Certificate: %t\n"+
   143  				"  SSH Agent: %t\n"+
   144  				"  Checking Host Key: %t\n"+
   145  				"  Target Platform: %s\n",
   146  			c.connInfo.Host, c.connInfo.User,
   147  			c.connInfo.Password != "",
   148  			c.connInfo.PrivateKey != "",
   149  			c.connInfo.Certificate != "",
   150  			c.connInfo.Agent,
   151  			c.connInfo.HostKey != "",
   152  			c.connInfo.TargetPlatform,
   153  		))
   154  
   155  		if c.connInfo.BastionHost != "" {
   156  			o.Output(fmt.Sprintf(
   157  				"Using configured bastion host...\n"+
   158  					"  Host: %s\n"+
   159  					"  User: %s\n"+
   160  					"  Password: %t\n"+
   161  					"  Private key: %t\n"+
   162  					"  Certificate: %t\n"+
   163  					"  SSH Agent: %t\n"+
   164  					"  Checking Host Key: %t",
   165  				c.connInfo.BastionHost, c.connInfo.BastionUser,
   166  				c.connInfo.BastionPassword != "",
   167  				c.connInfo.BastionPrivateKey != "",
   168  				c.connInfo.BastionCertificate != "",
   169  				c.connInfo.Agent,
   170  				c.connInfo.BastionHostKey != "",
   171  			))
   172  		}
   173  	}
   174  
   175  	hostAndPort := fmt.Sprintf("%s:%d", c.connInfo.Host, c.connInfo.Port)
   176  	log.Printf("[DEBUG] Connecting to %s for SSH", hostAndPort)
   177  	c.conn, err = c.config.connection()
   178  	if err != nil {
   179  		// Explicitly set this to the REAL nil. Connection() can return
   180  		// a nil implementation of net.Conn which will make the
   181  		// "if c.conn == nil" check fail above. Read here for more information
   182  		// on this psychotic language feature:
   183  		//
   184  		// http://golang.org/doc/faq#nil_error
   185  		c.conn = nil
   186  
   187  		log.Printf("[ERROR] connection error: %s", err)
   188  		return err
   189  	}
   190  
   191  	log.Printf("[DEBUG] Connection established. Handshaking for user %v", c.connInfo.User)
   192  	sshConn, sshChan, req, err := ssh.NewClientConn(c.conn, hostAndPort, c.config.config)
   193  	if err != nil {
   194  		err = fmt.Errorf("SSH authentication failed (%s@%s): %w", c.connInfo.User, hostAndPort, err)
   195  
   196  		// While in theory this should be a fatal error, some hosts may start
   197  		// the ssh service before it is properly configured, or before user
   198  		// authentication data is available.
   199  		// Log the error, and allow the provisioner to retry.
   200  		log.Printf("[WARN] %s", err)
   201  		return err
   202  	}
   203  
   204  	c.client = ssh.NewClient(sshConn, sshChan, req)
   205  
   206  	if c.config.sshAgent != nil {
   207  		log.Printf("[DEBUG] Telling SSH config to forward to agent")
   208  		if err := c.config.sshAgent.ForwardToAgent(c.client); err != nil {
   209  			return fatalError{err}
   210  		}
   211  
   212  		log.Printf("[DEBUG] Setting up a session to request agent forwarding")
   213  		session, err := c.client.NewSession()
   214  		if err != nil {
   215  			return err
   216  		}
   217  		defer session.Close()
   218  
   219  		err = agent.RequestAgentForwarding(session)
   220  
   221  		if err == nil {
   222  			log.Printf("[INFO] agent forwarding enabled")
   223  		} else {
   224  			log.Printf("[WARN] error forwarding agent: %s", err)
   225  		}
   226  	}
   227  
   228  	if err != nil {
   229  		return err
   230  	}
   231  
   232  	if o != nil {
   233  		o.Output("Connected!")
   234  	}
   235  
   236  	ctx, cancelKeepAlive := context.WithCancel(context.TODO())
   237  	c.cancelKeepAlive = cancelKeepAlive
   238  
   239  	// Start a keepalive goroutine to help maintain the connection for
   240  	// long-running commands.
   241  	log.Printf("[DEBUG] starting ssh KeepAlives")
   242  
   243  	// We want a local copy of the ssh client pointer, so that a reconnect
   244  	// doesn't race with the running keep-alive loop.
   245  	sshClient := c.client
   246  	go func() {
   247  		defer cancelKeepAlive()
   248  		// Along with the KeepAlives generating packets to keep the tcp
   249  		// connection open, we will use the replies to verify liveness of the
   250  		// connection. This will prevent dead connections from blocking the
   251  		// provisioner indefinitely.
   252  		respCh := make(chan error, 1)
   253  
   254  		go func() {
   255  			t := time.NewTicker(keepAliveInterval)
   256  			defer t.Stop()
   257  			for {
   258  				select {
   259  				case <-t.C:
   260  					_, _, err := sshClient.SendRequest("keepalive@terraform.io", true, nil)
   261  					respCh <- err
   262  				case <-ctx.Done():
   263  					return
   264  				}
   265  			}
   266  		}()
   267  
   268  		after := time.NewTimer(maxKeepAliveDelay)
   269  		defer after.Stop()
   270  
   271  		for {
   272  			select {
   273  			case err := <-respCh:
   274  				if err != nil {
   275  					log.Printf("[ERROR] ssh keepalive: %s", err)
   276  					sshConn.Close()
   277  					return
   278  				}
   279  			case <-after.C:
   280  				// abort after too many missed keepalives
   281  				log.Println("[ERROR] no reply from ssh server")
   282  				sshConn.Close()
   283  				return
   284  			case <-ctx.Done():
   285  				return
   286  			}
   287  			if !after.Stop() {
   288  				<-after.C
   289  			}
   290  			after.Reset(maxKeepAliveDelay)
   291  		}
   292  	}()
   293  
   294  	return nil
   295  }
   296  
   297  // Disconnect implementation of communicator.Communicator interface
   298  func (c *Communicator) Disconnect() error {
   299  	c.lock.Lock()
   300  	defer c.lock.Unlock()
   301  
   302  	if c.cancelKeepAlive != nil {
   303  		c.cancelKeepAlive()
   304  	}
   305  
   306  	if c.config.sshAgent != nil {
   307  		if err := c.config.sshAgent.Close(); err != nil {
   308  			return err
   309  		}
   310  	}
   311  
   312  	if c.conn != nil {
   313  		conn := c.conn
   314  		c.conn = nil
   315  		return conn.Close()
   316  	}
   317  
   318  	return nil
   319  }
   320  
   321  // Timeout implementation of communicator.Communicator interface
   322  func (c *Communicator) Timeout() time.Duration {
   323  	return c.connInfo.TimeoutVal
   324  }
   325  
   326  // ScriptPath implementation of communicator.Communicator interface
   327  func (c *Communicator) ScriptPath() string {
   328  	randLock.Lock()
   329  	defer randLock.Unlock()
   330  
   331  	return strings.Replace(
   332  		c.connInfo.ScriptPath, "%RAND%",
   333  		strconv.FormatInt(int64(randShared.Int31()), 10), -1)
   334  }
   335  
   336  // Start implementation of communicator.Communicator interface
   337  func (c *Communicator) Start(cmd *remote.Cmd) error {
   338  	cmd.Init()
   339  
   340  	session, err := c.newSession()
   341  	if err != nil {
   342  		return err
   343  	}
   344  
   345  	// Set up our session
   346  	session.Stdin = cmd.Stdin
   347  	session.Stdout = cmd.Stdout
   348  	session.Stderr = cmd.Stderr
   349  
   350  	if !c.config.noPty && c.connInfo.TargetPlatform != TargetPlatformWindows {
   351  		// Request a PTY
   352  		termModes := ssh.TerminalModes{
   353  			ssh.ECHO:          0,     // do not echo
   354  			ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
   355  			ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
   356  		}
   357  
   358  		if err := session.RequestPty("xterm", 80, 40, termModes); err != nil {
   359  			return err
   360  		}
   361  	}
   362  
   363  	log.Printf("[DEBUG] starting remote command: %s", cmd.Command)
   364  	err = session.Start(strings.TrimSpace(cmd.Command) + "\n")
   365  	if err != nil {
   366  		return err
   367  	}
   368  
   369  	// Start a goroutine to wait for the session to end and set the
   370  	// exit boolean and status.
   371  	go func() {
   372  		defer session.Close()
   373  
   374  		err := session.Wait()
   375  		exitStatus := 0
   376  		if err != nil {
   377  			exitErr, ok := err.(*ssh.ExitError)
   378  			if ok {
   379  				exitStatus = exitErr.ExitStatus()
   380  			}
   381  		}
   382  
   383  		cmd.SetExitStatus(exitStatus, err)
   384  		log.Printf("[DEBUG] remote command exited with '%d': %s", exitStatus, cmd.Command)
   385  	}()
   386  
   387  	return nil
   388  }
   389  
   390  // Upload implementation of communicator.Communicator interface
   391  func (c *Communicator) Upload(path string, input io.Reader) error {
   392  	// The target directory and file for talking the SCP protocol
   393  	targetDir := filepath.Dir(path)
   394  	targetFile := filepath.Base(path)
   395  
   396  	// On windows, filepath.Dir uses backslash separators (ie. "\tmp").
   397  	// This does not work when the target host is unix.  Switch to forward slash
   398  	// which works for unix and windows
   399  	targetDir = filepath.ToSlash(targetDir)
   400  
   401  	// Skip copying if we can get the file size directly from common io.Readers
   402  	size := int64(0)
   403  
   404  	switch src := input.(type) {
   405  	case *os.File:
   406  		fi, err := src.Stat()
   407  		if err != nil {
   408  			size = fi.Size()
   409  		}
   410  	case *bytes.Buffer:
   411  		size = int64(src.Len())
   412  	case *bytes.Reader:
   413  		size = int64(src.Len())
   414  	case *strings.Reader:
   415  		size = int64(src.Len())
   416  	}
   417  
   418  	scpFunc := func(w io.Writer, stdoutR *bufio.Reader) error {
   419  		return scpUploadFile(targetFile, input, w, stdoutR, size)
   420  	}
   421  
   422  	cmd, err := quoteShell([]string{"scp", "-vt", targetDir}, c.connInfo.TargetPlatform)
   423  	if err != nil {
   424  		return err
   425  	}
   426  	return c.scpSession(cmd, scpFunc)
   427  }
   428  
   429  // UploadScript implementation of communicator.Communicator interface
   430  func (c *Communicator) UploadScript(path string, input io.Reader) error {
   431  	reader := bufio.NewReader(input)
   432  	prefix, err := reader.Peek(2)
   433  	if err != nil {
   434  		return fmt.Errorf("Error reading script: %s", err)
   435  	}
   436  	var script bytes.Buffer
   437  
   438  	if string(prefix) != "#!" && c.connInfo.TargetPlatform != TargetPlatformWindows {
   439  		script.WriteString(DefaultShebang)
   440  	}
   441  	script.ReadFrom(reader)
   442  
   443  	if err := c.Upload(path, &script); err != nil {
   444  		return err
   445  	}
   446  	if c.connInfo.TargetPlatform != TargetPlatformWindows {
   447  		var stdout, stderr bytes.Buffer
   448  		cmd := &remote.Cmd{
   449  			Command: fmt.Sprintf("chmod 0777 %s", path),
   450  			Stdout:  &stdout,
   451  			Stderr:  &stderr,
   452  		}
   453  		if err := c.Start(cmd); err != nil {
   454  			return fmt.Errorf(
   455  				"Error chmodding script file to 0777 in remote "+
   456  					"machine: %s", err)
   457  		}
   458  
   459  		if err := cmd.Wait(); err != nil {
   460  			return fmt.Errorf(
   461  				"Error chmodding script file to 0777 in remote "+
   462  					"machine %v: %s %s", err, stdout.String(), stderr.String())
   463  		}
   464  	}
   465  	return nil
   466  }
   467  
   468  // UploadDir implementation of communicator.Communicator interface
   469  func (c *Communicator) UploadDir(dst string, src string) error {
   470  	log.Printf("[DEBUG] Uploading dir '%s' to '%s'", src, dst)
   471  	scpFunc := func(w io.Writer, r *bufio.Reader) error {
   472  		uploadEntries := func() error {
   473  			f, err := os.Open(src)
   474  			if err != nil {
   475  				return err
   476  			}
   477  			defer f.Close()
   478  
   479  			entries, err := f.Readdir(-1)
   480  			if err != nil {
   481  				return err
   482  			}
   483  
   484  			return scpUploadDir(src, entries, w, r)
   485  		}
   486  
   487  		if src[len(src)-1] != '/' {
   488  			log.Printf("[DEBUG] No trailing slash, creating the source directory name")
   489  			return scpUploadDirProtocol(filepath.Base(src), w, r, uploadEntries)
   490  		}
   491  		// Trailing slash, so only upload the contents
   492  		return uploadEntries()
   493  	}
   494  
   495  	cmd, err := quoteShell([]string{"scp", "-rvt", dst}, c.connInfo.TargetPlatform)
   496  	if err != nil {
   497  		return err
   498  	}
   499  	return c.scpSession(cmd, scpFunc)
   500  }
   501  
   502  func (c *Communicator) newSession() (session *ssh.Session, err error) {
   503  	log.Println("[DEBUG] opening new ssh session")
   504  	if c.client == nil {
   505  		err = errors.New("ssh client is not connected")
   506  	} else {
   507  		session, err = c.client.NewSession()
   508  	}
   509  
   510  	if err != nil {
   511  		log.Printf("[WARN] ssh session open error: '%s', attempting reconnect", err)
   512  		if err := c.Connect(nil); err != nil {
   513  			return nil, err
   514  		}
   515  
   516  		return c.client.NewSession()
   517  	}
   518  
   519  	return session, nil
   520  }
   521  
   522  func (c *Communicator) scpSession(scpCommand string, f func(io.Writer, *bufio.Reader) error) error {
   523  	session, err := c.newSession()
   524  	if err != nil {
   525  		return err
   526  	}
   527  	defer session.Close()
   528  
   529  	// Get a pipe to stdin so that we can send data down
   530  	stdinW, err := session.StdinPipe()
   531  	if err != nil {
   532  		return err
   533  	}
   534  
   535  	// We only want to close once, so we nil w after we close it,
   536  	// and only close in the defer if it hasn't been closed already.
   537  	defer func() {
   538  		if stdinW != nil {
   539  			stdinW.Close()
   540  		}
   541  	}()
   542  
   543  	// Get a pipe to stdout so that we can get responses back
   544  	stdoutPipe, err := session.StdoutPipe()
   545  	if err != nil {
   546  		return err
   547  	}
   548  	stdoutR := bufio.NewReader(stdoutPipe)
   549  
   550  	// Set stderr to a bytes buffer
   551  	stderr := new(bytes.Buffer)
   552  	session.Stderr = stderr
   553  
   554  	// Start the sink mode on the other side
   555  	// TODO(mitchellh): There are probably issues with shell escaping the path
   556  	log.Println("[DEBUG] Starting remote scp process: ", scpCommand)
   557  	if err := session.Start(scpCommand); err != nil {
   558  		return err
   559  	}
   560  
   561  	// Call our callback that executes in the context of SCP. We ignore
   562  	// EOF errors if they occur because it usually means that SCP prematurely
   563  	// ended on the other side.
   564  	log.Println("[DEBUG] Started SCP session, beginning transfers...")
   565  	if err := f(stdinW, stdoutR); err != nil && err != io.EOF {
   566  		return err
   567  	}
   568  
   569  	// Close the stdin, which sends an EOF, and then set w to nil so that
   570  	// our defer func doesn't close it again since that is unsafe with
   571  	// the Go SSH package.
   572  	log.Println("[DEBUG] SCP session complete, closing stdin pipe.")
   573  	stdinW.Close()
   574  	stdinW = nil
   575  
   576  	// Wait for the SCP connection to close, meaning it has consumed all
   577  	// our data and has completed. Or has errored.
   578  	log.Println("[DEBUG] Waiting for SSH session to complete.")
   579  	err = session.Wait()
   580  
   581  	// log any stderr before exiting on an error
   582  	scpErr := stderr.String()
   583  	if len(scpErr) > 0 {
   584  		log.Printf("[ERROR] scp stderr: %q", stderr)
   585  	}
   586  
   587  	if err != nil {
   588  		if exitErr, ok := err.(*ssh.ExitError); ok {
   589  			// Otherwise, we have an ExitErorr, meaning we can just read
   590  			// the exit status
   591  			log.Printf("[ERROR] %s", exitErr)
   592  
   593  			// If we exited with status 127, it means SCP isn't available.
   594  			// Return a more descriptive error for that.
   595  			if exitErr.ExitStatus() == 127 {
   596  				return errors.New(
   597  					"SCP failed to start. This usually means that SCP is not\n" +
   598  						"properly installed on the remote system.")
   599  			}
   600  		}
   601  
   602  		return err
   603  	}
   604  
   605  	return nil
   606  }
   607  
   608  // checkSCPStatus checks that a prior command sent to SCP completed
   609  // successfully. If it did not complete successfully, an error will
   610  // be returned.
   611  func checkSCPStatus(r *bufio.Reader) error {
   612  	code, err := r.ReadByte()
   613  	if err != nil {
   614  		return err
   615  	}
   616  
   617  	if code != 0 {
   618  		// Treat any non-zero (really 1 and 2) as fatal errors
   619  		message, _, err := r.ReadLine()
   620  		if err != nil {
   621  			return fmt.Errorf("Error reading error message: %s", err)
   622  		}
   623  
   624  		return errors.New(string(message))
   625  	}
   626  
   627  	return nil
   628  }
   629  
   630  func scpUploadFile(dst string, src io.Reader, w io.Writer, r *bufio.Reader, size int64) error {
   631  	if size == 0 {
   632  		// Create a temporary file where we can copy the contents of the src
   633  		// so that we can determine the length, since SCP is length-prefixed.
   634  		tf, err := ioutil.TempFile("", "terraform-upload")
   635  		if err != nil {
   636  			return fmt.Errorf("Error creating temporary file for upload: %s", err)
   637  		}
   638  		defer os.Remove(tf.Name())
   639  		defer tf.Close()
   640  
   641  		log.Println("[DEBUG] Copying input data into temporary file so we can read the length")
   642  		if _, err := io.Copy(tf, src); err != nil {
   643  			return err
   644  		}
   645  
   646  		// Sync the file so that the contents are definitely on disk, then
   647  		// read the length of it.
   648  		if err := tf.Sync(); err != nil {
   649  			return fmt.Errorf("Error creating temporary file for upload: %s", err)
   650  		}
   651  
   652  		// Seek the file to the beginning so we can re-read all of it
   653  		if _, err := tf.Seek(0, 0); err != nil {
   654  			return fmt.Errorf("Error creating temporary file for upload: %s", err)
   655  		}
   656  
   657  		fi, err := tf.Stat()
   658  		if err != nil {
   659  			return fmt.Errorf("Error creating temporary file for upload: %s", err)
   660  		}
   661  
   662  		src = tf
   663  		size = fi.Size()
   664  	}
   665  
   666  	// Start the protocol
   667  	log.Println("[DEBUG] Beginning file upload...")
   668  	fmt.Fprintln(w, "C0644", size, dst)
   669  	if err := checkSCPStatus(r); err != nil {
   670  		return err
   671  	}
   672  
   673  	if _, err := io.Copy(w, src); err != nil {
   674  		return err
   675  	}
   676  
   677  	fmt.Fprint(w, "\x00")
   678  	if err := checkSCPStatus(r); err != nil {
   679  		return err
   680  	}
   681  
   682  	return nil
   683  }
   684  
   685  func scpUploadDirProtocol(name string, w io.Writer, r *bufio.Reader, f func() error) error {
   686  	log.Printf("[DEBUG] SCP: starting directory upload: %s", name)
   687  	fmt.Fprintln(w, "D0755 0", name)
   688  	err := checkSCPStatus(r)
   689  	if err != nil {
   690  		return err
   691  	}
   692  
   693  	if err := f(); err != nil {
   694  		return err
   695  	}
   696  
   697  	fmt.Fprintln(w, "E")
   698  	if err != nil {
   699  		return err
   700  	}
   701  
   702  	return nil
   703  }
   704  
   705  func scpUploadDir(root string, fs []os.FileInfo, w io.Writer, r *bufio.Reader) error {
   706  	for _, fi := range fs {
   707  		realPath := filepath.Join(root, fi.Name())
   708  
   709  		// Track if this is actually a symlink to a directory. If it is
   710  		// a symlink to a file we don't do any special behavior because uploading
   711  		// a file just works. If it is a directory, we need to know so we
   712  		// treat it as such.
   713  		isSymlinkToDir := false
   714  		if fi.Mode()&os.ModeSymlink == os.ModeSymlink {
   715  			symPath, err := filepath.EvalSymlinks(realPath)
   716  			if err != nil {
   717  				return err
   718  			}
   719  
   720  			symFi, err := os.Lstat(symPath)
   721  			if err != nil {
   722  				return err
   723  			}
   724  
   725  			isSymlinkToDir = symFi.IsDir()
   726  		}
   727  
   728  		if !fi.IsDir() && !isSymlinkToDir {
   729  			// It is a regular file (or symlink to a file), just upload it
   730  			f, err := os.Open(realPath)
   731  			if err != nil {
   732  				return err
   733  			}
   734  
   735  			err = func() error {
   736  				defer f.Close()
   737  				return scpUploadFile(fi.Name(), f, w, r, fi.Size())
   738  			}()
   739  
   740  			if err != nil {
   741  				return err
   742  			}
   743  
   744  			continue
   745  		}
   746  
   747  		// It is a directory, recursively upload
   748  		err := scpUploadDirProtocol(fi.Name(), w, r, func() error {
   749  			f, err := os.Open(realPath)
   750  			if err != nil {
   751  				return err
   752  			}
   753  			defer f.Close()
   754  
   755  			entries, err := f.Readdir(-1)
   756  			if err != nil {
   757  				return err
   758  			}
   759  
   760  			return scpUploadDir(realPath, entries, w, r)
   761  		})
   762  		if err != nil {
   763  			return err
   764  		}
   765  	}
   766  
   767  	return nil
   768  }
   769  
   770  // ConnectFunc is a convenience method for returning a function
   771  // that just uses net.Dial to communicate with the remote end that
   772  // is suitable for use with the SSH communicator configuration.
   773  func ConnectFunc(network, addr string) func() (net.Conn, error) {
   774  	return func() (net.Conn, error) {
   775  		c, err := net.DialTimeout(network, addr, 15*time.Second)
   776  		if err != nil {
   777  			return nil, err
   778  		}
   779  
   780  		if tcpConn, ok := c.(*net.TCPConn); ok {
   781  			tcpConn.SetKeepAlive(true)
   782  		}
   783  
   784  		return c, nil
   785  	}
   786  }
   787  
   788  // BastionConnectFunc is a convenience method for returning a function
   789  // that connects to a host over a bastion connection.
   790  func BastionConnectFunc(
   791  	bProto string,
   792  	bAddr string,
   793  	bConf *ssh.ClientConfig,
   794  	proto string,
   795  	addr string) func() (net.Conn, error) {
   796  	return func() (net.Conn, error) {
   797  		log.Printf("[DEBUG] Connecting to bastion: %s", bAddr)
   798  		bastion, err := ssh.Dial(bProto, bAddr, bConf)
   799  		if err != nil {
   800  			return nil, fmt.Errorf("Error connecting to bastion: %s", err)
   801  		}
   802  
   803  		log.Printf("[DEBUG] Connecting via bastion (%s) to host: %s", bAddr, addr)
   804  		conn, err := bastion.Dial(proto, addr)
   805  		if err != nil {
   806  			bastion.Close()
   807  			return nil, err
   808  		}
   809  
   810  		// Wrap it up so we close both things properly
   811  		return &bastionConn{
   812  			Conn:    conn,
   813  			Bastion: bastion,
   814  		}, nil
   815  	}
   816  }
   817  
   818  type bastionConn struct {
   819  	net.Conn
   820  	Bastion *ssh.Client
   821  }
   822  
   823  func (c *bastionConn) Close() error {
   824  	c.Conn.Close()
   825  	return c.Bastion.Close()
   826  }
   827  
   828  func quoteShell(args []string, targetPlatform string) (string, error) {
   829  	if targetPlatform == TargetPlatformUnix {
   830  		return shquot.POSIXShell(args), nil
   831  	}
   832  	if targetPlatform == TargetPlatformWindows {
   833  		return shquot.WindowsArgv(args), nil
   834  	}
   835  
   836  	return "", fmt.Errorf("Cannot quote shell command, target platform unknown: %s", targetPlatform)
   837  
   838  }