Cosmos, Redis and Postgres show keyvault when using key access or passwords with emulator (#8406)

* Follow up keyvault changes
- Pass the IKeyVaultSecretReference to the SecretResolver
- Don't add the default keyvault when using the emulator.

* Add the KeyVault resource, but remove it from the model in BeforeStart if the Azure resource is emulated or container in run mode.

* Respond to PR feedback

Fix #8364
---------

Co-authored-by: Eric Erhardt <eric.erhardt@microsoft.com>

Author: davidfowl

src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs

@@ -342,6 +342,20 @@ public static IResourceBuilder<AzureCosmosDBResource> WithAccessKeyAuthenticatio
         var kv = builder.ApplicationBuilder.AddAzureKeyVault($"{builder.Resource.Name}-kv")
                                            .WithParentRelationship(builder.Resource);
 
+        // remove the KeyVault from the model if the emulator is used during run mode.
+        // need to do this later in case builder becomes an emulator after this method is called.
+        if (builder.ApplicationBuilder.ExecutionContext.IsRunMode)
+        {
+            builder.ApplicationBuilder.Eventing.Subscribe<BeforeStartEvent>((data, _) =>
+            {
+                if (builder.Resource.IsEmulator)
+                {
+                    data.Model.Resources.Remove(kv.Resource);
+                }
+                return Task.CompletedTask;
+            });
+        }
+
         return builder.WithAccessKeyAuthentication(kv);
     }
 

src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBResource.cs

@@ -134,7 +134,7 @@ internal ReferenceExpression GetChildConnectionString(string childResourceName,
 
         var builder = new ReferenceExpressionBuilder();
 
-        if (UseAccessKeyAuthentication)
+        if (UseAccessKeyAuthentication && !IsEmulator)
         {
             builder.AppendFormatted(ConnectionStringSecretOutput.Resource.GetSecretReference(GetKeyValueSecretName(childResourceName)));
         }

src/Aspire.Hosting.Azure.KeyVault/AzureKeyVaultResource.cs

@@ -34,9 +34,9 @@ public class AzureKeyVaultResource(string name, Action<AzureResourceInfrastructu
     BicepOutputReference IKeyVaultResource.VaultUriOutputReference => VaultUri;
 
     // In run mode, this is set to the secret client used to access the Azure Key Vault.
-    internal Func<string, CancellationToken, Task<string?>>? SecretResolver { get; set; }
+    internal Func<IKeyVaultSecretReference, CancellationToken, Task<string?>>? SecretResolver { get; set; }
 
-    Func<string, CancellationToken, Task<string?>>? IKeyVaultResource.SecretResolver
+    Func<IKeyVaultSecretReference, CancellationToken, Task<string?>>? IKeyVaultResource.SecretResolver
     {
         get => SecretResolver;
         set => SecretResolver = value;

src/Aspire.Hosting.Azure.KeyVault/AzureKeyVaultSecretReference.cs

@@ -28,7 +28,7 @@ internal sealed class AzureKeyVaultSecretReference(string secretName, AzureKeyVa
     {
         if (azureKeyVaultResource.SecretResolver is { } secretResolver)
         {
-            return await secretResolver(secretName, cancellationToken).ConfigureAwait(false);
+            return await secretResolver(this, cancellationToken).ConfigureAwait(false);
         }
 
         throw new InvalidOperationException($"Secret '{secretName}' not found in Key Vault '{azureKeyVaultResource.Name}'.");

src/Aspire.Hosting.Azure.PostgreSQL/AzurePostgresExtensions.cs

@@ -289,6 +289,20 @@ public static IResourceBuilder<AzurePostgresFlexibleServerResource> WithPassword
         var kv = builder.ApplicationBuilder.AddAzureKeyVault($"{builder.Resource.Name}-kv")
                                            .WithParentRelationship(builder.Resource);
 
+        // remove the KeyVault from the model if the emulator is used during run mode.
+        // need to do this later in case builder becomes an emulator after this method is called.
+        if (builder.ApplicationBuilder.ExecutionContext.IsRunMode)
+        {
+            builder.ApplicationBuilder.Eventing.Subscribe<BeforeStartEvent>((data, token) =>
+            {
+                if (builder.Resource.IsContainer())
+                {
+                    data.Model.Resources.Remove(kv.Resource);
+                }
+                return Task.CompletedTask;
+            });
+        }
+
         return builder.WithPasswordAuthentication(kv, userName, password);
     }
 

src/Aspire.Hosting.Azure.Redis/AzureRedisExtensions.cs

@@ -195,6 +195,20 @@ public static IResourceBuilder<AzureRedisCacheResource> WithAccessKeyAuthenticat
         var kv = builder.ApplicationBuilder.AddAzureKeyVault($"{builder.Resource.Name}-kv")
                                            .WithParentRelationship(builder.Resource);
 
+        // remove the KeyVault from the model if the emulator is used during run mode.
+        // need to do this later in case builder becomes an emulator after this method is called.
+        if (builder.ApplicationBuilder.ExecutionContext.IsRunMode)
+        {
+            builder.ApplicationBuilder.Eventing.Subscribe<BeforeStartEvent>((data, token) =>
+            {
+                if (builder.Resource.IsContainer())
+                {
+                    data.Model.Resources.Remove(kv.Resource);
+                }
+                return Task.CompletedTask;
+            });
+        }
+
         return builder.WithAccessKeyAuthentication(kv);
     }
 

src/Aspire.Hosting.Azure/IKeyVaultResource.cs

@@ -23,7 +23,7 @@ public interface IKeyVaultResource : IResource, IAzureResource
     /// <summary>
     /// Gets or sets the secret resolver function used to resolve secrets at runtime.
     /// </summary>
-    Func<string, CancellationToken, Task<string?>>? SecretResolver { get; set; }
+    Func<IKeyVaultSecretReference, CancellationToken, Task<string?>>? SecretResolver { get; set; }
 
     /// <summary>
     /// Gets a secret reference for the specified secret name.

src/Aspire.Hosting.Azure/Provisioning/Provisioners/BicepProvisioner.cs

@@ -284,9 +284,9 @@ await notificationService.PublishUpdateAsync(resource, state =>
 
             // Set the client for resolving secrets at runtime
             var client = new SecretClient(new(vaultUri), context.Credential);
-            kvr.SecretResolver = async (secretName, ct) =>
+            kvr.SecretResolver = async (secretRef, ct) =>
             {
-                var secret = await client.GetSecretAsync(secretName, cancellationToken: ct).ConfigureAwait(false);
+                var secret = await client.GetSecretAsync(secretRef.SecretName, cancellationToken: ct).ConfigureAwait(false);
                 return secret.Value.Value;
             };
         }

tests/Aspire.Hosting.Azure.Tests/AzureBicepResourceTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net.Sockets;
+using System.Runtime.CompilerServices;
 using System.Text.Json.Nodes;
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Lifecycle;
@@ -237,6 +238,24 @@ public async Task AddAzureCosmosDBEmulator()
         Assert.Equal(cs, await ((IResourceWithConnectionString)cosmos.Resource).GetConnectionStringAsync());
     }
 
+    [Fact]
+    public async Task AddAzureCosmosDB_WithAccessKeyAuthentication_NoKeyVaultWithEmulator()
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+
+        builder.AddAzureCosmosDB("cosmos").WithAccessKeyAuthentication().RunAsEmulator();
+
+#pragma warning disable ASPIRECOSMOSDB001
+        builder.AddAzureCosmosDB("cosmos2").WithAccessKeyAuthentication().RunAsPreviewEmulator();
+#pragma warning restore ASPIRECOSMOSDB001
+
+        var app = builder.Build();
+        var model = app.Services.GetRequiredService<DistributedApplicationModel>();
+        await ExecuteBeforeStartHooksAsync(app, CancellationToken.None);
+
+        Assert.Empty(model.Resources.OfType<AzureKeyVaultResource>());
+    }
+
     [Theory]
     [InlineData(null)]
     [InlineData("mykeyvault")]
@@ -264,16 +283,24 @@ public async Task AddAzureCosmosDBViaRunMode_WithAccessKeyAuthentication(string?
         var db = cosmos.AddCosmosDatabase("db", databaseName: "mydatabase");
         db.AddContainer("container", "mypartitionkeypath", containerName: "mycontainer");
 
-        var kv = builder.CreateResourceBuilder<AzureKeyVaultResource>(kvName);
+        var app = builder.Build();
+
+        await ExecuteBeforeStartHooksAsync(app, CancellationToken.None);
+
+        var model = app.Services.GetRequiredService<DistributedApplicationModel>();
+
+        var kv = model.Resources.OfType<AzureKeyVaultResource>().Single();
+
+        Assert.Equal(kvName, kv.Name);
 
         var secrets = new Dictionary<string, string>
         {
             ["connectionstrings--cosmos"] = "mycosmosconnectionstring"
         };
 
-        kv.Resource.SecretResolver = (name, _) =>
+        kv.SecretResolver = (secretRef, _) =>
         {
-            if (!secrets.TryGetValue(name, out var value))
+            if (!secrets.TryGetValue(secretRef.SecretName, out var value))
             {
                 return Task.FromResult<string?>(null);
             }
@@ -533,9 +560,9 @@ public async Task AddAzureCosmosDBViaPublishMode_WithAccessKeyAuthentication(str
             ["connectionstrings--cosmos"] = "mycosmosconnectionstring"
         };
 
-        kv.Resource.SecretResolver = (name, _) =>
+        kv.Resource.SecretResolver = (secretRef, _) =>
         {
-            if (!secrets.TryGetValue(name, out var value))
+            if (!secrets.TryGetValue(secretRef.SecretName, out var value))
             {
                 return Task.FromResult<string?>(null);
             }
@@ -3130,6 +3157,9 @@ public async Task InfrastructureCanBeMutatedAfterCreation()
         Assert.Equal(expectedBicep, bicep);
     }
 
+    [UnsafeAccessor(UnsafeAccessorKind.Method, Name = "ExecuteBeforeStartHooksAsync")]
+    private static extern Task ExecuteBeforeStartHooksAsync(DistributedApplication app, CancellationToken cancellationToken);
+
     private sealed class ProjectA : IProjectMetadata
     {
         public string ProjectPath => "projectA";

tests/Aspire.Hosting.Azure.Tests/AzureCosmosDBExtensionsTests.cs

@@ -106,16 +106,26 @@ public void AzureCosmosDBHasCorrectConnectionStrings_ForAccountEndpoint()
         Assert.Equal("AccountEndpoint={cosmos.outputs.connectionString};Database=db1;Container=container1", container1.Resource.ConnectionStringExpression.ValueExpression);
     }
 
-    [Fact]
-    public void AzureCosmosDBHasCorrectConnectionStrings()
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void AzureCosmosDBHasCorrectConnectionStrings(bool useAccessKeyAuth)
     {
         using var builder = TestDistributedApplicationBuilder.Create();
 
         var cosmos = builder.AddAzureCosmosDB("cosmos").RunAsEmulator();
+        if (useAccessKeyAuth)
+        {
+            cosmos.WithAccessKeyAuthentication();
+        }
         var db1 = cosmos.AddCosmosDatabase("db1");
         var container1 = db1.AddContainer("container1", "id");
 
         var cosmos1 = builder.AddAzureCosmosDB("cosmos1").RunAsEmulator();
+        if (useAccessKeyAuth)
+        {
+            cosmos1.WithAccessKeyAuthentication();
+        }
         var db2 = cosmos1.AddCosmosDatabase("db2", "db");
         var container2 = db2.AddContainer("container2", "id", "container");