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 hostAssignmentContext)
         {
-            _logger.LogDebug($"Starting container assignment for host : {Request?.Host}. ContextLength is: {encryptedAssignmentContext.EncryptedContext?.Length}");
+            if (string.IsNullOrEmpty(hostAssignmentContext.EncryptedContext) &&
+                hostAssignmentContext.AssignmentContext is null)
+            {
+                return BadRequest("At least one of 'assignmentContext' or 'encryptedContext' must be provided.");
+            }
+
+            if (!string.IsNullOrEmpty(hostAssignmentContext.EncryptedContext) &&
+                hostAssignmentContext.AssignmentContext is not null)
+            {
+                return BadRequest("Only one of 'assignmentContext' or 'encryptedContext' may be set.");
+            }
+
+            if (!string.IsNullOrEmpty(hostAssignmentContext.EncryptedContext))
+            {
+                _logger.LogDebug("Starting container assignment. ContextLength is {ContextLength}", hostAssignmentContext.EncryptedContext.Length);
+            }
+            else
+            {
+                if (!User.HasClaim(c => c.Type == SecurityConstants.AssignUnencryptedClaimType && c.Value == "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(hostAssignmentContext);
 
             // 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

@@ -71,6 +71,10 @@ 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 +153,7 @@ public static TokenValidationParameters CreateTokenValidationParameters()
                 result.ValidIssuers =
                 [
                     AppServiceCoreUri,
+                    LegionCoreUri, // ATOL: Specialization is invoked from Legion.
                     string.Format(ScmSiteUriFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName)),
                     string.Format(SiteUriFormat, ScriptSettingsManager.Instance.GetSetting(AzureWebsiteName))
                 ];

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

@@ -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 unencrypted payload.
+        /// </summary>
+        public const string AssignUnencryptedClaimType = "http://schemas.microsoft.com/2017/07/functions/claims/assign-unencrypted";
     }
 }

src/WebJobs.Script.WebHost/StartupContextProvider.cs

@@ -135,12 +135,17 @@ private StartupContext GetStartupContextOrNull()
         /// Decrypt and deserialize the specified context, and apply values from it to the
         /// startup cache context.
         /// </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 decrypted assignment context.</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 useEncyptedPayload)
         {
             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, useEncyptedPayload);
 
             // 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 };
         }
     }
 }

test/WebJobs.Script.Tests.Integration/Management/InstanceControllerTests.cs

@@ -1,24 +1,27 @@
-// 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;
-using System.Collections.Generic;
-using System.Net;
-using System.Net.Http;
-using System.Threading;
-using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.Azure.WebJobs.Script.WebHost;
 using Microsoft.Azure.WebJobs.Script.WebHost.Controllers;
 using Microsoft.Azure.WebJobs.Script.WebHost.Management;
 using Microsoft.Azure.WebJobs.Script.WebHost.Management.LinuxSpecialization;
 using Microsoft.Azure.WebJobs.Script.WebHost.Models;
 using Microsoft.Azure.WebJobs.Script.WebHost.Security;
+using Microsoft.Azure.WebJobs.Script.WebHost.Security.Authentication;
 using Microsoft.Extensions.Logging;
 using Microsoft.WebJobs.Script.Tests;
 using Moq;
 using Moq.Protected;
 using Newtonsoft.Json;
+using System;
+using System.Collections.Generic;
+using System.Net;
+using System.Net.Http;
+using System.Security.Claims;
+using System.Threading;
+using System.Threading.Tasks;
 using Xunit;
 
 namespace Microsoft.Azure.WebJobs.Script.Tests.Managment
@@ -80,14 +83,14 @@ public async Task Assign_MSISpecializationFailure_ReturnsError()
 
             var encryptedHostAssignmentValue = EncryptionHelper.Encrypt(JsonConvert.SerializeObject(hostAssignmentContext), TestHelpers.EncryptionKey.ToKeyBytes());
 
-            var encryptedHostAssignmentContext = new EncryptedHostAssignmentContext()
+            var hostAssignmentRequest = new HostAssignmentRequest()
             {
                 EncryptedContext = encryptedHostAssignmentValue
             };
 
             environment.SetEnvironmentVariable(EnvironmentSettingNames.ContainerEncryptionKey, TestHelpers.EncryptionKey);
 
-            IActionResult result = await instanceController.Assign(encryptedHostAssignmentContext);
+            IActionResult result = await instanceController.Assign(hostAssignmentRequest);
 
             var objectResult = result as ObjectResult;
 
@@ -156,14 +159,14 @@ public async Task Assignment_Sets_Secrets_Context()
 
             var encryptedHostAssignmentValue = EncryptionHelper.Encrypt(JsonConvert.SerializeObject(hostAssignmentContext), TestHelpers.EncryptionKey.ToKeyBytes());
 
-            var encryptedHostAssignmentContext = new EncryptedHostAssignmentContext()
+            var hostAssignmentRequest = new HostAssignmentRequest()
             {
                 EncryptedContext = encryptedHostAssignmentValue
             };
 
             environment.SetEnvironmentVariable(EnvironmentSettingNames.ContainerEncryptionKey, TestHelpers.EncryptionKey);
 
-            await instanceController.Assign(encryptedHostAssignmentContext);
+            await instanceController.Assign(hostAssignmentRequest);
             Assert.NotNull(startupContextProvider.Context);
         }
 
@@ -209,21 +212,23 @@ public async Task Assignment_Does_Not_Set_Secrets_Context_For_Warmup_Request()
 
             var encryptedHostAssignmentValue = EncryptionHelper.Encrypt(JsonConvert.SerializeObject(hostAssignmentContext), TestHelpers.EncryptionKey.ToKeyBytes());
 
-            var encryptedHostAssignmentContext = new EncryptedHostAssignmentContext()
+            var hostAssignmentRequest = new HostAssignmentRequest()
             {
                 EncryptedContext = encryptedHostAssignmentValue
             };
 
             environment.SetEnvironmentVariable(EnvironmentSettingNames.ContainerEncryptionKey, TestHelpers.EncryptionKey);
 
-            await instanceController.Assign(encryptedHostAssignmentContext);
+            await instanceController.Assign(hostAssignmentRequest);
             Assert.Null(startupContextProvider.Context);
         }
 
         [Theory]
-        [InlineData(true, true)]
-        [InlineData(false, true)]
-        public async Task Assignment_Invokes_InstanceManager_Methods_For_Warmup_Requests_Also(bool isWarmupRequest, bool shouldInvokeMethod)
+        [InlineData(true, true, true)]
+        [InlineData(false, true, true)]
+        [InlineData(true, true, false)]
+        [InlineData(false, true, false)]
+        public async Task Assignment_Invokes_InstanceManager_Methods_For_Warmup_Requests_Also(bool isWarmupRequest, bool shouldInvokeMethod, bool useEncryptedPayload)
         {
             var environment = new TestEnvironment();
             environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsitePlaceholderMode, "1");
@@ -239,6 +244,11 @@ public async Task Assignment_Invokes_InstanceManager_Methods_For_Warmup_Requests
 
             var instanceController = new InstanceController(environment, instanceManager.Object, loggerFactory,
                 startupContextProvider);
+            DefaultHttpContext context = new() { User = new ClaimsPrincipal(new ClaimsIdentity(new Claim[]
+            {
+                new Claim(SecurityConstants.AssignUnencryptedClaimType, "true")
+            })) };
+            instanceController.ControllerContext = new() { HttpContext = context };
 
             var hostAssignmentContext = new HostAssignmentContext
             {
@@ -250,14 +260,18 @@ public async Task Assignment_Invokes_InstanceManager_Methods_For_Warmup_Requests
                 EncryptionHelper.Encrypt(JsonConvert.SerializeObject(hostAssignmentContext),
                     TestHelpers.EncryptionKey.ToKeyBytes());
 
-            var encryptedHostAssignmentContext = new EncryptedHostAssignmentContext()
+            var hostAssignmentRequest = new HostAssignmentRequest() { };
+            if (useEncryptedPayload)
             {
-                EncryptedContext = encryptedHostAssignmentValue
-            };
-
+                hostAssignmentRequest.EncryptedContext = encryptedHostAssignmentValue;
+            }
+            else
+            {
+                hostAssignmentRequest.AssignmentContext = hostAssignmentContext;
+            }
             environment.SetEnvironmentVariable(EnvironmentSettingNames.ContainerEncryptionKey, TestHelpers.EncryptionKey);
 
-            await instanceController.Assign(encryptedHostAssignmentContext);
+            await instanceController.Assign(hostAssignmentRequest);
 
             instanceManager.Verify(i => i.ValidateContext(It.IsAny<HostAssignmentContext>()),
                 shouldInvokeMethod ? Times.Once() : Times.Never());
@@ -266,5 +280,36 @@ public async Task Assignment_Invokes_InstanceManager_Methods_For_Warmup_Requests
             instanceManager.Verify(i => i.StartAssignment(It.IsAny<HostAssignmentContext>()),
                 shouldInvokeMethod ? Times.Once() : Times.Never());
         }
+
+        [Fact]
+        public async Task Assignment_ErrorScenarios()
+        {
+            var environment = new TestEnvironment();
+            environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsitePlaceholderMode, "1");
+            var loggerFactory = new LoggerFactory();
+            var loggerProvider = new TestLoggerProvider();
+            loggerFactory.AddProvider(loggerProvider);
+            var instanceController = new InstanceController(environment, null, loggerFactory, null);
+
+            // Both encrypted and unencrypted context are null
+            var hostAssignmentRequest = new HostAssignmentRequest() { };
+            var result = await instanceController.Assign(hostAssignmentRequest);
+            var badRequestResult = result as BadRequestObjectResult;
+            Assert.NotNull(badRequestResult);
+            Assert.Equal(400, badRequestResult.StatusCode);
+            Assert.Equal("At least one of 'assignmentContext' or 'encryptedContext' must be provided.", badRequestResult.Value);
+
+            // Both encrypted and unencrypted context are set
+            hostAssignmentRequest = new HostAssignmentRequest()
+            {
+                EncryptedContext = "EncryptedContext",
+                AssignmentContext = new HostAssignmentContext()
+            };
+            result = await instanceController.Assign(hostAssignmentRequest);
+            badRequestResult = result as BadRequestObjectResult;
+            Assert.NotNull(badRequestResult);
+            Assert.Equal(400, badRequestResult.StatusCode);
+            Assert.Equal("Only one of 'assignmentContext' or 'encryptedContext' may be set.", badRequestResult.Value);
+        }
     }
 }