From b1962b9ae549c20ba0f7cd0d37db3587ff174bac Mon Sep 17 00:00:00 2001
From: Alex Goodman <wagoodman@gmail.com>
Date: Sat, 24 Nov 2018 21:02:58 -0500
Subject: [PATCH] ensure tmp image is cleaned up on crash (addresses part of
 #58)

---
 image/image.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/image/image.go b/image/image.go
index 38497ce7..3c8baafd 100644
--- 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)
 	}