From e53d6aea6623c2900e50b095257f23450dde623c Mon Sep 17 00:00:00 2001
From: Wilken Rivera <dev@wilkenrivera.com>
Date: Fri, 24 Apr 2020 09:05:13 -0400
Subject: [PATCH] amazon/ssm_driver: Update to support command context

* Defer the release of the found port to prevent possible Packer build collisions.
---
 builder/amazon/common/ssm_driver.go             | 6 ++++--
 builder/amazon/common/step_create_ssm_tunnel.go | 7 +++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/builder/amazon/common/ssm_driver.go b/builder/amazon/common/ssm_driver.go
index 789821ef8..82ad8887b 100644
--- a/builder/amazon/common/ssm_driver.go
+++ b/builder/amazon/common/ssm_driver.go
@@ -2,6 +2,7 @@ package common
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"log"
 	"os/exec"
@@ -12,7 +13,8 @@ import (
 const SessionManagerPluginName string = "session-manager-plugin"
 
 type SSMDriver struct {
-	Ui packer.Ui
+	Ui  packer.Ui
+	Ctx context.Context
 	// Provided for testing purposes; if not specified it defaults to SessionManagerPluginName
 	PluginName string
 }
@@ -40,7 +42,7 @@ func (s *SSMDriver) StartSession(sessionData, region, profile, params, endpoint
 	}
 
 	log.Printf("Attempting to start session with the following args: %v", args)
-	cmd := exec.Command(s.PluginName, args...)
+	cmd := exec.CommandContext(s.Ctx, s.PluginName, args...)
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
 
diff --git a/builder/amazon/common/step_create_ssm_tunnel.go b/builder/amazon/common/step_create_ssm_tunnel.go
index 922911cfc..d06feca14 100644
--- a/builder/amazon/common/step_create_ssm_tunnel.go
+++ b/builder/amazon/common/step_create_ssm_tunnel.go
@@ -50,7 +50,6 @@ func (s *StepCreateSSMTunnel) Run(ctx context.Context, state multistep.StateBag)
 		"portNumber":      []*string{aws.String(dst)},
 		"localPortNumber": []*string{aws.String(src)},
 	}
-	l.Close()
 
 	instance, ok := state.Get("instance").(*ec2.Instance)
 	if !ok {
@@ -101,7 +100,11 @@ func (s *StepCreateSSMTunnel) Run(ctx context.Context, state multistep.StateBag)
 		return multistep.ActionHalt
 	}
 
-	driver := SSMDriver{Ui: ui}
+	// Stop listening on selected port so that the AWS session-manager-plugin can use it.
+	// The port is closed right before we start the session to avoid two Packer builds from getting the same port - fingers-crossed
+	l.Close()
+
+	driver := SSMDriver{Ui: ui, Ctx: ctx}
 	// sessionDetails, region, "StartSession", profile, paramJson, endpoint
 	region := aws.StringValue(s.AWSSession.Config.Region)
 	// how to best get Profile name