From eeea8648e4cd9bc493a7ec6381a5022a71f5971b Mon Sep 17 00:00:00 2001
From: Radek Simko <radek.simko@gmail.com>
Date: Sun, 7 Feb 2016 11:47:09 +0000
Subject: [PATCH] provider/aws: Refactor logic for getting credentials
 providers

---
 builtin/providers/aws/auth_helpers.go      | 58 ++++++++++------------
 builtin/providers/aws/auth_helpers_test.go | 24 ++++++++-
 2 files changed, 50 insertions(+), 32 deletions(-)

diff --git a/builtin/providers/aws/auth_helpers.go b/builtin/providers/aws/auth_helpers.go
index 0a9c852e4370..ecb02e7dcde1 100644
--- a/builtin/providers/aws/auth_helpers.go
+++ b/builtin/providers/aws/auth_helpers.go
@@ -1,8 +1,8 @@
 package aws
 
 import (
+	"fmt"
 	"log"
-	"net/http"
 	"os"
 	"strings"
 	"time"
@@ -12,12 +12,25 @@ import (
 	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/hashicorp/go-cleanhttp"
 )
 
+const DefaultMetadataEndpoint = "http://169.254.169.254:80/latest"
+
 // This function is responsible for reading credentials from the
 // environment in the case that they're not explicitly specified
 // in the Terraform configuration.
 func getCreds(key, secret, token, profile, credsfile string) *awsCredentials.Credentials {
+	// Build isolated HTTP client to avoid issues with globally-shared settings
+	client := cleanhttp.DefaultClient()
+
+	// Keep the timeout low as we don't want to wait in non-EC2 environments
+	client.Timeout = 100 * time.Millisecond
+	cfg := &aws.Config{
+		Endpoint:   aws.String(getMetadataEndpoint()),
+		HTTPClient: client,
+	}
+
 	// build a chain provider, lazy-evaulated by aws-sdk
 	providers := []awsCredentials.Provider{
 		&awsCredentials.StaticProvider{Value: awsCredentials.Value{
@@ -30,39 +43,22 @@ func getCreds(key, secret, token, profile, credsfile string) *awsCredentials.Cre
 			Filename: credsfile,
 			Profile:  profile,
 		},
+		&ec2rolecreds.EC2RoleProvider{
+			Client: ec2metadata.New(session.New(cfg)),
+		},
 	}
 
-	// We only look in the EC2 metadata API if we can connect
-	// to the metadata service within a reasonable amount of time
-	metadataURL := os.Getenv("AWS_METADATA_URL")
-	if metadataURL == "" {
-		metadataURL = "http://169.254.169.254:80/latest"
-	}
-	c := http.Client{
-		Timeout: 100 * time.Millisecond,
-	}
+	return awsCredentials.NewChainCredentials(providers)
+}
 
-	r, err := c.Get(metadataURL)
-	// Flag to determine if we should add the EC2Meta data provider. Default false
-	var useIAM bool
-	if err == nil {
-		// AWS will add a "Server: EC2ws" header value for the metadata request. We
-		// check the headers for this value to ensure something else didn't just
-		// happent to be listening on that IP:Port
-		if r.Header["Server"] != nil && strings.Contains(r.Header["Server"][0], "EC2") {
-			useIAM = true
-		}
+// getMetadataEndpoint returns metadata endpoint URL
+// including the version path
+func getMetadataEndpoint() string {
+	endpoint := os.Getenv("AWS_METADATA_ENDPOINT")
+	if endpoint != "" {
+		log.Printf("[DEBUG] Using custom metadata endpoint: %q", endpoint)
+		return endpoint
 	}
 
-	if useIAM {
-		log.Printf("[DEBUG] EC2 Metadata service found, adding EC2 Role Credential Provider")
-		providers = append(providers, &ec2rolecreds.EC2RoleProvider{
-			Client: ec2metadata.New(session.New(&aws.Config{
-				Endpoint: aws.String(metadataURL),
-			})),
-		})
-	} else {
-		log.Printf("[DEBUG] EC2 Metadata service not found, not adding EC2 Role Credential Provider")
-	}
-	return awsCredentials.NewChainCredentials(providers)
+	return DefaultMetadataEndpoint
 }
diff --git a/builtin/providers/aws/auth_helpers_test.go b/builtin/providers/aws/auth_helpers_test.go
index 5c58a57290fb..9ef97e8cf231 100644
--- a/builtin/providers/aws/auth_helpers_test.go
+++ b/builtin/providers/aws/auth_helpers_test.go
@@ -154,6 +154,28 @@ func TestAWSConfig_shouldIgnoreIAM(t *testing.T) {
 	}
 }
 
+func TestAWSConfig_shouldCatchEC2RoleProvider(t *testing.T) {
+	resetEnv := unsetEnv(t)
+	defer resetEnv()
+	// capture the test server's close method, to call after the test returns
+	ts := awsEnv(t)
+	defer ts()
+
+	creds := getCreds("", "", "", "", "")
+	if creds == nil {
+		t.Fatalf("Expected an EC2Role creds provider to be returned")
+	}
+	v, err := creds.Get()
+	if err != nil {
+		t.Fatalf("Expected no error when getting creds: %s", err)
+	}
+	expectedProvider := "EC2RoleProvider"
+	if v.ProviderName != expectedProvider {
+		t.Fatal("Expected provider name to be %q, %q given",
+			expectedProvider, v.ProviderName)
+	}
+}
+
 var credentialsFileContents = `[myprofile]
 aws_access_key_id = accesskey
 aws_secret_access_key = secretkey
@@ -331,7 +353,7 @@ func awsEnv(t *testing.T) func() {
 		}
 	}))
 
-	os.Setenv("AWS_METADATA_URL", ts.URL+"/latest")
+	os.Setenv("AWS_METADATA_ENDPOINT", ts.URL+"/latest")
 	return ts.Close
 }