From a20ba021db3d994b5c3521c5d366bc70e27db5dc Mon Sep 17 00:00:00 2001
From: Coda Hale <coda@stripe.com>
Date: Mon, 15 Dec 2014 21:51:39 -0800
Subject: [PATCH] Added IAM credentials.

Closes #6. Closes #13.
---
 aws/auth.go      | 51 +++++++++++++++++++++++++++++++++++++
 aws/auth_test.go | 65 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)

diff --git a/aws/auth.go b/aws/auth.go
index 005b3752fe0..38f43287b11 100644
--- a/aws/auth.go
+++ b/aws/auth.go
@@ -1,7 +1,11 @@
 package aws
 
 import (
+	"encoding/json"
+	"net/http"
 	"os"
+	"sync"
+	"time"
 
 	"github.com/juju/errors"
 )
@@ -65,6 +69,53 @@ func Creds(accessKeyID, secretAccessKey, securityToken string) CredentialsProvid
 	}
 }
 
+// IAMCreds returns a provider which pulls credentials from the local EC2
+// instance's IAM roles.
+func IAMCreds() CredentialsProvider {
+	return &iamProvider{}
+}
+
+type iamProvider struct {
+	creds      Credentials
+	m          sync.Mutex
+	expiration time.Time
+}
+
+var metadataCredentialsEndpoint = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
+
+func (p *iamProvider) Credentials() (*Credentials, error) {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	if p.expiration.Before(currentTime()) {
+		var body struct {
+			Expiration      time.Time
+			AccessKeyID     string
+			SecretAccessKey string
+			Token           string
+		}
+
+		resp, err := http.Get(metadataCredentialsEndpoint)
+		if err != nil {
+			return nil, err
+		}
+		defer resp.Body.Close()
+
+		if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
+			return nil, err
+		}
+
+		p.creds = Credentials{
+			AccessKeyID:     body.AccessKeyID,
+			SecretAccessKey: body.SecretAccessKey,
+			SecurityToken:   body.Token,
+		}
+		p.expiration = body.Expiration
+	}
+
+	return &p.creds, nil
+}
+
 type staticCredentialsProvider struct {
 	creds Credentials
 }
diff --git a/aws/auth_test.go b/aws/auth_test.go
index 033ffb3ba3a..0f7cbf28589 100644
--- a/aws/auth_test.go
+++ b/aws/auth_test.go
@@ -1,8 +1,12 @@
 package aws
 
 import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"testing"
+	"time"
 )
 
 func TestEnvCreds(t *testing.T) {
@@ -77,3 +81,64 @@ func TestEnvCredsAlternateNames(t *testing.T) {
 		t.Errorf("Secret access key was %v, expected %v", v, want)
 	}
 }
+
+func TestIAMCreds(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, `{
+  "AccessKeyId" : "accessKey",
+  "SecretAccessKey" : "secret",
+  "Token" : "token",
+  "Expiration" : "2014-12-16T01:51:37Z"
+}`)
+	}))
+	defer server.Close()
+
+	defer func(s string) {
+		metadataCredentialsEndpoint = s
+	}(metadataCredentialsEndpoint)
+	metadataCredentialsEndpoint = server.URL
+
+	defer func() {
+		currentTime = time.Now
+	}()
+	currentTime = func() time.Time {
+		return time.Date(2014, 12, 15, 21, 26, 0, 0, time.UTC)
+	}
+
+	prov := IAMCreds()
+	t.Log(prov.Credentials())
+}
+
+func BenchmarkIAMCreds(b *testing.B) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, `{
+  "AccessKeyId" : "accessKey",
+  "SecretAccessKey" : "secret",
+  "Token" : "token",
+  "Expiration" : "2014-12-16T01:51:37Z"
+}`)
+	}))
+	defer server.Close()
+
+	defer func(s string) {
+		metadataCredentialsEndpoint = s
+	}(metadataCredentialsEndpoint)
+	metadataCredentialsEndpoint = server.URL
+
+	defer func() {
+		currentTime = time.Now
+	}()
+	currentTime = func() time.Time {
+		return time.Date(2014, 12, 15, 21, 26, 0, 0, time.UTC)
+	}
+
+	b.ResetTimer()
+
+	prov := IAMCreds()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			prov.Credentials()
+		}
+	})
+}