From 0653ef27f4f66b6d3c25d903b35ecd9941becfda Mon Sep 17 00:00:00 2001 From: Robbie Ginsburg Date: Thu, 28 May 2026 17:23:18 -0400 Subject: [PATCH] Fix #6024: probe IMDSv2 first and remove preview latch ManagedIdentityClient.GetManagedIdentitySourceAsync now probes the IMDSv2 CSR metadata endpoint before falling back to the IMDSv1 token endpoint. Previously v1 was probed first, and because the IMDS contract returns HTTP 400 (success signal) on every host for the v1 token path, v2-capable hosts were never detected and were always cached as ImdsV1. This broke the Azure SDK mTLS PoP discovery path end-to-end. Also removes the dormant s_imdsV1UsedForPreview latch and its associated CannotSwitchBetweenImdsVersionsForPreview throw block. With v2-first discovery, that latch would have started firing on any process that performs a non-PoP MI call before a PoP MI call (the exact Azure SDK call pattern). Per-request IMDSv1 fallback for non-PoP requests is preserved without any state being latched, so subsequent PoP requests still route to the cached ImdsV2 source. The MsalError.CannotSwitchBetweenImdsVersionsForPreview public constant is kept for backward compatibility (it ships in PublicAPI.Shipped.txt across all TFMs); only the throw site is removed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/mtlspop_managed_identity.md | 3 +- .../ManagedIdentity/ManagedIdentityClient.cs | 45 +++++------- .../ManagedIdentityTests/ImdsV2Tests.cs | 71 ++++++++++--------- .../ManagedIdentityTests.cs | 10 +-- 4 files changed, 60 insertions(+), 69 deletions(-) diff --git a/docs/mtlspop_managed_identity.md b/docs/mtlspop_managed_identity.md index 5d796c325a..9ec86f2bea 100644 --- a/docs/mtlspop_managed_identity.md +++ b/docs/mtlspop_managed_identity.md @@ -283,7 +283,7 @@ When attestation is configured: | **Framework** | .NET Core / .NET 5+ only. Not supported on .NET Framework 4.6.2. | | **IMDS version** | Requires IMDSv2. If the VM only has IMDSv1, throws `MtlsPopTokenNotSupportedinImdsV1`. | | **Key type** | KeyGuard RSA key required. Throws error code `mtls_pop_requires_keyguard` if not available (hardcoded string — not yet a constant in `MsalError`). | -| **Mixed usage** | Once IMDSv1 is used in a process while IMDSv2 is cached, switching to IMDSv2 PoP in the same process is blocked (preview behavior). Throws `CannotSwitchBetweenImdsVersionsForPreview`. | +| **Mixed usage** | Non-mTLS PoP calls in a process where IMDSv2 is cached transparently use IMDSv1 for that request; subsequent PoP calls still route to IMDSv2. | | **Experimental** | This feature is in preview. Not all regions and environments may be supported. | --- @@ -297,7 +297,6 @@ When attestation is configured: | `mtls_pop_requires_keyguard` | The managed identity key is not a KeyGuard key (hardcoded string — not yet a constant in `MsalError`) | Use a VM/VMSS with KeyGuard support enabled | | `MtlsCertificateNotProvided` | (CCA path) No certificate was found for binding | Pass a certificate via `.WithCertificate(cert, sendX5C: true)` | | `MtlsPopWithoutRegion` | (CCA path) Azure region not set | Add `.WithAzureRegion("region")` to the app builder | -| `CannotSwitchBetweenImdsVersionsForPreview` | Mixed IMDSv1/v2 usage in same process | Use a single IMDS version per process; restart the app | --- diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs b/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs index 6888e31bcd..005c70c5cd 100644 --- a/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs +++ b/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs @@ -22,9 +22,6 @@ internal class ManagedIdentityClient private const string WindowsHimdsFilePath = "%Programfiles%\\AzureConnectedMachineAgent\\himds.exe"; private const string LinuxHimdsFilePath = "/opt/azcmagent/bin/himds"; - // Preview guard: once we fall back to IMDSv1 while IMDSv2 is cached, - // disallow switching to IMDSv2 PoP in the same process (preview behavior). - internal static bool s_imdsV1UsedForPreview = false; // Non-null only after the explicit discovery API (GetManagedIdentitySourceAsync) runs. // Allows caching "NoneFound" (Source=None) without confusing it with "not discovered yet". private static ManagedIdentitySourceResult s_cachedSourceResult = null; @@ -36,7 +33,6 @@ internal class ManagedIdentityClient internal static void ResetSourceForTest() { s_cachedSourceResult = null; - s_imdsV1UsedForPreview = false; // Clear cert caches so each test starts fresh ImdsV2ManagedIdentitySource.ResetCertCacheForTest(); @@ -105,28 +101,17 @@ private Task GetOrSelectManagedIdentitySourceAsync( throw CreateManagedIdentityUnavailableException(s_cachedSourceResult); } - // Preview fallback: if ImdsV2 is cached but mTLS PoP not requested, fall back per-request to ImdsV1 + // Per-request fallback: if ImdsV2 is cached but mTLS PoP not requested, use ImdsV1 for this request only. + // We do NOT latch this state; future PoP requests can still leverage the cached ImdsV2 discovery. if (source == ManagedIdentitySource.ImdsV2 && !isMtlsPopRequested) { - requestContext.Logger.Info("[Managed Identity] ImdsV2 detected, but mTLS PoP was not requested. Falling back to ImdsV1 for this request only. Please use the \"WithMtlsProofOfPossession\" API to request a token via ImdsV2."); - - // Mark that we used IMDSv1 in this process while IMDSv2 is cached (preview behavior). - s_imdsV1UsedForPreview = true; + requestContext.Logger.Info("[Managed Identity] ImdsV2 detected, but mTLS PoP was not requested. Using ImdsV1 for this request only. Please use the \"WithMtlsProofOfPossession\" API to request a token via ImdsV2."); // Do NOT modify s_cachedSourceResult; keep cached ImdsV2 so future PoP // requests can leverage it. source = ManagedIdentitySource.Imds; } - // Preview behavior: once we've used IMDSv1 fallback while IMDSv2 is cached, - // we disallow switching back to IMDSv2 PoP in this process. - if (source == ManagedIdentitySource.ImdsV2 && isMtlsPopRequested && s_imdsV1UsedForPreview) - { - throw new MsalClientException( - MsalError.CannotSwitchBetweenImdsVersionsForPreview, - MsalErrorMessage.CannotSwitchBetweenImdsVersionsForPreview); - } - // If the source is determined to be ImdsV1 and mTLS PoP was requested, // throw an exception since ImdsV1 does not support mTLS PoP if (source == ManagedIdentitySource.Imds && isMtlsPopRequested) @@ -158,7 +143,7 @@ private static ManagedIdentitySourceResult CacheDiscoveryResult(ManagedIdentityS // Detect managed identity source by probing IMDS endpoints. // This method is called only by the explicit discovery path (GetManagedIdentitySourceAsync in ManagedIdentityApplication.cs). - // It probes IMDS v1 first, then v2 if v1 fails, and caches the result. + // It probes IMDS v2 first, then v1 if v2 fails, and caches the result. internal async Task GetManagedIdentitySourceAsync( RequestContext requestContext, CancellationToken cancellationToken) @@ -180,16 +165,9 @@ internal async Task GetManagedIdentitySourceAsync( string imdsV1FailureReason = null; string imdsV2FailureReason = null; - // Probe IMDS v1 first - var (imdsV1Success, imdsV1Failure) = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1, cancellationToken).ConfigureAwait(false); - if (imdsV1Success) - { - requestContext.Logger.Info("[Managed Identity] ImdsV1 detected."); - return CacheDiscoveryResult(new ManagedIdentitySourceResult(ManagedIdentitySource.Imds)); - } - imdsV1FailureReason = imdsV1Failure; - - // If v1 fails, probe IMDS v2 + // Probe IMDS v2 first. The v2 path (CSR metadata endpoint) only exists on hosts that + // actually support IMDSv2; on v1-only hosts it returns 404. Probing v2 first avoids + // the v1 success-on-400 contract masking a v2-capable host (see issue #6024). var (imdsV2Success, imdsV2Failure) = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2, cancellationToken).ConfigureAwait(false); if (imdsV2Success) { @@ -198,6 +176,15 @@ internal async Task GetManagedIdentitySourceAsync( } imdsV2FailureReason = imdsV2Failure; + // If v2 fails, fall back to probing IMDS v1. + var (imdsV1Success, imdsV1Failure) = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1, cancellationToken).ConfigureAwait(false); + if (imdsV1Success) + { + requestContext.Logger.Info("[Managed Identity] ImdsV1 detected."); + return CacheDiscoveryResult(new ManagedIdentitySourceResult(ManagedIdentitySource.Imds)); + } + imdsV1FailureReason = imdsV1Failure; + requestContext.Logger.Info($"[Managed Identity] {MsalErrorMessage.ManagedIdentityAllSourcesUnavailable}"); return CacheDiscoveryResult(new ManagedIdentitySourceResult(ManagedIdentitySource.None) { diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs index 40f38b5d9f..7f48193be8 100644 --- a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs +++ b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs @@ -126,7 +126,8 @@ private async Task CreateManagedIdentityAsync( if (imdsVersion == ImdsVersion.V1) { - // New discovery order: V1 probed first (succeeds) → ImdsV1 cached + // Discovery order: V2 probed first (fails), then V1 (succeeds) → ImdsV1 cached + httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2, userAssignedIdentityId, userAssignedId)); httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1, userAssignedIdentityId, userAssignedId)); if (addSourceCheck) @@ -140,8 +141,7 @@ private async Task CreateManagedIdentityAsync( if (addProbeMock) { - // New discovery order: V1 probed first (fails), then V2 (succeeds) - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1, userAssignedIdentityId, userAssignedId)); + // Discovery order: V2 probed first (succeeds) → ImdsV2 cached httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V2, userAssignedIdentityId, userAssignedId)); } @@ -433,12 +433,15 @@ public async Task MtlsPopWithoutPriorDiscovery_UsesImdsV2AndSucceeds( } } + // Verifies that after a non-mTLS request uses IMDSv1 (per-request fallback), a subsequent + // mTLS PoP request still succeeds against the cached IMDSv2 source. Previously this combination + // threw CannotSwitchBetweenImdsVersionsForPreview; the preview latch has been removed (issue #6024). [TestMethod] [DataRow(UserAssignedIdentityId.None, null)] // SAMI [DataRow(UserAssignedIdentityId.ClientId, TestConstants.ClientId)] // UAMI [DataRow(UserAssignedIdentityId.ResourceId, TestConstants.MiResourceId)] // UAMI [DataRow(UserAssignedIdentityId.ObjectId, TestConstants.ObjectId)] // UAMI - public async Task ApplicationsCannotSwitchBetweenImdsVersionsForPreview( + public async Task ApplicationsCanSwitchBetweenImdsVersions( UserAssignedIdentityId userAssignedIdentityId, string userAssignedId) { @@ -449,7 +452,7 @@ public async Task ApplicationsCannotSwitchBetweenImdsVersionsForPreview( var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false); - // IMDSv1 request mock + // Arrange: non-mTLS request will route to IMDSv1 per-request fallback httpManager.AddManagedIdentityMockHandler( ManagedIdentityTests.ImdsEndpoint, ManagedIdentityTests.Resource, @@ -458,28 +461,32 @@ public async Task ApplicationsCannotSwitchBetweenImdsVersionsForPreview( userAssignedId: userAssignedId, userAssignedIdentityId: userAssignedIdentityId); - var result = await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) - //.WithMtlsProofOfPossession() - excluding this will cause fallback to ImdsV1 - //.WithAttestationSupport() + // Act: first call without WithMtlsProofOfPossession → IMDSv1 fallback + var bearerResult = await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) .ExecuteAsync().ConfigureAwait(false); - Assert.IsNotNull(result); - Assert.AreEqual(Bearer, result.TokenType); - Assert.AreEqual(TokenSource.IdentityProvider, result.AuthenticationResultMetadata.TokenSource); + // Assert: Bearer token returned, but cached source remains ImdsV2 (no latching) + Assert.IsNotNull(bearerResult); + Assert.AreEqual(Bearer, bearerResult.TokenType); + Assert.AreEqual(TokenSource.IdentityProvider, bearerResult.AuthenticationResultMetadata.TokenSource); - // even though the app fell back to ImdsV1, the source should still be ImdsV2 var miSourceResult = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(ManagedIdentityTests.ImdsProbesCancellationToken).ConfigureAwait(false); Assert.AreEqual(ManagedIdentitySource.ImdsV2, miSourceResult.Source); - // none of the mocks from AddMocksToGetEntraToken are needed since checking the cache occurs before the network requests - var ex = await Assert.ThrowsAsync(async () => - await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) - .WithMtlsProofOfPossession() // this will cause an error to be thrown since the app already fell back to ImdsV1 + // Arrange: mocks for the IMDSv2 mTLS PoP token request + AddMocksToGetEntraToken(httpManager, userAssignedIdentityId, userAssignedId); + + // Act: second call WITH mTLS PoP → must succeed (no CannotSwitchBetweenImdsVersionsForPreview throw) + var popResult = await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) + .WithMtlsProofOfPossession() .WithAttestationSupport() - .ExecuteAsync().ConfigureAwait(false) - ).ConfigureAwait(false); + .ExecuteAsync().ConfigureAwait(false); - Assert.AreEqual(MsalError.CannotSwitchBetweenImdsVersionsForPreview, ex.ErrorCode); + // Assert: mTLS PoP token returned successfully + Assert.IsNotNull(popResult); + Assert.AreEqual(MTLSPoP, popResult.TokenType); + Assert.IsNotNull(popResult.BindingCertificate); + Assert.AreEqual(TokenSource.IdentityProvider, popResult.AuthenticationResultMetadata.TokenSource); } } #endregion Failure Tests @@ -493,8 +500,7 @@ public async Task ProbeImdsEndpointAsyncSucceeds() { SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint); - // New discovery order: V1 probed first (fails), then V2 (succeeds) - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1)); + // Discovery order: V2 probed first (succeeds) httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V2)); await CreateManagedIdentityAsync(httpManager, addProbeMock: false).ConfigureAwait(false); @@ -509,8 +515,7 @@ public async Task ProbeImdsEndpointAsyncSucceedsAfterRetry() { SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint); - // New discovery order: V1 probed first (fails), then V2 (first attempt fails with retry, second succeeds) - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1)); + // Discovery order: V2 probed first (first attempt fails with retry, second succeeds) // `retry: true` indicates a retriable status code will be returned httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2, retry: true)); // Second V2 attempt succeeds @@ -528,15 +533,15 @@ public async Task ProbeImdsEndpointAsyncFails404WhichIsNonRetriableAndRetryPolic { SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint); - // New discovery order: V1 probed first (fails with non-retriable 404), then V2 (succeeds) - // `retry: false` indicates a non-retriable status code (404) will be returned for V1 - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1, retry: false)); - httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V2)); + // Discovery order: V2 probed first (fails with non-retriable 404), then V1 (succeeds) + // `retry: false` indicates a non-retriable status code (404) will be returned for V2 + httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2, retry: false)); + httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1)); var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, addProbeMock: false, addSourceCheck: false).ConfigureAwait(false); var miSourceResult = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(ManagedIdentityTests.ImdsProbesCancellationToken).ConfigureAwait(false); - Assert.AreEqual(ManagedIdentitySource.ImdsV2, miSourceResult.Source); + Assert.AreEqual(ManagedIdentitySource.Imds, miSourceResult.Source); } } @@ -554,8 +559,8 @@ public async Task ImdsProbeEndpointAsync_TimeOutThrowsOperationCanceledException var managedIdentityApp = miBuilder.Build(); - // New discovery order: V1 is probed first - httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1)); + // Discovery order: V2 is probed first + httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V2)); var cts = new CancellationTokenSource(); cts.Cancel(); @@ -617,9 +622,7 @@ public async Task BothImdsProbesFailMaxRetries_ReturnsNoneFound() { SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint); - // New discovery order: V1 probed first (fails), then V2 fails with max retries → NoneFound - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1, retry: false)); - + // Discovery order: V2 probed first (fails with max retries), then V1 (fails) → NoneFound const int Num500Errors = 1 + TestImdsProbeRetryPolicy.ExponentialStrategyNumRetries; for (int i = 0; i < Num500Errors; i++) { @@ -627,6 +630,8 @@ public async Task BothImdsProbesFailMaxRetries_ReturnsNoneFound() httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2, retry: true)); } + httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1, retry: false)); + var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, addProbeMock: false, addSourceCheck: false).ConfigureAwait(false); var miSourceResult = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(ManagedIdentityTests.ImdsProbesCancellationToken).ConfigureAwait(false); diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs index f56f4f3f17..9d85275562 100644 --- a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs +++ b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs @@ -71,13 +71,13 @@ public async Task GetManagedIdentityTests( if (managedIdentitySource == ManagedIdentitySource.ImdsV2) { - // New discovery order: V1 probed first (fails), then V2 (succeeds) - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1)); + // Discovery order: V2 probed first (succeeds) httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V2)); } else if (managedIdentitySource == ManagedIdentitySource.Imds) { - // New discovery order: V1 probed first (succeeds) + // Discovery order: V2 probed first (fails), then V1 (succeeds) + httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2)); httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1)); } @@ -1128,9 +1128,9 @@ public async Task UnavailableManagedIdentitySource_ThrowsExceptionDuringTokenAcq var mi = miBuilder.Build() as ManagedIdentityApplication; Assert.IsNotNull(mi, "Build() should return a ManagedIdentityApplication instance."); - // Explicit discovery: V1 probe fails, then V2 probe also fails → NoneFound cached - httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1)); + // Explicit discovery: V2 probe fails, then V1 probe also fails → NoneFound cached httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2)); + httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V1)); var sourceResult = await mi.GetManagedIdentitySourceAsync(ImdsProbesCancellationToken).ConfigureAwait(false); Assert.AreEqual(ManagedIdentitySource.None, sourceResult.Source);