added ecs compatibility

iam.go

@@ -8,24 +8,24 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"time"
 )
 
 const (
-	imdsTokenHeader        = "X-aws-ec2-metadata-token"
-	imdsTokenTtlHeader     = "X-aws-ec2-metadata-token-ttl-seconds"
-	metadataBaseURL        = "http://169.254.169.254/latest"
-	securityCredentialsURI = "/meta-data/iam/security-credentials/"
-	imdsTokenURI           = "/api/token"
-	defaultIMDSTokenTTL    = "60"
+	imdsTokenHeader                = "X-aws-ec2-metadata-token"
+	imdsTokenTtlHeader             = "X-aws-ec2-metadata-token-ttl-seconds"
+	metadataBaseURL                = "http://169.254.169.254/latest"
+	securityCredentialsURI         = "/meta-data/iam/security-credentials/"
+	imdsTokenURI                   = "/api/token"
+	defaultIMDSTokenTTL            = "60"
+	ecsContainerCredentialsEnv     = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
+	ecsContainerCredentialsBaseUrl = "http://169.254.170.2"
 )
 
 // IAMResponse is used by NewUsingIAM to auto
 // detect the credentials.
 type IAMResponse struct {
-	Code            string    `json:"Code"`
-	LastUpdated     string    `json:"LastUpdated"`
-	Type            string    `json:"Type"`
 	AccessKeyID     string    `json:"AccessKeyId"`
 	SecretAccessKey string    `json:"SecretAccessKey"`
 	Token           string    `json:"Token"`
@@ -77,11 +77,51 @@ func fetchIMDSToken(cl *http.Client, baseURL string) (string, bool, error) {
 	return string(token), true, nil
 }
 
+// fetchIAMDataECS fetches the IAM credentials from the ECS default endpoint.
+func fetchIAMDataForEcs(cl *http.Client) (IAMResponse, error) {
+	env, isSet := os.LookupEnv(ecsContainerCredentialsEnv)
+
+	if !isSet {
+		return IAMResponse{}, fmt.Errorf("error getting ecs environment variable, most likely not using ecs.")
+	}
+
+	url := ecsContainerCredentialsBaseUrl + env
+
+	req, err := http.NewRequest(http.MethodGet, url, nil)
+	if err != nil {
+		return IAMResponse{}, fmt.Errorf("error creating IAM ECS request: %w", err)
+	}
+
+	resp, err := cl.Do(req)
+	if err != nil {
+		return IAMResponse{}, fmt.Errorf("error fetching IAM ECS request: %w", err)
+	}
+
+	var jResp IAMResponse
+	jsonString, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return IAMResponse{}, fmt.Errorf("error reading role data: %w", err)
+	}
+
+	if err := json.Unmarshal(jsonString, &jResp); err != nil {
+		return IAMResponse{}, fmt.Errorf("error unmarshalling role data: %w (%s)", err, jsonString)
+	}
+
+	return jResp, nil
+}
+
 // fetchIAMData fetches the IAM data from the given URL.
 // In case of a normal AWS setup, baseURL would be metadataBaseURL.
 // 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) {
+	response, err := fetchIAMDataForEcs(cl)
+
+	// If already have ECS response, skip the rest of the function and use it instead.
+	if err == nil {
+		return response, nil
+	}
+
 	token, useIMDSv2, err := fetchIMDSToken(cl, baseURL)
 	if err != nil {
 		return IAMResponse{}, fmt.Errorf("error fetching IMDSv2 token: %w", err)