From 9c2d6f93c6ad689708b5af978da6518343912a6f Mon Sep 17 00:00:00 2001
From: Maxim Hayes <hayemaxi@amazon.com>
Date: Thu, 1 Feb 2024 15:52:54 -0500
Subject: [PATCH] fix: use IMDSv2 for all metadata service calls.

Problem: We did not attach the token to IMDS calls, resulting in IMDSv1 being used. The SDK does support using tokens but only
for calls to the endpoint `/latest/meta-data/iam/security-credentials/`.
https://github.com/aws/aws-sdk-js/blob/3333f8b49283f5bbff823ab8a8469acedb7fe3d5/lib/metadata_service.js#L123-L235

Solution: Call the "private" sdk method to get the token so code isn't duplicated, and attach that in the header of our calls
just like the sdk does for the above endpoint.

AWS SDK v3 does not support token handling https://github.com/aws/aws-sdk-js-v3/issues/4004
---
 src/shared/clients/ec2MetadataClient.ts | 37 ++++++++++++++++++-------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/shared/clients/ec2MetadataClient.ts b/src/shared/clients/ec2MetadataClient.ts
index 391d9ec3e44..a611088e289 100644
--- a/src/shared/clients/ec2MetadataClient.ts
+++ b/src/shared/clients/ec2MetadataClient.ts
@@ -4,7 +4,7 @@
  */
 
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import { MetadataService } from 'aws-sdk'
+import { AWSError, MetadataService } from 'aws-sdk'
 
 export interface IamInfo {
     Code: string
@@ -33,17 +33,34 @@ export class DefaultEc2MetadataClient {
 
     public invoke<T>(path: string): Promise<T> {
         return new Promise((resolve, reject) => {
-            this.metadata.request(path, (err, response) => {
-                if (err) {
-                    reject(err)
+            // fetchMetadataToken is private for some reason, but has the exact token functionality
+            // that we want out of the metadata service.
+            ;(this.metadata as any).fetchMetadataToken((tokenErr: AWSError, token: string) => {
+                if (tokenErr) {
+                    reject(tokenErr)
                     return
                 }
-                try {
-                    const jsonResponse: T = JSON.parse(response)
-                    resolve(jsonResponse)
-                } catch (e) {
-                    reject(`Ec2MetadataClient: invalid response from "${path}": ${response}\nerror: ${e}`)
-                }
+
+                this.metadata.request(
+                    path,
+                    {
+                        // By attaching the token we force the use of IMDSv2.
+                        // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-v2-how-it-works.html
+                        headers: { 'x-aws-ec2-metadata-token': token },
+                    },
+                    (err, response) => {
+                        if (err) {
+                            reject(err)
+                            return
+                        }
+                        try {
+                            const jsonResponse: T = JSON.parse(response)
+                            resolve(jsonResponse)
+                        } catch (e) {
+                            reject(`Ec2MetadataClient: invalid response from "${path}": ${response}\nerror: ${e}`)
+                        }
+                    }
+                )
             })
         })
     }