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,12 +1,14 @@
-// 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;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
+using Microsoft.Azure.WebJobs.Script.Diagnostics;
 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;
 
@@ -21,47 +23,81 @@ public class InstanceController : Controller
     {
         private readonly IEnvironment _environment;
         private readonly IInstanceManager _instanceManager;
+        private readonly IMetricsLogger _metricsLogger;
         private readonly ILogger _logger;
         private readonly StartupContextProvider _startupContextProvider;
 
-        public InstanceController(IEnvironment environment, IInstanceManager instanceManager, ILoggerFactory loggerFactory, StartupContextProvider startupContextProvider)
+        public InstanceController(IEnvironment environment, IInstanceManager instanceManager, ILoggerFactory loggerFactory, StartupContextProvider startupContextProvider, IMetricsLogger metricsLogger)
         {
             _environment = environment;
             _instanceManager = instanceManager;
             _logger = loggerFactory.CreateLogger<InstanceController>();
             _startupContextProvider = startupContextProvider;
+            _metricsLogger = metricsLogger;
         }
 
         [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}");
+            using (_metricsLogger.LatencyEvent(MetricEventNames.LinuxContainerSpecializationMSIInit))
+            {
+                if (hostAssignmentRequest == null)
+                {
+                    return BadRequest($"{nameof(hostAssignmentRequest)} cannot be null.");
+                }
 
-            var assignmentContext = _startupContextProvider.SetContext(encryptedAssignmentContext);
+                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.");
+                }
 
-            // before starting the assignment we want to perform as much
-            // up front validation on the context as possible
-            string error = await _instanceManager.ValidateContext(assignmentContext);
-            if (error != null)
-            {
-                return StatusCode(StatusCodes.Status400BadRequest, error);
-            }
+                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.");
+                }
 
-            // Wait for Sidecar specialization to complete before returning ok.
-            // This shouldn't take too long so ok to do this sequentially.
-            error = await _instanceManager.SpecializeMSISidecar(assignmentContext);
-            if (error != null)
-            {
-                return StatusCode(StatusCodes.Status500InternalServerError, error);
-            }
+                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 succeeded = _instanceManager.StartAssignment(assignmentContext);
+                var assignmentContext = _startupContextProvider.SetContext(hostAssignmentRequest);
 
-            return succeeded
-                ? Accepted()
-                : StatusCode(StatusCodes.Status409Conflict, "Instance already assigned");
+                // before starting the assignment we want to perform as much
+                // up front validation on the context as possible
+                string error = await _instanceManager.ValidateContext(assignmentContext);
+                if (error != null)
+                {
+                    return StatusCode(StatusCodes.Status400BadRequest, error);
+                }
+
+                // Wait for Sidecar specialization to complete before returning ok.
+                // This shouldn't take too long so ok to do this sequentially.
+                error = await _instanceManager.SpecializeMSISidecar(assignmentContext);
+                if (error != null)
+                {
+                    return StatusCode(StatusCodes.Status500InternalServerError, error);
+                }
+
+                var succeeded = _instanceManager.StartAssignment(assignmentContext);
+
+                return succeeded
+                    ? Accepted()
+                    : StatusCode(StatusCodes.Status409Conflict, "Instance already assigned");
+            }
         }
 
         [HttpGet]

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

@@ -0,0 +1,16 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using Newtonsoft.Json;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Models
+{
+    public sealed class HostAssignmentRequest
+    {
+        [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/Diagnostics/MetricEventNames.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.
 
 namespace Microsoft.Azure.WebJobs.Script.Diagnostics
@@ -83,6 +83,7 @@ public static class MetricEventNames
         public const string SecretManagerPurgeOldSecrets = "secretmanager.purgeoldsecrets.{0}";
 
         // Linux container specialization events
+        public const string LinuxContainerSpecializationAssign = "linux.container.specialization.assign";
         public const string LinuxContainerSpecializationBindMount = "linux.container.specialization.bind.mount";
         public const string LinuxContainerSpecializationMountCifs = "linux.container.specialization.mount.cifs";
         public const string LinuxContainerSpecializationZipExtract = "linux.container.specialization.zip.extract";

src/WebJobs.Script/ScriptConstants.cs

@@ -152,6 +152,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 };
         }
     }
 }