From 71d833139263d2c9f6c0d94b04090011d0aa242d Mon Sep 17 00:00:00 2001
From: Stanislas Michalak <stanislas-m@users.noreply.github.com>
Date: Sun, 14 Jul 2019 16:35:16 +0200
Subject: [PATCH] Fix assets are not extracted (#1736)

Final zip was deleted by Cleanup routine, and the generated zip
was not valid.

Fixes #1717
---
 genny/build/archived_assets.go | 13 +++++++++----
 genny/build/cleanup.go         |  4 ++++
 genny/build/options.go         |  4 ++++
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/genny/build/archived_assets.go b/genny/build/archived_assets.go
index 0d7385569..65f35030b 100644
--- a/genny/build/archived_assets.go
+++ b/genny/build/archived_assets.go
@@ -72,14 +72,19 @@ func archivedAssets(opts *Options) (*genny.Generator, error) {
 				return nil
 			}
 
-			if _, err = io.Copy(writer, file); err != nil {
-				return err
-			}
-			return r.File(genny.NewFile(target, bb))
+			_, err = io.Copy(writer, file)
+			return err
 		})
 		if err != nil {
 			return err
 		}
+		// We need to close the archive before passing the buffer to genny, otherwise the zip
+		// will be corrupted.
+		archive.Close()
+		if err := r.File(genny.NewFile(target, bb)); err != nil {
+			return err
+		}
+		opts.keep.Store(target, struct{}{})
 		return nil
 	})
 
diff --git a/genny/build/cleanup.go b/genny/build/cleanup.go
index 90ca20bd4..cccc92648 100644
--- a/genny/build/cleanup.go
+++ b/genny/build/cleanup.go
@@ -32,6 +32,10 @@ func Cleanup(opts *Options) genny.RunFn {
 			return err
 		}
 		for _, f := range r.Disk.Files() {
+			if _, keep := opts.keep.Load(f.Name()); keep {
+				// Keep this file
+				continue
+			}
 			if err := r.Disk.Delete(f.Name()); err != nil {
 				return err
 			}
diff --git a/genny/build/options.go b/genny/build/options.go
index 7bb7495a3..492259119 100644
--- a/genny/build/options.go
+++ b/genny/build/options.go
@@ -44,6 +44,7 @@ type Options struct {
 	// GoCommand is the `go X` command to be used. Default is "build".
 	GoCommand string `json:"go_command"`
 	rollback  *sync.Map
+	keep      *sync.Map
 }
 
 // Validate that options are usuable
@@ -64,6 +65,9 @@ func (opts *Options) Validate() error {
 	if opts.rollback == nil {
 		opts.rollback = &sync.Map{}
 	}
+	if opts.keep == nil {
+		opts.keep = &sync.Map{}
+	}
 	if len(opts.GoCommand) == 0 {
 		opts.GoCommand = "build"
 	}