updated to IMDSv2 (#5783)

pytorch · Oct 30, 2024 · 29dc85d · 29dc85d
1 parent a728aa1
commit 29dc85d
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 3 deletions.
diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml
@@ -13,7 +13,7 @@ runs:
           # Pulled from instance metadata endpoint for EC2
           # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
           category=$1
-          curl -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
         }
         echo "ami-id: $(get_ec2_metadata ami-id)"
         echo "instance-id: $(get_ec2_metadata instance-id)"

diff --git a/.github/actions/setup-windows/action.yml b/.github/actions/setup-windows/action.yml
@@ -18,7 +18,7 @@ runs:
           # Pulled from instance metadata endpoint for EC2
           # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
           category=$1
-          curl -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
         }
         echo "ami-id: $(get_ec2_metadata ami-id)"
         echo "instance-id: $(get_ec2_metadata instance-id)"

diff --git a/setup-ssh/src/ec2-utils.ts b/setup-ssh/src/ec2-utils.ts
@@ -6,9 +6,28 @@ export async function getEC2Metadata(category: string): Promise<string> {
     allowRetries: true,
     maxRetries
   })
+  // convert these two curls:
+  // curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+    const tokenResponse = await http.put(
+    `http://169.254.169.254/latest/api/token`, undefined, {
+      headers: {
+        'X-aws-ec2-metadata-token-ttl-seconds': '30'
+      }
+    }
+  )
+
+  if (tokenResponse.message.statusCode !== 200) {
+    return ''
+  }
+
   const resp = await http.get(
-    `http://169.254.169.254/latest/meta-data/${category}`
+    `http://169.254.169.254/latest/meta-data/${category}`, {
+      headers: {
+        'X-aws-ec2-metadata-token': tokenResponse.result
+      }
+    }
   )
+
   if (resp.message.statusCode !== 200) {
     return ''
   }