Split test host controller into focused partial files

src/Platform/Microsoft.Testing.Platform/Hosts/TestHostControllersTestHost.Disposal.cs

@@ -0,0 +1,69 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.Testing.Platform.Extensions.TestHostControllers;
+using Microsoft.Testing.Platform.Helpers;
+using Microsoft.Testing.Platform.Messages;
+using Microsoft.Testing.Platform.Services;
+
+namespace Microsoft.Testing.Platform.Hosts;
+
+internal sealed partial class TestHostControllersTestHost
+{
+    private async Task DisposeServicesAsync()
+    {
+        // A CompositeExtensionFactory builds one object that is reused for every role it was registered
+        // under, so the lifetime handlers and environment-variable providers disposed below can be the very
+        // same instances that the message bus holds as IDataConsumer. Because we dispose them here, before
+        // DisposeServiceProviderAsync gets a chance to close the handshake, we have to disable the bus first
+        // or a consumer can be disposed while its ConsumeAsync is still running. The in-run disable is not
+        // enough: it is gated on there being lifetime handlers, and it is skipped entirely on every early-out
+        // path (an invalid platform setup, or any failure between building the bus and reaching it).
+        object[] consumersStillRunning = [];
+        if (ServiceProvider.GetService<BaseMessageBus>() is { } messageBus)
+        {
+            await EnsureMessageBusDisabledAsync(messageBus, ServiceProvider).ConfigureAwait(false);
+
+            // Disabling is bounded on an aborted run, so it can return while a consumer that ignores the
+            // cancellation token is still inside ConsumeAsync. Those instances must be skipped here too, for
+            // exactly the reason above: a multi-role extension is reached by these manual loops.
+            consumersStillRunning = [.. messageBus.ConsumersStillRunning];
+        }
+
+        ITestHostEnvironmentVariableProvider[] variableProviders = _testHostsInformation.EnvironmentVariableProviders;
+        ITestHostProcessLifetimeHandler[] lifetimeHandlers = _testHostsInformation.LifetimeHandlers;
+
+        List<object> alreadyDisposed = [with(lifetimeHandlers.Length + variableProviders.Length)];
+
+        // Recording them as already disposed is what keeps them from being disposed by the service-provider
+        // walk below either.
+        alreadyDisposed.AddRange(consumersStillRunning);
+
+        foreach (ITestHostProcessLifetimeHandler service in lifetimeHandlers)
+        {
+            if (alreadyDisposed.Contains(service))
+            {
+                continue;
+            }
+
+            await DisposeHelper.DisposeAsync(service).ConfigureAwait(false);
+            alreadyDisposed.Add(service);
+        }
+
+        foreach (ITestHostEnvironmentVariableProvider service in variableProviders)
+        {
+            if (alreadyDisposed.Contains(service))
+            {
+                continue;
+            }
+
+            await DisposeHelper.DisposeAsync(service).ConfigureAwait(false);
+            alreadyDisposed.Add(service);
+        }
+
+        await DisposeServiceProviderAsync(ServiceProvider, alreadyDisposed: alreadyDisposed).ConfigureAwait(false);
+    }
+
+    public void Dispose()
+        => _waitForPid.Dispose();
+}

src/Platform/Microsoft.Testing.Platform/Hosts/TestHostControllersTestHost.Ipc.cs

@@ -0,0 +1,122 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.Testing.Platform.Extensions.TestHostControllers;
+using Microsoft.Testing.Platform.IPC;
+using Microsoft.Testing.Platform.IPC.Models;
+using Microsoft.Testing.Platform.IPC.Serializers;
+using Microsoft.Testing.Platform.Logging;
+using Microsoft.Testing.Platform.Resources;
+using Microsoft.Testing.Platform.Services;
+
+namespace Microsoft.Testing.Platform.Hosts;
+
+internal sealed partial class TestHostControllersTestHost
+{
+    private async Task<NamedPipeServer> CreateTestHostControllerIpcAsync(ExecutableInfo executableInfo, CancellationToken cancellationToken)
+    {
+        NamedPipeServer testHostControllerIpc = new(
+            $"MONITORTOHOST_{Guid.NewGuid():N}",
+            HandleRequestAsync,
+            _environment,
+            _loggerFactory.CreateLogger<NamedPipeServer>(),
+            ServiceProvider.GetTask(),
+            await GetAuthorizedSecurityIdentitiesAsync(_testHostsInformation.TestHostLauncher, executableInfo.FilePath, cancellationToken).ConfigureAwait(false),
+            cancellationToken);
+        testHostControllerIpc.RegisterAllSerializers();
+        return testHostControllerIpc;
+    }
+
+    /// <summary>
+    /// Asks the registered launcher, when it implements
+    /// <see cref="ITestHostControllerConnectionAuthorizer"/>, for the security identities that must
+    /// additionally be authorized on the controller-to-host connection.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This runs before the connection is created, which is the only point where its access control can
+    /// still be composed: it has to be listening before the host is launched, so the launcher cannot
+    /// contribute this from <see cref="ITestHostLauncher.LaunchTestHostAsync"/>.
+    /// </para>
+    /// <para>
+    /// Every returned value is validated against the platform's least-privilege policy: only the identity of
+    /// a single sandboxed application may be authorized, never a user, a group, <c>Everyone</c>, or an
+    /// identity shared by every sandboxed application on the machine. An extension that asks for anything
+    /// else fails the run instead of silently getting a weaker connection, so a mistake in an extension
+    /// cannot degrade into an over-permissive access control list.
+    /// </para>
+    /// </remarks>
+    private async Task<IReadOnlyList<string>?> GetAuthorizedSecurityIdentitiesAsync(
+        ITestHostLauncher? testHostLauncher,
+        string testHostFileName,
+        CancellationToken cancellationToken)
+    {
+        if (testHostLauncher is not ITestHostControllerConnectionAuthorizer connectionAuthorizer)
+        {
+            return null;
+        }
+
+        IReadOnlyList<string> extensionResult = await connectionAuthorizer.GetAuthorizedSecurityIdentitiesAsync(testHostFileName, cancellationToken).ConfigureAwait(false);
+        if (extensionResult is null || extensionResult.Count == 0)
+        {
+            return null;
+        }
+
+        // Snapshot immediately. Everything below — validation, logging, and ultimately the security
+        // descriptor — must operate on one fixed set of values; the sequence itself is extension-supplied
+        // and re-enumerating it is not guaranteed to yield what was validated.
+        string[] securityIdentities = [.. extensionResult];
+
+        if (!NamedPipeServerSecurity.IsSupported)
+        {
+            // Sandboxed-application identities and connection access control lists are expressible only on
+            // Windows. Anywhere else the request is meaningless, so it is ignored rather than failing an
+            // otherwise valid run.
+            await _logger.LogDebugAsync($"'{testHostLauncher.Uid}' requested {securityIdentities.Length} connection authorization(s), ignored on this operating system.").ConfigureAwait(false);
+            return null;
+        }
+
+        foreach (string securityIdentity in securityIdentities.Where(static securityIdentity =>
+            !NamedPipeServerSecurity.IsAuthorizableSandboxedApplicationIdentity(securityIdentity)))
+        {
+            throw new InvalidOperationException(string.Format(
+                CultureInfo.InvariantCulture,
+                PlatformResources.TestHostControllerConnectionInvalidAuthorizedSecurityIdentityErrorMessage,
+                testHostLauncher.DisplayName,
+                testHostLauncher.Uid,
+                securityIdentity ?? "<null>",
+                NamedPipeServerSecurity.AllApplicationPackagesSid));
+        }
+
+        await _logger.LogDebugAsync($"'{testHostLauncher.Uid}' authorized the following security identity/identities on the test host controller connection: {string.Join(", ", securityIdentities)}").ConfigureAwait(false);
+        return securityIdentities;
+    }
+
+    private Task<IResponse> HandleRequestAsync(IRequest request)
+    {
+        try
+        {
+            switch (request)
+            {
+                case TestHostCompletedRequest testHostCompletedRequest:
+                    _testHostCompletedReceived = true;
+                    _testHostExitCodeReceived = testHostCompletedRequest.ExitCode;
+                    _testHostUnfilteredExitCodeReceived = testHostCompletedRequest.UnfilteredExitCode;
+                    return Task.FromResult<IResponse>(VoidResponse.CachedInstance);
+
+                case TestHostProcessPIDRequest testHostProcessPIDRequest:
+                    _testHostPID = testHostProcessPIDRequest.PID;
+                    _waitForPid.Set();
+                    return Task.FromResult<IResponse>(VoidResponse.CachedInstance);
+
+                default:
+                    throw new NotSupportedException($"Request '{request}' not supported");
+            }
+        }
+        catch (Exception ex)
+        {
+            _environment.FailFast($"[TestHostControllersTestHost] Unhandled exception:\n{ex}", ex);
+            throw;
+        }
+    }
+}

src/Platform/Microsoft.Testing.Platform/Hosts/TestHostControllersTestHost.ProcessConfiguration.cs

@@ -0,0 +1,174 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.Testing.Platform.CommandLine;
+using Microsoft.Testing.Platform.Extensions;
+using Microsoft.Testing.Platform.Extensions.TestHostControllers;
+using Microsoft.Testing.Platform.Helpers;
+using Microsoft.Testing.Platform.IPC;
+using Microsoft.Testing.Platform.Logging;
+using Microsoft.Testing.Platform.Messages;
+using Microsoft.Testing.Platform.OutputDevice;
+using Microsoft.Testing.Platform.Resources;
+using Microsoft.Testing.Platform.ServerMode;
+using Microsoft.Testing.Platform.Services;
+using Microsoft.Testing.Platform.TestHostControllers;
+
+namespace Microsoft.Testing.Platform.Hosts;
+
+internal sealed partial class TestHostControllersTestHost
+{
+    private async Task<(ProcessStartInfo ProcessStartInfo, IReadOnlyList<string> PartialCommandLine)?> PrepareProcessConfigurationAsync(
+        ExecutableInfo executableInfo,
+        int currentPid,
+        string processIdString,
+        string processCorrelationId,
+        NamedPipeServer testHostControllerIpc,
+        IEnvironment environment,
+        ProxyOutputDevice outputDevice,
+        CancellationToken cancellationToken)
+    {
+        List<string> partialCommandLine =
+        [
+            .. executableInfo.Arguments,
+            $"--{PlatformCommandLineProvider.TestHostControllerPIDOptionKey}",
+            processIdString
+        ];
+
+#if NET8_0_OR_GREATER
+        // On net8.0+, we can pass the arguments as a collection directly to ProcessStartInfo.
+        // When passing the collection, it's expected to be unescaped, so we pass what we have directly.
+        IEnumerable<string> arguments = partialCommandLine;
+#else
+        // Current target framework (.NET Framework and .NET Standard 2.0) only supports arguments as a single string.
+        // In this case, escaping is essential. For example, one of the arguments could already contain spaces.
+        // PasteArguments is borrowed from dotnet/runtime.
+        var builder = new StringBuilder();
+        foreach (string arg in partialCommandLine)
+        {
+            PasteArguments.AppendArgument(builder, arg);
+        }
+
+        string arguments = builder.ToString();
+#endif
+
+        ProcessStartInfo processStartInfo = new(
+            executableInfo.FilePath,
+            arguments)
+        {
+            EnvironmentVariables =
+            {
+                { $"{EnvironmentVariableConstants.TESTINGPLATFORM_TESTHOSTCONTROLLER_CORRELATIONID}_{currentPid}", processCorrelationId },
+                { $"{EnvironmentVariableConstants.TESTINGPLATFORM_TESTHOSTCONTROLLER_PARENTPID}_{currentPid}", processIdString },
+                { $"{EnvironmentVariableConstants.TESTINGPLATFORM_TESTHOSTCONTROLLER_SKIPEXTENSION}_{currentPid}", "1" },
+                { $"{EnvironmentVariableConstants.TESTINGPLATFORM_TESTHOSTCONTROLLER_PIPENAME}_{currentPid}", testHostControllerIpc.PipeName.Name },
+            },
+            UseShellExecute = false,
+        };
+
+        List<IDataConsumer> dataConsumersBuilder = [.. _testHostsInformation.DataConsumer];
+        if (ServiceProvider.GetService<TestCoverageCapabilities>() is { } coverageCapabilities)
+        {
+            coverageCapabilities.RegisterProducers(dataConsumersBuilder);
+        }
+
+        // Register the coverage result consumer so that coverage messages published by
+        // ITestHostProcessLifetimeHandler extensions in this (controller) process are tracked.
+        // This is the same instance later read by the coverage threshold exit code check.
+        if (ServiceProvider.GetService<TestCoverageResult>() is { } testCoverageResult)
+        {
+            dataConsumersBuilder.Add(testCoverageResult);
+        }
+
+        // We add the IPlatformOutputDevice after all users extensions.
+        IPlatformOutputDevice? display = ServiceProvider.GetServiceInternal<IPlatformOutputDevice>();
+        if (display is IDataConsumer dataConsumerDisplay)
+        {
+            dataConsumersBuilder.Add(dataConsumerDisplay);
+        }
+
+        // We register the DotnetTestDataConsumer as last to ensure that it will be the last one to consume the data.
+        IPushOnlyProtocol? pushOnlyProtocol = ServiceProvider.GetService<IPushOnlyProtocol>();
+        if (pushOnlyProtocol?.IsServerMode == true)
+        {
+            dataConsumersBuilder.Add(await pushOnlyProtocol.GetDataConsumerAsync().ConfigureAwait(false));
+        }
+
+        // If we're in server mode jsonrpc we add as last consumer the PassiveNodeDataConsumer for the attachments.
+        // Connect the passive node if it's available
+        if (_passiveNode is not null)
+        {
+            if (await _passiveNode.ConnectAsync().ConfigureAwait(false))
+            {
+                dataConsumersBuilder.Add(new PassiveNodeDataConsumer(_passiveNode));
+            }
+            else
+            {
+                await _logger.LogWarningAsync("PassiveNode was expected to connect but failed").ConfigureAwait(false);
+            }
+        }
+
+        var concreteMessageBusService = new AsynchronousMessageBus(
+            [.. dataConsumersBuilder],
+            ServiceProvider.GetTestApplicationCancellationTokenSource(),
+            ServiceProvider.GetTask(),
+            ServiceProvider.GetLoggerFactory(),
+            ServiceProvider.GetEnvironment(),
+            ServiceProvider.GetService<IShutdownProgressReporter>());
+        await concreteMessageBusService.InitAsync().ConfigureAwait(false);
+        ((MessageBusProxy)ServiceProvider.GetMessageBus()).SetBuiltMessageBus(concreteMessageBusService);
+
+        // Apply the ITestHostEnvironmentVariableProvider
+        if (_testHostsInformation.EnvironmentVariableProviders.Length > 0)
+        {
+            SystemEnvironmentVariableProvider systemEnvironmentVariableProvider = new(environment);
+            EnvironmentVariables environmentVariables = new(_loggerFactory)
+            {
+                CurrentProvider = systemEnvironmentVariableProvider,
+            };
+            await systemEnvironmentVariableProvider.UpdateAsync(environmentVariables).ConfigureAwait(false);
+
+            foreach (ITestHostEnvironmentVariableProvider environmentVariableProvider in _testHostsInformation.EnvironmentVariableProviders)
+            {
+                environmentVariables.CurrentProvider = environmentVariableProvider;
+                await environmentVariableProvider.UpdateAsync(environmentVariables).ConfigureAwait(false);
+            }
+
+            environmentVariables.CurrentProvider = null;
+
+            List<(IExtension, string)> failedValidations = [];
+            foreach (ITestHostEnvironmentVariableProvider hostEnvironmentVariableProvider in _testHostsInformation.EnvironmentVariableProviders)
+            {
+                ValidationResult variableResult = await hostEnvironmentVariableProvider.ValidateTestHostEnvironmentVariablesAsync(environmentVariables).ConfigureAwait(false);
+                if (!variableResult.IsValid)
+                {
+                    failedValidations.Add((hostEnvironmentVariableProvider, variableResult.ErrorMessage));
+                }
+            }
+
+            if (failedValidations.Count > 0)
+            {
+                StringBuilder displayErrorMessageBuilder = new();
+                StringBuilder logErrorMessageBuilder = new();
+                displayErrorMessageBuilder.AppendLine(PlatformResources.GlobalValidationOfTestHostEnvironmentVariablesFailedErrorMessage);
+                logErrorMessageBuilder.AppendLine("The following 'ITestHostEnvironmentVariableProvider' providers rejected the final environment variables setup:");
+                foreach ((IExtension extension, string errorMessage) in failedValidations)
+                {
+                    displayErrorMessageBuilder.AppendLine(string.Format(CultureInfo.InvariantCulture, PlatformResources.EnvironmentVariableProviderFailedWithError, extension.DisplayName, extension.Uid, errorMessage));
+                    displayErrorMessageBuilder.AppendLine(CultureInfo.InvariantCulture, $"Provider '{extension.DisplayName}' (UID: {extension.Uid}) failed with error: {errorMessage}");
+                }
+
+                await outputDevice.DisplayAsync(this, new ErrorMessageOutputDeviceData(displayErrorMessageBuilder.ToString()), cancellationToken).ConfigureAwait(false);
+                await _logger.LogErrorAsync(logErrorMessageBuilder.ToString()).ConfigureAwait(false);
+                return null;
+            }
+
+            foreach (EnvironmentVariable envVar in environmentVariables.GetAll())
+            {
+                processStartInfo.EnvironmentVariables[envVar.Variable] = envVar.Value;
+            }
+        }
+
+        return (processStartInfo, partialCommandLine);
+    }
+}