Merge pull request pnp#1053 from gautamdsheth/fix/mi-perf-issue

Fix: Improve perf for managed identity auth
gautamdsheth · Aug 28, 2024 · 8dc4d7c · 8dc4d7c
2 parents 7801b0d + 4e201eb
commit 8dc4d7c
Showing 1 changed file with 32 additions and 44 deletions.
diff --git a/src/lib/PnP.Framework/AuthenticationManager.cs b/src/lib/PnP.Framework/AuthenticationManager.cs
@@ -92,6 +92,7 @@ public class AuthenticationManager : IDisposable
 
         private readonly IPublicClientApplication publicClientApplication;
         private readonly IConfidentialClientApplication confidentialClientApplication;
+        private readonly IManagedIdentityApplication mi;
 
         // Azure environment setup
         private AzureEnvironment azureEnvironment;
@@ -347,13 +348,36 @@ public AuthenticationManager(string endpoint, string identityHeader, ManagedIden
             {
                 throw new ArgumentException($"When {nameof(managedIdentityType)} is not SystemAssigned, {nameof(managedIdentityUserAssignedIdentifier)} must be provided", nameof(managedIdentityType));
             }
-
-            this.accessToken = new NetworkCredential("", accessToken).SecurePassword;
-            this.managedIdendityEndpoint = endpoint;
-            this.managedIdentityHeader = identityHeader;
-            this.authenticationType = managedIdentityType == ManagedIdentityType.SystemAssigned ? ClientContextType.SystemAssignedManagedIdentity : ClientContextType.UserAssignedManagedIdentity;
+
+            authenticationType = managedIdentityType == ManagedIdentityType.SystemAssigned ? ClientContextType.SystemAssignedManagedIdentity : ClientContextType.UserAssignedManagedIdentity;
             this.managedIdentityType = managedIdentityType;            
-            this.managedIdentityUserAssignedIdentifier = managedIdentityUserAssignedIdentifier;
+            this.managedIdentityUserAssignedIdentifier = managedIdentityUserAssignedIdentifier;            
+
+            // Construct the URL to call to get the token based on the type of Managed Identity in use
+            switch (managedIdentityType)
+            {
+                case ManagedIdentityType.UserAssignedByClientId:
+                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with client ID: {managedIdentityUserAssignedIdentifier}");
+                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedClientId(managedIdentityUserAssignedIdentifier)).Build();
+                    break;
+
+                case ManagedIdentityType.UserAssignedByObjectId:
+                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with object/principal ID: {managedIdentityUserAssignedIdentifier}");
+                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedObjectId(managedIdentityUserAssignedIdentifier)).Build();
+                    break;
+
+
+                case ManagedIdentityType.UserAssignedByResourceId:
+                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with Azure Resource ID: {managedIdentityUserAssignedIdentifier}");
+                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedResourceId(managedIdentityUserAssignedIdentifier)).Build();
+                    break;
+
+                case ManagedIdentityType.SystemAssigned:
+                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, "Using the system assigned managed identity");
+                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned).Build();
+                    break;
+            }            
+
         }
 
         /// <summary>
@@ -1144,10 +1168,10 @@ public async Task<ClientContext> GetContextAsync(string siteUrl, CancellationTok
                         });
                         ClientContextSettings clientContextSettings = new ClientContextSettings()
                         {
-                            Type = ClientContextType.AccessToken,
+                            Type = managedIdentityType == ManagedIdentityType.SystemAssigned ? ClientContextType.SystemAssignedManagedIdentity : ClientContextType.UserAssignedManagedIdentity,
                             SiteUrl = siteUrl,
                             AuthenticationManager = this,
-                            Environment = this.azureEnvironment
+                            Environment = azureEnvironment
                         };
                         context.AddContextSettings(clientContextSettings);
 
@@ -1532,42 +1556,6 @@ public ClientContext GetAccessTokenContext(string siteUrl, string accessToken)
         /// <returns>Access token</returns>
         private string GetManagedIdentityToken(string audience)
         {
-            // Ensure our AuthenticationManager is set up to handle Managed Identities
-            if(!managedIdentityType.HasValue)
-            {
-                throw new InvalidOperationException("Trying to get a Managed Identity access token within a non Managed Identity authentication context is not possible");
-            }
-
-            IManagedIdentityApplication mi;
-
-            // Construct the URL to call to get the token based on the type of Managed Identity in use
-            switch(managedIdentityType.Value)
-            {
-                case ManagedIdentityType.UserAssignedByClientId:
-                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with client ID: {managedIdentityUserAssignedIdentifier}");
-                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedClientId(managedIdentityUserAssignedIdentifier)).Build();
-                    break;
-
-                case ManagedIdentityType.UserAssignedByObjectId:
-                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with object/principal ID: {managedIdentityUserAssignedIdentifier}");
-                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedObjectId(managedIdentityUserAssignedIdentifier)).Build();
-                    break;
-
-
-                case ManagedIdentityType.UserAssignedByResourceId:
-                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, $"Using the user assigned managed identity with Azure Resource ID: {managedIdentityUserAssignedIdentifier}");
-                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.WithUserAssignedResourceId(managedIdentityUserAssignedIdentifier)).Build();
-                    break;
-
-                case ManagedIdentityType.SystemAssigned:
-                    Diagnostics.Log.Debug(Constants.LOGGING_SOURCE, "Using the system assigned managed identity");
-                    mi = ManagedIdentityApplicationBuilder.Create(ManagedIdentityId.SystemAssigned).Build();
-                    break;
-
-                default:
-                    throw new ArgumentException("Using an unsupported type of Managed Identity", nameof(managedIdentityType));
-            }
-
             AuthenticationResult result = mi.AcquireTokenForManagedIdentity(audience).ExecuteAsync().GetAwaiter().GetResult();
             return result?.AccessToken;
         }