feat: add support for IMDSv2

simples3.go

@@ -23,6 +23,9 @@ import (
 )
 
 const (
+	imdsTokenURL           = "http://169.254.169.254/latest/api/token"
+	imdsTokenHeader        = "X-aws-ec2-metadata-token"
+	imdsTokenTtlHeader     = "X-aws-ec2-metadata-token-ttl-seconds"
 	securityCredentialsURL = "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
 
 	// AMZMetaPrefix to prefix metadata key.
@@ -89,12 +92,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 +153,56 @@ func NewUsingIAM(region string) (*S3, error) {
 		}, securityCredentialsURL, region)
 }
 
+// fetchIMDSToken retrieves an IMDSv2 token from the
+// EC2 instance metadata service. It returns a token and boolean,
+// only if IMDSv2 is enabled in the EC2 instance metadata
+// configuration, otherwise returns an error.
+func fetchIMDSToken(cl *http.Client) (string, bool, error) {
+	req, err := http.NewRequest(http.MethodPut, imdsTokenURL, nil)
+	if err != nil {
+		return "", false, err
+	}
+	req.Header.Set(imdsTokenTtlHeader, "21600")
+
+	resp, err := cl.Do(req)
+	if err != nil {
+		return "", false, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", false, fmt.Errorf("failed to request IMDSv2 token: %s", resp.Status)
+	}
+
+	token, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", false, err
+	}
+
+	return string(token), true, 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
+	var useIMDSv2 bool
+	token, useIMDSv2, err := fetchIMDSToken(cl)
+	if err != nil {
+		return IAMResponse{}, fmt.Errorf("Error fetching IMDSv2 token: %w", err)
+	}
+
+	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 +217,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
 	}