From 1fb6d8f3590e89ec658b70989ef013b795df92e6 Mon Sep 17 00:00:00 2001
From: abhilashra <abhilash.ramadugu@idt.net>
Date: Sun, 11 Jan 2026 23:28:28 +0530
Subject: [PATCH] fixed the imdsv2 changes

---
 src/erlcloud_ec2_meta.erl | 79 +++++++++++++++++++++++++++++++++++----
 1 file changed, 72 insertions(+), 7 deletions(-)

diff --git a/src/erlcloud_ec2_meta.erl b/src/erlcloud_ec2_meta.erl
index b40ed4422..9594d60e3 100644
--- a/src/erlcloud_ec2_meta.erl
+++ b/src/erlcloud_ec2_meta.erl
@@ -5,7 +5,13 @@
 
 -export([get_instance_metadata/0, get_instance_metadata/1, get_instance_metadata/2,
         get_instance_user_data/0, get_instance_user_data/1,
-        get_instance_dynamic_data/0, get_instance_dynamic_data/1, get_instance_dynamic_data/2]).
+        get_instance_dynamic_data/0, get_instance_dynamic_data/1, get_instance_dynamic_data/2,
+        get_imds_token/1]).
+
+-define(IMDS_BASE_URL, "http://169.254.169.254").
+-define(IMDS_TOKEN_TTL, 21600).  % 6 hours in seconds
+-define(IMDS_TOKEN_REFRESH_THRESHOLD, 10800).  % 3 hours in seconds
+-define(ERLCLOUD_TOKEN_CACHE_KEY, {erlcloud_ec2_meta, imds_token}).
 
 
 
@@ -30,8 +36,10 @@ get_instance_metadata(Config) ->
 %%
 %%
 get_instance_metadata(ItemPath, Config) ->
-    MetaDataPath = "http://169.254.169.254/latest/meta-data/" ++ ItemPath,
-    erlcloud_aws:http_body(erlcloud_httpc:request(MetaDataPath, get, [], <<>>, erlcloud_aws:get_timeout(Config), Config)).
+    MetaDataPath = ?IMDS_BASE_URL ++ "/latest/meta-data/" ++ ItemPath,
+    Headers = get_imds_headers(Config),
+    Timeout = erlcloud_aws:get_timeout(Config),
+    erlcloud_aws:http_body(erlcloud_httpc:request(MetaDataPath, get, Headers, <<>>, Timeout, Config)).
 
 
 -spec get_instance_user_data() -> {ok, binary()} | {error, erlcloud_aws:httpc_result_error()}.
@@ -48,8 +56,10 @@ get_instance_user_data() ->
 %%
 %%
 get_instance_user_data(Config) ->
-    UserDataPath = "http://169.254.169.254/latest/user-data/",
-    erlcloud_aws:http_body(erlcloud_httpc:request(UserDataPath, get, [], <<>>, erlcloud_aws:get_timeout(Config), Config)).
+    UserDataPath = ?IMDS_BASE_URL ++ "/latest/user-data/",
+    Headers = get_imds_headers(Config),
+    Timeout = erlcloud_aws:get_timeout(Config),
+    erlcloud_aws:http_body(erlcloud_httpc:request(UserDataPath, get, Headers, <<>>, Timeout, Config)).
 
 
 -spec get_instance_dynamic_data() -> {ok, binary()} | {error, erlcloud_aws:httpc_result_error()}.
@@ -66,6 +76,61 @@ get_instance_dynamic_data(Config) ->
 %%%---------------------------------------------------------------------------
 
 get_instance_dynamic_data(ItemPath, Config) ->
-    DynamicDataPath = "http://169.254.169.254/latest/dynamic/" ++ ItemPath,
-    erlcloud_aws:http_body(erlcloud_httpc:request(DynamicDataPath, get, [], <<>>, erlcloud_aws:get_timeout(Config), Config)).
+    DynamicDataPath = ?IMDS_BASE_URL ++ "/latest/dynamic/" ++ ItemPath,
+    Headers = get_imds_headers(Config),
+    Timeout = erlcloud_aws:get_timeout(Config),
+    erlcloud_aws:http_body(erlcloud_httpc:request(DynamicDataPath, get, Headers, <<>>, Timeout, Config)).
+
+%%%---------------------------------------------------------------------------
+%% @doc Get or refresh IMDSv2 token from metadata service
+%% Token is cached in application environment with timestamp
+%% Refreshes if token is older than 3 hours to prevent expiration
+%%%---------------------------------------------------------------------------
+-spec get_imds_token(aws_config()) -> {ok, string()} | {error, any()}.
+get_imds_token(Config) ->
+    case application:get_env(erlcloud, ?ERLCLOUD_TOKEN_CACHE_KEY) of
+        {ok, {Timestamp, Token}} ->
+            CurrentTime = erlang:system_time(second),
+            Age = CurrentTime - Timestamp,
+            if
+                Age < ?IMDS_TOKEN_REFRESH_THRESHOLD ->
+                    {ok, Token};
+                true ->
+                    fetch_new_imds_token(Config)
+            end;
+        undefined ->
+            fetch_new_imds_token(Config)
+    end.
+
+%%%---------------------------------------------------------------------------
+%% @doc Fetch a new IMDSv2 token and cache it
+%%%---------------------------------------------------------------------------
+-spec fetch_new_imds_token(aws_config()) -> {ok, string()} | {error, any()}.
+fetch_new_imds_token(Config) ->
+    URL = ?IMDS_BASE_URL ++ "/latest/api/token",
+    Headers = [{"X-aws-ec2-metadata-token-ttl-seconds", 
+                integer_to_list(?IMDS_TOKEN_TTL)}],
+    Timeout = erlcloud_aws:get_timeout(Config),
+    case erlcloud_httpc:request(URL, put, Headers, <<>>, Timeout, Config) of
+        {ok, {{200, _}, _, Body}} ->
+            Token = binary_to_list(Body),
+            CurrentTime = erlang:system_time(second),
+            application:set_env(erlcloud, ?ERLCLOUD_TOKEN_CACHE_KEY, {CurrentTime, Token}),
+            {ok, Token};
+        _ ->
+            {error, no_token}
+    end.
+
+%%%---------------------------------------------------------------------------
+%% @doc Get headers for IMDS requests with IMDSv2 token if available
+%% Falls back to empty headers (IMDSv1) if token fetch fails
+%%%---------------------------------------------------------------------------
+-spec get_imds_headers(aws_config()) -> [{string(), string()}].
+get_imds_headers(Config) ->
+    case get_imds_token(Config) of
+        {ok, Token} ->
+            [{"X-aws-ec2-metadata-token", Token}];
+        {error, _} ->
+            []  % Fallback to IMDSv1
+    end.