Skip to content

Commit dde608f

Browse files
committed
add support for AWS IMDSv2
Signed-off-by: Dmitry Shmulevich <[email protected]>
1 parent 524efe2 commit dde608f

File tree

1 file changed

+49
-8
lines changed

1 file changed

+49
-8
lines changed

pkg/providers/aws/imds.go

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import (
2121
"context"
2222
"encoding/json"
2323
"fmt"
24-
"io"
2524
"net/http"
2625
"strings"
2726
"time"
@@ -32,7 +31,9 @@ import (
3231
)
3332

3433
const (
35-
IMDSURL = "http://169.254.169.254/latest/meta-data/"
34+
IMDS = "http://169.254.169.254"
35+
IMDS_TOKEN_URL = IMDS + "/latest/api/token"
36+
IMDS_URL = IMDS + "/latest/meta-data"
3637

3738
tokenTimeDelay = 15 * time.Second
3839
)
@@ -45,20 +46,59 @@ type Creds struct {
4546
Expiration string `json:"Expiration"`
4647
}
4748

49+
func getToken() (string, error) {
50+
req, err := http.NewRequest("PUT", IMDS_TOKEN_URL, nil)
51+
if err != nil {
52+
return "", fmt.Errorf("failed to create HTTP request: %v", err)
53+
}
54+
55+
req.Header.Add("X-aws-ec2-metadata-token-ttl-seconds", "21600")
56+
57+
_, data, err := utils.HttpRequest(req)
58+
if err != nil {
59+
return "", fmt.Errorf("failed to send HTTP request: %v", err)
60+
}
61+
62+
return string(data), nil
63+
}
64+
65+
func addToken(req *http.Request) error {
66+
token, err := getToken()
67+
if err != nil {
68+
return err
69+
}
70+
71+
if len(token) != 0 {
72+
req.Header.Add("X-aws-ec2-metadata-token", token)
73+
}
74+
75+
return nil
76+
}
77+
4878
func getMetadata(path string) ([]byte, error) {
49-
url := IMDSURL + path
79+
url := fmt.Sprintf("%s/%s", IMDS_URL, path)
80+
klog.V(4).Infof("Requesting URL %s", url)
5081

51-
resp, err := http.Get(url)
82+
req, err := http.NewRequest("GET", url, nil)
83+
if err != nil {
84+
return nil, fmt.Errorf("failed to create HTTP request: %v", err)
85+
}
86+
87+
err = addToken(req)
5288
if err != nil {
5389
return nil, err
5490
}
55-
defer func() { _ = resp.Body.Close() }()
91+
92+
resp, data, err := utils.HttpRequest(req)
93+
if err != nil {
94+
return nil, fmt.Errorf("failed to send HTTP request: %v", err)
95+
}
5696

5797
if resp.StatusCode != http.StatusOK {
5898
return nil, fmt.Errorf("HTTP status: %s", resp.Status)
5999
}
60100

61-
return io.ReadAll(resp.Body)
101+
return data, nil
62102
}
63103

64104
func GetRegion() (string, error) {
@@ -85,7 +125,7 @@ func GetCredentials() (*Creds, error) {
85125

86126
// ensure the credentials remain valid for at least the next tokenTimeDelay
87127
for {
88-
klog.V(4).Infof("Getting credentials from path %s", path)
128+
klog.V(4).Infof("Getting credentials from %s", path)
89129
data, err = getMetadata(path)
90130
if err != nil {
91131
return nil, err
@@ -114,7 +154,8 @@ func GetCredentials() (*Creds, error) {
114154
}
115155

116156
func Instance2NodeMap(ctx context.Context, nodes []string) (map[string]string, error) {
117-
args := []string{"-w", strings.Join(nodes, ","), fmt.Sprintf("echo $(curl -s %s/instance-id)", IMDSURL)}
157+
args := []string{"-w", strings.Join(nodes, ","),
158+
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/instance-id)", IMDS_TOKEN_URL, IMDS_URL)}
118159

119160
stdout, err := utils.Exec(ctx, "pdsh", args, nil)
120161
if err != nil {

0 commit comments

Comments
 (0)