Merge pull request #954 from AzureAD/avdunn/imds-fix

Correct IMDS resource ID query parameter

Author: Avery-Dunn

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/Constants.java

@@ -16,6 +16,7 @@ final class Constants {
 
     public static final String MANAGED_IDENTITY_CLIENT_ID = "client_id";
     public static final String MANAGED_IDENTITY_RESOURCE_ID = "mi_res_id";
+    public static final String MANAGED_IDENTITY_RESOURCE_ID_IMDS = "msi_res_id";
     public static final String MANAGED_IDENTITY_OBJECT_ID = "object_id";
     public static final String MANAGED_IDENTITY_DEFAULT_TENTANT = "managed_identity";
 

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityRequest.java

@@ -68,7 +68,13 @@ void addUserAssignedIdToQuery(ManagedIdentityIdType idType, String userAssignedI
                 break;
             case RESOURCE_ID:
                 LOG.info("[Managed Identity] Adding user assigned resource id to the request.");
-                queryParameters.put(Constants.MANAGED_IDENTITY_RESOURCE_ID, Collections.singletonList(userAssignedId));
+                if (ManagedIdentityClient.getManagedIdentitySource() == ManagedIdentitySourceType.IMDS) {
+                    // IMDS seems to accept both mi_res_id and msi_res_id but their API only documents msi_res_id,
+                    // and using mi_res_id leads to issues in some scenarios that use the IMDS code path.
+                    queryParameters.put(Constants.MANAGED_IDENTITY_RESOURCE_ID_IMDS, Collections.singletonList(userAssignedId));
+                } else {
+                    queryParameters.put(Constants.MANAGED_IDENTITY_RESOURCE_ID, Collections.singletonList(userAssignedId));
+                }
                 break;
             case OBJECT_ID:
                 LOG.info("[Managed Identity] Adding user assigned object id to the request.");

msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java

@@ -138,7 +138,11 @@ private HttpRequest expectedRequest(ManagedIdentitySourceType source, String res
                 queryParameters.put("client_id", Collections.singletonList(id.getUserAssignedId()));
                 break;
             case RESOURCE_ID:
-                queryParameters.put("mi_res_id", Collections.singletonList(id.getUserAssignedId()));
+                if (ManagedIdentityClient.getManagedIdentitySource() == ManagedIdentitySourceType.IMDS) {
+                    queryParameters.put(Constants.MANAGED_IDENTITY_RESOURCE_ID_IMDS, Collections.singletonList(id.getUserAssignedId()));
+                } else {
+                    queryParameters.put(Constants.MANAGED_IDENTITY_RESOURCE_ID, Collections.singletonList(id.getUserAssignedId()));
+                }
                 break;
             case OBJECT_ID:
                 queryParameters.put("object_id", singletonList(id.getUserAssignedId()));