Use context for timeouts, interruption in ssh and winrm communicators

Also don't waste 5 seconds waiting before first winrm connection attempt
Minor code cleanup as well

helper/communicator/step_connect_ssh.go

@@ -36,17 +36,17 @@ func (s *StepConnectSSH) Run(ctx context.Context, state multistep.StateBag) mult
 	var comm packer.Communicator
 	var err error
 
-	cancel := make(chan struct{})
+	subCtx, cancel := context.WithCancel(ctx)
 	waitDone := make(chan bool, 1)
 	go func() {
 		ui.Say("Waiting for SSH to become available...")
-		comm, err = s.waitForSSH(state, cancel)
+		comm, err = s.waitForSSH(state, subCtx)
+		cancel() // just to make 'possible context leak' analysis happy
 		waitDone <- true
 	}()
 
 	log.Printf("[INFO] Waiting for SSH, up to timeout: %s", s.Config.SSHTimeout)
 	timeout := time.After(s.Config.SSHTimeout)
-WaitLoop:
 	for {
 		// Wait for either SSH to become available, a timeout to occur,
 		// or an interrupt to come through.
@@ -60,31 +60,28 @@ WaitLoop:
 
 			ui.Say("Connected to SSH!")
 			state.Put("communicator", comm)
-			break WaitLoop
+			return multistep.ActionContinue
 		case <-timeout:
 			err := fmt.Errorf("Timeout waiting for SSH.")
 			state.Put("error", err)
 			ui.Error(err.Error())
-			close(cancel)
+			cancel()
+			return multistep.ActionHalt
+		case <-ctx.Done():
+			// The step sequence was cancelled, so cancel waiting for SSH
+			// and just start the halting process.
+			cancel()
+			log.Println("[WARN] Interrupt detected, quitting waiting for SSH.")
 			return multistep.ActionHalt
 		case <-time.After(1 * time.Second):
-			if _, ok := state.GetOk(multistep.StateCancelled); ok {
-				// The step sequence was cancelled, so cancel waiting for SSH
-				// and just start the halting process.
-				close(cancel)
-				log.Println("[WARN] Interrupt detected, quitting waiting for SSH.")
-				return multistep.ActionHalt
-			}
 		}
 	}
-
-	return multistep.ActionContinue
 }
 
 func (s *StepConnectSSH) Cleanup(multistep.StateBag) {
 }
 
-func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, cancel <-chan struct{}) (packer.Communicator, error) {
+func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, ctx context.Context) (packer.Communicator, error) {
 	// Determine if we're using a bastion host, and if so, retrieve
 	// that configuration. This configuration doesn't change so we
 	// do this one before entering the retry loop.
@@ -123,7 +120,7 @@ func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, cancel <-chan stru
 		// Don't check for cancel or wait on first iteration
 		if !first {
 			select {
-			case <-cancel:
+			case <-ctx.Done():
 				log.Println("[DEBUG] SSH wait cancelled. Exiting loop.")
 				return nil, errors.New("SSH wait cancelled")
 			case <-time.After(5 * time.Second):

helper/communicator/step_connect_winrm.go

@@ -33,23 +33,23 @@ type StepConnectWinRM struct {
 	WinRMPort   func(multistep.StateBag) (int, error)
 }
 
-func (s *StepConnectWinRM) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
+func (s *StepConnectWinRM) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
 	ui := state.Get("ui").(packer.Ui)
 
 	var comm packer.Communicator
 	var err error
 
-	cancel := make(chan struct{})
+	subCtx, cancel := context.WithCancel(ctx)
 	waitDone := make(chan bool, 1)
 	go func() {
 		ui.Say("Waiting for WinRM to become available...")
-		comm, err = s.waitForWinRM(state, cancel)
+		comm, err = s.waitForWinRM(state, subCtx)
+		cancel() // just to make 'possible context leak' analysis happy
 		waitDone <- true
 	}()
 
 	log.Printf("Waiting for WinRM, up to timeout: %s", s.Config.WinRMTimeout)
 	timeout := time.After(s.Config.WinRMTimeout)
-WaitLoop:
 	for {
 		// Wait for either WinRM to become available, a timeout to occur,
 		// or an interrupt to come through.
@@ -62,39 +62,40 @@ WaitLoop:
 
 			ui.Say("Connected to WinRM!")
 			state.Put("communicator", comm)
-			break WaitLoop
+			return multistep.ActionContinue
 		case <-timeout:
 			err := fmt.Errorf("Timeout waiting for WinRM.")
 			state.Put("error", err)
 			ui.Error(err.Error())
-			close(cancel)
+			cancel()
+			return multistep.ActionHalt
+		case <-ctx.Done():
+			// The step sequence was cancelled, so cancel waiting for WinRM
+			// and just start the halting process.
+			cancel()
+			log.Println("Interrupt detected, quitting waiting for WinRM.")
 			return multistep.ActionHalt
 		case <-time.After(1 * time.Second):
-			if _, ok := state.GetOk(multistep.StateCancelled); ok {
-				// The step sequence was cancelled, so cancel waiting for WinRM
-				// and just start the halting process.
-				close(cancel)
-				log.Println("Interrupt detected, quitting waiting for WinRM.")
-				return multistep.ActionHalt
-			}
 		}
 	}
-
-	return multistep.ActionContinue
 }
 
 func (s *StepConnectWinRM) Cleanup(multistep.StateBag) {
 }
 
-func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, cancel <-chan struct{}) (packer.Communicator, error) {
+func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, ctx context.Context) (packer.Communicator, error) {
-	ctx := context.TODO()
 	var comm packer.Communicator
+	first := true
 	for {
-		select {
+		// Don't check for cancel or wait on first iteration
-		case <-cancel:
+		if !first {
-			log.Println("[INFO] WinRM wait cancelled. Exiting loop.")
+			select {
-			return nil, errors.New("WinRM wait cancelled")
+			case <-ctx.Done():
-		case <-time.After(5 * time.Second):
+				log.Println("[INFO] WinRM wait cancelled. Exiting loop.")
+				return nil, errors.New("WinRM wait cancelled")
+			case <-time.After(5 * time.Second):
+			}
+			first = false
 		}
 
 		host, err := s.Host(state)
@@ -157,7 +158,7 @@ func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, cancel <-chan
 		cmd.Stdout = &buf
 		cmd.Stdout = io.MultiWriter(cmd.Stdout, &buf2)
 		select {
-		case <-cancel:
+		case <-ctx.Done():
 			log.Println("WinRM wait canceled, exiting loop")
 			return comm, fmt.Errorf("WinRM wait canceled")
 		case <-time.After(retryableSleep):