From b199ce9a222bf40a1b6bb4e28478efc12b9db0dd Mon Sep 17 00:00:00 2001
From: Tim Dawson <timdawson264@gmail.com>
Date: Thu, 1 Oct 2020 17:23:09 +1300
Subject: [PATCH] Added snapshots to artifacts Increased Snapshot wait timeout

---
 builder/amazon/common/state.go                   |  9 ++++++++-
 builder/amazon/ebsvolume/artifact.go             | 16 ++++++++++++++--
 builder/amazon/ebsvolume/builder.go              |  1 +
 .../ebsvolume/step_snapshot_ebs_volumes.go       | 10 ++++++++++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/builder/amazon/common/state.go b/builder/amazon/common/state.go
index 1ceacef84..a00443e64 100644
--- a/builder/amazon/common/state.go
+++ b/builder/amazon/common/state.go
@@ -154,10 +154,17 @@ func (w *AWSPollingConfig) WaitUntilSnapshotDone(ctx aws.Context, conn *ec2.EC2,
 		SnapshotIds: []*string{&snapshotID},
 	}
 
+	waitOpts := w.getWaiterOptions()
+	if len(waitOpts) == 0 {
+		// Bump this default to 30 minutes.
+		// Large snapshots can take a long time for the copy to s3
+		waitOpts = append(waitOpts, request.WithWaiterMaxAttempts(120))
+	}
+
 	err := conn.WaitUntilSnapshotCompletedWithContext(
 		ctx,
 		&snapInput,
-		w.getWaiterOptions()...)
+		waitOpts...)
 	return err
 }
 
diff --git a/builder/amazon/ebsvolume/artifact.go b/builder/amazon/ebsvolume/artifact.go
index 683bcbbd3..a6e56060c 100644
--- a/builder/amazon/ebsvolume/artifact.go
+++ b/builder/amazon/ebsvolume/artifact.go
@@ -13,11 +13,15 @@ import (
 // map of region to list of volume IDs
 type EbsVolumes map[string][]string
 
+// map of region to list of snapshot IDs
+type EbsSnapshots map[string][]string
+
 // Artifact is an artifact implementation that contains built AMIs.
 type Artifact struct {
 	// A map of regions to EBS Volume IDs.
 	Volumes EbsVolumes
-
+	// A map of regions to EBS Snapshot IDs.
+	Snapshots EbsSnapshots
 	// BuilderId is the unique ID for the builder that created this AMI
 	BuilderIdValue string
 
@@ -40,13 +44,21 @@ func (*Artifact) Files() []string {
 
 // returns a sorted list of region:ID pairs
 func (a *Artifact) idList() []string {
-	parts := make([]string, 0, len(a.Volumes))
+
+	parts := make([]string, 0, len(a.Volumes)+len(a.Snapshots))
+
 	for region, volumeIDs := range a.Volumes {
 		for _, volumeID := range volumeIDs {
 			parts = append(parts, fmt.Sprintf("%s:%s", region, volumeID))
 		}
 	}
 
+	for region, snapshotIDs := range a.Snapshots {
+		for _, snapshotID := range snapshotIDs {
+			parts = append(parts, fmt.Sprintf("%s:%s", region, snapshotID))
+		}
+	}
+
 	sort.Strings(parts)
 	return parts
 }
diff --git a/builder/amazon/ebsvolume/builder.go b/builder/amazon/ebsvolume/builder.go
index 1984f4cb0..6e43ad67d 100644
--- a/builder/amazon/ebsvolume/builder.go
+++ b/builder/amazon/ebsvolume/builder.go
@@ -337,6 +337,7 @@ func (b *Builder) Run(ctx context.Context, ui packersdk.Ui, hook packersdk.Hook)
 	// Build the artifact and return it
 	artifact := &Artifact{
 		Volumes:        state.Get("ebsvolumes").(EbsVolumes),
+		Snapshots:      state.Get("ebssnapshots").(EbsSnapshots),
 		BuilderIdValue: BuilderId,
 		Conn:           ec2conn,
 		StateData:      map[string]interface{}{"generated_data": state.Get("generated_data")},
diff --git a/builder/amazon/ebsvolume/step_snapshot_ebs_volumes.go b/builder/amazon/ebsvolume/step_snapshot_ebs_volumes.go
index ad3a00331..e234013c0 100644
--- a/builder/amazon/ebsvolume/step_snapshot_ebs_volumes.go
+++ b/builder/amazon/ebsvolume/step_snapshot_ebs_volumes.go
@@ -144,6 +144,16 @@ func (s *stepSnapshotEBSVolumes) Run(ctx context.Context, state multistep.StateB
 		}
 	}
 
+	//Record all snapshots in current Region.
+	snapshots := make(EbsSnapshots)
+	for snapID := range s.SnapshotMap {
+		snapshots[*ec2conn.Config.Region] = append(
+			snapshots[*ec2conn.Config.Region],
+			snapID)
+	}
+	//Records artifacts
+	state.Put("ebssnapshots", snapshots)
+
 	return multistep.ActionContinue
 }