Cache IMDS tokens per-user (#1386)

awslabs · Aug 15, 2023 · 95f1141 · 95f1141
1 parent c2cf65d
commit 95f1141
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/files/bin/imds b/files/bin/imds
@@ -49,8 +49,8 @@ function imdscurl() {
 }
 
 function get-token() {
-  local TOKEN_DIR=/tmp/imds-tokens
-  mkdir -p -m a+wrx $TOKEN_DIR
+  local TOKEN_DIR="/tmp/imds-tokens/$(whoami)"
+  mkdir -p -m 0600 $TOKEN_DIR
 
   # cleanup expired tokens
   local DELETED_TOKENS=0
@@ -71,8 +71,7 @@ function get-token() {
       -X PUT \
       "http://$IMDS_ENDPOINT/latest/api/token")
     echo "$TOKEN" > "$TOKEN_DIR/$TOKEN_FILE"
-    # make sure any user can utilize (and clean up) these tokens
-    chmod a+rwx $TOKEN_DIR/$TOKEN_FILE
+    chmod 0600 "$TOKEN_DIR/$TOKEN_FILE"
     log "🔑 Retrieved a fresh IMDS token that will expire in $IMDS_TOKEN_TTL_SECONDS seconds."
   else
     log "ℹ️ Using cached IMDS token that expires in $(($TOKEN_FILE - $CURRENT_TIME)) seconds."

diff --git a/test/cases/imds-token-refresh.sh b/test/cases/imds-token-refresh.sh
@@ -6,7 +6,7 @@ set -o pipefail
 
 echo "--> Should refresh IMDS token on configured interval"
 exit_code=0
-TOKEN_DIR=/tmp/imds-tokens
+TOKEN_DIR=/tmp/imds-tokens/$(whoami)
 TTL=5
 export IMDS_TOKEN_TTL_SECONDS=$TTL
 export IMDS_DEBUG=true