From 335f42b38c47bdee069984499ac3d25af55b9d05 Mon Sep 17 00:00:00 2001
From: Justin Santa Barbara <justin@fathomdb.com>
Date: Sun, 19 Apr 2015 21:00:18 -0700
Subject: [PATCH] AWS: Return public & private addresses in GetNodeAddresses

---
 pkg/cloudprovider/aws/aws.go      | 30 +++++++++++++++++++++++++-----
 pkg/cloudprovider/aws/aws_test.go | 20 +++++++++++++++-----
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/pkg/cloudprovider/aws/aws.go b/pkg/cloudprovider/aws/aws.go
index fa2a42b4650..63eabd22a63 100644
--- a/pkg/cloudprovider/aws/aws.go
+++ b/pkg/cloudprovider/aws/aws.go
@@ -275,16 +275,36 @@ func (aws *AWSCloud) Zones() (cloudprovider.Zones, bool) {
 
 // NodeAddresses is an implementation of Instances.NodeAddresses.
 func (aws *AWSCloud) NodeAddresses(name string) ([]api.NodeAddress, error) {
-	inst, err := aws.getInstancesByDnsName(name)
+	instance, err := aws.getInstancesByDnsName(name)
 	if err != nil {
 		return nil, err
 	}
-	ip := net.ParseIP(inst.PrivateIpAddress)
-	if ip == nil {
-		return nil, fmt.Errorf("invalid network IP: %s", inst.PrivateIpAddress)
+
+	addresses := []api.NodeAddress{}
+
+	if instance.PrivateIpAddress != "" {
+		ipAddress := instance.PrivateIpAddress
+		ip := net.ParseIP(ipAddress)
+		if ip == nil {
+			return nil, fmt.Errorf("EC2 instance had invalid private address: %s (%s)", instance.InstanceId, ipAddress)
+		}
+		addresses = append(addresses, api.NodeAddress{Type: api.NodeInternalIP, Address: ip.String()})
+
+		// Legacy compatibility: the private ip was the legacy host ip
+		addresses = append(addresses, api.NodeAddress{Type: api.NodeLegacyHostIP, Address: ip.String()})
 	}
 
-	return []api.NodeAddress{{Type: api.NodeLegacyHostIP, Address: ip.String()}}, nil
+	// TODO: Other IP addresses (multiple ips)?
+	if instance.PublicIpAddress != "" {
+		ipAddress := instance.PublicIpAddress
+		ip := net.ParseIP(ipAddress)
+		if ip == nil {
+			return nil, fmt.Errorf("EC2 instance had invalid public address: %s (%s)", instance.InstanceId, ipAddress)
+		}
+		addresses = append(addresses, api.NodeAddress{Type: api.NodeExternalIP, Address: ip.String()})
+	}
+
+	return addresses, nil
 }
 
 // ExternalID returns the cloud provider ID of the specified instance.
diff --git a/pkg/cloudprovider/aws/aws_test.go b/pkg/cloudprovider/aws/aws_test.go
index 79c89ea5cf0..21ffc227a3e 100644
--- a/pkg/cloudprovider/aws/aws_test.go
+++ b/pkg/cloudprovider/aws/aws_test.go
@@ -267,12 +267,22 @@ func TestList(t *testing.T) {
 	}
 }
 
+func testHasNodeAddress(t *testing.T, addrs []api.NodeAddress, addressType api.NodeAddressType, address string) {
+	for _, addr := range addrs {
+		if addr.Type == addressType && addr.Address == address {
+			return
+		}
+	}
+	t.Errorf("Did not find expected address: %s:%s in %v", addressType, address, addrs)
+}
+
 func TestNodeAddresses(t *testing.T) {
 	// Note these instances have the same name
 	// (we test that this produces an error)
 	instances := make([]ec2.Instance, 2)
 	instances[0].PrivateDNSName = "instance1"
 	instances[0].PrivateIpAddress = "192.168.0.1"
+	instances[0].PublicIpAddress = "1.2.3.4"
 	instances[0].State.Name = "running"
 	instances[1].PrivateDNSName = "instance1"
 	instances[1].PrivateIpAddress = "192.168.0.2"
@@ -295,12 +305,12 @@ func TestNodeAddresses(t *testing.T) {
 	if err3 != nil {
 		t.Errorf("Should not error when instance found")
 	}
-	if len(addrs3) != 1 {
-		t.Errorf("Should return exactly one NodeAddress")
-	}
-	if e, a := instances[0].PrivateIpAddress, addrs3[0].Address; e != a {
-		t.Errorf("Expected %v, got %v", e, a)
+	if len(addrs3) != 3 {
+		t.Errorf("Should return exactly 3 NodeAddresses")
 	}
+	testHasNodeAddress(t, addrs3, api.NodeInternalIP, "192.168.0.1")
+	testHasNodeAddress(t, addrs3, api.NodeLegacyHostIP, "192.168.0.1")
+	testHasNodeAddress(t, addrs3, api.NodeExternalIP, "1.2.3.4")
 }
 
 func TestGetRegion(t *testing.T) {