From 54c94c3335dc6790d98fcc6e7517af6723e3ffb7 Mon Sep 17 00:00:00 2001 From: Robbie Ginsburg Date: Tue, 9 Dec 2025 12:21:19 -0500 Subject: [PATCH] Implemented Service Exception for Imds Probe --- .../ImdsManagedIdentitySource.cs | 5 ++-- .../Microsoft.Identity.Client/MsalError.cs | 5 ++++ .../PublicApi/net462/PublicAPI.Unshipped.txt | 3 +- .../PublicApi/net472/PublicAPI.Unshipped.txt | 3 +- .../net8.0-android/PublicAPI.Unshipped.txt | 3 +- .../net8.0-ios/PublicAPI.Unshipped.txt | 3 +- .../PublicApi/net8.0/PublicAPI.Unshipped.txt | 3 +- .../netstandard2.0/PublicAPI.Unshipped.txt | 3 +- .../ManagedIdentityTests/ImdsTests.cs | 23 +++++++------- .../ManagedIdentityTests/ImdsV2Tests.cs | 30 +++++++++++++++++++ 10 files changed, 61 insertions(+), 20 deletions(-) diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs b/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs index eb322a08a8..ef767f2afd 100644 --- a/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs +++ b/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs @@ -305,8 +305,9 @@ public static async Task ProbeImdsEndpointAsync( } catch (Exception ex) { - requestContext.Logger.Info($"[Managed Identity] {imdsStringHelper} probe endpoint failure. Exception occurred while sending request to probe endpoint: {ex}"); - return false; + throw new MsalServiceException( + MsalError.ImdsServiceError, + $"[Managed Identity] {imdsStringHelper} probe endpoint failure. Exception occurred while sending request to probe endpoint: {ex}"); } // probe omits the "Metadata: true" header and then treats 400 Bad Request as success diff --git a/src/client/Microsoft.Identity.Client/MsalError.cs b/src/client/Microsoft.Identity.Client/MsalError.cs index 688458a3cf..221042e069 100644 --- a/src/client/Microsoft.Identity.Client/MsalError.cs +++ b/src/client/Microsoft.Identity.Client/MsalError.cs @@ -1232,5 +1232,10 @@ public static class MsalError /// All managed identity sources are unavailable. /// public const string ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable"; + + /// + /// Represents the error code returned when an IMDS operation fails. + /// + public const string ImdsServiceError = "imds_service_error"; } } diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt index 3241ccd9cd..692bd0feeb 100644 --- a/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt +++ b/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt @@ -1,2 +1,3 @@ -const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string +const Microsoft.Identity.Client.MsalError.ImdsServiceError = "imds_service_error" -> string +const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs index 204074482d..d9b4b91ab3 100644 --- a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs +++ b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs @@ -3,7 +3,6 @@ using System; using System.Net; -using System.Net.Http; using System.Threading; using System.Threading.Tasks; using Microsoft.Identity.Client; @@ -436,7 +435,7 @@ await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) } [TestMethod] - public async Task ProbeImdsEndpointAsync_TimesOutAfterOneSecond() + public async Task ProbeImdsV1EndpointAsync_TimesOutAfterOneSecond() { using (new EnvVariableContext()) using (var httpManager = new MockHttpManager()) @@ -450,19 +449,19 @@ public async Task ProbeImdsEndpointAsync_TimesOutAfterOneSecond() var managedIdentityApp = miBuilder.Build(); httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2)); - httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1)); + // ImdsV1 mock is not needed, as the request will not be sent due to cancellation token being cancelled - var imdsProbesCancellationToken = new CancellationTokenSource(TimeSpan.FromSeconds(0)).Token; // timeout immediately - - var miSource = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(imdsProbesCancellationToken).ConfigureAwait(false); - Assert.AreEqual(ManagedIdentitySource.None, miSource); // Probe timed out, no source available + var cts = new CancellationTokenSource(); + cts.Cancel(); + var imdsProbesCancellationToken = cts.Token; - var ex = await Assert.ThrowsExceptionAsync(async () => - await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource) - .ExecuteAsync().ConfigureAwait(false) - ).ConfigureAwait(false); + var ex = + await Assert.ThrowsExceptionAsync(async () => + await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(imdsProbesCancellationToken) + .ConfigureAwait(false)) + .ConfigureAwait(false); - Assert.AreEqual(MsalError.ManagedIdentityAllSourcesUnavailable, ex.ErrorCode); + Assert.AreEqual(MsalError.ImdsServiceError, ex.ErrorCode); } } } diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs index 7e61d7d908..7d5f00b00a 100644 --- a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs +++ b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs @@ -456,6 +456,36 @@ public async Task ProbeImdsEndpointAsyncFails404WhichIsNonRetriableAndRetryPolic Assert.AreEqual(ManagedIdentitySource.Imds, miSource); } } + + [TestMethod] + public async Task ProbeImdsV2EndpointAsync_TimesOutAfterOneSecond() + { + using (new EnvVariableContext()) + using (var httpManager = new MockHttpManager()) + { + var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned); + + miBuilder + .WithHttpManager(httpManager) + .WithRetryPolicyFactory(_testRetryPolicyFactory); + + var managedIdentityApp = miBuilder.Build(); + + // ImdsV2 mock is not needed, as the request will not be sent due to cancellation token being cancelled + + var cts = new CancellationTokenSource(); + cts.Cancel(); + var imdsProbesCancellationToken = cts.Token; + + var ex = + await Assert.ThrowsExceptionAsync(async () => + await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(imdsProbesCancellationToken) + .ConfigureAwait(false)) + .ConfigureAwait(false); + + Assert.AreEqual(MsalError.ImdsServiceError, ex.ErrorCode); + } + } #endregion Probe Tests #region Fallback Behavior Tests