Support using IMDSv2 to get availability zone (#253)

Author: yi-tseng-intel

mage/deploy.go

@@ -1538,6 +1538,26 @@ func (d Deploy) orch(targetEnv string) error {
 	return err
 }
 
+// getAWSAvailabilityZone retrieves the AWS availability zone using IMDSv2 with fallback to IMDSv1
+func getAWSAvailabilityZone() (string, error) {
+	// Try IMDSv2 first - requires getting a token
+	tokenCmd := "curl -s -X PUT \"http://169.254.169.254/latest/api/token\" -H \"X-aws-ec2-metadata-token-ttl-seconds: 60\""
+	token, err := script.Exec(tokenCmd).String()
+
+	if err == nil && token != "" {
+		// Use the token to get the AZ with IMDSv2
+		azCmd := fmt.Sprintf("curl -s -H \"X-aws-ec2-metadata-token: %s\" http://169.254.169.254/latest/meta-data/placement/availability-zone", strings.TrimSpace(token))
+		az, err := script.Exec(azCmd).String()
+		if err == nil && az != "" {
+			return strings.TrimSpace(az), nil
+		}
+	}
+
+	// Fall back to IMDSv1 if IMDSv2 fails
+	az, err := script.Exec("curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone").String()
+	return strings.TrimSpace(az), err
+}
+
 func (d Deploy) orchLocal(targetEnv string) error {
 	targetConfig := getTargetConfig(targetEnv)
 
@@ -1578,7 +1598,7 @@ func (d Deploy) orchLocal(targetEnv string) error {
 		cmd = cmd + " " + fmt.Sprintf("--set-string argo.aws.account=%s", strings.Trim(awsAccountID, "\n"))
 
 		// Get AWS region of this VM
-		az, err := script.Exec("curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone").String()
+		az, err := getAWSAvailabilityZone()
 		if err != nil || az == "" {
 			return fmt.Errorf("error retrieving the AWS AZ: %w", err)
 		}