cleanup on exit of non-gocui methods (fixes #42)

cmd/analyze.go

@@ -2,17 +2,18 @@ package cmd
 
 import (
 	"fmt"
-	"os"
 
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/wagoodman/dive/image"
 	"github.com/wagoodman/dive/ui"
+	"github.com/wagoodman/dive/utils"
 )
 
 // analyze takes a docker image tag, digest, or id and displayes the
 // image analysis to the screen
 func analyze(cmd *cobra.Command, args []string) {
+	defer utils.Cleanup()
 	if len(args) == 0 {
 		printVersionFlag, err := cmd.PersistentFlags().GetBool("version")
 		if err == nil && printVersionFlag {
@@ -22,14 +23,14 @@ func analyze(cmd *cobra.Command, args []string) {
 
 		fmt.Println("No image argument given")
 		cmd.Help()
-		os.Exit(1)
+		utils.Exit(1)
 	}
 
 	userImage := args[0]
 	if userImage == "" {
 		fmt.Println("No image argument given")
 		cmd.Help()
-		os.Exit(1)
+		utils.Exit(1)
 	}
 	color.New(color.Bold).Println("Analyzing Image")
 	manifest, refTrees, efficiency, inefficiencies := image.InitializeData(userImage)

cmd/build.go

@@ -24,6 +24,7 @@ func init() {
 
 // doBuild implements the steps taken for the build command
 func doBuild(cmd *cobra.Command, args []string) {
+	defer utils.Cleanup()
 	iidfile, err := ioutil.TempFile("/tmp", "dive.*.iid")
 	if err != nil {
 		log.Fatal(err)

cmd/root.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"github.com/wagoodman/dive/utils"
 	"os"
 
 	"github.com/k0kubun/go-ansi"
@@ -9,7 +10,6 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"github.com/tebeka/atexit"
 )
 
 var cfgFile string
@@ -28,17 +28,12 @@ the amount of wasted space and identifies the offending files from the image.`,
 func Execute() {
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)
-		os.Exit(1)
+		utils.Exit(1)
 	}
 }
 
-func exitHandler() {
-	ansi.CursorShow()
-}
-
 func init() {
 	ansi.CursorHide()
-	atexit.Register(exitHandler)
 
 	cobra.OnInitialize(initConfig)
 	cobra.OnInitialize(initLogging)
@@ -59,7 +54,7 @@ func initConfig() {
 		home, err := homedir.Dir()
 		if err != nil {
 			fmt.Println(err)
-			os.Exit(1)
+			utils.Exit(1)
 		}
 
 		// Search config in home directory with name ".dive" (without extension).

filetree/data.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"crypto/md5"
 	"fmt"
+	"github.com/sirupsen/logrus"
 	"io"
 )
 
@@ -85,7 +86,7 @@ func NewFileInfo(reader *tar.Reader, header *tar.Header, path string) FileInfo {
 	fileBytes := make([]byte, header.Size)
 	_, err := reader.Read(fileBytes)
 	if err != nil && err != io.EOF {
-		panic(err)
+		logrus.Panic(err)
 	}
 
 	return FileInfo{

image/image.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/sirupsen/logrus"
 	"io"
 	"io/ioutil"
 	"os"
@@ -100,12 +101,12 @@ func NewImageManifest(reader *tar.Reader, header *tar.Header) ImageManifest {
 	manifestBytes := make([]byte, size)
 	_, err := reader.Read(manifestBytes)
 	if err != nil && err != io.EOF {
-		panic(err)
+		logrus.Panic(err)
 	}
 	var manifest []ImageManifest
 	err = json.Unmarshal(manifestBytes, &manifest)
 	if err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 	return manifest[0]
 }
@@ -115,12 +116,12 @@ func NewImageConfig(reader *tar.Reader, header *tar.Header) ImageConfig {
 	configBytes := make([]byte, size)
 	_, err := reader.Read(configBytes)
 	if err != nil && err != io.EOF {
-		panic(err)
+		logrus.Panic(err)
 	}
 	var imageConfig ImageConfig
 	err = json.Unmarshal(configBytes, &imageConfig)
 	if err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 
 	layerIdx := 0
@@ -143,7 +144,7 @@ func GetImageConfig(imageTarPath string, manifest ImageManifest) ImageConfig {
 	tarFile, err := os.Open(imageTarPath)
 	if err != nil {
 		fmt.Println(err)
-		os.Exit(1)
+		utils.Exit(1)
 	}
 	defer tarFile.Close()
 
@@ -157,7 +158,7 @@ func GetImageConfig(imageTarPath string, manifest ImageManifest) ImageConfig {
 
 		if err != nil {
 			fmt.Println(err)
-			os.Exit(1)
+			utils.Exit(1)
 		}
 
 		name := header.Name
@@ -203,7 +204,7 @@ func InitializeData(imageID string) ([]*Layer, []*filetree.FileTree, float64, fi
 	dockerClient, err := client.NewClientWithOpts(client.WithVersion(dockerVersion))
 	if err != nil {
 		fmt.Println("Could not connect to the Docker daemon:" + err.Error())
-		os.Exit(1)
+		utils.Exit(1)
 	}
 	_, _, err = dockerClient.ImageInspectWithRaw(ctx, imageID)
 	if err != nil {
@@ -222,13 +223,13 @@ func InitializeData(imageID string) ([]*Layer, []*filetree.FileTree, float64, fi
 	tarFile, err := os.Open(imageTarPath)
 	if err != nil {
 		fmt.Println(err)
-		os.Exit(1)
+		utils.Exit(1)
 	}
 	defer tarFile.Close()
 
 	fi, err := tarFile.Stat()
 	if err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 	totalSize := fi.Size()
 	var observedBytes int64
@@ -250,7 +251,7 @@ func InitializeData(imageID string) ([]*Layer, []*filetree.FileTree, float64, fi
 
 		if err != nil {
 			fmt.Println(err)
-			os.Exit(1)
+			utils.Exit(1)
 		}
 
 		observedBytes += header.Size
@@ -265,7 +266,7 @@ func InitializeData(imageID string) ([]*Layer, []*filetree.FileTree, float64, fi
 			if strings.HasSuffix(name, "layer.tar") {
 				line, err := frame.Prepend()
 				if err != nil {
-					panic(err)
+					logrus.Panic(err)
 				}
 				shortName := name[:15]
 				io.WriteString(line, "    ├─ "+shortName+" : loading...")
@@ -274,7 +275,7 @@ func InitializeData(imageID string) ([]*Layer, []*filetree.FileTree, float64, fi
 
 				_, err = tarReader.Read(tarredBytes)
 				if err != nil && err != io.EOF {
-					panic(err)
+					logrus.Panic(err)
 				}
 
 				go processLayerTar(line, layerMap, name, tarredBytes)
@@ -336,7 +337,7 @@ func saveImage(imageID string) (string, string) {
 	dockerClient, err := client.NewClientWithOpts(client.WithVersion(dockerVersion))
 	if err != nil {
 		fmt.Println("Could not connect to the Docker daemon:" + err.Error())
-		os.Exit(1)
+		utils.Exit(1)
 	}
 
 	frame := jotframe.NewFixedFrame(0, false, false, true)
@@ -365,7 +366,7 @@ func saveImage(imageID string) (string, string) {
 
 	defer func() {
 		if err := imageFile.Close(); err != nil {
-			panic(err)
+			logrus.Panic(err)
 		}
 	}()
 	imageWriter := bufio.NewWriter(imageFile)
@@ -377,7 +378,7 @@ func saveImage(imageID string) (string, string) {
 	for {
 		n, err := readCloser.Read(buf)
 		if err != nil && err != io.EOF {
-			panic(err)
+			logrus.Panic(err)
 		}
 		if n == 0 {
 			break
@@ -390,12 +391,12 @@ func saveImage(imageID string) (string, string) {
 		}
 
 		if _, err := imageWriter.Write(buf[:n]); err != nil {
-			panic(err)
+			logrus.Panic(err)
 		}
 	}
 
 	if err = imageWriter.Flush(); err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 
 	pb.Done()
@@ -419,7 +420,7 @@ func getFileList(tarredBytes []byte) []filetree.FileInfo {
 
 		if err != nil {
 			fmt.Println(err)
-			os.Exit(1)
+			utils.Exit(1)
 		}
 
 		name := header.Name

ui/filetreeview.go

@@ -2,6 +2,7 @@ package ui
 
 import (
 	"fmt"
+	"github.com/sirupsen/logrus"
 	"regexp"
 	"strings"
 
@@ -231,7 +232,7 @@ func (view *FileTreeView) CursorLeft() error {
 
 	err = view.ModelTree.VisitDepthParentFirst(visitor, evaluator)
 	if err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 
 	view.TreeIndex = newIndex
@@ -285,7 +286,7 @@ func (view *FileTreeView) getAbsPositionNode() (node *filetree.FileNode) {
 
 	err = view.ModelTree.VisitDepthParentFirst(visitor, evaluator)
 	if err != nil {
-		panic(err)
+		logrus.Panic(err)
 	}
 
 	return node

utils/exit.go

@@ -0,0 +1,16 @@
+package utils
+
+import (
+	"github.com/k0kubun/go-ansi"
+	"os"
+)
+
+// Note: this should only be used when exiting from non-gocui code
+func Exit(rc int) {
+	Cleanup()
+	os.Exit(rc)
+}
+
+func Cleanup() {
+	ansi.CursorShow()
+}