Remove output file when save/export fail

Signed-off-by: Daehyeok Mun <daehyeok@gmail.com>

api/client/export.go

@@ -3,7 +3,6 @@ package client
 import (
 	"errors"
 	"io"
-	"os"
 
 	Cli "github.com/docker/docker/cli"
 	flag "github.com/docker/docker/pkg/mflag"
@@ -21,16 +20,7 @@ func (cli *DockerCli) CmdExport(args ...string) error {
 
 	cmd.ParseFlags(args, true)
 
-	var (
-		output = cli.out
-		err    error
-	)
-	if *outfile != "" {
-		output, err = os.Create(*outfile)
-		if err != nil {
-			return err
-		}
-	} else if cli.isTerminalOut {
+	if *outfile == "" && cli.isTerminalOut {
 		return errors.New("Cowardly refusing to save to a terminal. Use the -o flag or redirect.")
 	}
 
@@ -40,6 +30,11 @@ func (cli *DockerCli) CmdExport(args ...string) error {
 	}
 	defer responseBody.Close()
 
-	_, err = io.Copy(output, responseBody)
-	return err
+	if *outfile == "" {
+		_, err := io.Copy(cli.out, responseBody)
+		return err
+	}
+
+	return copyToFile(*outfile, responseBody)
+
 }

api/client/save.go

@@ -3,7 +3,6 @@ package client
 import (
 	"errors"
 	"io"
-	"os"
 
 	Cli "github.com/docker/docker/cli"
 	flag "github.com/docker/docker/pkg/mflag"
@@ -21,19 +20,9 @@ func (cli *DockerCli) CmdSave(args ...string) error {
 
 	cmd.ParseFlags(args, true)
 
-	var (
-		output = cli.out
-		err    error
-	)
-
 	if *outfile == "" && cli.isTerminalOut {
 		return errors.New("Cowardly refusing to save to a terminal. Use the -o flag or redirect.")
 	}
-	if *outfile != "" {
-		if output, err = os.Create(*outfile); err != nil {
-			return err
-		}
-	}
 
 	responseBody, err := cli.client.ImageSave(cmd.Args())
 	if err != nil {
@@ -41,6 +30,11 @@ func (cli *DockerCli) CmdSave(args ...string) error {
 	}
 	defer responseBody.Close()
 
-	_, err = io.Copy(output, responseBody)
-	return err
+	if *outfile == "" {
+		_, err := io.Copy(cli.out, responseBody)
+		return err
+	}
+
+	return copyToFile(*outfile, responseBody)
+
 }

api/client/utils.go

@@ -4,8 +4,11 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"os"
 	gosignal "os/signal"
+	"path/filepath"
 	"runtime"
 	"time"
 
@@ -138,3 +141,27 @@ func (cli *DockerCli) getTtySize() (int, int) {
 	}
 	return int(ws.Height), int(ws.Width)
 }
+
+func copyToFile(outfile string, r io.Reader) error {
+	tmpFile, err := ioutil.TempFile(filepath.Dir(outfile), ".docker_temp_")
+	if err != nil {
+		return err
+	}
+
+	tmpPath := tmpFile.Name()
+
+	_, err = io.Copy(tmpFile, r)
+	tmpFile.Close()
+
+	if err != nil {
+		os.Remove(tmpPath)
+		return err
+	}
+
+	if err = os.Rename(tmpPath, outfile); err != nil {
+		os.Remove(tmpPath)
+		return err
+	}
+
+	return nil
+}