Skip to content

Commit 09155f3

Browse files
Robbie-MicrosoftRobbie-Microsoft
authored
ImdsV2 probe does not fire when .WithMtlsProofOfPossesstion is not used (#5579)
ImdsV2 probe does not fire when .WithMtlsProofOfPosesstion is not used
1 parent 8e0a65a commit 09155f3

File tree

4 files changed

+68
-33
lines changed

4 files changed

+68
-33
lines changed

‎src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
5858
if (s_sourceName == ManagedIdentitySource.None)
5959
{
6060
// First invocation: detect and cache
61-
source = await GetManagedIdentitySourceAsync(requestContext).ConfigureAwait(false);
61+
source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested).ConfigureAwait(false);
6262
}
6363
else
6464
{
@@ -101,13 +101,18 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
101101

102102
// Detect managed identity source based on the availability of environment variables and csr metadata probe request.
103103
// This method is perf sensitive any changes should be benchmarked.
104-
internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(RequestContext requestContext)
104+
internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
105+
RequestContext requestContext,
106+
bool isMtlsPopRequested)
105107
{
106108
// First check env vars to avoid the probe if possible
107109
ManagedIdentitySource source = GetManagedIdentitySourceNoImdsV2(requestContext.Logger);
108110

109-
// If a source is detected via env vars, use it
110-
if (source != ManagedIdentitySource.DefaultToImds)
111+
// If a source is detected via env vars, or
112+
// a source wasn't detected (it defaulted to ImdsV1) and MtlsPop was NOT requested,
113+
// use the source.
114+
// (don't trigger the ImdsV2 probe endpoint if MtlsPop was NOT requested)
115+
if (source != ManagedIdentitySource.DefaultToImds || !isMtlsPopRequested)
111116
{
112117
s_sourceName = source;
113118
return source;

‎src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync()
6666
// Create a temporary RequestContext for the CSR metadata probe request.
6767
var csrMetadataProbeRequestContext = new RequestContext(this.ServiceBundle, Guid.NewGuid(), null, CancellationToken.None);
6868

69-
return await ManagedIdentityClient.GetManagedIdentitySourceAsync(csrMetadataProbeRequestContext).ConfigureAwait(false);
69+
// GetManagedIdentitySourceAsync might return ImdsV2 = true, but it still requires .WithMtlsProofOfPossesion on the Managed Identity Application object to hit the ImdsV2 flow
70+
return await ManagedIdentityClient.GetManagedIdentitySourceAsync(csrMetadataProbeRequestContext, isMtlsPopRequested: true).ConfigureAwait(false);
7071
}
7172

7273
/// <summary>

‎tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,6 +506,62 @@ public async Task GetCsrMetadataAsyncFails404WhichIsNonRetriableAndRetryPolicyIs
506506
}
507507
}
508508

509+
[DataTestMethod]
510+
[DataRow(UserAssignedIdentityId.None, null)] // SAMI
511+
[DataRow(UserAssignedIdentityId.ClientId, TestConstants.ClientId)] // UAMI
512+
[DataRow(UserAssignedIdentityId.ResourceId, TestConstants.MiResourceId)] // UAMI
513+
[DataRow(UserAssignedIdentityId.ObjectId, TestConstants.ObjectId)] // UAMI
514+
public async Task ProbeDoesNotFireWhenMtlsPopNotRequested(
515+
UserAssignedIdentityId userAssignedIdentityId,
516+
string userAssignedId)
517+
{
518+
using (new EnvVariableContext())
519+
using (var httpManager = new MockHttpManager())
520+
{
521+
SetEnvironmentVariables(ManagedIdentitySource.Imds, TestConstants.ImdsEndpoint);
522+
523+
ManagedIdentityApplicationBuilder miBuilder = null;
524+
525+
var uami = userAssignedIdentityId != UserAssignedIdentityId.None && userAssignedId != null;
526+
if (uami)
527+
{
528+
miBuilder = CreateMIABuilder(userAssignedId, userAssignedIdentityId);
529+
}
530+
else
531+
{
532+
miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned);
533+
}
534+
535+
miBuilder
536+
.WithHttpManager(httpManager)
537+
.WithRetryPolicyFactory(_testRetryPolicyFactory);
538+
539+
var managedIdentityApp = miBuilder.Build();
540+
541+
// mock probe to show ImdsV2 is available
542+
httpManager.AddMockHandler(MockHelpers.MockCsrResponse(userAssignedIdentityId: userAssignedIdentityId, userAssignedId: userAssignedId));
543+
544+
var miSource = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync().ConfigureAwait(false);
545+
// this indicates ImdsV2 is available
546+
Assert.AreEqual(ManagedIdentitySource.ImdsV2, miSource);
547+
548+
httpManager.AddManagedIdentityMockHandler(
549+
ManagedIdentityTests.ImdsEndpoint,
550+
ManagedIdentityTests.Resource,
551+
MockHelpers.GetMsiSuccessfulResponse(),
552+
ManagedIdentitySource.Imds,
553+
userAssignedId: userAssignedId,
554+
userAssignedIdentityId: userAssignedIdentityId);
555+
556+
// ImdsV1 flow will be used since .WithMtlsProofOfPossession() is not used here
557+
var result = await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource).ExecuteAsync().ConfigureAwait(false);
558+
559+
Assert.IsNotNull(result);
560+
Assert.IsNotNull(result.AccessToken);
561+
Assert.AreEqual(TokenSource.IdentityProvider, result.AuthenticationResultMetadata.TokenSource);
562+
}
563+
}
564+
509565
#region Cuid Tests
510566
[TestMethod]
511567
public void TestCsrGeneration_OnlyVmId()
@@ -684,7 +740,7 @@ await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Res
684740
}
685741
#endregion
686742

687-
#region cached certificate tests
743+
#region Cached certificate tests
688744
[TestMethod]
689745
public async Task mTLSPop_ForceRefresh_UsesCachedCert_NoIssueCredential_PostsCanonicalClientId_AndSkipsAttestation()
690746
{

‎tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -42,16 +42,6 @@ public class ManagedIdentityTests : TestBase
4242

4343
private readonly TestRetryPolicyFactory _testRetryPolicyFactory = new TestRetryPolicyFactory();
4444

45-
private void AddImdsV2CsrMockHandlerIfNeeded(
46-
ManagedIdentitySource managedIdentitySource,
47-
MockHttpManager httpManager)
48-
{
49-
if (managedIdentitySource == ManagedIdentitySource.Imds)
50-
{
51-
httpManager.AddMockHandler(MockHelpers.MockCsrResponseFailure());
52-
}
53-
}
54-
5545
[DataTestMethod]
5646
[DataRow("http://127.0.0.1:41564/msi/token/", ManagedIdentitySource.AppService, ManagedIdentitySource.AppService)]
5747
[DataRow(AppServiceEndpoint, ManagedIdentitySource.AppService, ManagedIdentitySource.AppService)]
@@ -69,15 +59,11 @@ public async Task GetManagedIdentityTests(
6959
using (new EnvVariableContext())
7060
using (var httpManager = new MockHttpManager())
7161
{
72-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
7362
SetEnvironmentVariables(managedIdentitySource, endpoint);
7463

7564
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
7665
.WithHttpManager(httpManager);
7766

78-
79-
80-
8167
ManagedIdentityApplication mi = miBuilder.Build() as ManagedIdentityApplication;
8268

8369
Assert.AreEqual(expectedManagedIdentitySource, await mi.GetManagedIdentitySourceAsync().ConfigureAwait(false));
@@ -106,7 +92,6 @@ public async Task SAMIHappyPathAsync(
10692
using (new EnvVariableContext())
10793
using (var httpManager = new MockHttpManager())
10894
{
109-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
11095
SetEnvironmentVariables(managedIdentitySource, endpoint);
11196

11297
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -155,7 +140,6 @@ public async Task UAMIHappyPathAsync(
155140
using (new EnvVariableContext())
156141
using (var httpManager = new MockHttpManager())
157142
{
158-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
159143
SetEnvironmentVariables(managedIdentitySource, endpoint);
160144

161145
var miBuilder = CreateMIABuilder(userAssignedId, userAssignedIdentityId);
@@ -203,7 +187,6 @@ public async Task ManagedIdentityDifferentScopesTestAsync(
203187
using (new EnvVariableContext())
204188
using (var httpManager = new MockHttpManager())
205189
{
206-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
207190
SetEnvironmentVariables(managedIdentitySource, endpoint);
208191

209192
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -261,7 +244,6 @@ public async Task ManagedIdentityForceRefreshTestAsync(
261244
using (new EnvVariableContext())
262245
using (var httpManager = new MockHttpManager())
263246
{
264-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
265247
SetEnvironmentVariables(managedIdentitySource, endpoint);
266248

267249
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -320,7 +302,6 @@ public async Task ManagedIdentityWithClaimsAndCapabilitiesTestAsync(
320302
using (new EnvVariableContext())
321303
using (var httpManager = new MockHttpManager())
322304
{
323-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
324305
SetEnvironmentVariables(managedIdentitySource, endpoint);
325306

326307
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -384,7 +365,6 @@ public async Task ManagedIdentityWithClaimsTestAsync(
384365
using (new EnvVariableContext())
385366
using (var httpManager = new MockHttpManager())
386367
{
387-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
388368
SetEnvironmentVariables(managedIdentitySource, endpoint);
389369

390370
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -456,7 +436,6 @@ public async Task ManagedIdentityTestWrongScopeAsync(string resource, ManagedIde
456436
using (new EnvVariableContext())
457437
using (var httpManager = new MockHttpManager())
458438
{
459-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
460439
SetEnvironmentVariables(managedIdentitySource, endpoint);
461440

462441
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -560,7 +539,6 @@ public async Task ManagedIdentityErrorResponseNoPayloadTestAsync(ManagedIdentity
560539
using (new EnvVariableContext())
561540
using (var httpManager = new MockHttpManager())
562541
{
563-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
564542
SetEnvironmentVariables(managedIdentitySource, endpoint);
565543

566544
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -600,7 +578,6 @@ public async Task ManagedIdentityNullResponseAsync(ManagedIdentitySource managed
600578
using (new EnvVariableContext())
601579
using (var httpManager = new MockHttpManager())
602580
{
603-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
604581
SetEnvironmentVariables(managedIdentitySource, endpoint);
605582

606583
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -638,7 +615,6 @@ public async Task ManagedIdentityUnreachableNetworkAsync(ManagedIdentitySource m
638615
using (new EnvVariableContext())
639616
using (var httpManager = new MockHttpManager())
640617
{
641-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
642618
SetEnvironmentVariables(managedIdentitySource, endpoint);
643619

644620
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)
@@ -1036,7 +1012,6 @@ public async Task InvalidJsonResponseHandling(ManagedIdentitySource managedIdent
10361012
using (new EnvVariableContext())
10371013
using (var httpManager = new MockHttpManager())
10381014
{
1039-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
10401015
SetEnvironmentVariables(managedIdentitySource, endpoint);
10411016

10421017
var miBuilder = ManagedIdentityApplicationBuilder
@@ -1077,7 +1052,6 @@ public async Task ManagedIdentityRequestTokensForDifferentScopesTestAsync(
10771052
using (new EnvVariableContext())
10781053
using (var httpManager = new MockHttpManager())
10791054
{
1080-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
10811055
SetEnvironmentVariables(managedIdentitySource, endpoint);
10821056

10831057
var miBuilder = ManagedIdentityApplicationBuilder
@@ -1332,7 +1306,6 @@ public async Task ManagedIdentityWithCapabilitiesTestAsync(
13321306
using (new EnvVariableContext())
13331307
using (var httpManager = new MockHttpManager())
13341308
{
1335-
AddImdsV2CsrMockHandlerIfNeeded(managedIdentitySource, httpManager);
13361309
SetEnvironmentVariables(managedIdentitySource, endpoint);
13371310

13381311
var miBuilder = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned)

0 commit comments

Comments
 (0)