builder/amazon/chroot: perform early cleanup

builder/amazon/chroot/builder.go

@@ -18,6 +18,10 @@ import (
 // The unique ID for this builder
 const BuilderId = "mitchellh.amazon.chroot"
 
+// CleanupFunc is a type that is strung throughout the state bag in
+// order to perform cleanup at earlier points.
+type CleanupFunc func(map[string]interface{}) error
+
 // Config is the configuration that is chained through the steps and
 // settable from the template.
 type Config struct {
@@ -138,6 +142,7 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		&StepMountExtra{},
 		&StepCopyFiles{},
 		&StepChrootProvision{},
+		&StepEarlyCleanup{},
 	}
 
 	// Run!

builder/amazon/chroot/step_attach_volume.go

@@ -14,6 +14,7 @@ import (
 //
 // Produces:
 //   device string - The location where the volume was attached.
+//   attach_cleanup CleanupFunc
 type StepAttachVolume struct {
 	attached bool
 	volumeId string
@@ -70,12 +71,20 @@ func (s *StepAttachVolume) Run(state map[string]interface{}) multistep.StepActio
 	}
 
 	state["device"] = config.AttachedDevicePath
+	state["attach_cleanup"] = s.CleanupFunc
 	return multistep.ActionContinue
 }
 
 func (s *StepAttachVolume) Cleanup(state map[string]interface{}) {
+	ui := state["ui"].(packer.Ui)
+	if err := s.CleanupFunc(state); err != nil {
+		ui.Error(err.Error())
+	}
+}
+
+func (s *StepAttachVolume) CleanupFunc(state map[string]interface{}) error {
 	if !s.attached {
-		return
+		return nil
 	}
 
 	ec2conn := state["ec2"].(*ec2.EC2)
@@ -84,10 +93,11 @@ func (s *StepAttachVolume) Cleanup(state map[string]interface{}) {
 	ui.Say("Detaching EBS volume...")
 	_, err := ec2conn.DetachVolume(s.volumeId)
 	if err != nil {
-		ui.Error(fmt.Sprintf("Error detaching EBS volume: %s", err))
-		return
+		return fmt.Errorf("Error detaching EBS volume: %s", err)
 	}
 
+	s.attached = false
+
 	// Wait for the volume to detach
 	stateChange := awscommon.StateChangeConf{
 		Conn:      ec2conn,
@@ -111,7 +121,8 @@ func (s *StepAttachVolume) Cleanup(state map[string]interface{}) {
 
 	_, err = awscommon.WaitForState(&stateChange)
 	if err != nil {
-		ui.Error(fmt.Sprintf("Error waiting for volume: %s", err))
-		return
+		return fmt.Errorf("Error waiting for volume: %s", err)
 	}
+
+	return nil
 }

builder/amazon/chroot/step_copy_files.go

@@ -11,8 +11,12 @@ import (
 )
 
 // StepCopyFiles copies some files from the host into the chroot environment.
+//
+// Produces:
+//   copy_files_cleanup CleanupFunc - A function to clean up the copied files
+//   early.
 type StepCopyFiles struct {
-	mounts []string
+	files []string
 }
 
 func (s *StepCopyFiles) Run(state map[string]interface{}) multistep.StepAction {
@@ -20,6 +24,7 @@ func (s *StepCopyFiles) Run(state map[string]interface{}) multistep.StepAction {
 	mountPath := state["mount_path"].(string)
 	ui := state["ui"].(packer.Ui)
 
+	s.files = make([]string, len(config.CopyFiles))
 	if len(config.CopyFiles) > 0 {
 		ui.Say("Copying files from host to chroot...")
 		for _, path := range config.CopyFiles {
@@ -33,13 +38,35 @@ func (s *StepCopyFiles) Run(state map[string]interface{}) multistep.StepAction {
 				ui.Error(err.Error())
 				return multistep.ActionHalt
 			}
+
+			s.files = append(s.files, chrootPath)
 		}
 	}
 
+	state["copy_files_cleanup"] = s.CleanupFunc
 	return multistep.ActionContinue
 }
 
-func (s *StepCopyFiles) Cleanup(state map[string]interface{}) {}
+func (s *StepCopyFiles) Cleanup(state map[string]interface{}) {
+	ui := state["ui"].(packer.Ui)
+	if err := s.CleanupFunc(state); err != nil {
+		ui.Error(err.Error())
+	}
+}
+
+func (s *StepCopyFiles) CleanupFunc(map[string]interface{}) error {
+	if s.files != nil {
+		for _, file := range s.files {
+			log.Printf("Removing: %s", file)
+			if err := os.Remove(file); err != nil {
+				return err
+			}
+		}
+	}
+
+	s.files = nil
+	return nil
+}
 
 func (s *StepCopyFiles) copySingle(dst, src string) error {
 	// Stat the src file so we can copy the mode later

builder/amazon/chroot/step_early_cleanup.go

@@ -0,0 +1,37 @@
+package chroot
+
+import (
+	"fmt"
+	"github.com/mitchellh/multistep"
+	"github.com/mitchellh/packer/packer"
+	"log"
+)
+
+// StepEarlyCleanup performs some of the cleanup steps early in order to
+// prepare for snapshotting and creating an AMI.
+type StepEarlyCleanup struct{}
+
+func (s *StepEarlyCleanup) Run(state map[string]interface{}) multistep.StepAction {
+	ui := state["ui"].(packer.Ui)
+	cleanupKeys := []string{
+		"copy_files_cleanup",
+		"mount_extra_cleanup",
+		"mount_device_cleanup",
+		"attach_cleanup",
+	}
+
+	for _, key := range cleanupKeys {
+		f := state[key].(CleanupFunc)
+		log.Printf("Running cleanup func: %s", key)
+		if err := f(state); err != nil {
+			err := fmt.Errorf("Error cleaning up: %s", err)
+			state["error"] = err
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+	}
+
+	return multistep.ActionContinue
+}
+
+func (s *StepEarlyCleanup) Cleanup(state map[string]interface{}) {}

builder/amazon/chroot/step_mount_device.go

@@ -20,6 +20,7 @@ type mountPathData struct {
 //
 // Produces:
 //   mount_path string - The location where the volume was mounted.
+//   mount_device_cleanup CleanupFunc - To perform early cleanup
 type StepMountDevice struct {
 	mountPath string
 }
@@ -61,13 +62,21 @@ func (s *StepMountDevice) Run(state map[string]interface{}) multistep.StepAction
 	// Set the mount path so we remember to unmount it later
 	s.mountPath = mountPath
 	state["mount_path"] = s.mountPath
+	state["mount_device_cleanup"] = s.CleanupFunc
 
 	return multistep.ActionContinue
 }
 
 func (s *StepMountDevice) Cleanup(state map[string]interface{}) {
+	ui := state["ui"].(packer.Ui)
+	if err := s.CleanupFunc(state); err != nil {
+		ui.Error(err.Error())
+	}
+}
+
+func (s *StepMountDevice) CleanupFunc(state map[string]interface{}) error {
 	if s.mountPath == "" {
-		return
+		return nil
 	}
 
 	config := state["config"].(*Config)
@@ -77,8 +86,9 @@ func (s *StepMountDevice) Cleanup(state map[string]interface{}) {
 	unmountCommand := fmt.Sprintf("%s %s", config.UnmountCommand, s.mountPath)
 	cmd := exec.Command("/bin/sh", "-c", unmountCommand)
 	if err := cmd.Run(); err != nil {
-		ui.Error(fmt.Sprintf(
-			"Error unmounting root device: %s", err))
-		return
+		return fmt.Errorf("Error unmounting root device: %s", err)
 	}
+
+	s.mountPath = ""
+	return nil
 }

builder/amazon/chroot/step_mount_extra.go

@@ -12,7 +12,7 @@ import (
 // StepMountExtra mounts the attached device.
 //
 // Produces:
-//   mount_path string - The location where the volume was mounted.
+//   mount_extra_cleanup CleanupFunc - To perform early cleanup
 type StepMountExtra struct {
 	mounts []string
 }
@@ -61,28 +61,40 @@ func (s *StepMountExtra) Run(state map[string]interface{}) multistep.StepAction
 		s.mounts = append(s.mounts, innerPath)
 	}
 
+	state["mount_extra_cleanup"] = s.CleanupFunc
 	return multistep.ActionContinue
 }
 
 func (s *StepMountExtra) Cleanup(state map[string]interface{}) {
-	if s.mounts == nil {
+	ui := state["ui"].(packer.Ui)
+
+	if err := s.CleanupFunc(state); err != nil {
+		ui.Error(err.Error())
 		return
 	}
+}
 
-	config := state["config"].(*Config)
-	ui := state["ui"].(packer.Ui)
+func (s *StepMountExtra) CleanupFunc(state map[string]interface{}) error {
+	if s.mounts == nil {
+		return nil
+	}
 
-	for i := len(s.mounts) - 1; i >= 0; i-- {
-		path := s.mounts[i]
+	config := state["config"].(*Config)
+	for len(s.mounts) > 0 {
+		var path string
+		lastIndex := len(s.mounts) - 1
+		path, s.mounts = s.mounts[lastIndex], s.mounts[:lastIndex]
 		unmountCommand := fmt.Sprintf("%s %s", config.UnmountCommand, path)
 
 		stderr := new(bytes.Buffer)
 		cmd := exec.Command("/bin/sh", "-c", unmountCommand)
 		cmd.Stderr = stderr
 		if err := cmd.Run(); err != nil {
-			ui.Error(fmt.Sprintf(
-				"Error unmounting device: %s\nStderr: %s", err, stderr.String()))
-			return
+			return fmt.Errorf(
+				"Error unmounting device: %s\nStderr: %s", err, stderr.String())
 		}
 	}
+
+	s.mounts = nil
+	return nil
 }