Skip to content

Fix #6024: probe IMDSv2 first and remove preview latch #6041

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Robbie-Microsoft merged 1 commit into from
May 29, 2026
+60 −69
Merged

Fix #6024: probe IMDSv2 first and remove preview latch #6041

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions docs/mtlspop_managed_identity.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. |

---
Expand All @@ -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 |

---

Expand Down
45 changes: 16 additions & 29 deletions src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
Expand Down Expand Up @@ -105,28 +101,17 @@ private Task<AbstractManagedIdentity> 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)
Expand Down Expand Up @@ -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<ManagedIdentitySourceResult> GetManagedIdentitySourceAsync(
RequestContext requestContext,
CancellationToken cancellationToken)
Expand All @@ -180,16 +165,9 @@ internal async Task<ManagedIdentitySourceResult> 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)
{
Expand All @@ -198,6 +176,15 @@ internal async Task<ManagedIdentitySourceResult> 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)
{
Expand Down
71 changes: 38 additions & 33 deletions tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,8 @@ private async Task<IManagedIdentityApplication> 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)
Expand All @@ -140,8 +141,7 @@ private async Task<IManagedIdentityApplication> 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));
}

Expand Down Expand Up @@ -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)
{
Expand All @@ -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,
Expand All @@ -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<MsalClientException>(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
Expand All @@ -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);
Expand All @@ -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
Expand All @@ -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);
}
}

Expand All @@ -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();
Expand Down Expand Up @@ -617,16 +622,16 @@ 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++)
{
// `retry: true` indicates a retriable status code will be returned
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);
Expand Down
Loading
Loading