Support managed identity for Event Grid

src/WebJobs.Extensions.DurableTask/DurableTaskExtension.cs

@@ -508,10 +508,17 @@ private void TraceConfigurationSettings()
         private ILifeCycleNotificationHelper CreateLifeCycleNotificationHelper()
         {
             // First: EventGrid
-            if (this.Options.Notifications.EventGrid != null
-                && (!string.IsNullOrEmpty(this.Options.Notifications.EventGrid.TopicEndpoint) || !string.IsNullOrEmpty(this.Options.Notifications.EventGrid.KeySettingName)))
+
+            EventGridNotificationOptions eventGridOptions = this.Options.Notifications.EventGrid;
+            bool topicKeySettingOrKeySettingNameConfigured = eventGridOptions != null
+                && (!string.IsNullOrEmpty(eventGridOptions.TopicEndpoint)
+                || !string.IsNullOrEmpty(eventGridOptions.KeySettingName));
+            bool usingManagedIdentity = !string.IsNullOrEmpty(this.nameResolver.Resolve(EventGridLifeCycleNotificationHelper.TopicEndpointKey));
+
+            if (topicKeySettingOrKeySettingNameConfigured || usingManagedIdentity)
             {
-                return new EventGridLifeCycleNotificationHelper(this.Options, this.nameResolver, this.TraceHelper);
+                var notificationHelper = new EventGridLifeCycleNotificationHelper(this.Options, this.nameResolver, this.TraceHelper);
+                return notificationHelper;
             }
 
             // Fallback: Disable Notification
@@ -1325,6 +1332,11 @@ internal async Task<bool> StartTaskHubWorkerIfNotStartedAsync()
                             this.taskHubWorker.TaskOrchestrationDispatcher.IncludeDetails = true;
                         }
 
+                        if (this.LifeCycleNotificationHelper is EventGridLifeCycleNotificationHelper lifeCycleNotificationHelper)
+                        {
+                            await lifeCycleNotificationHelper.SetUpAuthenticationAsync();
+                        }
+
                         this.isTaskHubWorkerStarted = true;
                         return true;
                     }

src/WebJobs.Extensions.DurableTask/EventGridLifeCycleNotificationHelper.cs

@@ -17,6 +17,11 @@ namespace Microsoft.Azure.WebJobs.Extensions.DurableTask
 {
     internal class EventGridLifeCycleNotificationHelper : ILifeCycleNotificationHelper
     {
+        internal const string TopicEndpointKey = "EventGrid:topicEndpoint";
+        private const string CredentialKey = "EventGrid:credential";
+        private const string ClientIdKey = "EventGrid:clientId";
+        private const string ManagedIdentityValue = "managedidentity";
+
         private readonly DurableTaskOptions options;
         private readonly EndToEndTraceHelper traceHelper;
         private readonly bool useTrace;
@@ -37,96 +42,141 @@ public EventGridLifeCycleNotificationHelper(
             if (nameResolver == null)
             {
                 throw new ArgumentNullException(nameof(nameResolver));
-            }
-
-            if (options.Notifications == null)
-            {
-                throw new ArgumentNullException(nameof(options.Notifications));
-            }
-
-            var eventGridNotificationsConfig = options.Notifications.EventGrid ?? throw new ArgumentNullException(nameof(options.Notifications.EventGrid));
+            }
+
+            EventGridNotificationOptions eventGridNotificationsConfig = null;
 
-            this.eventGridKeyValue = nameResolver.Resolve(eventGridNotificationsConfig.KeySettingName);
-            this.eventGridTopicEndpoint = eventGridNotificationsConfig.TopicEndpoint;
+            this.UseManagedIdentity = false;
 
-            if (nameResolver.TryResolveWholeString(eventGridNotificationsConfig.TopicEndpoint, out var endpoint))
-            {
-                this.eventGridTopicEndpoint = endpoint;
-            }
+            // Check to see if we have a topic name app setting defined. If so, we will use managed identity to authenticate.
+            string topicEndpoint = nameResolver.Resolve(TopicEndpointKey);
+            if (!string.IsNullOrEmpty(topicEndpoint))
+            {
+                this.UseManagedIdentity = true;
+                this.eventGridTopicEndpoint = topicEndpoint;
+
+                string clientId = nameResolver.Resolve(ClientIdKey);
+                if (string.Equals(nameResolver.Resolve(CredentialKey), ManagedIdentityValue) &&
+                    !string.IsNullOrEmpty(clientId))
+                {
+                    // Use user assigned managed identity
+                    this.ManagedIdentityTokenSource = new ManagedIdentityTokenSource("https://eventgrid.azure.net/.default", new ManagedIdentityOptions(null, null, clientId));
+                }
+                else
+                {
+                    // Use system assigned managed identity
+                    this.ManagedIdentityTokenSource = new ManagedIdentityTokenSource("https://eventgrid.azure.net/.default");
+                }
+            }
+            else
+            {
+                // Set up configuration for key based authentication
+                eventGridNotificationsConfig = options.Notifications.EventGrid ?? throw new ArgumentNullException(nameof(options.Notifications.EventGrid));
+
+                this.eventGridKeyValue = nameResolver.Resolve(eventGridNotificationsConfig.KeySettingName);
+
+                this.eventGridTopicEndpoint = eventGridNotificationsConfig.TopicEndpoint;
+
+                if (nameResolver.TryResolveWholeString(eventGridNotificationsConfig.TopicEndpoint, out var endpoint))
+                {
+                    this.eventGridTopicEndpoint = endpoint;
+                }
+            }
+
+            // Log warning if both managed identity and key based authentication are configured.
+            if (this.UseManagedIdentity && options.Notifications != null && options.Notifications.EventGrid != null && (!string.IsNullOrEmpty(options.Notifications.EventGrid.TopicEndpoint) || !string.IsNullOrEmpty(options.Notifications.EventGrid.KeySettingName)))
+            {
+                this.traceHelper.ExtensionWarningEvent(
+                        hubName: this.options.HubName,
+                        functionName: "",
+                        instanceId: "",
+                        "Both managed identity and key based authentication are configured for Event Grid notifications. Managed Identity will be used for authentication. Please configure either managed identity or key based authentication for best results.");
+            }
+
+            // Check if we have the minimum required settings to enable Event Grid notifications with key based authentication.
+            bool eventGridNotificationSettingsConfigured = false;
+
+            if (eventGridNotificationsConfig != null)
+            {
+                if (!string.IsNullOrEmpty(eventGridNotificationsConfig.KeySettingName) && string.IsNullOrEmpty(this.eventGridKeyValue))
+                {
+                    throw new ArgumentException($"Failed to start lifecycle notification feature. Please check the configuration values for {eventGridNotificationsConfig.KeySettingName} on AppSettings.");
+                }
+
+                if (string.IsNullOrEmpty(eventGridNotificationsConfig.KeySettingName) || string.IsNullOrEmpty(eventGridNotificationsConfig.TopicEndpoint))
+                {
+                    throw new ArgumentException($"Failed to start lifecycle notification feature. Please check the configuration values for {eventGridNotificationsConfig.TopicEndpoint} and {eventGridNotificationsConfig.KeySettingName}.");
+                }
+
+                eventGridNotificationSettingsConfigured = true;
+            }
 
-            if (!string.IsNullOrEmpty(this.eventGridTopicEndpoint))
-            {
-                if (!string.IsNullOrEmpty(eventGridNotificationsConfig.KeySettingName))
-                {
-                    this.useTrace = true;
+            if (this.UseManagedIdentity || eventGridNotificationSettingsConfigured)
+            {
+                this.useTrace = true;
 
-                    var retryStatusCode = eventGridNotificationsConfig.PublishRetryHttpStatus?
-                                              .Where(x => Enum.IsDefined(typeof(HttpStatusCode), x))
-                                              .Select(x => (HttpStatusCode)x)
-                                              .ToArray()
-                                          ?? Array.Empty<HttpStatusCode>();
-
-                    if (eventGridNotificationsConfig.PublishEventTypes == null || eventGridNotificationsConfig.PublishEventTypes.Length == 0)
+                var retryStatusCode = (eventGridNotificationsConfig != null && eventGridNotificationsConfig.PublishRetryHttpStatus != null) ?
+                                        eventGridNotificationsConfig.PublishRetryHttpStatus?
+                                            .Where(x => Enum.IsDefined(typeof(HttpStatusCode), x))
+                                            .Select(x => (HttpStatusCode)x)
+                                            .ToArray()
+                                        : Array.Empty<HttpStatusCode>();
+
+                if (eventGridNotificationsConfig == null || eventGridNotificationsConfig.PublishEventTypes == null || eventGridNotificationsConfig.PublishEventTypes.Length == 0)
+                {
+                    this.eventGridPublishEventTypes = (OrchestrationRuntimeStatus[])Enum.GetValues(typeof(OrchestrationRuntimeStatus));
+                }
+                else
+                {
+                    var startedIndex = Array.FindIndex(eventGridNotificationsConfig.PublishEventTypes, x => x == "Started");
+                    if (startedIndex > -1)
                     {
-                        this.eventGridPublishEventTypes = (OrchestrationRuntimeStatus[])Enum.GetValues(typeof(OrchestrationRuntimeStatus));
+                        eventGridNotificationsConfig.PublishEventTypes[startedIndex] = OrchestrationRuntimeStatus.Running.ToString();
                     }
-                    else
-                    {
-                        var startedIndex = Array.FindIndex(eventGridNotificationsConfig.PublishEventTypes, x => x == "Started");
-                        if (startedIndex > -1)
-                        {
-                            eventGridNotificationsConfig.PublishEventTypes[startedIndex] = OrchestrationRuntimeStatus.Running.ToString();
-                        }
 
-                        OrchestrationRuntimeStatus ParseAndvalidateEvents(string @event)
+                    OrchestrationRuntimeStatus ParseAndvalidateEvents(string @event)
+                    {
+                        var success = Enum.TryParse(@event, out OrchestrationRuntimeStatus @enum);
+                        if (success)
                         {
-                            var success = Enum.TryParse(@event, out OrchestrationRuntimeStatus @enum);
-                            if (success)
-                            {
-                                switch (@enum)
-                                {
-                                    case OrchestrationRuntimeStatus.Canceled:
-                                    case OrchestrationRuntimeStatus.ContinuedAsNew:
-                                    case OrchestrationRuntimeStatus.Pending:
-                                        success = false;
-                                        break;
-                                    default:
-                                        break;
-                                }
-                            }
-
-                            if (!success)
+                            switch (@enum)
                             {
-                                throw new ArgumentException("Failed to start lifecycle notification feature. Unsupported event types detected in 'EventGridPublishEventTypes'. You may only specify one or more of the following 'Started', 'Completed', 'Failed', 'Terminated'.");
+                                case OrchestrationRuntimeStatus.Canceled:
+                                case OrchestrationRuntimeStatus.ContinuedAsNew:
+                                case OrchestrationRuntimeStatus.Pending:
+                                    success = false;
+                                    break;
+                                default:
+                                    break;
                             }
+                        }
 
-                            return @enum;
+                        if (!success)
+                        {
+                            throw new ArgumentException("Failed to start lifecycle notification feature. Unsupported event types detected in 'EventGridPublishEventTypes'. You may only specify one or more of the following 'Started', 'Completed', 'Failed', 'Terminated'.");
                         }
 
-                        this.eventGridPublishEventTypes = eventGridNotificationsConfig.PublishEventTypes.Select(x => ParseAndvalidateEvents(x)).ToArray();
-                    }
-
-                    // Currently, we support Event Grid Custom Topic for notify the lifecycle event of an orchestrator.
-                    // For more detail about the Event Grid, please refer this document.
-                    // Post to custom topic for Azure Event Grid
-                    // https://docs.microsoft.com/en-us/azure/event-grid/post-to-custom-topic
-                    this.HttpMessageHandler = options.NotificationHandler ?? new HttpRetryMessageHandler(
-                        new HttpClientHandler(),
-                        eventGridNotificationsConfig.PublishRetryCount,
-                        eventGridNotificationsConfig.PublishRetryInterval,
-                        retryStatusCode);
-
-                    if (string.IsNullOrEmpty(this.eventGridKeyValue))
-                    {
-                        throw new ArgumentException($"Failed to start lifecycle notification feature. Please check the configuration values for {eventGridNotificationsConfig.KeySettingName} on AppSettings.");
-                    }
-                }
-                else
-                {
-                    throw new ArgumentException($"Failed to start lifecycle notification feature. Please check the configuration values for {eventGridNotificationsConfig.TopicEndpoint} and {eventGridNotificationsConfig.KeySettingName}.");
+                        return @enum;
+                    }
+
+                    this.eventGridPublishEventTypes = eventGridNotificationsConfig.PublishEventTypes.Select(x => ParseAndvalidateEvents(x)).ToArray();
                 }
-            }
-        }
+
+                // Currently, we support Event Grid Custom Topic for notify the lifecycle event of an orchestrator.
+                // For more detail about the Event Grid, please refer this document.
+                // Post to custom topic for Azure Event Grid
+                // https://docs.microsoft.com/en-us/azure/event-grid/post-to-custom-topic
+                this.HttpMessageHandler = options.NotificationHandler ?? new HttpRetryMessageHandler(
+                    new HttpClientHandler(),
+                    eventGridNotificationsConfig != null ? eventGridNotificationsConfig.PublishRetryCount : 0,
+                    eventGridNotificationsConfig != null ? eventGridNotificationsConfig.PublishRetryInterval : TimeSpan.FromMinutes(5),
+                    retryStatusCode);
+            }
+        }
+
+        internal bool UseManagedIdentity { get; private set; }
+
+        internal ManagedIdentityTokenSource ManagedIdentityTokenSource { get; private set; }
 
         public string EventGridKeyValue => this.eventGridKeyValue;
 
@@ -139,11 +189,33 @@ public HttpMessageHandler HttpMessageHandler
             {
                 httpClient?.Dispose();
                 httpMessageHandler = value;
-                httpClient = new HttpClient(httpMessageHandler);
-                httpClient.DefaultRequestHeaders.Add("aeg-sas-key", this.eventGridKeyValue);
+                httpClient = new HttpClient(httpMessageHandler);
             }
         }
 
+        internal async Task SetUpAuthenticationAsync()
+        {
+            if (this.UseManagedIdentity)
+            {
+                // Use Bearer token for Managed Identity
+                await this.RefreshAccessTokenAsync();
+            }
+            else
+            {
+                // Use key-based authentication
+                if (httpClient != null)
+                {
+                    httpClient.DefaultRequestHeaders.Add("aeg-sas-key", this.eventGridKeyValue);
+                }
+            }
+        }
+
+        private async Task RefreshAccessTokenAsync()
+        {
+            string accessToken = await this.ManagedIdentityTokenSource.GetTokenAsync();
+            httpClient.DefaultRequestHeaders.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", accessToken);
+        }
+
         private async Task SendNotificationAsync(
             EventGridEvent[] eventGridEventArray,
             string hubName,
@@ -161,6 +233,11 @@ private async Task SendNotificationAsync(
             HttpResponseMessage result = null;
             try
             {
+                if (this.UseManagedIdentity)
+                {
+                    await this.RefreshAccessTokenAsync();
+                }
+
                 result = await httpClient.PostAsync(this.eventGridTopicEndpoint, content);
             }
             catch (Exception e)

src/WebJobs.Extensions.DurableTask/ManagedIdentityOptions.cs

@@ -11,15 +11,26 @@ namespace Microsoft.Azure.WebJobs.Extensions.DurableTask
     /// </summary>
     public class ManagedIdentityOptions
     {
-       /// <summary>
+        /// <summary>
         /// Initializes a new instance of the <see cref="ManagedIdentityOptions"/> class.
         /// </summary>
         /// <param name="authorityHost">The host of the Azure Active Directory authority.</param>
         /// <param name="tenantId">The tenant id of the user to authenticate.</param>
         public ManagedIdentityOptions(Uri authorityHost = null, string tenantId = null)
+            : this(authorityHost, tenantId, null)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ManagedIdentityOptions"/> class.
+        /// </summary>
+        /// <param name="clientId">The client id of the user assigned managed identity.</param>
+        [JsonConstructor]
+        public ManagedIdentityOptions(Uri authorityHost, string tenantId, string clientId)
         {
             this.AuthorityHost = authorityHost;
             this.TenantId = tenantId;
+            this.ClientId = clientId;
         }
 
         /// <summary>
@@ -33,5 +44,11 @@ public ManagedIdentityOptions(Uri authorityHost = null, string tenantId = null)
         /// </summary>
         [JsonProperty("tenantid")]
         public string TenantId { get; set; }
+
+        /// <summary>
+        /// The client id of the user assigned managed identity.
+        /// </summary>
+        [JsonProperty("clientid")]
+        public string ClientId { get; set; }
     }
 }

src/WebJobs.Extensions.DurableTask/ManagedIdentityTokenSource.cs

@@ -70,6 +70,11 @@ public async Task<string> GetTokenAsync()
                 defaultAzureCredentialOptions.InteractiveBrowserTenantId = this.Options.TenantId;
             }
 
+            if (!string.IsNullOrEmpty(this.Options?.ClientId))
+            {
+                defaultAzureCredentialOptions.ManagedIdentityClientId = this.Options.ClientId;
+            }
+
             defaultCredential = this.Options == null ? new DefaultAzureCredential() : new DefaultAzureCredential(defaultAzureCredentialOptions); // CodeQL [SM05137] Use DefaultAzureCredential explicitly for local development and is decided by the user
 
             AccessToken defaultToken = await defaultCredential.GetTokenAsync(context);