From af0f21d7440551f4b57c2d7effc5b639af3efff6 Mon Sep 17 00:00:00 2001
From: Josh Hogle <josh@hysolate.com>
Date: Tue, 19 May 2020 18:42:12 -0400
Subject: [PATCH] added support for IMDSv2 API

---
 authority/provisioner/aws.go | 40 +++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/authority/provisioner/aws.go b/authority/provisioner/aws.go
index 3f2c8873..221634a4 100644
--- a/authority/provisioner/aws.go
+++ b/authority/provisioner/aws.go
@@ -29,6 +29,19 @@ const awsIdentityURL = "http://169.254.169.254/latest/dynamic/instance-identity/
 // awsSignatureURL is the url used to retrieve the instance identity signature.
 const awsSignatureURL = "http://169.254.169.254/latest/dynamic/instance-identity/signature"
 
+// awsAPITokenURL is the url used to get the IMDSv2 API token
+const awsAPITokenURL = "http://169.254.169.254/latest/api/token"
+
+// awsAPITokenTTL is the default TTL to use when requesting IMDSv2 API tokens
+// -- we keep this short-lived since we get a new token with every call to readURL()
+const awsAPITokenTTL = "30"
+
+// awsMetadataTokenHeader is the header that must be passed with every IMDSv2 request
+const awsMetadataTokenHeader = "X-aws-ec2-metadata-token"
+
+// awsMetadataTokenTTLHeader is the header used to indicate the token TTL requested
+const awsMetadataTokenTTLHeader = "X-aws-ec2-metadata-token-ttl-seconds"
+
 // awsCertificate is the certificate used to validate the instance identity
 // signature.
 const awsCertificate = `-----BEGIN CERTIFICATE-----
@@ -332,7 +345,15 @@ func (p *AWS) checkSignature(signed, signature []byte) error {
 // using pkg/errors to avoid verbose errors, the caller should use it and write
 // the appropriate error.
 func (p *AWS) readURL(url string) ([]byte, error) {
-	r, err := http.Get(url)
+	client := &http.Client{}
+
+	// get authorization token
+	req, err := http.NewRequest(http.MethodPut, awsAPITokenURL, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set(awsMetadataTokenTTLHeader, awsAPITokenTTL)
+	r, err := client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -341,6 +362,23 @@ func (p *AWS) readURL(url string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	token := string(b)
+
+	// now get the data
+	req, err = http.NewRequest(http.MethodGet, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set(awsMetadataTokenHeader, token)
+	r, err = client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer r.Body.Close()
+	b, err = ioutil.ReadAll(r.Body)
+	if err != nil {
+		return nil, err
+	}
 	return b, nil
 }