From 9bca9db6677902848458b0cc7375fd7e26a185fd Mon Sep 17 00:00:00 2001
From: Radek Simko <radek.simko@gmail.com>
Date: Sun, 7 Feb 2016 11:41:14 +0000
Subject: [PATCH] provider/aws: Make account ID validation work in EC2
 instances

---
 builtin/providers/aws/auth_helpers.go      | 19 +++++++-
 builtin/providers/aws/auth_helpers_test.go | 54 ++++++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/builtin/providers/aws/auth_helpers.go b/builtin/providers/aws/auth_helpers.go
index 0e5a87387440..914c7e97174d 100644
--- a/builtin/providers/aws/auth_helpers.go
+++ b/builtin/providers/aws/auth_helpers.go
@@ -18,7 +18,24 @@ import (
 )
 
 func GetAccountId(iamconn *iam.IAM, authProviderName string) (string, error) {
-	// Try IAM GetUser
+	// If we have creds from instance profile, we can use metadata API
+	if authProviderName == ec2rolecreds.ProviderName {
+		log.Println("[DEBUG] Trying to get account ID via AWS Metadata API")
+
+		cfg := &aws.Config{}
+		setOptionalEndpoint(cfg)
+		metadataClient := ec2metadata.New(session.New(cfg))
+		info, err := metadataClient.IAMInfo()
+		if err != nil {
+			// This can be triggered when no IAM Role is assigned
+			// or AWS just happens to return invalid response
+			return "", fmt.Errorf("Failed getting EC2 IAM info: %s", err)
+		}
+
+		return parseAccountIdFromArn(info.InstanceProfileArn)
+	}
+
+	// Then try IAM GetUser
 	log.Println("[DEBUG] Trying to get account ID via iam:GetUser")
 	outUser, err := iamconn.GetUser(nil)
 	if err == nil {
diff --git a/builtin/providers/aws/auth_helpers_test.go b/builtin/providers/aws/auth_helpers_test.go
index 6c69ec02bd30..b3f13403923e 100644
--- a/builtin/providers/aws/auth_helpers_test.go
+++ b/builtin/providers/aws/auth_helpers_test.go
@@ -15,10 +15,60 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	awsCredentials "github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/iam"
 )
 
+func TestAWSGetAccountId_shouldBeValid_fromEC2Role(t *testing.T) {
+	resetEnv := unsetEnv(t)
+	defer resetEnv()
+	// capture the test server's close method, to call after the test returns
+	awsTs := awsEnv(t)
+	defer awsTs()
+
+	iamEndpoints := []*iamEndpoint{}
+	ts, iamConn := getMockedAwsIamApi(iamEndpoints)
+	defer ts()
+
+	id, err := GetAccountId(iamConn, ec2rolecreds.ProviderName)
+	if err != nil {
+		t.Fatalf("Getting account ID from EC2 metadata API failed: %s", err)
+	}
+
+	expectedAccountId := "123456789013"
+	if id != expectedAccountId {
+		t.Fatalf("Expected account ID: %s, given: %s", expectedAccountId, id)
+	}
+}
+
+func TestAWSGetAccountId_shouldBeValid_EC2RoleHasPriority(t *testing.T) {
+	resetEnv := unsetEnv(t)
+	defer resetEnv()
+	// capture the test server's close method, to call after the test returns
+	awsTs := awsEnv(t)
+	defer awsTs()
+
+	iamEndpoints := []*iamEndpoint{
+		&iamEndpoint{
+			Request:  &iamRequest{"POST", "/", "Action=GetUser&Version=2010-05-08"},
+			Response: &iamResponse{200, iamResponse_GetUser_valid, "text/xml"},
+		},
+	}
+	ts, iamConn := getMockedAwsIamApi(iamEndpoints)
+	defer ts()
+
+	id, err := GetAccountId(iamConn, ec2rolecreds.ProviderName)
+	if err != nil {
+		t.Fatalf("Getting account ID from EC2 metadata API failed: %s", err)
+	}
+
+	expectedAccountId := "123456789013"
+	if id != expectedAccountId {
+		t.Fatalf("Expected account ID: %s, given: %s", expectedAccountId, id)
+	}
+}
+
 func TestAWSGetAccountId_shouldBeValid_fromIamUser(t *testing.T) {
 	iamEndpoints := []*iamEndpoint{
 		&iamEndpoint{
@@ -610,6 +660,10 @@ const metadataApiRoutes = `
       "uri": "/latest/meta-data/instance-id",
       "body": "mock-instance-id"
     },
+    {
+      "uri": "/latest/meta-data/iam/info",
+      "body": "{\"Code\": \"Success\",\"LastUpdated\": \"2016-03-17T12:27:32Z\",\"InstanceProfileArn\": \"arn:aws:iam::123456789013:instance-profile/my-instance-profile\",\"InstanceProfileId\": \"AIPAABCDEFGHIJKLMN123\"}"
+    },
     {
       "uri": "/latest/meta-data/iam/security-credentials",
       "body": "test_role"