Tactical fix to update JWT validation parameters after specialization

Author: fabiocav

src/WebJobs.Script.WebHost/Security/Authentication/Jwt/ScriptJwtBearerExtensions.cs

@@ -1,14 +1,17 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System;
 using System.Security.Claims;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Authentication;
 using Microsoft.AspNetCore.Authentication.JwtBearer;
 using Microsoft.Azure.Web.DataProtection;
 using Microsoft.Azure.WebJobs.Extensions.Http;
 using Microsoft.Azure.WebJobs.Script.Config;
+using Microsoft.Azure.WebJobs.Script.WebHost;
 using Microsoft.Azure.WebJobs.Script.WebHost.Security.Authentication;
 using Microsoft.IdentityModel.Tokens;
 using static Microsoft.Azure.WebJobs.Script.EnvironmentSettingNames;
@@ -18,11 +21,23 @@ namespace Microsoft.Extensions.DependencyInjection
 {
     public static class ScriptJwtBearerExtensions
     {
+        private static double _specialized = 0;
+
         public static AuthenticationBuilder AddScriptJwtBearer(this AuthenticationBuilder builder)
             => builder.AddJwtBearer(o =>
                         {
                             o.Events = new JwtBearerEvents()
                             {
+                                OnMessageReceived = c =>
+                                {
+                                    // Temporary: Tactical fix to address specialization issues. This should likely be moved to a token validator
+                                    if (_specialized == 0 && !WebScriptHostManager.InStandbyMode && Interlocked.CompareExchange(ref _specialized, 1, 0) == 0)
+                                    {
+                                        o.TokenValidationParameters = CreateTokenValidationParameters();
+                                    }
+
+                                    return Task.CompletedTask;
+                                },
                                 OnTokenValidated = c =>
                                 {
                                     c.Principal.AddIdentity(new ClaimsIdentity(new Claim[]
@@ -35,19 +50,32 @@ public static AuthenticationBuilder AddScriptJwtBearer(this AuthenticationBuilde
                                     return Task.CompletedTask;
                                 }
                             };
-                            string defaultKey = Util.GetDefaultKeyValue();
-                            if (defaultKey != null)
+
+                            o.TokenValidationParameters = CreateTokenValidationParameters();
+
+                            if (!WebScriptHostManager.InStandbyMode)
                             {
-                                // TODO: Once ScriptSettingsManager is gone, Audience and Issuer shouold be pulled from configuration.
-                                o.TokenValidationParameters = new TokenValidationParameters()
-                                {
-                                    IssuerSigningKey = new SymmetricSecurityKey(Encoding.UTF8.GetBytes(defaultKey)),
-                                    ValidateAudience = true,
-                                    ValidateIssuer = true,
-                                    ValidAudience = string.Format(AdminJwtValidAudienceFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName)),
-                                    ValidIssuer = string.Format(AdminJwtValidIssuerFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName))
-                                };
+                                // We're not in standby mode, so flag as specialized
+                                _specialized = 1;
                             }
                         });
+
+        private static TokenValidationParameters CreateTokenValidationParameters()
+        {
+            string defaultKey = Util.GetDefaultKeyValue();
+
+            var result = new TokenValidationParameters();
+            if (defaultKey != null)
+            {
+                // TODO: Once ScriptSettingsManager is gone, Audience and Issuer shouold be pulled from configuration.
+                result.IssuerSigningKey = new SymmetricSecurityKey(Encoding.UTF8.GetBytes(defaultKey));
+                result.ValidateAudience = true;
+                result.ValidateIssuer = true;
+                result.ValidAudience = string.Format(AdminJwtValidAudienceFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName));
+                result.ValidIssuer = string.Format(AdminJwtValidIssuerFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName));
+            }
+
+            return result;
+        }
     }
 }