Support setting existing app identity on compute resources

src/Aspire.Hosting.Azure/AzureResourcePreparer.cs

@@ -29,11 +29,11 @@ public async Task BeforeStartAsync(DistributedApplicationModel appModel, Cancell
         }
 
         var options = provisioningOptions.Value;
-        if (!EnvironmentSupportsTargetedRoleAssignments(options))
+        if (!EnvironmentSupportsIdentitiesAndAssignments(options))
         {
-            // If the app infrastructure does not support targeted role assignments, then we need to ensure that
-            // there are no role assignment annotations in the app model because they won't be honored otherwise.
-            EnsureNoRoleAssignmentAnnotations(appModel);
+            // If the app infrastructure does not support targeted identities and role assignments, then we need to ensure that
+            // there are no identity or role assignment annotations in the app model because they won't be honored otherwise.
+            EnsureNoIdentityOrRoleAssignmentAnnotations(appModel);
         }
 
         await BuildRoleAssignmentAnnotations(appModel, azureResources, options, cancellationToken).ConfigureAwait(false);
@@ -78,29 +78,34 @@ public async Task BeforeStartAsync(DistributedApplicationModel appModel, Cancell
         return azureResources;
     }
 
-    private bool EnvironmentSupportsTargetedRoleAssignments(AzureProvisioningOptions options)
+    private bool EnvironmentSupportsIdentitiesAndAssignments(AzureProvisioningOptions options)
     {
         // run mode always supports targeted role assignments
         // publish mode only supports targeted role assignments if the environment supports it
         return executionContext.IsRunMode || options.SupportsTargetedRoleAssignments;
     }
 
-    private static void EnsureNoRoleAssignmentAnnotations(DistributedApplicationModel appModel)
+    private static void EnsureNoIdentityOrRoleAssignmentAnnotations(DistributedApplicationModel appModel)
     {
         foreach (var resource in appModel.Resources)
         {
             if (resource.HasAnnotationOfType<RoleAssignmentAnnotation>())
             {
                 throw new InvalidOperationException("The application model does not support role assignments. Ensure you are using an environment that supports role assignments, for example AddAzureContainerAppEnvironment.");
             }
+
+            if (resource.HasAnnotationOfType<AppIdentityAnnotation>())
+            {
+                throw new InvalidOperationException("The application model does not support using explicit managed identities. Ensure you are using an environment that supports managed identities, for example AddAzureContainerAppEnvironment.");
+            }
         }
     }
 
     private async Task BuildRoleAssignmentAnnotations(DistributedApplicationModel appModel, List<(IResource Resource, IAzureResource AzureResource)> azureResources, AzureProvisioningOptions options, CancellationToken cancellationToken)
     {
         var globalRoleAssignments = new Dictionary<AzureProvisioningResource, HashSet<RoleDefinition>>();
 
-        if (!EnvironmentSupportsTargetedRoleAssignments(options))
+        if (!EnvironmentSupportsIdentitiesAndAssignments(options))
         {
             // when the app infrastructure doesn't support targeted role assignments, just copy all the default role assignments to applied role assignments
             foreach (var resource in azureResources.Select(r => r.AzureResource).OfType<AzureProvisioningResource>())
@@ -180,11 +185,21 @@ private async Task BuildRoleAssignmentAnnotations(DistributedApplicationModel ap
 
                         if (resource != identityResource)
                         {
-                            // attach the identity resource to compute resource so it can be used by the compute environment
-                            resource.Annotations.Add(new AppIdentityAnnotation(identityResource));
+                            // Only add the AppIdentityAnnotation if the resource doesn't already have one
+                            if (!resource.TryGetLastAnnotation<AppIdentityAnnotation>(out var existingAppIdentityAnnotation) ||
+                                existingAppIdentityAnnotation.IdentityResource != identityResource)
+                            {
+                                resource.Annotations.Add(new AppIdentityAnnotation(identityResource));
+                            }
+
                             // add the identity resource to the resource collection so it can be provisioned
-                            appModel.Resources.Add(identityResource);
+                            // but only if it's not already there
+                            if (!appModel.Resources.Contains(identityResource))
+                            {
+                                appModel.Resources.Add(identityResource);
+                            }
                         }
+
                         foreach (var roleAssignmentResource in roleAssignmentResources)
                         {
                             appModel.Resources.Add(roleAssignmentResource);
@@ -239,12 +254,27 @@ private static (AzureUserAssignedIdentityResource IdentityResource, List<AzureBi
         Dictionary<AzureProvisioningResource, IEnumerable<RoleDefinition>> roleAssignments,
         DistributedApplicationExecutionContext executionContext)
     {
-        var identityResource = resource is AzureUserAssignedIdentityResource existingIdentityResource
-            ? existingIdentityResource
-            : new AzureUserAssignedIdentityResource($"{resource.Name}-identity")
+        AzureUserAssignedIdentityResource identityResource;
+
+        // If we're currently targeting an AzureUserAssignedIdentityResource, we can use it as the identity resource
+        // for the role assignments. If we are targeting a compute resource that has an AppIdentityAnnotation, we can
+        // use the identity resource from that annotation. Otherwise, create a new identity resource to use for role assignments.
+        if (resource is AzureUserAssignedIdentityResource existingIdentityResource)
+        {
+            identityResource = existingIdentityResource;
+        }
+        else if (resource.TryGetLastAnnotation<AppIdentityAnnotation>(out var appIdentityAnnotation) &&
+                appIdentityAnnotation.IdentityResource is AzureUserAssignedIdentityResource existingAppIdentity)
+        {
+            identityResource = existingAppIdentity;
+        }
+        else
+        {
+            identityResource = new AzureUserAssignedIdentityResource($"{resource.Name}-identity")
             {
                 ProvisioningBuildOptions = provisioningOptions.ProvisioningBuildOptions
             };
+        }
 
         var roleAssignmentResources = CreateRoleAssignmentsResources(provisioningOptions, resource, roleAssignments, identityResource, executionContext);
         return (identityResource, roleAssignmentResources);
@@ -287,7 +317,7 @@ private static void AddRoleAssignmentsInfrastructure(
         IEnumerable<RoleDefinition> roles,
         AzureUserAssignedIdentityResource appIdentityResource)
     {
-        
+
         var context = new AddRoleAssignmentsContext(
             infra,
             executionContext,

src/Aspire.Hosting.Azure/AzureUserAssignedIdentityExtensions.cs

@@ -1,3 +1,5 @@
+#pragma warning disable ASPIRECOMPUTE001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
@@ -41,4 +43,31 @@ public static IResourceBuilder<AzureUserAssignedIdentityResource> AddAzureUserAs
 
         return builder.AddResource(resource);
     }
+
+    /// <summary>
+    /// Attaches an existing <see cref="AzureUserAssignedIdentityResource"/> to a compute resource, 
+    /// setting it as the target identity for the builder.
+    /// </summary>
+    /// <param name="builder">The builder for the <see cref="IComputeResource"/> the identity will be associated with.</param>
+    /// <param name="identityResourceBuilder">The builder for the <see cref="AzureUserAssignedIdentityResource"/>.</param>
+    /// <returns>A reference to the <see cref="IResourceBuilder{IComputeResource}"/> builder.</returns>
+    /// <example>
+    /// <code>
+    /// var identity = builder.AddAzureUserAssignedIdentity("myIdentity");
+    /// var app = builder.AddProject("myApp")
+    ///     .WithAzureUserAssignedIdentity(identity);
+    /// </code>
+    /// </example>
+    public static IResourceBuilder<T> WithAzureUserAssignedIdentity<T>(
+        this IResourceBuilder<T> builder,
+        IResourceBuilder<AzureUserAssignedIdentityResource> identityResourceBuilder)
+        where T : IComputeResource
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+        ArgumentNullException.ThrowIfNull(identityResourceBuilder);
+
+        builder.WithAnnotation(new AppIdentityAnnotation(identityResourceBuilder.Resource));
+
+        return builder;
+    }
 }