WindowsSdkHelper for https://github.com/InfectedLibraries/Biohazrd/issues/199

WindowsSdkHelper.cs

using Biohazrd;
using Microsoft.VisualStudio.Setup.Configuration;
using Microsoft.Win32;
using System;
using System.IO;
using System.Runtime.InteropServices;

namespace InfectedWin32.Generator
{
    /// <summary>This class helps configure Biohazrd to include a specific version of the Windows SDK and reference files from it</summary>
    /// <remarks>
    /// Currently only supports the Windows 10 SDK since it relies on the UniversalCRT
    /// </remarks>
    public sealed class WindowsSdkHelper
    {
        public string SdkVersion { get; }
        public string SdkRootPath { get; }
        private string MsvcToolchainPath { get; }
        private string IncludeRootPath { get; }
        private string LibraryRootPath { get; }

        // These environment variables are the same ones set by a Visual Studio command prompt
        private const string Windows10SdkRootEnvironmentVariable = "WindowsSdkDir";
        private const string MsvcToolchainRootEnvironmentVariable = "VCToolsInstallDir";

        public WindowsSdkHelper(string sdkVersion)
        {
            SdkVersion = sdkVersion;
            SdkRootPath = LocateWindowsSdkRoot()
                ?? throw new Exception
                (
                    "Could not locate Windows 10 SDK root path! " +
                    // Using the environment variable is allowed on Windows, but ideally the SDK should be properly installed.
                    (RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "Is it installed?" : "Specify by setting " + Windows10SdkRootEnvironmentVariable)
                );
            MsvcToolchainPath = LocateToolchainRoot()
                ?? throw new Exception
                (
                    "Could not locate MSVC toolchain path!" +
                    // Using the environment variable is allowed on Windows, but ideally the SDK should be properly installed.
                    (RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "Is Visual Studio 2019 or newer installed with desktop development with C++?" : "Specify by setting " + MsvcToolchainRootEnvironmentVariable)
                );
            IncludeRootPath = Path.Combine(SdkRootPath, "Include", SdkVersion);
            LibraryRootPath = Path.Combine(SdkRootPath, "Lib", SdkVersion, "um", "x64");

            if (!Directory.Exists(IncludeRootPath))
            { throw new DirectoryNotFoundException($"Could not locate Windows SDK {SdkVersion} include root at '{IncludeRootPath}'"); }

            if (!Directory.Exists(LibraryRootPath))
            { throw new DirectoryNotFoundException($"Could not locate Windows SDK {SdkVersion} library root at '{LibraryRootPath}' (Is the x64 SDK not installed?)"); }

            string userModeIncludeDirectoryPath = GetIncludeDirectory("um");
            if (!Directory.Exists(userModeIncludeDirectoryPath))
            { throw new DirectoryNotFoundException($"Could not locate Windows SDK {SdkVersion} user mode header files directory '{userModeIncludeDirectoryPath}' (Is the UCRT installed but not the Windows SDK?)"); }

            string windowHeaderPath = GetHeaderFilePath("um", "Windows.h");
            if (!File.Exists(windowHeaderPath))
            { throw new DirectoryNotFoundException($"Could not locate Windows.h in Windows SDK {SdkVersion} '{windowHeaderPath}' (Is the UCRT installed but not the Windows SDK?)"); }
        }

        public string GetIncludeDirectory(string categoryName)
            => Path.Combine(IncludeRootPath, categoryName);

        private string GetMsvcToolchainIncludeDirectory(string? subdirectory = null)
        {
            string root = MsvcToolchainPath;

            if (subdirectory is not null)
            { root = Path.Combine(root, subdirectory); }

            return Path.Combine(root, "include");
        }

        public string GetHeaderFilePath(string categoryName, string headerFileName)
            => Path.Combine(IncludeRootPath, categoryName, headerFileName);

        public string GetLibraryFilePath(string libraryFileName)
            => Path.Combine(LibraryRootPath, libraryFileName);

        /// <summary>Configures the specified builder to use this version of the Windows SDK</summary>
        /// <remarks>
        /// In many ways, this method manually implements this Clang functionality:
        /// https://github.com/InfectedLibraries/llvm-project/blob/6d5c430eb3c0bd49f6f5bda4b0d2d8aa79b0fa3f/clang/lib/Driver/ToolChains/MSVC.cpp#L1228-L1313
        /// </remarks>
        public void ConfigureBuilder(TranslatedLibraryBuilder builder)
        {
            // Disable standard includes (otherwise Clang will just include the latest SDK that's installed)
            // You *can* avoid doing this by adding the include directories with -I (--include-directory) since it is search first
            // but this has two issues:
            // 1) The system includes are still there just asking for accidental usage
            // 2) You loose out on Clang's special treatment of warnings in system headers
            builder.AddCommandLineArgument("--no-standard-includes");

            //TODO: Do we need to replicate Clang's builtin includes?
            // https://github.com/InfectedLibraries/llvm-project/blob/6d5c430eb3c0bd49f6f5bda4b0d2d8aa79b0fa3f/clang/lib/Driver/ToolChains/MSVC.cpp#L1233-L1236
            // This option introduced in Clang 11 would be helpful: https://releases.llvm.org/11.0.0/tools/clang/docs/ClangCommandLineReference.html#cmdoption-clang-ibuiltininc

            // Include the MSVC toolchain headers
            builder.AddCommandLineArgument($"-isystem{GetMsvcToolchainIncludeDirectory()}");
            // We shouldn't need ATL or MFC so we leave them out
            //builder.AddCommandLineArgument($"-isystem{GetMsvcToolchainIncludeDirectory("atlmfc")}");

            // Include the Universal CRT
            builder.AddCommandLineArgument($"-isystem{GetIncludeDirectory("ucrt")}");

            // Include the Windows SDK
            builder.AddCommandLineArgument($"-isystem{GetIncludeDirectory("shared")}");
            builder.AddCommandLineArgument($"-isystem{GetIncludeDirectory("um")}");
            // We don't intend to use WinRT stuff so we leave it out
            //builder.AddCommandLineArgument($"-isystem{GetIncludeDirectory("winrt")}");
        }

        private static string? LocateWindowsSdkRoot()
        {
            if (Environment.GetEnvironmentVariable(Windows10SdkRootEnvironmentVariable) is string pathFromEnvironment)
            { return pathFromEnvironment; }

            // On Linux the SDK path must be specified via environment variable
            if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            { return null; }

            // Query registry for the Windows 10 SDK root //HKEY_LOCAL_MACHINE\SOFTWARE\WOW6432Node\Microsoft\Windows Kits\Installed Roots\KitsRoot10
            // This registry key does not seem to be properly documented, but it is used by Clang so it's probably safe:
            // https://github.com/InfectedLibraries/llvm-project/blob/6d5c430eb3c0bd49f6f5bda4b0d2d8aa79b0fa3f/clang/lib/Driver/ToolChains/MSVC.cpp#L1148-L1150
            // It's also randomly mentioned in the HoloLens documentation:
            // https://docs.microsoft.com/en-us/windows/mixed-reality/develop/platform-capabilities-and-apis/using-the-hololens-emulator#troubleshooting
            if (Registry.GetValue(@"HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows Kits\Installed Roots", "KitsRoot10", null) is string pathFromRegistry)
            { return pathFromRegistry; }

            // Try default install directory as a fallback
            string? TryDefault(string root)
            {
                string path = Path.Combine(root, "Windows Kits", "10");
                return Directory.Exists(path) ? path : null;
            }

            return TryDefault(Environment.GetFolderPath(Environment.SpecialFolder.ProgramFilesX86))
                ?? TryDefault(Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles));
        }

        private static string? LocateToolchainRoot()
        {
            if (Environment.GetEnvironmentVariable(MsvcToolchainRootEnvironmentVariable) is string pathFromEnvironment)
            { return pathFromEnvironment; }

            // On Linux the SDK path must be specified via environment variable
            if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            { return null; }

            if (LocateToolchainRootUsingSetupApi() is string pathFromSetupApi)
            { return pathFromSetupApi; }

            return null;
        }

        private static string? LocateToolchainRootUsingSetupApi()
        {
            static string? GetToolchainRoot(ISetupInstance visualStudio)
            {
                string vcPath = visualStudio.ResolvePath("VC");
                string defaultToolsVersionPath = Path.Combine(vcPath, "Auxiliary", "Build", "Microsoft.VCToolsVersion.default.txt");

                if (!File.Exists(defaultToolsVersionPath))
                { return null; }

                string defaultToolsVersion = File.ReadAllText(defaultToolsVersionPath).Trim();
                string toolchainRootPath = Path.Combine(vcPath, "Tools", "MSVC", defaultToolsVersion);
                return Directory.Exists(toolchainRootPath) ? toolchainRootPath : null;
            }

            // The following uses logic similar to Clang
            // https://github.com/InfectedLibraries/llvm-project/blob/6d5c430eb3c0bd49f6f5bda4b0d2d8aa79b0fa3f/clang/lib/Driver/ToolChains/MSVC.cpp#L178-L268
            // Unlike Clang, we check if the desktop development with C++ workload is installed and restrict to Visual Studio 2019 or newer
            string? result = null;
            try
            {
                SetupConfiguration query = new();
                IEnumSetupInstances enumInstances = query.EnumAllInstances();
                ISetupHelper helper = (ISetupHelper)query;

                ISetupInstance[] instance = new ISetupInstance[1];
                ulong newestVersionNum = 0;

                while (true)
                {
                    int fetched;
                    enumInstances.Next(1, instance, out fetched);

                    if (fetched == 0)
                    { break; }

                    string versionString = instance[0].GetInstallationVersion();

                    // Skip versions prior to Visual Studio 2019
                    if (new Version(versionString).Major < 16)
                    { continue; }

                    // Skip versions without the desktop development with C++ workload
                    string? toolchainRootPath = GetToolchainRoot(instance[0]);
                    if (toolchainRootPath is null)
                    { continue; }

                    // Look for the newest version available
                    ulong versionNum = helper.ParseVersion(versionString);
                    if (newestVersionNum == 0 || versionNum > newestVersionNum)
                    {
                        result = toolchainRootPath;
                        newestVersionNum = versionNum;
                    }
                }

                return result;
            }
            catch (COMException ex) when (ex.HResult == unchecked((int)0x80040154)) // REGDB_E_CLASSNOTREG
            { throw new Exception("Could not locate the Visual Studio setup configuration COM service. (Is Visual Studio installed?)", ex); }
        }
    }
}