From b08287f46922d0fdb41549649b178905127d95c8 Mon Sep 17 00:00:00 2001
From: Karol Kokoszka <karol.kokoszka@scylladb.com>
Date: Wed, 28 Feb 2024 15:23:38 +0100
Subject: [PATCH] feat(aws): fallback to IMDSv1 if session token cannot be
 retrieved

---
 pkg/rclone/aws.go | 68 ++++++++++++++++++++++++++---------------------
 1 file changed, 38 insertions(+), 30 deletions(-)

diff --git a/pkg/rclone/aws.go b/pkg/rclone/aws.go
index 83f106134c..ec8bf47ca2 100644
--- a/pkg/rclone/aws.go
+++ b/pkg/rclone/aws.go
@@ -9,53 +9,34 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/pkg/errors"
 	"github.com/rclone/rclone/fs"
 )
 
 // awsRegionFromMetadataAPI uses instance metadata API v2 to fetch region of the
 // running instance see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
 // Returns empty string if region can't be obtained for whatever reason.
+// Fallbacks to IMDSv1 if the session token cannot be obtained.
 func awsRegionFromMetadataAPI() string {
-	const (
-		tokenUrl = "http://169.254.169.254/latest/api/token"
-		docURL   = "http://169.254.169.254/latest/dynamic/instance-identity/document"
-	)
+	const docURL = "http://169.254.169.254/latest/dynamic/instance-identity/document"
 
 	// Step 1: Request an IMDSv2 session token
-	reqToken, err := http.NewRequestWithContext(context.Background(), http.MethodPut, tokenUrl, nil)
+	token, err := awsAPIToken()
+	// fallback to IMDSv1 when error on token retrieval
 	if err != nil {
-		fs.Errorf(nil, "create token request: %+v", err)
-		return ""
-	}
-	reqToken.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
-	tokenClient := http.Client{
-		Timeout: 2 * time.Second,
-	}
-	resToken, err := tokenClient.Do(reqToken)
-	if err != nil {
-		fs.Errorf(nil, "IMDSv2 failed to fetch session token: %+v", err)
-		return ""
-	}
-	defer resToken.Body.Close()
-
-	if resToken.StatusCode != http.StatusOK {
-		fs.Errorf(nil, "failed to retrieve session token: %s", resToken.Status)
-		return ""
-	}
-
-	token, err := io.ReadAll(resToken.Body)
-	if err != nil {
-		fs.Errorf(nil, "Failed to read session token: %+v", err)
-		return ""
+		fs.Errorf(nil, "%+v", err)
+		token = ""
 	}
 
 	// Step 2: Use the session token to retrieve instance metadata
-	reqMetadata, err := http.NewRequestWithContext(context.Background(), http.MethodGet, docURL, nil)
+	reqMetadata, err := http.NewRequestWithContext(context.Background(), http.MethodGet, docURL, http.NoBody)
 	if err != nil {
 		fs.Errorf(nil, "create metadata request: %+v", err)
 		return ""
 	}
-	reqMetadata.Header.Set("X-aws-ec2-metadata-token", string(token))
+	if token != "" {
+		reqMetadata.Header.Set("X-aws-ec2-metadata-token", token)
+	}
 
 	metadataClient := http.Client{
 		Timeout: 2 * time.Second,
@@ -77,3 +58,30 @@ func awsRegionFromMetadataAPI() string {
 
 	return metadata.Region
 }
+
+func awsAPIToken() (string, error) {
+	const tokenURL = "http://169.254.169.254/latest/api/token"
+
+	reqToken, err := http.NewRequestWithContext(context.Background(), http.MethodPut, tokenURL, http.NoBody)
+	if err != nil {
+		return "", errors.Wrap(err, "create token request")
+	}
+	reqToken.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+	tokenClient := http.Client{
+		Timeout: 2 * time.Second,
+	}
+	resToken, err := tokenClient.Do(reqToken)
+	if err != nil {
+		return "", errors.Wrap(err, "IMDSv2 failed to fetch session token")
+	}
+	defer resToken.Body.Close()
+
+	if resToken.StatusCode != http.StatusOK {
+		return "", errors.Wrap(err, "failed to retrieve session token")
+	}
+	token, err := io.ReadAll(resToken.Body)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to read session token")
+	}
+	return string(token), nil
+}