Supporting new platform specialization payload (#9302)

Co-authored-by: anandagopal <[email protected]>

Author: anandagopal6

src/WebJobs.Script.WebHost/Management/AtlasInstanceManager.cs

@@ -15,6 +15,7 @@
 using Microsoft.Azure.WebJobs.Script.WebHost.Models;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
+using Microsoft.IdentityModel.Tokens;
 using Newtonsoft.Json;
 
 namespace Microsoft.Azure.WebJobs.Script.WebHost.Management
@@ -63,25 +64,36 @@ public override async Task<string> SpecializeMSISidecar(HostAssignmentContext co
 
             if (msiEnabled)
             {
-                if (context.MSIContext == null)
+                if (context.MSIContext == null && context.EncryptedTokenServiceSpecializationPayload == null)
                 {
-                    _logger.LogWarning("Skipping specialization of MSI sidecar since MSIContext was absent");
+                    _logger.LogWarning("Skipping specialization of MSI sidecar since MSIContext and EncryptedTokenServiceSpecializationPayload were absent");
                     await _meshServiceClient.NotifyHealthEvent(ContainerHealthEventType.Fatal, this.GetType(),
-                        "Could not specialize MSI sidecar since MSIContext was empty");
+                        "Could not specialize MSI sidecar since MSIContext and EncryptedTokenServiceSpecializationPayload were empty");
                 }
                 else
                 {
                     using (_metricsLogger.LatencyEvent(MetricEventNames.LinuxContainerSpecializationMSIInit))
                     {
                         var uri = new Uri(endpoint);
-                        var address = $"http://{uri.Host}:{uri.Port}{ScriptConstants.LinuxMSISpecializationStem}";
+                        var addressStem = GetMsiSpecializationRequestAddressStem(context);
 
+                        var address = $"http://{uri.Host}:{uri.Port}{addressStem}";
                         _logger.LogDebug($"Specializing sidecar at {address}");
 
+                        StringContent payload;
+                        if (string.IsNullOrEmpty(context.EncryptedTokenServiceSpecializationPayload))
+                        {
+                            payload = new StringContent(JsonConvert.SerializeObject(context.MSIContext),
+                                    Encoding.UTF8, "application/json");
+                        }
+                        else
+                        {
+                            payload = new StringContent(context.EncryptedTokenServiceSpecializationPayload, Encoding.UTF8);
+                        }
+
                         var requestMessage = new HttpRequestMessage(HttpMethod.Post, address)
                         {
-                            Content = new StringContent(JsonConvert.SerializeObject(context.MSIContext),
-                                Encoding.UTF8, "application/json")
+                            Content = payload
                         };
 
                         var response = await _client.SendAsync(requestMessage);
@@ -364,5 +376,22 @@ await Utility.InvokeWithRetriesAsync(async () =>
                 return false;
             }
         }
+
+        private string GetMsiSpecializationRequestAddressStem(HostAssignmentContext context)
+        {
+            var stem = ScriptConstants.LinuxMSISpecializationStem;
+
+            if (!string.IsNullOrEmpty(context.EncryptedTokenServiceSpecializationPayload))
+            {
+                _logger.LogDebug("Using encrypted TokenService payload format");
+
+                // use default encrypted API endpoint if endpoint not provided in context
+                stem = string.IsNullOrEmpty(context.TokenServiceApiEndpoint)
+                    ? ScriptConstants.LinuxEncryptedTokenServiceSpecializationStem
+                    : context.TokenServiceApiEndpoint;
+            }
+
+            return stem;
+        }
     }
 }

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

@@ -27,6 +27,12 @@ public class HostAssignmentContext
         [JsonProperty("MSISpecializationPayload")]
         public MSIContext MSIContext { get; set; }
 
+        [JsonProperty("EncryptedTokenServiceSpecializationPayload")]
+        public string EncryptedTokenServiceSpecializationPayload { get; set; }
+
+        [JsonProperty("TokenServiceApiEndpoint")]
+        public string TokenServiceApiEndpoint { get; set; }
+
         [JsonProperty("CorsSpecializationPayload")]
         public CorsSettings CorsSettings { get; set; }
 

src/WebJobs.Script/ScriptConstants.cs

@@ -172,6 +172,7 @@ public static class ScriptConstants
         public const string LinuxFunctionDetailsEventStreamName = "MS_FUNCTION_DETAILS";
         public const string LinuxAzureMonitorEventStreamName = "MS_FUNCTION_AZURE_MONITOR_EVENT";
         public const string LinuxMSISpecializationStem = "/api/specialize?api-version=2022-01-01";
+        public const string LinuxEncryptedTokenServiceSpecializationStem = "/api/specialize?api-version=2023-05-01";
 
         public const string DurableTaskPropertyName = "durableTask";
         public const string DurableTaskHubName = "HubName";

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

@@ -618,6 +618,69 @@ public async Task SpecializeMSISidecar_Succeeds()
                 p => Assert.StartsWith("Specialize MSI sidecar returned OK", p));
         }
 
+        [Fact]
+        public async Task SpecializeMSISidecar_Succeeds_EncryptedMSIContextWithoutProvidedEndpoint()
+        {
+            var environment = new Dictionary<string, string>()
+            {
+                { EnvironmentSettingNames.MsiEndpoint, "http://localhost:8081" },
+                { EnvironmentSettingNames.MsiSecret, "secret" }
+            };
+            var assignmentContext = new HostAssignmentContext
+            {
+                SiteId = 1234,
+                SiteName = "TestSite",
+                Environment = environment,
+                IsWarmupRequest = false,
+                MSIContext = new MSIContext(),
+                EncryptedTokenServiceSpecializationPayload = "TestContext"
+            };
+
+            var instanceManager = GetInstanceManagerForMSISpecialization(assignmentContext, HttpStatusCode.OK, null);
+
+            string error = await instanceManager.SpecializeMSISidecar(assignmentContext);
+            Assert.Null(error);
+
+            var logs = _loggerProvider.GetAllLogMessages().Select(p => p.FormattedMessage).ToArray();
+            Assert.Collection(logs,
+                p => Assert.StartsWith("MSI enabled status: True", p),
+                p => Assert.StartsWith("Using encrypted TokenService payload format", p),
+                p => Assert.Equal($"Specializing sidecar at http://localhost:8081{ScriptConstants.LinuxEncryptedTokenServiceSpecializationStem}", p),
+                p => Assert.StartsWith("Specialize MSI sidecar returned OK", p));
+        }
+
+        [Fact]
+        public async Task SpecializeMSISidecar_Succeeds_EncryptedMSIContextWithProvidedEndpoint()
+        {
+            var environment = new Dictionary<string, string>()
+            {
+                { EnvironmentSettingNames.MsiEndpoint, "http://localhost:8081" },
+                { EnvironmentSettingNames.MsiSecret, "secret" }
+            };
+            var assignmentContext = new HostAssignmentContext
+            {
+                SiteId = 1234,
+                SiteName = "TestSite",
+                Environment = environment,
+                IsWarmupRequest = false,
+                MSIContext = new MSIContext(),
+                EncryptedTokenServiceSpecializationPayload = "TestContext",
+                TokenServiceApiEndpoint = "/api/TestEndpoint"
+            };
+
+            var instanceManager = GetInstanceManagerForMSISpecialization(assignmentContext, HttpStatusCode.OK, null);
+
+            string error = await instanceManager.SpecializeMSISidecar(assignmentContext);
+            Assert.Null(error);
+
+            var logs = _loggerProvider.GetAllLogMessages().Select(p => p.FormattedMessage).ToArray();
+            Assert.Collection(logs,
+                p => Assert.StartsWith("MSI enabled status: True", p),
+                p => Assert.StartsWith("Using encrypted TokenService payload format", p),
+                p => Assert.Equal($"Specializing sidecar at http://localhost:8081{assignmentContext.TokenServiceApiEndpoint}", p),
+                p => Assert.StartsWith("Specialize MSI sidecar returned OK", p));
+        }
+
         [Fact]
         public async Task SpecializeMsiSidecar_RequiredPropertiesInPayload()
         {
@@ -751,7 +814,7 @@ public async Task DoesNotSpecializeMSISidecar_WhenMSIContextNull()
 
             var meshServiceClient = new Mock<IMeshServiceClient>(MockBehavior.Strict);
             meshServiceClient.Setup(c => c.NotifyHealthEvent(ContainerHealthEventType.Fatal,
-                It.Is<Type>(t => t == typeof(AtlasInstanceManager)), "Could not specialize MSI sidecar since MSIContext was empty")).Returns(Task.CompletedTask);
+                It.Is<Type>(t => t == typeof(AtlasInstanceManager)), "Could not specialize MSI sidecar since MSIContext and EncryptedTokenServiceSpecializationPayload were empty")).Returns(Task.CompletedTask);
 
             var instanceManager = GetInstanceManagerForMSISpecialization(assignmentContext, HttpStatusCode.BadRequest, meshServiceClient.Object);
 
@@ -761,10 +824,10 @@ public async Task DoesNotSpecializeMSISidecar_WhenMSIContextNull()
             var logs = _loggerProvider.GetAllLogMessages().Select(p => p.FormattedMessage).ToArray();
             Assert.Collection(logs,
                 p => Assert.StartsWith("MSI enabled status: True", p),
-                p => Assert.StartsWith("Skipping specialization of MSI sidecar since MSIContext was absent", p));
+                p => Assert.StartsWith("Skipping specialization of MSI sidecar since MSIContext and EncryptedTokenServiceSpecializationPayload were absent", p));
 
             meshServiceClient.Verify(c => c.NotifyHealthEvent(ContainerHealthEventType.Fatal,
-                It.Is<Type>(t => t == typeof(AtlasInstanceManager)), "Could not specialize MSI sidecar since MSIContext was empty"), Times.Once);
+                It.Is<Type>(t => t == typeof(AtlasInstanceManager)), "Could not specialize MSI sidecar since MSIContext and EncryptedTokenServiceSpecializationPayload were empty"), Times.Once);
         }
 
         [Fact]
@@ -1308,9 +1371,15 @@ private AtlasInstanceManager GetInstanceManagerForMSISpecialization(HostAssignme
 
             var msiEndpoint = hostAssignmentContext.Environment[EnvironmentSettingNames.MsiEndpoint] + ScriptConstants.LinuxMSISpecializationStem;
 
+            var defaultEncryptedMsiEndpoint = hostAssignmentContext.Environment[EnvironmentSettingNames.MsiEndpoint] + ScriptConstants.LinuxEncryptedTokenServiceSpecializationStem;
+
+            var providedEncryptedMsiEndpoint = hostAssignmentContext.Environment[EnvironmentSettingNames.MsiEndpoint] + hostAssignmentContext.TokenServiceApiEndpoint;
+
             handlerMock.Protected().Setup<Task<HttpResponseMessage>>("SendAsync",
                 ItExpr.Is<HttpRequestMessage>(request => request.Method == HttpMethod.Post
-                                                         && request.RequestUri.AbsoluteUri.Equals(msiEndpoint)
+                                                         && (request.RequestUri.AbsoluteUri.Equals(msiEndpoint) 
+                                                            || request.RequestUri.AbsoluteUri.Equals(defaultEncryptedMsiEndpoint)
+                                                            || request.RequestUri.AbsoluteUri.Equals(providedEncryptedMsiEndpoint))
                                                          && request.Content != null),
                 ItExpr.IsAny<CancellationToken>())
                 .Callback<HttpRequestMessage, CancellationToken>((request, token) => customAction?.Invoke(request, token))