@@ -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
3433const (
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+
4878func 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
64104func 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
116156func 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