From a1061ed9311487d0bb16357eb5dd7ed770cc7389 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Thu, 15 Mar 2018 12:55:57 -0400
Subject: [PATCH] update the chef and habitat error handling

Use the new ExitStatus method, and also check the cmd.Err() method for
errors.

Remove leaks from the output goroutines in both provisioners by
deferring their cleanup, and returning early on all error conditions.
---
 .../provisioners/chef/resource_provisioner.go | 22 ++++++++++--------
 .../habitat/resource_provisioner.go           | 23 +++++++++++--------
 2 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/builtin/provisioners/chef/resource_provisioner.go b/builtin/provisioners/chef/resource_provisioner.go
index 6edf973f3..ea44e9fd0 100644
--- a/builtin/provisioners/chef/resource_provisioner.go
+++ b/builtin/provisioners/chef/resource_provisioner.go
@@ -684,6 +684,13 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 	errDoneCh := make(chan struct{})
 	go p.copyOutput(o, outR, outDoneCh)
 	go p.copyOutput(o, errR, errDoneCh)
+	go func() {
+		// Wait for output to clean up
+		outW.Close()
+		errW.Close()
+		<-outDoneCh
+		<-errDoneCh
+	}()
 
 	cmd := &remote.Cmd{
 		Command: command,
@@ -697,18 +704,15 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 	}
 
 	cmd.Wait()
-	if cmd.ExitStatus != 0 {
-		err = fmt.Errorf(
-			"Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus)
+	if cmd.Err() != nil {
+		return cmd.Err()
 	}
 
-	// Wait for output to clean up
-	outW.Close()
-	errW.Close()
-	<-outDoneCh
-	<-errDoneCh
+	if cmd.ExitStatus() != 0 {
+		return fmt.Errorf("Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus())
+	}
 
-	return err
+	return nil
 }
 
 func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader, doneCh chan<- struct{}) {
diff --git a/builtin/provisioners/habitat/resource_provisioner.go b/builtin/provisioners/habitat/resource_provisioner.go
index aa404dae1..ada06a88c 100644
--- a/builtin/provisioners/habitat/resource_provisioner.go
+++ b/builtin/provisioners/habitat/resource_provisioner.go
@@ -740,12 +740,17 @@ func (p *provisioner) copyOutput(o terraform.UIOutput, r io.Reader, doneCh chan<
 func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communicator, command string) error {
 	outR, outW := io.Pipe()
 	errR, errW := io.Pipe()
-	var err error
 
 	outDoneCh := make(chan struct{})
 	errDoneCh := make(chan struct{})
 	go p.copyOutput(o, outR, outDoneCh)
 	go p.copyOutput(o, errR, errDoneCh)
+	defer func() {
+		outW.Close()
+		errW.Close()
+		<-outDoneCh
+		<-errDoneCh
+	}()
 
 	cmd := &remote.Cmd{
 		Command: command,
@@ -753,22 +758,20 @@ func (p *provisioner) runCommand(o terraform.UIOutput, comm communicator.Communi
 		Stderr:  errW,
 	}
 
-	if err = comm.Start(cmd); err != nil {
+	if err := comm.Start(cmd); err != nil {
 		return fmt.Errorf("Error executing command %q: %v", cmd.Command, err)
 	}
 
 	cmd.Wait()
-	if cmd.ExitStatus != 0 {
-		err = fmt.Errorf(
-			"Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus)
+	if cmd.Err() != nil {
+		return cmd.Err()
 	}
 
-	outW.Close()
-	errW.Close()
-	<-outDoneCh
-	<-errDoneCh
+	if cmd.ExitStatus() != 0 {
+		return fmt.Errorf("Command %q exited with non-zero exit status: %d", cmd.Command, cmd.ExitStatus())
+	}
 
-	return err
+	return nil
 }
 
 func getBindFromString(bind string) (Bind, error) {