WorkloadIdentityCredential in KVSecretsRepository (#10310)

Azure · Jul 23, 2024 · cd592ff · cd592ff
1 parent 6aa4638
commit cd592ff
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git a/src/WebJobs.Script.WebHost/Security/KeyManagement/KeyVaultSecretsRepository.cs b/src/WebJobs.Script.WebHost/Security/KeyManagement/KeyVaultSecretsRepository.cs
@@ -26,6 +26,7 @@ public class KeyVaultSecretsRepository : BaseSecretsRepository
         private const string FunctionKeyPrefix = "host--functionKey--";
         private const string SystemKeyPrefix = "host--systemKey--";
 
+        private readonly Lazy<TokenCredential> _tokenCredential;
         private readonly Lazy<SecretClient> _secretClient;
         private readonly IEnvironment _environment;
 
@@ -38,19 +39,37 @@ public KeyVaultSecretsRepository(string secretsSentinelFilePath, string vaultUri
 
             Uri keyVaultUri = string.IsNullOrEmpty(vaultUri) ? throw new ArgumentException(nameof(vaultUri)) : new Uri(vaultUri);
 
-            _secretClient = new Lazy<SecretClient>(() =>
+            _tokenCredential = new Lazy<TokenCredential>(() =>
             {
-                // If clientSecret and tenantId are provided, use ClientSecret credential; otherwise use managed identity
-                TokenCredential credential = !string.IsNullOrEmpty(clientSecret) && !string.IsNullOrEmpty(tenantId)
-                    ? new ClientSecretCredential(tenantId, clientId, clientSecret)
-                    : new ChainedTokenCredential(new ManagedIdentityCredential(clientId), new ManagedIdentityCredential());
+                if (!TryCreateTokenCredential(clientId, clientSecret, tenantId, out TokenCredential credential))
+                {
+                    throw new InvalidOperationException("Failed to create token credential for KeyVaultSecretsRepository");
+                }
 
-                return new SecretClient(keyVaultUri, credential);
+                return credential;
+            });
+
+            _secretClient = new Lazy<SecretClient>(() =>
+            {
+                return new SecretClient(keyVaultUri, _tokenCredential.Value);
             });
 
             _environment = environment ?? throw new ArgumentNullException(nameof(environment));
         }
 
+        internal KeyVaultSecretsRepository(string secretsSentinelFilePath, string vaultUri, string clientId, string clientSecret, string tenantId, ILogger logger, IEnvironment environment, TokenCredential testEnvironmentTokenCredential) : this(secretsSentinelFilePath, vaultUri, clientId, clientSecret, tenantId, logger, environment)
+        {
+            _tokenCredential = new Lazy<TokenCredential>(() =>
+            {
+                if (!TryCreateTokenCredential(clientId, clientSecret, tenantId, out TokenCredential credential))
+                {
+                    throw new InvalidOperationException("Failed to create token credential for KeyVaultSecretsRepository");
+                }
+
+                return new ChainedTokenCredential(testEnvironmentTokenCredential, credential);
+            });
+        }
+
         // For testing
         internal KeyVaultSecretsRepository(SecretClient secretClient, string secretsSentinelFilePath, ILogger logger, IEnvironment environment) : base(secretsSentinelFilePath, logger, environment)
         {
@@ -72,6 +91,25 @@ public override async Task<ScriptSecrets> ReadAsync(ScriptSecretsType type, stri
             return type == ScriptSecretsType.Host ? await ReadHostSecrets() : await ReadFunctionSecrets(functionName);
         }
 
+        private static bool TryCreateTokenCredential(string clientId, string clientSecret, string tenantId, out TokenCredential credential)
+        {
+            if (!string.IsNullOrEmpty(clientId) && !string.IsNullOrEmpty(clientSecret) && !string.IsNullOrEmpty(tenantId))
+            {
+                credential = new ClientSecretCredential(tenantId, clientId, clientSecret);
+                return true;
+            }
+            else if (!string.IsNullOrEmpty(clientId))
+            {
+                credential = new ManagedIdentityCredential(clientId);
+                return true;
+            }
+            else
+            {
+                credential = new ManagedIdentityCredential();
+                return true;
+            }
+        }
+
         public override async Task WriteAsync(ScriptSecretsType type, string functionName, ScriptSecrets secrets)
         {
             // Get secrets as dictionary

diff --git a/test/WebJobs.Script.Tests.Integration/Host/SecretsRepositoryTests.cs b/test/WebJobs.Script.Tests.Integration/Host/SecretsRepositoryTests.cs
@@ -626,7 +626,7 @@ public ISecretsRepository GetNewSecretRepository()
                 }
                 else
                 {
-                    return new KeyVaultSecretsRepository(SecretsDirectory, KeyVaultUri, KeyVaultClientId, KeyVaultClientSecret, KeyVaultTenantId, logger, Environment);
+                    return new KeyVaultSecretsRepository(SecretsDirectory, KeyVaultUri, KeyVaultClientId, KeyVaultClientSecret, KeyVaultTenantId, logger, Environment, new WorkloadIdentityCredential());
                 }
             }