ensure tmp image is cleaned up on crash (addresses part of #58)

wagoodman · Nov 25, 2018 · b1962b9 · b1962b9
1 parent 7d86809
commit b1962b9
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/image/image.go b/image/image.go
@@ -361,12 +361,17 @@ func saveImage(imageID string) (string, string) {
 	tmpDir, err := ioutil.TempDir("", "dive")
 	check(err)
 
+	cleanUpTmp := func() {
+		os.RemoveAll(tmpDir)
+	}
+
 	imageTarPath := filepath.Join(tmpDir, "image.tar")
 	imageFile, err := os.Create(imageTarPath)
 	check(err)
 
 	defer func() {
 		if err := imageFile.Close(); err != nil {
+			cleanUpTmp()
 			logrus.Panic(err)
 		}
 	}()
@@ -379,6 +384,7 @@ func saveImage(imageID string) (string, string) {
 	for {
 		n, err := readCloser.Read(buf)
 		if err != nil && err != io.EOF {
+			cleanUpTmp()
 			logrus.Panic(err)
 		}
 		if n == 0 {
@@ -392,11 +398,13 @@ func saveImage(imageID string) (string, string) {
 		}
 
 		if _, err := imageWriter.Write(buf[:n]); err != nil {
+			cleanUpTmp()
 			logrus.Panic(err)
 		}
 	}
 
 	if err = imageWriter.Flush(); err != nil {
+		cleanUpTmp()
 		logrus.Panic(err)
 	}