From 74ffe3749d70b9d115a9b9790969b8dcb4a76e12 Mon Sep 17 00:00:00 2001
From: Stephen Lowrie <stephen.lowrie@gmail.com>
Date: Mon, 11 Jan 2021 11:27:12 -0600
Subject: [PATCH] internal/providers/aws: probe the IMDS token URL

Probing the `/latest` path causes a 401 Unauthorized when running with
IMDSv2 only. Instead ping the token URL.
---
 internal/providers/aws/aws.go | 37 +++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/internal/providers/aws/aws.go b/internal/providers/aws/aws.go
index 54373dbb..4a6655f0 100644
--- a/internal/providers/aws/aws.go
+++ b/internal/providers/aws/aws.go
@@ -40,11 +40,6 @@ var (
 		Host:   "169.254.169.254",
 		Path:   "2019-10-01/user-data",
 	}
-	metadataServiceProbeURL = url.URL{
-		Scheme: "http",
-		Host:   "169.254.169.254",
-		Path:   "latest",
-	}
 	imdsTokenURL = url.URL{
 		Scheme: "http",
 		Host:   "169.254.169.254",
@@ -78,17 +73,17 @@ func NewFetcher(l *log.Logger) (resource.Fetcher, error) {
 // Init prepares the fetcher for this platform
 func Init(f *resource.Fetcher) error {
 	// During the fetch stage we might be running before the networking
-	// is fully ready. Perform an HTTP fetch against the metadata probe
-	// URL to ensure that networking is up before we attempt to fetch
-	// the region hint from ec2metadata.
+	// is fully ready. Perform an HTTP fetch against the IMDS token URL
+	// to ensure that networking is up before we attempt to fetch the
+	// region hint from ec2metadata.
 	//
-	// NOTE: the FetchToBuffer call against the metadata service probe
-	// URL is a temporary solution to handle waiting for networking
-	// before fetching from the AWS API. We do this instead of an
-	// infinite retry loop on the API call because, without a clear
-	// understanding of the failure cases, that would risk provisioning
-	// failures due to quirks of the ec2metadata API.  Additionally a
-	// finite retry loop would have to time out quickly enough to avoid
+	// NOTE: the FetchToBuffer call against the IMDS token URL is a
+	// temporary solution to handle waiting for networking before
+	// fetching from the AWS API. We do this instead of an infinite
+	// retry loop on the API call because, without a clear understanding
+	// of the failure cases, that would risk provisioning failures due
+	// to quirks of the ec2metadata API.  Additionally a finite retry
+	// loop would have to time out quickly enough to avoid
 	// extraordinarily long boots on failure (since this code runs in
 	// every stage) but that would risk premature timeouts if the
 	// network takes a while to come up.
@@ -102,8 +97,16 @@ func Init(f *resource.Fetcher) error {
 	// NOTE: FetchToBuffer is handling the ErrNeedNet case.  If we move
 	// to an alternative method, we will need to handle the detection in
 	// this function.
-	_, err := f.FetchToBuffer(metadataServiceProbeURL, resource.FetchOptions{})
-	if err != nil {
+	opts := resource.FetchOptions{
+		Headers: http.Header{
+			"x-aws-ec2-metadata-token-ttl-seconds": []string{"21600"},
+		},
+		HTTPVerb: "PUT",
+	}
+	_, err := f.FetchToBuffer(imdsTokenURL, opts)
+	// ErrNotFound would just mean that the instance might not have
+	// IMDSv2 enabled
+	if err != nil && err != resource.ErrNotFound {
 		return err
 	}
 
-- 
2.29.2
