From a70e8462cbe7aa294dd2a51d054beaab0e83afd4 Mon Sep 17 00:00:00 2001
From: Simon Thulbourn <simon+github@thulbourn.com>
Date: Thu, 8 Jan 2015 21:51:34 +0000
Subject: [PATCH] adds support for AWS IAM roles

Signed-off-by: Simon Thulbourn <simon+github@thulbourn.com>
---
 README.md                      |  9 +++++----
 drivers/amazonec2/amazonec2.go | 11 ++++++++++-
 drivers/amazonec2/amz/auth.go  |  6 +++---
 drivers/amazonec2/amz/ec2.go   |  1 +
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index fad0248b..bdc05dea 100644
--- a/README.md
+++ b/README.md
@@ -107,6 +107,7 @@ Options:
  - `--amazonec2-region`: The region to use when launching the instance.  Default: `us-east-1`
  - `--amazonec2-root-size`: The root disk size of the instance (in GB).  Default: `16`
  - `--amazonec2-secret-key`: Your secret access key for the Amazon Web Services API.
+ - `--amazonec2-session-token`: Your session token for the Amazon Web Services API.
  - `--amazonec2-vpc-id`: Your VPC ID to launch the instance in.
  - `--amazonec2-zone`: The AWS zone launch the instance in (i.e. one of a,b,c,d,e).
 
@@ -128,7 +129,7 @@ Options:
 
 Creates machines locally on [VMware Fusion](http://www.vmware.com/products/fusion). Requires VMware Fusion to be installed.
 
-Options: 
+Options:
 
  - `--vmwarefusion-boot2docker-url`: URL for boot2docker image.
  - `--vmwarefusion-disk-size`: Size of disk for host VM (in MB). Default: `20000`
@@ -155,7 +156,7 @@ Options:
  - `--vmwarevcloudair-publicip`: Org Public IP to use.
  - `--vmwarevcloudair-ssh-port`: SSH port. Default: `22`
  - `--vmwarevcloudair-vdcid`: Virtual Data Center ID.
-   
+
 ### VMware vSphere
 
 Creates machines on a [VMware vSphere](http://www.vmware.com/products/vsphere) Virtual Infrastructure. Requires a working vSphere (ESXi and optionally vCenter) installation. The vSphere driver depends on [`govc`](https://github.com/vmware/govmomi/tree/master/govc) (must be in path) and has been tested with [vmware/govmomi@`c848630`](https://github.com/vmware/govmomi/commit/c8486300bfe19427e4f3226e3b3eac067717ef17).
@@ -215,9 +216,9 @@ There is a suite of integration tests that will run for the drivers.  In order
 to use these you must export the corresponding environment variables for each
 driver as these perform the actual actions (start, stop, restart, kill, etc).
 
-By default, the suite will run tests against all drivers in master.  You can 
+By default, the suite will run tests against all drivers in master.  You can
 override this by setting the environment variable `MACHINE_TESTS`.  For example,
-`MACHINE_TESTS="virtualbox" ./script/run-integration-tests` will only run the 
+`MACHINE_TESTS="virtualbox" ./script/run-integration-tests` will only run the
 virtualbox driver integration tests.
 
 To run, use the helper script `./script/run-integration-tests`.
diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
index ad302850..bca4db00 100644
--- a/drivers/amazonec2/amazonec2.go
+++ b/drivers/amazonec2/amazonec2.go
@@ -33,6 +33,7 @@ type Driver struct {
 	Id              string
 	AccessKey       string
 	SecretKey       string
+	SessionToken    string
 	Region          string
 	AMI             string
 	SSHKeyID        int
@@ -40,6 +41,7 @@ type Driver struct {
 	InstanceId      string
 	InstanceType    string
 	IPAddress       string
+	MachineName     string
 	SecurityGroupId string
 	ReservationId   string
 	RootSize        int64
@@ -81,6 +83,12 @@ func GetCreateFlags() []cli.Flag {
 			Value:  "",
 			EnvVar: "AWS_SECRET_ACCESS_KEY",
 		},
+		cli.StringFlag{
+			Name:   "amazonec2-session-token",
+			Usage:  "AWS Session Token",
+			Value:  "",
+			EnvVar: "AWS_SESSION_TOKEN",
+		},
 		cli.StringFlag{
 			Name:   "amazonec2-ami",
 			Usage:  "AWS machine image",
@@ -134,6 +142,7 @@ func NewDriver(storePath string) (drivers.Driver, error) {
 func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 	d.AccessKey = flags.String("amazonec2-access-key")
 	d.SecretKey = flags.String("amazonec2-secret-key")
+	d.SessionToken = flags.String("amazonec2-session-token")
 	d.AMI = flags.String("amazonec2-ami")
 	d.Region = flags.String("amazonec2-region")
 	d.InstanceType = flags.String("amazonec2-instance-type")
@@ -422,7 +431,7 @@ func (d *Driver) GetSSHCommand(args ...string) (*exec.Cmd, error) {
 }
 
 func (d *Driver) getClient() *amz.EC2 {
-	auth := amz.GetAuth(d.AccessKey, d.SecretKey)
+	auth := amz.GetAuth(d.AccessKey, d.SecretKey, d.SessionToken)
 	return amz.NewEC2(auth, d.Region)
 }
 
diff --git a/drivers/amazonec2/amz/auth.go b/drivers/amazonec2/amz/auth.go
index c307bc4f..22291b53 100644
--- a/drivers/amazonec2/amz/auth.go
+++ b/drivers/amazonec2/amz/auth.go
@@ -1,9 +1,9 @@
 package amz
 
 type Auth struct {
-	AccessKey, SecretKey string
+	AccessKey, SecretKey, SessionToken string
 }
 
-func GetAuth(accessKey, secretKey string) Auth {
-	return Auth{accessKey, secretKey}
+func GetAuth(accessKey, secretKey, sessionToken string) Auth {
+	return Auth{accessKey, secretKey, sessionToken}
 }
diff --git a/drivers/amazonec2/amz/ec2.go b/drivers/amazonec2/amz/ec2.go
index 33cd34b3..13c617ac 100644
--- a/drivers/amazonec2/amz/ec2.go
+++ b/drivers/amazonec2/amz/ec2.go
@@ -152,6 +152,7 @@ func (e *EC2) awsApiCall(v url.Values) (http.Response, error) {
 	awsauth.Sign(req, awsauth.Credentials{
 		AccessKeyID:     e.Auth.AccessKey,
 		SecretAccessKey: e.Auth.SecretKey,
+		SecurityToken:   e.Auth.SessionToken,
 	})
 	resp, err := client.Do(req)
 	if err != nil {