From cd592ff400a6c26f8fe33959e85f8c02ac51842b Mon Sep 17 00:00:00 2001
From: Matthew Henderson <mattchenderson@users.noreply.github.com>
Date: Tue, 23 Jul 2024 14:09:38 -0700
Subject: [PATCH] WorkloadIdentityCredential in KVSecretsRepository (#10310)

---
 .../KeyVaultSecretsRepository.cs              | 50 ++++++++++++++++---
 .../Host/SecretsRepositoryTests.cs            |  2 +-
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/src/WebJobs.Script.WebHost/Security/KeyManagement/KeyVaultSecretsRepository.cs b/src/WebJobs.Script.WebHost/Security/KeyManagement/KeyVaultSecretsRepository.cs
index bf19b00520..8341b3d91e 100644
--- 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
index 07fa9fd6b7..2681dc4674 100644
--- 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());
                 }
             }