From 547dfa87563685890e3c9a90bf5b75ac7de8ef53 Mon Sep 17 00:00:00 2001
From: "Chinmay D. Pai" <chinmaydpai@gmail.com>
Date: Sat, 29 Jun 2024 00:52:18 +0530
Subject: [PATCH] feat: add support for IMDSv2

Currently, IMDSv2 is not supported by simple-s3's `NewUsingIAM` method, causing it
to fail on EC2 instances with IMDSv2 enabled in the EC2 metadata configuration.

This patch adds a failsafe method that executes PUT request on the IMDSv2 token endpoint,
and based on the response decides whether IMDSv2 support needs to be enabled. Doing so
will not break or require any code-level changes to existing deployments, and allows the
same binaries to run on both IMDSv2 enabled and optional instances.

ref: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html

Signed-off-by: Chinmay D. Pai <chinmaydpai@gmail.com>
---
 simples3.go | 70 +++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 62 insertions(+), 8 deletions(-)

diff --git a/simples3.go b/simples3.go
index 3d30d2d..e5e41a6 100644
--- a/simples3.go
+++ b/simples3.go
@@ -23,6 +23,8 @@ import (
 )
 
 const (
+	imdsTokenURL           = "http://169.254.169.254/latest/api/token"
+	imdsTokenHeader        = "X-aws-ec2-metadata-token"
 	securityCredentialsURL = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
 
 	// AMZMetaPrefix to prefix metadata key.
@@ -89,12 +91,13 @@ type UploadInput struct {
 // UploadResponse receives the following XML
 // in case of success, since we set a 201 response from S3.
 // Sample response:
-//     <PostResponse>
-//       <Location>https://s3.amazonaws.com/link-to-the-file</Location>
-//       <Bucket>s3-bucket</Bucket>
-//       <Key>development/8614bd40-691b-4668-9241-3b342c6cf429/image.jpg</Key>
-//       <ETag>"32-bit-tag"</ETag>
-//     </PostResponse>
+//
+//	<PostResponse>
+//	  <Location>https://s3.amazonaws.com/link-to-the-file</Location>
+//	  <Bucket>s3-bucket</Bucket>
+//	  <Key>development/8614bd40-691b-4668-9241-3b342c6cf429/image.jpg</Key>
+//	  <ETag>"32-bit-tag"</ETag>
+//	</PostResponse>
 type UploadResponse struct {
 	Location string `xml:"Location"`
 	Bucket   string `xml:"Bucket"`
@@ -149,12 +152,55 @@ func NewUsingIAM(region string) (*S3, error) {
 		}, securityCredentialsURL, region)
 }
 
+// fetchIMDSToken retrieves an IMDSv2 token from the
+// EC2 instance metadata service. It returns a token
+// only if IMDSv2 is enabled in the EC2 instance metadata
+// configuration, otherwise returns an error.
+func fetchIMDSToken(cl *http.Client) (string, error) {
+	req, err := http.NewRequest(http.MethodPut, imdsTokenURL, nil)
+	if err != nil {
+		return "", err
+	}
+	req.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+
+	resp, err := cl.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("failed to fetch IMDSv2 token: %s", resp.Status)
+	}
+
+	token, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	return string(token), nil
+}
+
 // fetchIAMData fetches the IAM data from the given URL.
 // In case of a normal AWS setup, baseURL would be securityCredentialsURL.
 // You can use this method, to manually fetch IAM data from a custom
 // endpoint and pass it to SetIAMData.
 func fetchIAMData(cl *http.Client, baseURL string) (IAMResponse, error) {
-	resp, err := cl.Get(baseURL)
+	var token string
+	token, err := fetchIMDSToken(cl)
+	// useIMDSv2 if no errors were returned from
+	// the fetchIMDSToken method call.
+	useIMDSv2 := err == nil
+
+	req, err := http.NewRequest(http.MethodGet, baseURL, nil)
+	if err != nil {
+		return IAMResponse{}, err
+	}
+	if useIMDSv2 {
+		req.Header.Set(imdsTokenHeader, token)
+	}
+
+	resp, err := cl.Do(req)
 	if err != nil {
 		return IAMResponse{}, err
 	}
@@ -169,7 +215,15 @@ func fetchIAMData(cl *http.Client, baseURL string) (IAMResponse, error) {
 		return IAMResponse{}, err
 	}
 
-	resp, err = http.Get(baseURL + "/" + string(role))
+	req, err = http.NewRequest(http.MethodGet, baseURL+"/"+string(role), nil)
+	if err != nil {
+		return IAMResponse{}, err
+	}
+	if useIMDSv2 {
+		req.Header.Set(imdsTokenHeader, token)
+	}
+
+	resp, err = cl.Do(req)
 	if err != nil {
 		return IAMResponse{}, err
 	}