github.com/kaixiang/packer@v0.5.2-0.20140114230416-1f5786b0d7f1/common/step_connect_ssh.go (about)

     1  package common
     2  
     3  import (
     4  	gossh "code.google.com/p/go.crypto/ssh"
     5  	"errors"
     6  	"fmt"
     7  	"github.com/mitchellh/multistep"
     8  	"github.com/mitchellh/packer/communicator/ssh"
     9  	"github.com/mitchellh/packer/packer"
    10  	"log"
    11  	"strings"
    12  	"time"
    13  )
    14  
    15  // StepConnectSSH is a multistep Step implementation that waits for SSH
    16  // to become available. It gets the connection information from a single
    17  // configuration when creating the step.
    18  //
    19  // Uses:
    20  //   ui packer.Ui
    21  //
    22  // Produces:
    23  //   communicator packer.Communicator
    24  type StepConnectSSH struct {
    25  	// SSHAddress is a function that returns the TCP address to connect to
    26  	// for SSH. This is a function so that you can query information
    27  	// if necessary for this address.
    28  	SSHAddress func(multistep.StateBag) (string, error)
    29  
    30  	// SSHConfig is a function that returns the proper client configuration
    31  	// for SSH access.
    32  	SSHConfig func(multistep.StateBag) (*gossh.ClientConfig, error)
    33  
    34  	// SSHWaitTimeout is the total timeout to wait for SSH to become available.
    35  	SSHWaitTimeout time.Duration
    36  
    37  	// NoPty, if true, will not request a Pty from the remote end.
    38  	NoPty bool
    39  
    40  	comm packer.Communicator
    41  }
    42  
    43  func (s *StepConnectSSH) Run(state multistep.StateBag) multistep.StepAction {
    44  	ui := state.Get("ui").(packer.Ui)
    45  
    46  	var comm packer.Communicator
    47  	var err error
    48  
    49  	cancel := make(chan struct{})
    50  	waitDone := make(chan bool, 1)
    51  	go func() {
    52  		ui.Say("Waiting for SSH to become available...")
    53  		comm, err = s.waitForSSH(state, cancel)
    54  		waitDone <- true
    55  	}()
    56  
    57  	log.Printf("Waiting for SSH, up to timeout: %s", s.SSHWaitTimeout)
    58  	timeout := time.After(s.SSHWaitTimeout)
    59  WaitLoop:
    60  	for {
    61  		// Wait for either SSH to become available, a timeout to occur,
    62  		// or an interrupt to come through.
    63  		select {
    64  		case <-waitDone:
    65  			if err != nil {
    66  				ui.Error(fmt.Sprintf("Error waiting for SSH: %s", err))
    67  				return multistep.ActionHalt
    68  			}
    69  
    70  			ui.Say("Connected to SSH!")
    71  			s.comm = comm
    72  			state.Put("communicator", comm)
    73  			break WaitLoop
    74  		case <-timeout:
    75  			err := fmt.Errorf("Timeout waiting for SSH.")
    76  			state.Put("error", err)
    77  			ui.Error(err.Error())
    78  			close(cancel)
    79  			return multistep.ActionHalt
    80  		case <-time.After(1 * time.Second):
    81  			if _, ok := state.GetOk(multistep.StateCancelled); ok {
    82  				// The step sequence was cancelled, so cancel waiting for SSH
    83  				// and just start the halting process.
    84  				close(cancel)
    85  				log.Println("Interrupt detected, quitting waiting for SSH.")
    86  				return multistep.ActionHalt
    87  			}
    88  		}
    89  	}
    90  
    91  	return multistep.ActionContinue
    92  }
    93  
    94  func (s *StepConnectSSH) Cleanup(multistep.StateBag) {
    95  }
    96  
    97  func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, cancel <-chan struct{}) (packer.Communicator, error) {
    98  	handshakeAttempts := 0
    99  
   100  	var comm packer.Communicator
   101  	for {
   102  		select {
   103  		case <-cancel:
   104  			log.Println("SSH wait cancelled. Exiting loop.")
   105  			return nil, errors.New("SSH wait cancelled")
   106  		case <-time.After(5 * time.Second):
   107  		}
   108  
   109  		// First we request the TCP connection information
   110  		address, err := s.SSHAddress(state)
   111  		if err != nil {
   112  			log.Printf("Error getting SSH address: %s", err)
   113  			continue
   114  		}
   115  
   116  		// Retrieve the SSH configuration
   117  		sshConfig, err := s.SSHConfig(state)
   118  		if err != nil {
   119  			log.Printf("Error getting SSH config: %s", err)
   120  			continue
   121  		}
   122  
   123  		// Attempt to connect to SSH port
   124  		connFunc := ssh.ConnectFunc("tcp", address)
   125  		nc, err := connFunc()
   126  		if err != nil {
   127  			log.Printf("TCP connection to SSH ip/port failed: %s", err)
   128  			continue
   129  		}
   130  		nc.Close()
   131  
   132  		// Then we attempt to connect via SSH
   133  		config := &ssh.Config{
   134  			Connection: connFunc,
   135  			SSHConfig:  sshConfig,
   136  			NoPty:      s.NoPty,
   137  		}
   138  
   139  		log.Println("Attempting SSH connection...")
   140  		comm, err = ssh.New(config)
   141  		if err != nil {
   142  			log.Printf("SSH handshake err: %s", err)
   143  
   144  			// Only count this as an attempt if we were able to attempt
   145  			// to authenticate. Note this is very brittle since it depends
   146  			// on the string of the error... but I don't see any other way.
   147  			if strings.Contains(err.Error(), "authenticate") {
   148  				log.Printf("Detected authentication error. Increasing handshake attempts.")
   149  				handshakeAttempts += 1
   150  			}
   151  
   152  			if handshakeAttempts < 10 {
   153  				// Try to connect via SSH a handful of times
   154  				continue
   155  			}
   156  
   157  			return nil, err
   158  		}
   159  
   160  		break
   161  	}
   162  
   163  	return comm, nil
   164  }