From af840fd630befedca80a291e2402e723f374e9d5 Mon Sep 17 00:00:00 2001
From: Cosmos Nicolaou <cosmos.nicolaou@gmail.com>
Date: Thu, 17 Oct 2019 14:49:31 -0700
Subject: [PATCH] ec2machine: allow for an instances ssh key name to be
 specified directly (#11)

---
 ec2system/ec2machine.go | 23 +++++++++++++++++++++++
 go.mod                  |  4 +---
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/ec2system/ec2machine.go b/ec2system/ec2machine.go
index 117320d..05b7fa0 100644
--- a/ec2system/ec2machine.go
+++ b/ec2system/ec2machine.go
@@ -93,6 +93,13 @@ var (
 
 var immortal = flag.Bool("ec2machineimmortal", false, "make immortal EC2 instances (debugging only)")
 
+// SetMortality conrols the mortality of EC2 instances for help with debugging
+// low level boot time issues. It is equivalent to the 'ec2machineimmportal' flag
+// for configurations where flags cannot be used.
+func SetMortality(v bool) {
+	*immortal = v
+}
+
 var (
 	// InstanceTypes stores metadata for each known EC2 instance type.
 	// TODO(marius): generate this from the the EC2 inventory JSON instead.
@@ -188,6 +195,11 @@ type System struct {
 	// one exists.
 	SshKeys []string
 
+	// The EC2 key pair name to associate with the created instances when
+	// the instance this launched. This key name will appear in the EC2
+	// instance's metadata.
+	EC2KeyName string
+
 	// The user running the application. For tagging.
 	Username string
 
@@ -304,6 +316,9 @@ func (s *System) Init(b *bigmachine.B) error {
 		return err
 	}
 	s.authorityContents, err = ioutil.ReadFile(authorityPath)
+	if err != nil {
+		return err
+	}
 	return err
 }
 
@@ -402,6 +417,12 @@ func (s *System) Start(ctx context.Context, count int) ([]*bigmachine.Machine, e
 			securityGroups[i] = aws.String(s.SecurityGroups[i])
 		}
 	}
+
+	var ec2KeyName *string
+	if len(s.EC2KeyName) > 0 {
+		ec2KeyName = aws.String(s.EC2KeyName)
+	}
+
 	if s.OnDemand {
 		run = func() ([]string, error) {
 			resv, err := s.ec2.RunInstances(&ec2.RunInstancesInput{
@@ -423,6 +444,7 @@ func (s *System) Start(ctx context.Context, count int) ([]*bigmachine.Machine, e
 				},
 				UserData:         aws.String(base64.StdEncoding.EncodeToString(userData)),
 				SecurityGroupIds: securityGroups,
+				KeyName:          ec2KeyName,
 			})
 			if err != nil {
 				return nil, err
@@ -455,6 +477,7 @@ func (s *System) Start(ctx context.Context, count int) ([]*bigmachine.Machine, e
 						Arn: aws.String(s.InstanceProfile),
 					},
 					SecurityGroupIds: securityGroups,
+					KeyName:          ec2KeyName,
 				},
 			})
 			if err != nil {
diff --git a/go.mod b/go.mod
index 9239ce5..79a51dd 100644
--- a/go.mod
+++ b/go.mod
@@ -4,8 +4,8 @@ go 1.12
 
 require (
 	github.com/aws/aws-sdk-go v1.25.13
+	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/google/pprof v0.0.0-20190930153522-6ce02741cba3
-	github.com/google/uuid v1.1.1 // indirect
 	github.com/grailbio/base v0.0.4
 	github.com/grailbio/testutil v0.0.3
 	github.com/shirou/gopsutil v2.19.9+incompatible
@@ -13,8 +13,6 @@ require (
 	golang.org/x/net v0.0.0-20191007182048-72f939374954
 	golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e
 	golang.org/x/time v0.0.0-20190921001708-c4c64cad1fd0
-	gonum.org/v1/gonum v0.0.0-20190902003836-43865b531bee // indirect
-	google.golang.org/api v0.11.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/yaml.v2 v2.2.4
 )