Add MSAL client metadata headers (x-client-SKU, x-client-VER, x-ms-client-request-id) to IMDS managed identity requests

Agent-Logs-Url: https://github.com/AzureAD/microsoft-authentication-library-for-java/sessions/a002e01d-801c-4ba9-8cf4-0d0ad3b9a0c9

Co-authored-by: gladjohn <90415114+gladjohn@users.noreply.github.com>

Author: Copilot

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

@@ -77,6 +77,9 @@ public void createManagedIdentityRequest(String resource) {
 
         managedIdentityRequest.headers = new HashMap<>();
         managedIdentityRequest.headers.put("Metadata", "true");
+        managedIdentityRequest.headers.put(HttpHeaders.PRODUCT_HEADER_NAME, HttpHeaders.PRODUCT_HEADER_VALUE);
+        managedIdentityRequest.headers.put(HttpHeaders.PRODUCT_VERSION_HEADER_NAME, HttpHeaders.PRODUCT_VERSION_HEADER_VALUE);
+        managedIdentityRequest.headers.put("x-ms-client-request-id", managedIdentityRequest.requestContext().correlationId());
 
         managedIdentityRequest.queryParameters = new HashMap<>();
         managedIdentityRequest.queryParameters.put("api-version", IMDS_API_VERSION);

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

@@ -14,6 +14,8 @@
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import org.mockito.ArgumentCaptor;
+
 import java.net.SocketException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -37,6 +39,9 @@
 @TestInstance(TestInstance.Lifecycle.PER_METHOD)
 class ManagedIdentityTests {
 
+    private static final String EXPECTED_SKU = "MSAL.Java";
+    private static final String TEST_CORRELATION_ID = "00000000-0000-0000-0000-000000000001";
+
     @BeforeAll
     static void setupRetryPolicies() {
         // Set retry delays to 1ms for faster test execution
@@ -145,6 +150,9 @@ private String configureSourceSpecificParameters(ManagedIdentitySourceType sourc
             default:
                 queryParameters.put("api-version", "2018-02-01");
                 headers.put("Metadata", "true");
+                headers.put("x-client-SKU", EXPECTED_SKU);
+                headers.put("x-client-VER", HttpHeaders.PRODUCT_VERSION_HEADER_VALUE);
+                headers.put("x-ms-client-request-id", TEST_CORRELATION_ID);
                 return ManagedIdentityTestConstants.IMDS_ENDPOINT;
         }
     }
@@ -206,6 +214,7 @@ void initManagedIdentityApplication(ManagedIdentityId idType) {
             miApp = ManagedIdentityApplication
                     .builder(idType)
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // ManagedIdentityApplication uses a static token cache, avoid cross test pollution by clearing it
@@ -396,6 +405,7 @@ void managedIdentityTest_WithCapabilitiesOnly(ManagedIdentitySourceType source,
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
                     .clientCapabilities(singletonList("cp1"))
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             miApp.tokenCache.accessTokens.clear();
@@ -425,6 +435,7 @@ void managedIdentity_ClaimsAndCapabilities(ManagedIdentitySourceType source, Str
                     .builder(ManagedIdentityId.systemAssigned())
                     .clientCapabilities(singletonList("cp1"))
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // First call, get the token from the identity provider.
@@ -554,6 +565,7 @@ void managedIdentityTest_Retry(ManagedIdentitySourceType source, String endpoint
             miApp = ManagedIdentityApplication
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             //Several specific 4xx and 5xx errors, such as 500, should trigger MSAL's retry logic
@@ -602,6 +614,7 @@ void managedIdentityTest_RetriesDisabled(ManagedIdentitySourceType source, Strin
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
                     .disableInternalRetries()
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             //Several specific 4xx and 5xx errors, such as 500, should trigger MSAL's retry logic
@@ -631,6 +644,7 @@ void managedIdentityTest_IMDSRetry() throws Exception {
             miApp = ManagedIdentityApplication
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // IMDS has different retry logic for certain status codes, such as 410
@@ -672,6 +686,7 @@ void managedIdentityTest_RetrySucceedsAfterFailure() throws Exception {
             miApp = ManagedIdentityApplication
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // First call returns 500, subsequent calls return 200
@@ -685,7 +700,14 @@ void managedIdentityTest_RetrySucceedsAfterFailure() throws Exception {
             assertNotNull(result.accessToken());
 
             // Verify that the client was called exactly twice (first attempt + one retry)
-            verify(httpClientMock, times(2)).send(any());
+            ArgumentCaptor<HttpRequest> captor = ArgumentCaptor.forClass(HttpRequest.class);
+            verify(httpClientMock, times(2)).send(captor.capture());
+
+            // Verify IMDS client metadata headers on the captured request
+            HttpRequest capturedRequest = captor.getValue();
+            assertEquals(EXPECTED_SKU, capturedRequest.headers().get("x-client-SKU"));
+            assertNotNull(capturedRequest.headers().get("x-client-VER"));
+            assertDoesNotThrow(() -> UUID.fromString(capturedRequest.headers().get("x-ms-client-request-id")));
         }
 
         @Test
@@ -698,6 +720,7 @@ void managedIdentityTest_NonRetryableError() throws Exception {
             miApp = ManagedIdentityApplication
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // Use a status code that doesn't trigger retries (499)
@@ -727,6 +750,7 @@ void managedIdentityTest_RetriesARePerRequest() throws Exception {
             miApp = ManagedIdentityApplication
                     .builder(ManagedIdentityId.systemAssigned())
                     .httpClient(httpClientMock)
+                    .correlationId(TEST_CORRELATION_ID)
                     .build();
 
             // First call returns 500, subsequent calls return 200