From 521c7632bb9a740294951297e69c9acb536965ad Mon Sep 17 00:00:00 2001
From: Peter Engelbert <pmengelbert@gmail.com>
Date: Fri, 5 Jun 2020 14:25:21 -0500
Subject: [PATCH] Add SaveWithOpts and abstract code for reuse

Signed-off-by: Peter Engelbert <pmengelbert@gmail.com>
---
 internal/experimental/registry/cache.go |  2 +-
 pkg/chartutil/save.go                   | 99 +++++++++++++++++++------
 2 files changed, 76 insertions(+), 25 deletions(-)

diff --git a/internal/experimental/registry/cache.go b/internal/experimental/registry/cache.go
index 2be1214f5..991b12bf4 100644
--- a/internal/experimental/registry/cache.go
+++ b/internal/experimental/registry/cache.go
@@ -288,7 +288,7 @@ func (cache *Cache) saveChartConfig(ch *chart.Chart) (*ocispec.Descriptor, bool,
 func (cache *Cache) saveChartContentLayer(ch *chart.Chart) (*ocispec.Descriptor, bool, error) {
 	destDir := filepath.Join(cache.rootDir, ".build")
 	os.MkdirAll(destDir, 0755)
-	tmpFile, err := chartutil.Save(ch, destDir, chartutil.WithoutTimestamp())
+	tmpFile, err := chartutil.SaveWithOpts(ch, destDir, chartutil.WithoutTimestamp())
 	defer os.Remove(tmpFile)
 	if err != nil {
 		return nil, false, errors.Wrap(err, "failed to save")
diff --git a/pkg/chartutil/save.go b/pkg/chartutil/save.go
index 8e34367dd..5a253aeba 100644
--- a/pkg/chartutil/save.go
+++ b/pkg/chartutil/save.go
@@ -32,6 +32,12 @@ import (
 )
 
 type (
+	tarGzipWriter struct {
+		TarWriter  *tar.Writer
+		GzipWriter *gzip.Writer
+		File       *os.File
+	}
+
 	saveConfig struct {
 		noTimestamp bool
 	}
@@ -47,6 +53,12 @@ func WithoutTimestamp() SaveOpt {
 	}
 }
 
+func (w *tarGzipWriter) Close() {
+	w.TarWriter.Close()
+	w.GzipWriter.Close()
+	w.File.Close()
+}
+
 var headerBytes = []byte("+aHR0cHM6Ly95b3V0dS5iZS96OVV6MWljandyTQo=")
 
 // SaveDir saves a chart as files in a directory.
@@ -115,14 +127,65 @@ func SaveDir(c *chart.Chart, dest string) error {
 // will generate /foo/bar-1.0.0.tgz.
 //
 // This returns the absolute path to the chart archive file.
-func Save(c *chart.Chart, outDir string, saveOpts ...SaveOpt) (string, error) {
+func Save(c *chart.Chart, outDir string) (string, error) {
+	tgzWriter, err := createTarGzipWriter(c, outDir)
+	if err != nil {
+		return "", err
+	}
+
+	filename := tgzWriter.File.Name()
+	rollback := false
+
+	defer func() {
+		tgzWriter.Close()
+		if rollback {
+			os.Remove(filename)
+		}
+	}()
+
+	if err := writeTarContents(tgzWriter.TarWriter, c, "", time.Now()); err != nil {
+		rollback = true
+		return filename, err
+	}
+	return filename, nil
+}
+
+func SaveWithOpts(c *chart.Chart, outDir string, opts ...SaveOpt) (string, error) {
 	config := new(saveConfig)
-	for _, fn := range saveOpts {
+	for _, fn := range opts {
 		fn(config)
 	}
 
+	tgzWriter, err := createTarGzipWriter(c, outDir)
+	if err != nil {
+		return "", err
+	}
+
+	filename := tgzWriter.File.Name()
+	rollback := false
+
+	defer func() {
+		tgzWriter.Close()
+		if rollback {
+			os.Remove(filename)
+		}
+	}()
+
+	var timestamp time.Time
+	if !config.noTimestamp {
+		timestamp = time.Now()
+	}
+
+	if err := writeTarContents(tgzWriter.TarWriter, c, "", timestamp); err != nil {
+		rollback = true
+		return filename, err
+	}
+	return filename, nil
+}
+
+func createTarGzipWriter(c *chart.Chart, outDir string) (*tarGzipWriter, error) {
 	if err := c.Validate(); err != nil {
-		return "", errors.Wrap(err, "chart validation")
+		return nil, errors.Wrap(err, "chart validation")
 	}
 
 	filename := fmt.Sprintf("%s-%s.tgz", c.Name(), c.Metadata.Version)
@@ -131,18 +194,18 @@ func Save(c *chart.Chart, outDir string, saveOpts ...SaveOpt) (string, error) {
 	if stat, err := os.Stat(dir); err != nil {
 		if os.IsNotExist(err) {
 			if err2 := os.MkdirAll(dir, 0755); err2 != nil {
-				return "", err2
+				return nil, err2
 			}
 		} else {
-			return "", errors.Wrapf(err, "stat %s", dir)
+			return nil, errors.Wrapf(err, "stat %s", dir)
 		}
 	} else if !stat.IsDir() {
-		return "", errors.Errorf("is not a directory: %s", dir)
+		return nil, errors.Errorf("is not a directory: %s", dir)
 	}
 
 	f, err := os.Create(filename)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	// Wrap in gzip writer
@@ -152,26 +215,14 @@ func Save(c *chart.Chart, outDir string, saveOpts ...SaveOpt) (string, error) {
 
 	// Wrap in tar writer
 	twriter := tar.NewWriter(zipper)
-	rollback := false
-	defer func() {
-		twriter.Close()
-		zipper.Close()
-		f.Close()
-		if rollback {
-			os.Remove(filename)
-		}
-	}()
 
-	var timestamp time.Time
-	if !config.noTimestamp {
-		timestamp = time.Now()
+	cxvf := &tarGzipWriter{
+		TarWriter:  twriter,
+		GzipWriter: zipper,
+		File:       f,
 	}
 
-	if err := writeTarContents(twriter, c, "", timestamp); err != nil {
-		rollback = true
-		return filename, err
-	}
-	return filename, nil
+	return cxvf, nil
 }
 
 func writeTarContents(out *tar.Writer, c *chart.Chart, prefix string, timestamp time.Time) error {