From 3afd77a5e4473a2b9cd9055323579e4b08822a7f Mon Sep 17 00:00:00 2001
From: Luke Farnell <lfarnell@outlook.com>
Date: Tue, 16 May 2017 17:05:40 -0400
Subject: [PATCH] fixed logic error and fixed suggestion

---
 builder/digitalocean/builder.go       | 45 ++++++++++++++++++++++-----
 builder/digitalocean/step_snapshot.go |  9 +++---
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/builder/digitalocean/builder.go b/builder/digitalocean/builder.go
index 8d653ec13..66054f67f 100644
--- a/builder/digitalocean/builder.go
+++ b/builder/digitalocean/builder.go
@@ -35,6 +35,16 @@ func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
 	return nil, nil
 }
 
+func contains(slice []string, item string) bool {
+	set := make(map[string]struct{}, len(slice))
+	for _, s := range slice {
+		set[s] = struct{}{}
+	}
+
+	_, ok := set[item]
+	return ok
+}
+
 func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
 	client := godo.NewClient(oauth2.NewClient(oauth2.NoContext, &apiTokenSource{
 		AccessToken: b.config.APIToken,
@@ -54,15 +64,34 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		}
 		regions, _, err := client.Regions.List(context.TODO(), opt)
 		if err != nil {
-			return nil, fmt.Errorf("DigitalOcean: Unable to get regions, %s.", err)
+			return nil, fmt.Errorf("DigitalOcean: Unable to get regions, %s", err)
 		}
-		for _, snapShotRegion := range b.config.SnapshotRegions {
-			for _, region := range regions {
-				if snapShotRegion == region.Slug {
-					continue
-				} else {
-					return nil, fmt.Errorf("DigitalOcean: Invalid region, %s.", snapShotRegion)
-				}
+
+		var dcs []string
+		for _, val := range regions {
+			dcs = append(dcs, val.Slug)
+		}
+
+		regionSet := make(map[string]struct{})
+		regionsMap := make([]string, 0, len(b.config.SnapshotRegions))
+		regionSet[b.config.Region] = struct{}{}
+		for _, region := range b.config.SnapshotRegions {
+			// If we already saw the region, then don't look again
+			if _, ok := regionSet[region]; ok {
+				continue
+			}
+
+			// Mark that we saw the region
+			regionSet[region] = struct{}{}
+
+			regionsMap = append(regionsMap, region)
+		}
+
+		for _, val := range regionsMap {
+			if contains(dcs, val) {
+				continue
+			} else {
+				return nil, fmt.Errorf("DigitalOcean: Invalid region, %s", val)
 			}
 		}
 	}
diff --git a/builder/digitalocean/step_snapshot.go b/builder/digitalocean/step_snapshot.go
index 180cc8fcb..bd5298137 100644
--- a/builder/digitalocean/step_snapshot.go
+++ b/builder/digitalocean/step_snapshot.go
@@ -19,6 +19,7 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 	ui := state.Get("ui").(packer.Ui)
 	c := state.Get("config").(Config)
 	dropletId := state.Get("droplet_id").(int)
+	var snapshotRegions []string
 
 	ui.Say(fmt.Sprintf("Creating snapshot: %v", c.SnapshotName))
 	action, _, err := client.DropletActions.Snapshot(context.TODO(), dropletId, c.SnapshotName)
@@ -75,12 +76,12 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 
 			regions = append(regions, region)
 		}
-		c.SnapshotRegions = regions
+		snapshotRegions = regions
 
-		for transfer := range c.SnapshotRegions {
+		for transfer := range snapshotRegions {
 			transferRequest := &godo.ActionRequest{
 				"type":   "transfer",
-				"region": c.SnapshotRegions[transfer],
+				"region": snapshotRegions[transfer],
 			}
 			imageTransfer, _, err := client.ImageActions.Transfer(context.TODO(), images[0].ID, transferRequest)
 			if err != nil {
@@ -114,7 +115,7 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 	log.Printf("Snapshot image ID: %d", imageId)
 	state.Put("snapshot_image_id", imageId)
 	state.Put("snapshot_name", c.SnapshotName)
-	state.Put("regions", c.SnapshotRegions)
+	state.Put("regions", snapshotRegions)
 
 	return multistep.ActionContinue
 }