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;
@@ -35,11 +35,28 @@ 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 == null)
+            {
+                return BadRequest("Atleast one of Assignment context and EncryptedContext needs to be set.");
+            }
+            if (!string.IsNullOrEmpty(hostAssignmentContext.EncryptedContext) &&
+                !(hostAssignmentContext.AssignmentContext == null))
+            {
+                return BadRequest("Only one of Assignment context and EncryptedContext needs to be set.");
+            }
+            if (!string.IsNullOrEmpty(hostAssignmentContext.EncryptedContext))
+            {
+                _logger.LogDebug($"Starting container assignment for host : {Request?.Host}. ContextLength is {hostAssignmentContext.EncryptedContext.Length}");
+            }
+            else
+            {
+                _logger.LogDebug($"Starting container assignment for host : {Request?.Host}. Using unencrypted assignment context");
+            }
 
-            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

@@ -149,6 +149,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/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" },
-                    { RpcWorkerConstants.FunctionWorkerRuntimeSettingName, "node" }
-                };
+            {
+                { EnvironmentSettingNames.AzureWebsiteZipDeployment, sasUri.ToString() },
+                { RpcWorkerConstants.FunctionWorkerRuntimeVersionSettingName, "~2" },
+                { RpcWorkerConstants.FunctionWorkerRuntimeSettingName, "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)
         {
             string json = JsonConvert.SerializeObject(context);
             var encryptionKey = Convert.FromBase64String(key);
             string encrypted = EncryptionHelper.Encrypt(json, encryptionKey);
-
-            return new EncryptedHostAssignmentContext { EncryptedContext = encrypted };
+
+            if (useEncryptedPayload)
+            {
+                return new HostAssignmentRequest { EncryptedContext = encrypted };
+            }
+            return new HostAssignmentRequest { AssignmentContext = context };
         }
     }
 }

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

@@ -80,14 +80,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 +156,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 +209,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");
@@ -250,14 +252,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 +272,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("Atleast one of Assignment context and EncryptedContext needs to be set.", 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 Assignment context and EncryptedContext needs to be set.", badRequestResult.Value);
+        }
     }
 }

test/WebJobs.Script.Tests/ContainerManagment/LinuxContainerInitializationHostedServiceTests.cs

@@ -37,7 +37,7 @@ public async Task AssignInstanceAsyncIsAwaitedTest()
             _mockInstanceManager.Setup(m => m.AssignInstanceAsync(It.IsAny<HostAssignmentContext>())).Returns(tcs.Task);
 
             var assignmentContext = new HostAssignmentContext();
-            _mockStartupContextProvider.Setup(p => p.SetContext(It.IsAny<EncryptedHostAssignmentContext>())).Returns(assignmentContext);
+            _mockStartupContextProvider.Setup(p => p.SetContext(It.IsAny<HostAssignmentRequest>())).Returns(assignmentContext);
 
             var service = new TestLinuxContainerInitializationHostedService(GetTestEnvironment(), _mockInstanceManager.Object,
                 _mockLogger.Object, _mockStartupContextProvider.Object);
@@ -65,9 +65,9 @@ public TestLinuxContainerInitializationHostedService(IEnvironment environment, I
 
             protected override Task<(bool HasStartContext, string StartContext)> TryGetStartContextOrNullAsync(CancellationToken cancellationToken)
             {
-                var encryptedAssignmentContext = new EncryptedHostAssignmentContext { EncryptedContext = "test" };
+                var hostAssignmentRequest = new HostAssignmentRequest { EncryptedContext = "test" };
 
-                return Task.FromResult((true, JsonConvert.SerializeObject(encryptedAssignmentContext)));
+                return Task.FromResult((true, JsonConvert.SerializeObject(hostAssignmentRequest)));
             }
 
             protected override Task SpecializeMSISideCar(HostAssignmentContext assignmentContext)

test/WebJobs.Script.Tests/StartupContextProviderTests.cs

@@ -188,7 +188,7 @@ public void SetContext_AppliesHostAssignmentContext()
             };
             string json = JsonConvert.SerializeObject(context);
             string encrypted = EncryptionHelper.Encrypt(json, environment: _environment);
-            var encryptedContext = new EncryptedHostAssignmentContext { EncryptedContext = encrypted };
+            var encryptedContext = new HostAssignmentRequest { EncryptedContext = encrypted };
 
             var result = _startupContextProvider.SetContext(encryptedContext);
             Assert.Equal(context.SiteName, result.SiteName);
@@ -212,9 +212,9 @@ public void Does_Not_SetContext_AppliesHostAssignmentContext_For_Warmup_Request(
             };
             string json = JsonConvert.SerializeObject(context);
             string encrypted = EncryptionHelper.Encrypt(json, environment: _environment);
-            var encryptedContext = new EncryptedHostAssignmentContext { EncryptedContext = encrypted };
+            var hostAssignmentRequest = new HostAssignmentRequest { EncryptedContext = encrypted };
 
-            var result = _startupContextProvider.SetContext(encryptedContext);
+            var result = _startupContextProvider.SetContext(hostAssignmentRequest);
             Assert.Equal(context.SiteName, result.SiteName);
             Assert.Equal(_secrets.Host.Master, result.Secrets.Host.Master);