7 files changed, 565 insertions, 0 deletions

diff --git a/src/WixToolset.Core.ExtensionCache/CachedExtension.cs b/src/WixToolset.Core.ExtensionCache/CachedExtension.cs
new file mode 100644
index 00000000..9ed874d9
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/CachedExtension.cs
@@ -0,0 +1,20 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    public class CachedExtension
+    {
+        internal CachedExtension(string id, string version, bool damaged)
+        {
+            this.Id = id;
+            this.Version = version;
+            this.Damaged = damaged;
+        }
+
+        public string Id { get; }
+
+        public string Version { get; }
+
+        public bool Damaged { get; }
+    }
+}
diff --git a/src/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs
new file mode 100644
index 00000000..3ec6451e
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs
@@ -0,0 +1,252 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    using System;
+    using System.Collections.Generic;
+    using System.IO;
+    using System.Linq;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using NuGet.Common;
+    using NuGet.Configuration;
+    using NuGet.Credentials;
+    using NuGet.Packaging;
+    using NuGet.Protocol;
+    using NuGet.Protocol.Core.Types;
+    using NuGet.Versioning;
+
+    /// <summary>
+    /// Extension cache manager.
+    /// </summary>
+    public class ExtensionCacheManager
+    {
+        public string CacheFolder(bool global) => global ? this.GlobalCacheFolder() : this.LocalCacheFolder();
+
+        public string LocalCacheFolder() => Path.Combine(Environment.CurrentDirectory, @".wix\extensions\");
+
+        public string GlobalCacheFolder()
+        {
+            var baseFolder = Environment.GetEnvironmentVariable("WIX_EXTENSIONS") ?? Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
+            return Path.Combine(baseFolder, @".wix\extensions\");
+        }
+
+        public async Task<bool> AddAsync(bool global, string extension, CancellationToken cancellationToken)
+        {
+            if (String.IsNullOrEmpty(extension))
+            {
+                throw new ArgumentNullException(nameof(extension));
+            }
+
+            (var extensionId, var extensionVersion) = ParseExtensionReference(extension);
+
+            var result = await this.DownloadAndExtractAsync(global, extensionId, extensionVersion, cancellationToken);
+
+            return result;
+        }
+
+        public Task<bool> RemoveAsync(bool global, string extension, CancellationToken cancellationToken)
+        {
+            if (String.IsNullOrEmpty(extension))
+            {
+                throw new ArgumentNullException(nameof(extension));
+            }
+
+            (var extensionId, var extensionVersion) = ParseExtensionReference(extension);
+
+            var cacheFolder = this.CacheFolder(global);
+
+            cacheFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
+
+            if (Directory.Exists(cacheFolder))
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                Directory.Delete(cacheFolder, true);
+                return Task.FromResult(true);
+            }
+
+            return Task.FromResult(false);
+        }
+
+        public Task<IEnumerable<CachedExtension>> ListAsync(bool global, string extension, CancellationToken cancellationToken)
+        {
+            var found = new List<CachedExtension>();
+
+            (var extensionId, var extensionVersion) = ParseExtensionReference(extension);
+
+            var cacheFolder = this.CacheFolder(global);
+
+            var searchFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
+
+            if (!Directory.Exists(searchFolder))
+            {
+            }
+            else if (!String.IsNullOrEmpty(extensionVersion)) // looking for an explicit version of an extension.
+            {
+                var extensionFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
+                if (Directory.Exists(extensionFolder))
+                {
+                    var present = ExtensionFileExists(cacheFolder, extensionId, extensionVersion);
+                    found.Add(new CachedExtension(extensionId, extensionVersion, !present));
+                }
+            }
+            else // looking for all versions of an extension or all versions of all extensions.
+            {
+                IEnumerable<string> foundExtensionIds;
+
+                if (String.IsNullOrEmpty(extensionId))
+                {
+                    // Looking for all versions of all extensions.
+                    foundExtensionIds = Directory.GetDirectories(cacheFolder).Select(folder => Path.GetFileName(folder)).ToList();
+                }
+                else
+                {
+                    // Looking for all versions of a single extension.
+                    var extensionFolder = Path.Combine(cacheFolder, extensionId);
+                    foundExtensionIds = Directory.Exists(extensionFolder) ? new[] { extensionId } : Array.Empty<string>();
+                }
+
+                foreach (var foundExtensionId in foundExtensionIds)
+                {
+                    var extensionFolder = Path.Combine(cacheFolder, foundExtensionId);
+
+                    foreach (var folder in Directory.GetDirectories(extensionFolder))
+                    {
+                        cancellationToken.ThrowIfCancellationRequested();
+
+                        var foundExtensionVersion = Path.GetFileName(folder);
+
+                        if (!NuGetVersion.TryParse(foundExtensionVersion, out _))
+                        {
+                            continue;
+                        }
+
+                        var present = ExtensionFileExists(cacheFolder, foundExtensionId, foundExtensionVersion);
+                        found.Add(new CachedExtension(foundExtensionId, foundExtensionVersion, !present));
+                    }
+                }
+            }
+
+            return Task.FromResult((IEnumerable<CachedExtension>)found);
+        }
+
+        private async Task<bool> DownloadAndExtractAsync(bool global, string id, string version, CancellationToken cancellationToken)
+        {
+            var logger = NullLogger.Instance;
+
+            DefaultCredentialServiceUtility.SetupDefaultCredentialService(logger, nonInteractive: false);
+
+            var settings = Settings.LoadDefaultSettings(root: Environment.CurrentDirectory);
+            var sources = PackageSourceProvider.LoadPackageSources(settings).Where(s => s.IsEnabled);
+
+            using (var cache = new SourceCacheContext())
+            {
+                PackageSource versionSource = null;
+
+                var nugetVersion = String.IsNullOrEmpty(version) ? null : new NuGetVersion(version);
+
+                if (nugetVersion is null)
+                {
+                    foreach (var source in sources)
+                    {
+                        var repository = Repository.Factory.GetCoreV3(source.Source);
+                        var resource = await repository.GetResourceAsync<FindPackageByIdResource>();
+
+                        var availableVersions = await resource.GetAllVersionsAsync(id, cache, logger, cancellationToken);
+                        foreach (var availableVersion in availableVersions)
+                        {
+                            if (nugetVersion is null || nugetVersion < availableVersion)
+                            {
+                                nugetVersion = availableVersion;
+                                versionSource = source;
+                            }
+                        }
+                    }
+
+                    if (nugetVersion is null)
+                    {
+                        return false;
+                    }
+                }
+
+                var searchSources = versionSource is null ? sources : new[] { versionSource };
+
+                var extensionFolder = Path.Combine(this.CacheFolder(global), id, nugetVersion.ToString());
+
+                foreach (var source in searchSources)
+                {
+                    var repository = Repository.Factory.GetCoreV3(source.Source);
+                    var resource = await repository.GetResourceAsync<FindPackageByIdResource>();
+
+                    using (var stream = new MemoryStream())
+                    {
+                        var downloaded = await resource.CopyNupkgToStreamAsync(id, nugetVersion, stream, cache, logger, cancellationToken);
+
+                        if (downloaded)
+                        {
+                            stream.Position = 0;
+
+                            using (var archive = new PackageArchiveReader(stream))
+                            {
+                                var files = PackagingConstants.Folders.Known.SelectMany(folder => archive.GetFiles(folder)).Distinct(StringComparer.OrdinalIgnoreCase);
+                                await archive.CopyFilesAsync(extensionFolder, files, this.ExtractProgress, logger, cancellationToken);
+                            }
+
+                            return true;
+                        }
+                    }
+                }
+            }
+
+            return false;
+        }
+
+        private string ExtractProgress(string sourceFile, string targetPath, Stream fileStream) => fileStream.CopyToFile(targetPath);
+
+        private static (string extensionId, string extensionVersion) ParseExtensionReference(string extensionReference)
+        {
+            var extensionId = extensionReference ?? String.Empty;
+            var extensionVersion = String.Empty;
+
+            var index = extensionId.LastIndexOf('/');
+            if (index > 0)
+            {
+                extensionVersion = extensionReference.Substring(index + 1);
+                extensionId = extensionReference.Substring(0, index);
+
+                if (!NuGetVersion.TryParse(extensionVersion, out _))
+                {
+                    throw new ArgumentException($"Invalid extension version in {extensionReference}");
+                }
+
+                if (String.IsNullOrEmpty(extensionId))
+                {
+                    throw new ArgumentException($"Invalid extension id in {extensionReference}");
+                }
+            }
+
+            return (extensionId, extensionVersion);
+        }
+
+        private static bool ExtensionFileExists(string baseFolder, string extensionId, string extensionVersion)
+        {
+            var toolsFolder = Path.Combine(baseFolder, extensionId, extensionVersion, "tools");
+            if (!Directory.Exists(toolsFolder))
+            {
+                return false;
+            }
+
+            var extensionAssembly = Path.Combine(toolsFolder, extensionId + ".dll");
+
+            var present = File.Exists(extensionAssembly);
+            if (!present)
+            {
+                extensionAssembly = Path.Combine(toolsFolder, extensionId + ".exe");
+                present = File.Exists(extensionAssembly);
+            }
+
+            return present;
+        }
+    }
+}
diff --git a/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs
new file mode 100644
index 00000000..5016f430
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs
@@ -0,0 +1,170 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using WixToolset.Extensibility.Data;
+    using WixToolset.Extensibility.Services;
+
+    /// <summary>
+    /// Extension cache manager command.
+    /// </summary>
+    internal class ExtensionCacheManagerCommand : ICommandLineCommand
+    {
+        private enum CacheSubcommand
+        {
+            Add,
+            Remove,
+            List
+        }
+
+        public ExtensionCacheManagerCommand(IWixToolsetServiceProvider serviceProvider)
+        {
+            this.Messaging = serviceProvider.GetService<IMessaging>();
+            this.ExtensionReferences = new List<string>();
+        }
+
+        private IMessaging Messaging { get; }
+
+        public bool ShowLogo { get; private set; }
+
+        public bool StopParsing { get; private set; }
+
+        private bool ShowHelp { get; set; }
+
+        private bool Global { get; set; }
+
+        private CacheSubcommand? Subcommand { get; set; }
+
+        private List<string> ExtensionReferences { get; }
+
+        public async Task<int> ExecuteAsync(CancellationToken cancellationToken)
+        {
+            if (this.ShowHelp || !this.Subcommand.HasValue)
+            {
+                DisplayHelp();
+                return 1;
+            }
+
+            var success = false;
+            var cacheManager = new ExtensionCacheManager();
+
+            switch (this.Subcommand)
+            {
+                case CacheSubcommand.Add:
+                    success = await this.AddExtensions(cacheManager, cancellationToken);
+                    break;
+
+                case CacheSubcommand.Remove:
+                    success = await this.RemoveExtensions(cacheManager, cancellationToken);
+                    break;
+
+                case CacheSubcommand.List:
+                    success = await this.ListExtensions(cacheManager, cancellationToken);
+                    break;
+            }
+
+            return success ? 0 : 2;
+        }
+
+        public bool TryParseArgument(ICommandLineParser parser, string argument)
+        {
+            if (!parser.IsSwitch(argument))
+            {
+                if (!this.Subcommand.HasValue)
+                {
+                    if (!Enum.TryParse(argument, true, out CacheSubcommand subcommand))
+                    {
+                        return false;
+                    }
+
+                    this.Subcommand = subcommand;
+                }
+                else
+                {
+                    this.ExtensionReferences.Add(argument);
+                }
+
+                return true;
+            }
+
+            var parameter = argument.Substring(1);
+            switch (parameter.ToLowerInvariant())
+            {
+                case "?":
+                    this.ShowHelp = true;
+                    this.ShowLogo = true;
+                    this.StopParsing = true;
+                    return true;
+
+                case "nologo":
+                    this.ShowLogo = false;
+                    return true;
+
+                case "g":
+                case "-global":
+                    this.Global = true;
+                    return true;
+            }
+
+            return false;
+        }
+
+        private async Task<bool> AddExtensions(ExtensionCacheManager cacheManager, CancellationToken cancellationToken)
+        {
+            var success = true;
+
+            foreach (var extensionRef in this.ExtensionReferences)
+            {
+                var added = await cacheManager.AddAsync(this.Global, extensionRef, cancellationToken);
+                success |= added;
+            }
+
+            return success;
+        }
+
+        private async Task<bool> RemoveExtensions(ExtensionCacheManager cacheManager, CancellationToken cancellationToken)
+        {
+            var success = true;
+
+            foreach (var extensionRef in this.ExtensionReferences)
+            {
+                var removed = await cacheManager.RemoveAsync(this.Global, extensionRef, cancellationToken);
+                success |= removed;
+            }
+
+            return success;
+        }
+
+        private async Task<bool> ListExtensions(ExtensionCacheManager cacheManager, CancellationToken cancellationToken)
+        {
+            var found = false;
+            var extensionRef = this.ExtensionReferences.FirstOrDefault();
+
+            var extensions = await cacheManager.ListAsync(this.Global, extensionRef, cancellationToken);
+
+            foreach (var extension in extensions)
+            {
+                this.Messaging.Write($"{extension.Id} {extension.Version}{(extension.Damaged ? " (damaged)" : String.Empty)}");
+                found = true;
+            }
+
+            return found;
+        }
+
+        private static void DisplayHelp()
+        {
+            Console.WriteLine(" usage:  wix.exe extension add|remove|list [extensionRef]");
+            Console.WriteLine();
+            Console.WriteLine("   -g       add/remove the extension for the current user");
+            Console.WriteLine("   -nologo  suppress displaying the logo information");
+            Console.WriteLine("   -?       this help information");
+            Console.WriteLine();
+            Console.WriteLine("   extensionRef format: extensionId/version (the version is optional)");
+        }
+    }
+}
diff --git a/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionCommandLine.cs b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionCommandLine.cs
new file mode 100644
index 00000000..81e96718
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionCommandLine.cs
@@ -0,0 +1,39 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    using System;
+    using System.Collections.Generic;
+    using WixToolset.Extensibility;
+    using WixToolset.Extensibility.Data;
+    using WixToolset.Extensibility.Services;
+
+    /// <summary>
+    /// Parses the "extension" command-line command. See <c>ExtensionCacheManagerCommand</c>
+    /// for the bulk of the command-line processing.
+    /// </summary>
+    internal class ExtensionCacheManagerExtensionCommandLine : BaseExtensionCommandLine
+    {
+        public ExtensionCacheManagerExtensionCommandLine(IWixToolsetServiceProvider serviceProvider)
+        {
+            this.ServiceProvider = serviceProvider;
+        }
+
+        private IWixToolsetServiceProvider ServiceProvider { get; }
+
+        // TODO: Do something with CommandLineSwitches
+        public override IEnumerable<ExtensionCommandLineSwitch> CommandLineSwitches => base.CommandLineSwitches;
+
+        public override bool TryParseCommand(ICommandLineParser parser, string argument, out ICommandLineCommand command)
+        {
+            command = null;
+
+            if ("extension".Equals(argument, StringComparison.OrdinalIgnoreCase))
+            {
+                command = new ExtensionCacheManagerCommand(this.ServiceProvider);
+            }
+
+            return command != null;
+        }
+    }
+}
diff --git a/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionFactory.cs b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionFactory.cs
new file mode 100644
index 00000000..44fc4b86
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/ExtensionCacheManagerExtensionFactory.cs
@@ -0,0 +1,30 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    using System;
+    using WixToolset.Extensibility;
+    using WixToolset.Extensibility.Services;
+
+    internal class ExtensionCacheManagerExtensionFactory : IExtensionFactory
+    {
+        public ExtensionCacheManagerExtensionFactory(IWixToolsetServiceProvider serviceProvider)
+        {
+            this.ServiceProvider = serviceProvider;
+        }
+
+        private IWixToolsetServiceProvider ServiceProvider { get; }
+
+        public bool TryCreateExtension(Type extensionType, out object extension)
+        {
+            extension = null;
+
+            if (extensionType == typeof(IExtensionCommandLine))
+            {
+                extension = new ExtensionCacheManagerExtensionCommandLine(this.ServiceProvider);
+            }
+
+            return extension != null;
+        }
+    }
+}
diff --git a/src/WixToolset.Core.ExtensionCache/WixToolset.Core.ExtensionCache.csproj b/src/WixToolset.Core.ExtensionCache/WixToolset.Core.ExtensionCache.csproj
new file mode 100644
index 00000000..7ae5cdbb
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/WixToolset.Core.ExtensionCache.csproj
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFrameworks>netstandard2.0;net461;net472</TargetFrameworks>
+    <Description>Extension Cache</Description>
+    <Title>WiX Toolset Extension Cache</Title>
+    <DebugType>embedded</DebugType>
+    <PublishRepositoryUrl>true</PublishRepositoryUrl>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\WixToolset.Core\WixToolset.Core.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <PackageReference Include="NuGet.Credentials" Version="5.6.0" />
+    <PackageReference Include="NuGet.Protocol" Version="5.6.0" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Microsoft.SourceLink.GitHub" Version="1.0.0" PrivateAssets="All" />
+    <PackageReference Include="Nerdbank.GitVersioning" Version="2.1.65" PrivateAssets="All" />
+  </ItemGroup>
+</Project>
diff --git a/src/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs b/src/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs
new file mode 100644
index 00000000..c1579330
--- /dev/null
+++ b/src/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs
@@ -0,0 +1,28 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.ExtensionCache
+{
+    using System;
+    using System.Collections.Generic;
+    using WixToolset.Extensibility.Services;
+
+    public static class WixToolsetCoreServiceProviderExtensions
+    {
+        public static IWixToolsetCoreServiceProvider AddExtensionCacheManager(this IWixToolsetCoreServiceProvider serviceProvider)
+        {
+            var extensionManager = serviceProvider.GetService<IExtensionManager>();
+            extensionManager.Add(typeof(ExtensionCacheManagerExtensionFactory).Assembly);
+
+            serviceProvider.AddService(CreateExtensionCacheManager);
+            return serviceProvider;
+        }
+
+        private static ExtensionCacheManager CreateExtensionCacheManager(IWixToolsetCoreServiceProvider provider, Dictionary<Type, object> singletons)
+        {
+            var extensionCacheManager = new ExtensionCacheManager();
+            singletons.Add(typeof(ExtensionCacheManager), extensionCacheManager);
+
+            return extensionCacheManager;
+        }
+    }
+}