Implement Amazon's IMDSv2

CHANGELOG.md

@@ -1,5 +1,10 @@
 # ChangeLog
 
+## Unreleased
+
+* To keep up with the latest security protocols implemented by Amazon Web
+  Services, the agent now uses AWS IMDSv2 to find utilization data.
+
 ## 3.9.0
 
 ### Changes

v3/internal/utilization/aws.go

@@ -12,9 +12,12 @@ import (
 )
 
 const (
-	awsHostname     = "169.254.169.254"
-	awsEndpointPath = "/2016-09-02/dynamic/instance-identity/document"
-	awsEndpoint     = "http://" + awsHostname + awsEndpointPath
+	awsHostname          = "169.254.169.254"
+	awsEndpointPath      = "/2016-09-02/dynamic/instance-identity/document"
+	awsTokenEndpointPath = "/latest/api/token"
+	awsEndpoint          = "http://" + awsHostname + awsEndpointPath
+	awsTokenEndpoint     = "http://" + awsHostname + awsTokenEndpointPath
+	awsTokenTTL          = "60" // seconds this AWS utilization session will last
 )
 
 type aws struct {
@@ -44,6 +47,24 @@ func (e unexpectedAWSErr) Error() string {
 	return fmt.Sprintf("unexpected AWS error: %v", e.e)
 }
 
+// getAWSToken attempts to get the IMDSv2 token within the providerTimeout set
+// provider.go.
+func getAWSToken(client *http.Client) (token string, err error) {
+	request, err := http.NewRequest("PUT", awsTokenEndpoint, nil)
+	request.Header.Add("X-aws-ec2-metadata-token-ttl-seconds", awsTokenTTL)
+	response, err := client.Do(request)
+	if err != nil {
+		return "", err
+	}
+	defer response.Body.Close()
+	body, err := ioutil.ReadAll(response.Body)
+	if err != nil {
+		return "", err
+	}
+
+	return string(body), nil
+}
+
 func getAWS(client *http.Client) (ret *aws, err error) {
 	// In some cases, 3rd party providers might block requests to metadata
 	// endpoints in such a way that causes a panic in the underlying
@@ -56,7 +77,19 @@ func getAWS(client *http.Client) (ret *aws, err error) {
 		}
 	}()
 
-	response, err := client.Get(awsEndpoint)
+	// AWS' IMDSv2 requires us to get a token before requesting metadata.
+	awsToken, err := getAWSToken(client)
+	if err != nil {
+		// No unexpectedAWSErr here: A timeout is usually going to
+		// happen.
+		return nil, err
+	}
+
+	//Add the header to the outbound request.
+	request, err := http.NewRequest("GET", awsEndpoint, nil)
+	request.Header.Add("X-aws-ec2-metadata-token", awsToken)
+
+	response, err := client.Do(request)
 	if err != nil {
 		// No unexpectedAWSErr here: A timeout is usually going to
 		// happen.

v3/internal/utilization/provider.go

@@ -14,8 +14,8 @@ import (
 
 // Constants from the spec.
 const (
-	maxFieldValueSize = 255             // The maximum value size, in bytes.
-	providerTimeout   = 1 * time.Second // The maximum time a HTTP provider may block.
+	maxFieldValueSize = 255                    // The maximum value size, in bytes.
+	providerTimeout   = 500 * time.Millisecond // The maximum time a HTTP provider may block.
 	lookupAddrTimeout = 500 * time.Millisecond
 )
 

v3/internal/utilization/provider_test.go

@@ -48,8 +48,12 @@ type mockBody struct {
 }
 
 func (m *mockTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+
+	// Half the requests are going to the test's endpoint, while the other half
+	// are going to the AWS IMDSv2 token endpoint. Accept both.
 	for match, response := range m.responses {
-		if r.URL.String() == match {
+		if (r.URL.String() == match) ||
+			(r.URL.String() == awsTokenEndpoint) {
 			return m.respond(response)
 		}
 	}