Add support for host specialization with unencrypted payload

src/WebJobs.Script.WebHost/ContainerManagement/LinuxContainerInitializationHostedService.cs

@@ -48,8 +48,8 @@ private async Task ApplyStartContextIfPresent(CancellationToken cancellationToke
             {
                 _logger.LogDebug("Applying host context");
 
-                var encryptedAssignmentContext = JsonConvert.DeserializeObject<EncryptedHostAssignmentContext>(startContext);
-                var assignmentContext = _startupContextProvider.SetContext(encryptedAssignmentContext);
+                var hostAssignmentRequest = JsonConvert.DeserializeObject<HostAssignmentRequest>(startContext);
+                var assignmentContext = _startupContextProvider.SetContext(hostAssignmentRequest);
                 await SpecializeMSISideCar(assignmentContext);
 
                 try

src/WebJobs.Script.WebHost/Controllers/InstanceController.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System.Threading.Tasks;
@@ -7,6 +7,7 @@
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.Azure.WebJobs.Script.WebHost.Management;
 using Microsoft.Azure.WebJobs.Script.WebHost.Models;
+using Microsoft.Azure.WebJobs.Script.WebHost.Security.Authentication;
 using Microsoft.Azure.WebJobs.Script.WebHost.Security.Authorization.Policies;
 using Microsoft.Extensions.Logging;
 
@@ -35,11 +36,35 @@ public InstanceController(IEnvironment environment, IInstanceManager instanceMan
         [HttpPost]
         [Route("admin/instance/assign")]
         [Authorize(Policy = PolicyNames.AdminAuthLevel)]
-        public async Task<IActionResult> Assign([FromBody] EncryptedHostAssignmentContext encryptedAssignmentContext)
+        public async Task<IActionResult> Assign([FromBody] HostAssignmentRequest hostAssignmentRequest)
         {
-            _logger.LogDebug($"Starting container assignment for host : {Request?.Host}. ContextLength is: {encryptedAssignmentContext.EncryptedContext?.Length}");
+            if (string.IsNullOrEmpty(hostAssignmentRequest.EncryptedContext) &&
+                hostAssignmentRequest.AssignmentContext is null)
+            {
+                return BadRequest($"At least one of {nameof(HostAssignmentRequest.AssignmentContext)} or {nameof(HostAssignmentRequest.EncryptedContext)} must be provided.");
+            }
+
+            if (!string.IsNullOrEmpty(hostAssignmentRequest.EncryptedContext) &&
+                hostAssignmentRequest.AssignmentContext is not null)
+            {
+                return BadRequest($"Only one of {nameof(HostAssignmentRequest.AssignmentContext)} or {nameof(HostAssignmentRequest.EncryptedContext)} may be set.");
+            }
+
+            if (!string.IsNullOrEmpty(hostAssignmentRequest.EncryptedContext))
+            {
+                _logger.LogDebug("Starting container assignment. ContextLength is {ContextLength}", hostAssignmentRequest.EncryptedContext.Length);
+            }
+            else
+            {
+                if (!User.HasClaim(SecurityConstants.AssignUnencryptedClaimType, "true"))
+                {
+                    _logger.LogWarning("Required claims missing for invoking unencrypted assignment");
+                    return Forbid();
+                }
+                _logger.LogDebug("Starting container assignment.");
+            }
 
-            var assignmentContext = _startupContextProvider.SetContext(encryptedAssignmentContext);
+            var assignmentContext = _startupContextProvider.SetContext(hostAssignmentRequest);
 
             // before starting the assignment we want to perform as much
             // up front validation on the context as possible

src/WebJobs.Script.WebHost/Models/EncryptedHostAssignmentContext.cs → src/WebJobs.Script.WebHost/Models/HostAssignmentRequest.cs

@@ -5,9 +5,12 @@
 
 namespace Microsoft.Azure.WebJobs.Script.WebHost.Models
 {
-    public class EncryptedHostAssignmentContext
+    public class HostAssignmentRequest
     {
-        [JsonProperty("encryptedContext")]
+        [JsonProperty("encryptedContext", DefaultValueHandling = DefaultValueHandling.Ignore)]
         public string EncryptedContext { get; set; }
+
+        [JsonProperty("assignmentContext", DefaultValueHandling = DefaultValueHandling.Ignore)]
+        public HostAssignmentContext AssignmentContext { get; set; }
     }
 }

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

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -72,6 +72,11 @@ private static void ConfigureOptions(JwtBearerOptions options)
                         claims.Add(new Claim(SecurityConstants.InvokeClaimType, "true"));
                     }
 
+                    if (string.Equals(c.SecurityToken.Issuer, ScriptConstants.LegionCoreUri, StringComparison.OrdinalIgnoreCase))
+                    {
+                        claims.Add(new Claim(SecurityConstants.AssignUnencryptedClaimType, "true"));
+                    }
+
                     c.Principal.AddIdentity(new ClaimsIdentity(claims));
                     c.Success();
 
@@ -149,6 +154,7 @@ public static TokenValidationParameters CreateTokenValidationParameters()
                 result.ValidIssuers =
                 [
                     AppServiceCoreUri,
+                    LegionCoreUri,
                     string.Format(ScmSiteUriFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName)),
                     string.Format(SiteUriFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName))
                 ];

src/WebJobs.Script.WebHost/Security/Authentication/SecurityConstants.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -20,5 +20,10 @@ public class SecurityConstants
         /// /admin/functions/{function} API.
         /// </summary>
         public const string InvokeClaimType = "http://schemas.microsoft.com/2017/07/functions/claims/invoke";
+
+        /// <summary>
+        /// Claim indicating whether a principal is authorized to invoke assign with an unencrypted payload.
+        /// </summary>
+        public const string AssignUnencryptedClaimType = "http://schemas.microsoft.com/2017/07/functions/claims/assign-unencrypted";
     }
 }

src/WebJobs.Script.WebHost/StartupContextProvider.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -132,15 +132,20 @@ private StartupContext GetStartupContextOrNull()
         }
 
         /// <summary>
-        /// Decrypt and deserialize the specified context, and apply values from it to the
-        /// startup cache context.
+        /// Applies the values from the specified assignment request to the startup cache context,
+        /// performing any required decryption.
         /// </summary>
-        /// <param name="encryptedContext">The encrypted assignment context.</param>
-        /// <returns>The decrypted assignment context</returns>
-        public virtual HostAssignmentContext SetContext(EncryptedHostAssignmentContext encryptedContext)
+        /// <param name="hostAssignmentRequest">The Host assignment request.</param>
+        /// <returns>The assignment context applied.</returns>
+        public virtual HostAssignmentContext SetContext(HostAssignmentRequest hostAssignmentRequest)
         {
-            string decryptedContext = EncryptionHelper.Decrypt(encryptedContext.EncryptedContext, environment: _environment);
-            var hostAssignmentContext = JsonConvert.DeserializeObject<HostAssignmentContext>(decryptedContext);
+            var hostAssignmentContext = hostAssignmentRequest.AssignmentContext;
+
+            if (!string.IsNullOrEmpty(hostAssignmentRequest.EncryptedContext))
+            {
+                string decryptedContext = EncryptionHelper.Decrypt(hostAssignmentRequest.EncryptedContext, environment: _environment);
+                hostAssignmentContext = JsonConvert.DeserializeObject<HostAssignmentContext>(decryptedContext);
+            }
 
             // Don't update StartupContext for warmup requests
             if (!hostAssignmentContext.IsWarmupRequest)

src/WebJobs.Script/ScriptConstants.cs

@@ -151,6 +151,7 @@ public static class ScriptConstants
         public const string ScmSiteUriFormat = "https://{0}.scm.azurewebsites.net";
         public const string SiteUriFormat = "https://{0}.azurewebsites.net";
         public const string AppServiceCoreUri = "https://appservice.core.azurewebsites.net";
+        public const string LegionCoreUri = "https://legion.core.azurewebsites.net";
 
         public const string AzureFunctionsSystemDirectoryName = ".azurefunctions";
         public const string HttpMethodConstraintName = "httpMethod";

test/WebJobs.Script.Tests.Integration/Host/StandbyManager/StandbyManagerE2ETests_Linux.cs

@@ -70,8 +70,10 @@ public async Task StandbyModeE2E_LinuxConsumptionOnLegion()
             Assert.Equal(typeof(LinuxContainerLegionMetricsPublisher), webHost.Services.GetRequiredService<IMetricsPublisher>().GetType());
         }
 
-        [Fact]
-        public async Task StandbyModeE2E_LinuxContainer()
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public async Task StandbyModeE2E_LinuxContainer(bool useEncryptedPayload)
         {
             byte[] bytes = TestHelpers.GenerateKeyBytes();
             var encryptionKey = Convert.ToBase64String(bytes);
@@ -114,7 +116,7 @@ public async Task StandbyModeE2E_LinuxContainer()
             await VerifyWarmupSucceeds(restart: true);
 
             // now specialize the site
-            await Assign(encryptionKey);
+            await Assign(encryptionKey, useEncryptedPayload);
 
             // immediately call a function - expect the call to block until
             // the host is fully specialized
@@ -240,7 +242,7 @@ public async Task LinuxContainer_TimeZoneEnvVariableE2E()
 
         }
 
-        private async Task Assign(string encryptionKey)
+        private async Task Assign(string encryptionKey, bool useEncryptedPayload)
         {
             // create a zip package
             var testFunctionPath = Path.Combine("TestScripts", "Node", "HttpTrigger");
@@ -260,32 +262,36 @@ private async Task Assign(string encryptionKey)
             string uri = "admin/instance/assign";
             var request = new HttpRequestMessage(HttpMethod.Post, uri);
             var environment = new Dictionary<string, string>()
-                {
-                    { EnvironmentSettingNames.AzureWebsiteZipDeployment, sasUri.ToString() },
-                    { RpcWorkerConstants.FunctionWorkerRuntimeVersionSettingName, "~2" },
-                    { EnvironmentSettingNames.FunctionWorkerRuntime, "node" }
-                };
+            {
+                { EnvironmentSettingNames.AzureWebsiteZipDeployment, sasUri.ToString() },
+                { RpcWorkerConstants.FunctionWorkerRuntimeVersionSettingName, "~2" },
+                { EnvironmentSettingNames.FunctionWorkerRuntime, "node" }
+            };
             var assignmentContext = new HostAssignmentContext
             {
                 SiteId = 1234,
                 SiteName = "TestApp",
                 Environment = environment
             };
-            var encryptedAssignmentContext = CreateEncryptedContext(assignmentContext, encryptionKey);
+            var encryptedAssignmentContext = CreateHostAssignmentRequest(assignmentContext, encryptionKey, useEncryptedPayload);
             string json = JsonConvert.SerializeObject(encryptedAssignmentContext);
             request.Content = new StringContent(json, Encoding.UTF8, "application/json");
             request.Headers.Add(AuthenticationLevelHandler.FunctionsKeyHeaderName, masterKey);
             var response = await _httpClient.SendAsync(request);
             Assert.Equal(HttpStatusCode.Accepted, response.StatusCode);
         }
 
-        private static EncryptedHostAssignmentContext CreateEncryptedContext(HostAssignmentContext context, string key)
+        private static HostAssignmentRequest CreateHostAssignmentRequest(HostAssignmentContext context, string key, bool useEncryptedPayload)
         {
+            if (!useEncryptedPayload)
+            {
+                return new HostAssignmentRequest { AssignmentContext = context };
+            }
+
             string json = JsonConvert.SerializeObject(context);
             var encryptionKey = Convert.FromBase64String(key);
             string encrypted = EncryptionHelper.Encrypt(json, encryptionKey);
-
-            return new EncryptedHostAssignmentContext { EncryptedContext = encrypted };
+            return new HostAssignmentRequest { EncryptedContext = encrypted };
         }
     }
 }