Remove static caching of managed identity source (#4840)

* Remove static caching of managed identity source * Add comments * Update ManagedIdentityClient.cs
AzureAD · Jul 12, 2024 · 4419b3f · 4419b3f
1 parent 5c7c527
commit 4419b3f
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 27 deletions.
diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs b/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs
@@ -16,13 +16,6 @@ namespace Microsoft.Identity.Client.ManagedIdentity
     internal class ManagedIdentityClient
     {
         private readonly AbstractManagedIdentity _identitySource;
-        internal static Lazy<ManagedIdentitySource> s_managedIdentitySourceDetected = new Lazy<ManagedIdentitySource>(() => GetManagedIdentitySource());
-
-        // To reset the cached source for testing purposes.
-        internal static void resetCachedSource()
-        {
-            s_managedIdentitySourceDetected = new Lazy<ManagedIdentitySource>(() => GetManagedIdentitySource());
-        }
 
         public ManagedIdentityClient(RequestContext requestContext)
         {
@@ -40,7 +33,7 @@ internal Task<ManagedIdentityResponse> SendTokenRequestForManagedIdentityAsync(A
         // This method tries to create managed identity source for different sources, if none is created then defaults to IMDS.
         private static AbstractManagedIdentity SelectManagedIdentitySource(RequestContext requestContext)
         {
-            return s_managedIdentitySourceDetected.Value switch
+            return GetManagedIdentitySource() switch
             {
                 ManagedIdentitySource.ServiceFabric => ServiceFabricManagedIdentitySource.Create(requestContext),
                 ManagedIdentitySource.AppService => AppServiceManagedIdentitySource.Create(requestContext),
@@ -51,7 +44,9 @@ private static AbstractManagedIdentity SelectManagedIdentitySource(RequestContex
         }
 
         // Detect managed identity source based on the availability of environment variables.
-        private static ManagedIdentitySource GetManagedIdentitySource()
+        // The result of this method is not cached because reading environment variables is cheap. 
+        // This method is perf sensitive any changes should be benchmarked.
+        internal static ManagedIdentitySource GetManagedIdentitySource()
         {
             string identityEndpoint = EnvironmentVariables.IdentityEndpoint;
             string identityHeader = EnvironmentVariables.IdentityHeader;

diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs b/src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs
@@ -61,7 +61,7 @@ public AcquireTokenForManagedIdentityParameterBuilder AcquireTokenForManagedIden
         /// <returns>Managed identity source detected on the environment if any.</returns>
         public static ManagedIdentitySource GetManagedIdentitySource()
         {
-            return ManagedIdentityClient.s_managedIdentitySourceDetected.Value;
+            return ManagedIdentityClient.GetManagedIdentitySource();
         }
     }
 }
diff --git a/tests/Microsoft.Identity.Test.Common/Core/Helpers/ManagedIdentityTestUtil.cs b/tests/Microsoft.Identity.Test.Common/Core/Helpers/ManagedIdentityTestUtil.cs
@@ -60,8 +60,6 @@ public static void SetEnvironmentVariables(ManagedIdentitySource managedIdentity
                     Environment.SetEnvironmentVariable("IDENTITY_SERVER_THUMBPRINT", "thumbprint");
                     break;
             }
-
-            ManagedIdentityClient.resetCachedSource();
         }
 
         /// <summary>

diff --git a/.../Microsoft.Identity.Test.Integration.netcore/HeadlessTests/ManagedIdentityTests.NetFwk.cs b/.../Microsoft.Identity.Test.Integration.netcore/HeadlessTests/ManagedIdentityTests.NetFwk.cs
@@ -88,9 +88,6 @@ public async Task AcquireMSITokenAsync(MsiAzureResource azureResource, string us
                 //Set the Environment Variables
                 SetEnvironmentVariables(envVariables);
 
-                //Reset cached source with update in environment variables
-                ManagedIdentityClient.resetCachedSource();
-
                 //form the http proxy URI 
                 string uri = s_baseURL + $"MSIToken?" +
                     $"azureresource={azureResource}&uri=";
@@ -146,9 +143,6 @@ public async Task AcquireMsiToken_ForTokenExchangeResource_Successfully()
                 //Set the Environment Variables
                 SetEnvironmentVariables(envVariables);
 
-                //Reset cached source with update in environment variables
-                ManagedIdentityClient.resetCachedSource();
-
                 //form the http proxy URI 
                 string uri = s_baseURL + $"MSIToken?" +
                     $"azureresource={MsiAzureResource.WebApp}&uri=";
@@ -204,9 +198,6 @@ public async Task AcquireMsiToken_ExchangeForEstsToken_Successfully()
                 //Set the Environment Variables
                 SetEnvironmentVariables(envVariables);
 
-                //Reset cached source with update in environment variables
-                ManagedIdentityClient.resetCachedSource();
-
                 //form the http proxy URI 
                 string uri = s_baseURL + $"MSIToken?" +
                     $"azureresource={MsiAzureResource.WebApp}&uri=";
@@ -318,9 +309,6 @@ public async Task ManagedIdentityRequestFailureCheckAsync(MsiAzureResource azure
                 //Set the Environment Variables
                 SetEnvironmentVariables(envVariables);
 
-                //Reset cached source with update in environment variables
-                ManagedIdentityClient.resetCachedSource();
-
                 //form the http proxy URI 
                 string uri = s_baseURL + $"MSIToken?" +
                     $"azureresource={azureResource}&uri=";
@@ -363,9 +351,6 @@ public async Task MSIWrongScopesAsync(MsiAzureResource azureResource, string use
                 //Set the Environment Variables
                 SetEnvironmentVariables(envVariables);
 
-                //Reset cached source with update in environment variables
-                ManagedIdentityClient.resetCachedSource();
-
                 //form the http proxy URI 
                 string uri = s_baseURL + $"MSIToken?" +
                     $"azureresource={azureResource}&uri=";