Skip to content

Commit 732c9d6

Browse files
Robbie-MicrosoftRobbie-Microsoft
authored
ImdsV2: Added Caching for the IMDS Endpoint Env Variable + Improved Unit Tests (#5514)
1 parent 1c399c3 commit 732c9d6

File tree

3 files changed

+71
-15
lines changed

3 files changed

+71
-15
lines changed

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

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ internal class ImdsManagedIdentitySource : AbstractManagedIdentity
3434

3535
private readonly Uri _imdsEndpoint;
3636

37+
private static string s_cachedBaseEndpoint = null;
38+
3739
internal ImdsManagedIdentitySource(RequestContext requestContext) :
3840
base(requestContext, ManagedIdentitySource.Imds)
3941
{
@@ -181,25 +183,25 @@ public static Uri GetValidatedEndpoint(
181183
string queryParams = null
182184
)
183185
{
184-
UriBuilder builder;
185-
186-
if (!string.IsNullOrEmpty(EnvironmentVariables.PodIdentityEndpoint))
186+
if (s_cachedBaseEndpoint == null)
187187
{
188-
logger.Verbose(() => "[Managed Identity] Environment variable AZURE_POD_IDENTITY_AUTHORITY_HOST for IMDS returned endpoint: " + EnvironmentVariables.PodIdentityEndpoint);
189-
builder = new UriBuilder(EnvironmentVariables.PodIdentityEndpoint)
188+
if (!string.IsNullOrEmpty(EnvironmentVariables.PodIdentityEndpoint))
190189
{
191-
Path = subPath
192-
};
193-
}
194-
else
195-
{
196-
logger.Verbose(() => "[Managed Identity] Unable to find AZURE_POD_IDENTITY_AUTHORITY_HOST environment variable for IMDS, using the default endpoint.");
197-
builder = new UriBuilder(DefaultImdsBaseEndpoint)
190+
logger.Verbose(() => "[Managed Identity] Environment variable AZURE_POD_IDENTITY_AUTHORITY_HOST for IMDS returned endpoint: " + EnvironmentVariables.PodIdentityEndpoint);
191+
s_cachedBaseEndpoint = EnvironmentVariables.PodIdentityEndpoint;
192+
}
193+
else
198194
{
199-
Path = subPath
200-
};
195+
logger.Verbose(() => "[Managed Identity] Unable to find AZURE_POD_IDENTITY_AUTHORITY_HOST environment variable for IMDS, using the default endpoint.");
196+
s_cachedBaseEndpoint = DefaultImdsBaseEndpoint;
197+
}
201198
}
202-
199+
200+
UriBuilder builder = new UriBuilder(s_cachedBaseEndpoint)
201+
{
202+
Path = subPath
203+
};
204+
203205
if (!string.IsNullOrEmpty(queryParams))
204206
{
205207
builder.Query = queryParams;

‎tests/Microsoft.Identity.Test.Common/Core/Helpers/ManagedIdentityTestUtil.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public static void SetEnvironmentVariables(ManagedIdentitySource managedIdentity
4242
break;
4343

4444
case ManagedIdentitySource.Imds:
45+
case ManagedIdentitySource.ImdsV2:
4546
Environment.SetEnvironmentVariable("AZURE_POD_IDENTITY_AUTHORITY_HOST", endpoint);
4647
break;
4748

@@ -59,11 +60,15 @@ public static void SetEnvironmentVariables(ManagedIdentitySource managedIdentity
5960
Environment.SetEnvironmentVariable("IDENTITY_HEADER", secret);
6061
Environment.SetEnvironmentVariable("IDENTITY_SERVER_THUMBPRINT", thumbprint);
6162
break;
63+
6264
case ManagedIdentitySource.MachineLearning:
6365
Environment.SetEnvironmentVariable("MSI_ENDPOINT", endpoint);
6466
Environment.SetEnvironmentVariable("MSI_SECRET", secret);
6567
Environment.SetEnvironmentVariable("DEFAULT_IDENTITY_CLIENT_ID", "fake_DEFAULT_IDENTITY_CLIENT_ID");
6668
break;
69+
70+
default:
71+
throw new NotImplementedException($"Setting environment variables for {managedIdentitySource} is not implemented.");
6772
}
6873
}
6974

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
using Microsoft.Identity.Client.MtlsPop;
1717
using Microsoft.Identity.Client.PlatformsCommon.Interfaces;
1818
using Microsoft.Identity.Client.PlatformsCommon.Shared;
19+
using Microsoft.Identity.Test.Common.Core.Helpers;
1920
using Microsoft.Identity.Test.Common.Core.Mocks;
2021
using Microsoft.Identity.Test.Unit.Helpers;
2122
using Microsoft.Identity.Test.Unit.PublicApiTests;
@@ -152,8 +153,11 @@ public async Task BearerTokenHappyPath(
152153
UserAssignedIdentityId userAssignedIdentityId,
153154
string userAssignedId)
154155
{
156+
using (new EnvVariableContext())
155157
using (var httpManager = new MockHttpManager())
156158
{
159+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
160+
157161
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, managedIdentityKeyType: ManagedIdentityKeyType.InMemory).ConfigureAwait(false);
158162

159163
AddMocksToGetEntraToken(httpManager, userAssignedIdentityId, userAssignedId);
@@ -185,8 +189,11 @@ public async Task BearerTokenIsPerIdentity(
185189
string userAssignedId,
186190
string userAssignedId2)
187191
{
192+
using (new EnvVariableContext())
188193
using (var httpManager = new MockHttpManager())
189194
{
195+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
196+
190197
#region Identity 1
191198
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId).ConfigureAwait(false);
192199

@@ -245,8 +252,11 @@ public async Task BearerTokenIsReAcquiredWhenCertificatIsExpired(
245252
UserAssignedIdentityId userAssignedIdentityId,
246253
string userAssignedId)
247254
{
255+
using (new EnvVariableContext())
248256
using (var httpManager = new MockHttpManager())
249257
{
258+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
259+
250260
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId).ConfigureAwait(false);
251261

252262
AddMocksToGetEntraToken(httpManager, userAssignedIdentityId, userAssignedId, TestConstants.ExpiredRawCertificate); // cert will be expired on second request
@@ -287,8 +297,11 @@ public async Task mTLSPopTokenHappyPath(
287297
UserAssignedIdentityId userAssignedIdentityId,
288298
string userAssignedId)
289299
{
300+
using (new EnvVariableContext())
290301
using (var httpManager = new MockHttpManager())
291302
{
303+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
304+
292305
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
293306

294307
AddMocksToGetEntraToken(httpManager, userAssignedIdentityId, userAssignedId, mTLSPop: true);
@@ -327,8 +340,11 @@ public async Task mTLSPopTokenIsPerIdentity(
327340
string userAssignedId,
328341
string userAssignedId2)
329342
{
343+
using (new EnvVariableContext())
330344
using (var httpManager = new MockHttpManager())
331345
{
346+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
347+
332348
#region Identity 1
333349
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
334350

@@ -406,8 +422,11 @@ public async Task mTLSPopTokenIsReAcquiredWhenCertificatIsExpired(
406422
UserAssignedIdentityId userAssignedIdentityId,
407423
string userAssignedId)
408424
{
425+
using (new EnvVariableContext())
409426
using (var httpManager = new MockHttpManager())
410427
{
428+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
429+
411430
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
412431

413432
AddMocksToGetEntraToken(httpManager, userAssignedIdentityId, userAssignedId, TestConstants.ExpiredRawCertificate, mTLSPop: true);
@@ -448,8 +467,11 @@ public async Task mTLSPopTokenIsReAcquiredWhenCertificatIsExpired(
448467
[TestMethod]
449468
public async Task GetCsrMetadataAsyncSucceeds()
450469
{
470+
using (new EnvVariableContext())
451471
using (var httpManager = new MockHttpManager())
452472
{
473+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
474+
453475
var handler = httpManager.AddMockHandler(MockHelpers.MockCsrResponse());
454476

455477
await CreateManagedIdentityAsync(httpManager, addProbeMock: false).ConfigureAwait(false);
@@ -459,8 +481,11 @@ public async Task GetCsrMetadataAsyncSucceeds()
459481
[TestMethod]
460482
public async Task GetCsrMetadataAsyncSucceedsAfterRetry()
461483
{
484+
using (new EnvVariableContext())
462485
using (var httpManager = new MockHttpManager())
463486
{
487+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
488+
464489
// First attempt fails with INTERNAL_SERVER_ERROR (500)
465490
httpManager.AddMockHandler(MockHelpers.MockCsrResponse(HttpStatusCode.InternalServerError));
466491

@@ -472,8 +497,11 @@ public async Task GetCsrMetadataAsyncSucceedsAfterRetry()
472497
[TestMethod]
473498
public async Task GetCsrMetadataAsyncFailsWithMissingServerHeader()
474499
{
500+
using (new EnvVariableContext())
475501
using (var httpManager = new MockHttpManager())
476502
{
503+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
504+
477505
httpManager.AddMockHandler(MockHelpers.MockCsrResponse(responseServerHeader: null));
478506

479507
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, addProbeMock: false, addSourceCheck: false).ConfigureAwait(false);
@@ -486,8 +514,11 @@ public async Task GetCsrMetadataAsyncFailsWithMissingServerHeader()
486514
[TestMethod]
487515
public async Task GetCsrMetadataAsyncFailsWithInvalidFormat()
488516
{
517+
using (new EnvVariableContext())
489518
using (var httpManager = new MockHttpManager())
490519
{
520+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
521+
491522
httpManager.AddMockHandler(MockHelpers.MockCsrResponse(responseServerHeader: "I_MDS/150.870.65.1854"));
492523

493524
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, addProbeMock: false, addSourceCheck: false).ConfigureAwait(false);
@@ -500,8 +531,11 @@ public async Task GetCsrMetadataAsyncFailsWithInvalidFormat()
500531
[TestMethod]
501532
public async Task GetCsrMetadataAsyncFailsAfterMaxRetries()
502533
{
534+
using (new EnvVariableContext())
503535
using (var httpManager = new MockHttpManager())
504536
{
537+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
538+
505539
const int Num500Errors = 1 + TestCsrMetadataProbeRetryPolicy.ExponentialStrategyNumRetries;
506540
for (int i = 0; i < Num500Errors; i++)
507541
{
@@ -518,8 +552,11 @@ public async Task GetCsrMetadataAsyncFailsAfterMaxRetries()
518552
[TestMethod]
519553
public async Task GetCsrMetadataAsyncFails404WhichIsNonRetriableAndRetryPolicyIsNotTriggeredAsync()
520554
{
555+
using (new EnvVariableContext())
521556
using (var httpManager = new MockHttpManager())
522557
{
558+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
559+
523560
httpManager.AddMockHandler(MockHelpers.MockCsrResponse(HttpStatusCode.NotFound));
524561

525562
var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, addProbeMock: false, addSourceCheck: false).ConfigureAwait(false);
@@ -604,8 +641,11 @@ public void AttachPrivateKeyToCert_NullPrivateKey_ThrowsArgumentNullException()
604641
[TestMethod]
605642
public async Task MtlsPop_AttestationProviderMissing_ThrowsClientException()
606643
{
644+
using (new EnvVariableContext())
607645
using (var httpManager = new MockHttpManager())
608646
{
647+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
648+
609649
var mi = await CreateManagedIdentityAsync(httpManager, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
610650

611651
// CreateManagedIdentityAsync does a probe; Add one more CSR response for the actual acquire.
@@ -625,8 +665,11 @@ await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)
625665
[TestMethod]
626666
public async Task MtlsPop_AttestationProviderReturnsNull_ThrowsClientException()
627667
{
668+
using (new EnvVariableContext())
628669
using (var httpManager = new MockHttpManager())
629670
{
671+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
672+
630673
var mi = await CreateManagedIdentityAsync(httpManager, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
631674

632675
// CreateManagedIdentityAsync does a probe; Add one more CSR response for the actual acquire.
@@ -649,8 +692,11 @@ await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)
649692
[TestMethod]
650693
public async Task MtlsPop_AttestationProviderReturnsEmptyToken_ThrowsClientException()
651694
{
695+
using (new EnvVariableContext())
652696
using (var httpManager = new MockHttpManager())
653697
{
698+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
699+
654700
var mi = await CreateManagedIdentityAsync(httpManager, managedIdentityKeyType: ManagedIdentityKeyType.KeyGuard).ConfigureAwait(false);
655701

656702
// CreateManagedIdentityAsync does a probe; Add one more CSR response for the actual acquire.
@@ -673,8 +719,11 @@ await mi.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)
673719
[TestMethod]
674720
public async Task mTLSPop_RequestedWithoutKeyGuard_ThrowsClientException()
675721
{
722+
using (new EnvVariableContext())
676723
using (var httpManager = new MockHttpManager())
677724
{
725+
SetEnvironmentVariables(ManagedIdentitySource.ImdsV2, TestConstants.ImdsEndpoint);
726+
678727
// Force in-memory keys (i.e., not KeyGuard)
679728
var managedIdentityApp = await CreateManagedIdentityAsync(
680729
httpManager,

0 commit comments

Comments
 (0)