From d0651d0273152d6360734eef9e273bbf271b54f2 Mon Sep 17 00:00:00 2001
From: Simon Thulbourn <simon+github@thulbourn.com>
Date: Tue, 28 Apr 2015 16:10:37 +0100
Subject: [PATCH] Add --amazonec2-private-address-only

Signed-off-by: Simon Thulbourn <simon+github@thulbourn.com>
---
 docs/index.md                       |  1 +
 drivers/amazonec2/amazonec2.go      | 12 +++++++++++-
 drivers/amazonec2/amazonec2_test.go |  1 +
 drivers/amazonec2/amz/ec2.go        |  8 ++++++--
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/docs/index.md b/docs/index.md
index dafa687297..958f1932f5 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -739,6 +739,7 @@ Options:
  - `--amazonec2-subnet-id`: AWS VPC subnet id
  - `--amazonec2-vpc-id`: **required** 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). Default: `a`
+ - `--amazonec2-private-address-only`: Use the private IP address only
 
 By default, the Amazon EC2 driver will use a daily image of Ubuntu 14.04 LTS.
 
diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
index 1e91eff0c4..e9781acc61 100644
--- a/drivers/amazonec2/amazonec2.go
+++ b/drivers/amazonec2/amazonec2.go
@@ -69,6 +69,7 @@ type Driver struct {
 	keyPath             string
 	RequestSpotInstance bool
 	SpotPrice           string
+	PrivateIPOnly       bool
 }
 
 func init() {
@@ -164,6 +165,10 @@ func GetCreateFlags() []cli.Flag {
 			Usage: "AWS spot instance bid price (in dollar)",
 			Value: "0.50",
 		},
+		cli.BoolFlag{
+			Name:  "amazonec2-private-address-only",
+			Usage: "Only use a private IP address",
+		},
 	}
 }
 
@@ -221,6 +226,7 @@ func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 	d.SwarmDiscovery = flags.String("swarm-discovery")
 	d.SSHUser = flags.String("amazonec2-ssh-user")
 	d.SSHPort = 22
+	d.PrivateIPOnly = flags.Bool("amazonec2-private-address-only")
 
 	if d.AccessKey == "" {
 		return fmt.Errorf("amazonec2 driver requires the --amazonec2-access-key option")
@@ -372,7 +378,7 @@ func (d *Driver) Create() error {
 			return fmt.Errorf("Error get instance: %s", err)
 		}
 	} else {
-		inst, err := d.getClient().RunInstance(d.AMI, d.InstanceType, d.Zone, 1, 1, d.SecurityGroupId, d.KeyName, d.SubnetId, bdm, d.IamInstanceProfile)
+		inst, err := d.getClient().RunInstance(d.AMI, d.InstanceType, d.Zone, 1, 1, d.SecurityGroupId, d.KeyName, d.SubnetId, bdm, d.IamInstanceProfile, d.PrivateIPOnly)
 		if err != nil {
 			return fmt.Errorf("Error launching instance: %s", err)
 		}
@@ -427,6 +433,10 @@ func (d *Driver) GetIP() (string, error) {
 		return "", err
 	}
 
+	if d.PrivateIPOnly {
+		return inst.PrivateIpAddress, nil
+	}
+
 	return inst.IpAddress, nil
 }
 
diff --git a/drivers/amazonec2/amazonec2_test.go b/drivers/amazonec2/amazonec2_test.go
index a2fcf2f404..d94737eeb9 100644
--- a/drivers/amazonec2/amazonec2_test.go
+++ b/drivers/amazonec2/amazonec2_test.go
@@ -80,6 +80,7 @@ func getDefaultTestDriverFlags() *DriverOptionsMock {
 			"amazonec2-ssh-user":              "ubuntu",
 			"amazonec2-request-spot-instance": false,
 			"amazonec2-spot-price":            "",
+			"amazonec2-private-address-only":  false,
 		},
 	}
 }
diff --git a/drivers/amazonec2/amz/ec2.go b/drivers/amazonec2/amz/ec2.go
index 3250dcecb0..e1ddf8191d 100644
--- a/drivers/amazonec2/amz/ec2.go
+++ b/drivers/amazonec2/amz/ec2.go
@@ -173,7 +173,7 @@ func (e *EC2) awsApiCall(v url.Values) (*http.Response, error) {
 	return resp, nil
 }
 
-func (e *EC2) RunInstance(amiId string, instanceType string, zone string, minCount int, maxCount int, securityGroup string, keyName string, subnetId string, bdm *BlockDeviceMapping, role string) (EC2Instance, error) {
+func (e *EC2) RunInstance(amiId string, instanceType string, zone string, minCount int, maxCount int, securityGroup string, keyName string, subnetId string, bdm *BlockDeviceMapping, role string, privateIPOnly bool) (EC2Instance, error) {
 	instance := Instance{}
 	v := url.Values{}
 	v.Set("Action", "RunInstances")
@@ -186,7 +186,11 @@ func (e *EC2) RunInstance(amiId string, instanceType string, zone string, minCou
 	v.Set("NetworkInterface.0.DeviceIndex", "0")
 	v.Set("NetworkInterface.0.SecurityGroupId.0", securityGroup)
 	v.Set("NetworkInterface.0.SubnetId", subnetId)
-	v.Set("NetworkInterface.0.AssociatePublicIpAddress", "1")
+	if privateIPOnly {
+		v.Set("NetworkInterface.0.AssociatePublicIpAddress", "0")
+	} else {
+		v.Set("NetworkInterface.0.AssociatePublicIpAddress", "1")
+	}
 
 	if len(role) > 0 {
 		v.Set("IamInstanceProfile.Name", role)