add support for AWS IMDSv2

Signed-off-by: Dmitry Shmulevich <[email protected]>

Author: dmitsh

pkg/providers/aws/imds.go

@@ -21,8 +21,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net/http"
+	"path"
 	"strings"
 	"time"
 
@@ -32,7 +32,9 @@ import (
 )
 
 const (
-	IMDSURL = "http://169.254.169.254/latest/meta-data/"
+	IMDS           = "http://169.254.169.254"
+	IMDS_TOKEN_URL = IMDS + "/latest/api/token"
+	IMDS_URL       = IMDS + "/latest/meta-data"
 
 	tokenTimeDelay = 15 * time.Second
 )
@@ -45,20 +47,59 @@ type Creds struct {
 	Expiration      string `json:"Expiration"`
 }
 
-func getMetadata(path string) ([]byte, error) {
-	url := IMDSURL + path
+func getToken() (string, error) {
+	req, err := http.NewRequest("PUT", IMDS_TOKEN_URL, nil)
+	if err != nil {
+		return "", fmt.Errorf("failed to create HTTP request: %v", err)
+	}
+
+	req.Header.Add("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+
+	_, data, err := utils.HttpRequest(req)
+	if err != nil {
+		return "", fmt.Errorf("failed to send HTTP request: %v", err)
+	}
 
-	resp, err := http.Get(url)
+	return string(data), nil
+}
+
+func addToken(req *http.Request) error {
+	token, err := getToken()
+	if err != nil {
+		return err
+	}
+
+	if len(token) != 0 {
+		req.Header.Add("X-aws-ec2-metadata-token", token)
+	}
+
+	return nil
+}
+
+func getMetadata(route string) ([]byte, error) {
+	url := path.Join(IMDS_URL, route)
+	klog.V(4).Infof("Requesting URL %s", url)
+
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP request: %v", err)
+	}
+
+	err = addToken(req)
 	if err != nil {
 		return nil, err
 	}
-	defer func() { _ = resp.Body.Close() }()
+
+	resp, data, err := utils.HttpRequest(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to send HTTP request: %v", err)
+	}
 
 	if resp.StatusCode != http.StatusOK {
 		return nil, fmt.Errorf("HTTP status: %s", resp.Status)
 	}
 
-	return io.ReadAll(resp.Body)
+	return data, nil
 }
 
 func GetRegion() (string, error) {
@@ -71,22 +112,22 @@ func GetRegion() (string, error) {
 }
 
 func GetCredentials() (*Creds, error) {
-	path := "iam/security-credentials"
-	data, err := getMetadata(path)
+	route := "iam/security-credentials"
+	data, err := getMetadata(route)
 	if err != nil {
 		return nil, err
 	}
 
 	lines := strings.Split(string(data), "\n")
 	for _, line := range lines {
-		path = fmt.Sprintf("%s/%s", path, line)
+		route = path.Join(route, line)
 		break
 	}
 
 	// ensure the credentials remain valid for at least the next tokenTimeDelay
 	for {
-		klog.V(4).Infof("Getting credentials from path %s", path)
-		data, err = getMetadata(path)
+		klog.V(4).Infof("Getting credentials from %s", route)
+		data, err = getMetadata(route)
 		if err != nil {
 			return nil, err
 		}
@@ -114,7 +155,8 @@ func GetCredentials() (*Creds, error) {
 }
 
 func Instance2NodeMap(ctx context.Context, nodes []string) (map[string]string, error) {
-	args := []string{"-w", strings.Join(nodes, ","), fmt.Sprintf("echo $(curl -s %s/instance-id)", IMDSURL)}
+	args := []string{"-w", strings.Join(nodes, ","),
+		fmt.Sprintf("TOKEN=$(curl -s -X PUT -H \"X-aws-ec2-metadata-token-ttl-seconds: 21600\" %s); echo $(curl -s -H \"X-aws-ec2-metadata-token: $TOKEN\" %s)", IMDS_TOKEN_URL, path.Join(IMDS_URL, "instance-id"))}
 
 	stdout, err := utils.Exec(ctx, "pdsh", args, nil)
 	if err != nil {