11 files changed, 235 insertions, 61 deletions

diff --git a/src/api/wix/WixToolset.Extensibility/Data/IExtensionCacheLocation.cs b/src/api/wix/WixToolset.Extensibility/Data/IExtensionCacheLocation.cs
new file mode 100644
index 00000000..53158875
--- /dev/null
+++ b/src/api/wix/WixToolset.Extensibility/Data/IExtensionCacheLocation.cs
@@ -0,0 +1,41 @@
+// 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.Extensibility.Data
+{
+    /// <summary>
+    /// Extension cache location scope.
+    /// </summary>
+    public enum ExtensionCacheLocationScope
+    {
+        /// <summary>
+        /// Project extension cache location.
+        /// </summary>
+        Project,
+
+        /// <summary>
+        /// User extension cache location.
+        /// </summary>
+        User,
+
+        /// <summary>
+        /// Machine extension cache location.
+        /// </summary>
+        Machine,
+    }
+
+    /// <summary>
+    /// Location where extensions may be cached.
+    /// </summary>
+    public interface IExtensionCacheLocation
+    {
+        /// <summary>
+        /// Path for  the extension cache location.
+        /// </summary>
+        string Path { get; }
+
+        /// <summary>
+        /// Scope for the extension cache location.
+        /// </summary>
+        ExtensionCacheLocationScope Scope { get; }
+    }
+}
diff --git a/src/api/wix/WixToolset.Extensibility/Services/IExtensionManager.cs b/src/api/wix/WixToolset.Extensibility/Services/IExtensionManager.cs
index 8e49c38d..fe939a59 100644
--- a/src/api/wix/WixToolset.Extensibility/Services/IExtensionManager.cs
+++ b/src/api/wix/WixToolset.Extensibility/Services/IExtensionManager.cs
@@ -4,6 +4,7 @@ namespace WixToolset.Extensibility.Services
 {
     using System.Collections.Generic;
     using System.Reflection;
+    using WixToolset.Extensibility.Data;
 
     /// <summary>
     /// Loads extensions and uses the extensions' factories to provide services.
@@ -33,6 +34,12 @@ namespace WixToolset.Extensibility.Services
         void Load(string extensionReference);
 
         /// <summary>
+        /// Gets extensions cache locations.
+        /// </summary>
+        /// <returns>List of cache locations where extensions may be found.</returns>
+        IReadOnlyCollection<IExtensionCacheLocation> GetCacheLocations();
+
+        /// <summary>
         /// Gets extensions of specified type from factories loaded into the extension manager.
         /// </summary>
         /// <typeparam name="T">Type of extension to get.</typeparam>
diff --git a/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs
index 256eeb0b..b50ad6e8 100644
--- a/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs
+++ b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManager.cs
@@ -15,22 +15,26 @@ namespace WixToolset.Core.ExtensionCache
     using NuGet.Protocol;
     using NuGet.Protocol.Core.Types;
     using NuGet.Versioning;
+    using WixToolset.Extensibility.Data;
+    using WixToolset.Extensibility.Services;
 
     /// <summary>
     /// Extension cache manager.
     /// </summary>
     internal class ExtensionCacheManager
     {
-        public string CacheFolder(bool global) => global ? this.GlobalCacheFolder() : this.LocalCacheFolder();
+        private IReadOnlyCollection<IExtensionCacheLocation> cacheLocations;
 
-        public string LocalCacheFolder() => Path.Combine(Environment.CurrentDirectory, ".wix", "extensions");
-
-        public string GlobalCacheFolder()
+        public ExtensionCacheManager(IMessaging messaging, IExtensionManager extensionManager)
         {
-            var baseFolder = Environment.GetEnvironmentVariable("WIX_EXTENSIONS") ?? Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
-            return Path.Combine(baseFolder, ".wix", "extensions");
+            this.Messaging = messaging;
+            this.ExtensionManager = extensionManager;
         }
 
+        private IMessaging Messaging { get; }
+
+        private IExtensionManager ExtensionManager { get; }
+
         public async Task<bool> AddAsync(bool global, string extension, CancellationToken cancellationToken)
         {
             if (String.IsNullOrEmpty(extension))
@@ -54,11 +58,11 @@ namespace WixToolset.Core.ExtensionCache
 
             (var extensionId, var extensionVersion) = ParseExtensionReference(extension);
 
-            var cacheFolder = this.CacheFolder(global);
+            var cacheFolder = this.GetCacheFolder(global);
 
-            cacheFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
+            var extensionFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
 
-            if (Directory.Exists(cacheFolder))
+            if (Directory.Exists(extensionFolder))
             {
                 cancellationToken.ThrowIfCancellationRequested();
 
@@ -75,7 +79,7 @@ namespace WixToolset.Core.ExtensionCache
 
             (var extensionId, var extensionVersion) = ParseExtensionReference(extension);
 
-            var cacheFolder = this.CacheFolder(global);
+            var cacheFolder = this.GetCacheFolder(global);
 
             var searchFolder = Path.Combine(cacheFolder, extensionId, extensionVersion);
 
@@ -127,6 +131,20 @@ namespace WixToolset.Core.ExtensionCache
             return Task.FromResult((IEnumerable<CachedExtension>)found);
         }
 
+        private string GetCacheFolder(bool global)
+        {
+            if (this.cacheLocations == null)
+            {
+                this.cacheLocations = this.ExtensionManager.GetCacheLocations();
+            }
+
+            var requestedScope = global ? ExtensionCacheLocationScope.User : ExtensionCacheLocationScope.Project;
+
+            var cacheLocation = this.cacheLocations.First(l => l.Scope == requestedScope);
+
+            return cacheLocation.Path;
+        }
+
         private async Task<bool> DownloadAndExtractAsync(bool global, string id, string version, CancellationToken cancellationToken)
         {
             var logger = NullLogger.Instance;
@@ -149,15 +167,22 @@ namespace WixToolset.Core.ExtensionCache
                         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)
+                        try
                         {
-                            if (nugetVersion is null || nugetVersion < availableVersion)
+                            var availableVersions = await resource.GetAllVersionsAsync(id, cache, logger, cancellationToken);
+                            foreach (var availableVersion in availableVersions)
                             {
-                                nugetVersion = availableVersion;
-                                versionSource = source;
+                                if (nugetVersion is null || nugetVersion < availableVersion)
+                                {
+                                    nugetVersion = availableVersion;
+                                    versionSource = source;
+                                }
                             }
                         }
+                        catch (FatalProtocolException e)
+                        {
+                            this.Messaging.Write(ExtensionCacheWarnings.NugetException(id, e.Message));
+                        }
                     }
 
                     if (nugetVersion is null)
@@ -168,7 +193,9 @@ namespace WixToolset.Core.ExtensionCache
 
                 var searchSources = versionSource is null ? sources : new[] { versionSource };
 
-                var extensionFolder = Path.Combine(this.CacheFolder(global), id, nugetVersion.ToString());
+                var cacheFolder = this.GetCacheFolder(global);
+
+                var extensionFolder = Path.Combine(cacheFolder, id, nugetVersion.ToString());
 
                 foreach (var source in searchSources)
                 {
@@ -183,6 +210,8 @@ namespace WixToolset.Core.ExtensionCache
                         {
                             stream.Position = 0;
 
+                            Directory.CreateDirectory(extensionFolder);
+
                             using (var archive = new PackageArchiveReader(stream))
                             {
                                 var files = PackagingConstants.Folders.Known.SelectMany(folder => archive.GetFiles(folder)).Distinct(StringComparer.OrdinalIgnoreCase);
@@ -198,7 +227,10 @@ namespace WixToolset.Core.ExtensionCache
             return false;
         }
 
-        private string ExtractProgress(string sourceFile, string targetPath, Stream fileStream) => fileStream.CopyToFile(targetPath);
+        private string ExtractProgress(string sourceFile, string targetPath, Stream fileStream)
+        {
+            return fileStream.CopyToFile(targetPath);
+        }
 
         private static (string extensionId, string extensionVersion) ParseExtensionReference(string extensionReference)
         {
diff --git a/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs
index d37ee341..008fcebd 100644
--- a/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs
+++ b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheManagerCommand.cs
@@ -25,17 +25,20 @@ namespace WixToolset.Core.ExtensionCache
         public ExtensionCacheManagerCommand(IServiceProvider serviceProvider)
         {
             this.Messaging = serviceProvider.GetService<IMessaging>();
+            this.ExtensionManager = serviceProvider.GetService<IExtensionManager>();
             this.ExtensionReferences = new List<string>();
         }
 
-        private IMessaging Messaging { get; }
-
         public bool ShowHelp { get; set; }
 
         public bool ShowLogo { get; set; }
 
         public bool StopParsing { get; set; }
 
+        private IMessaging Messaging { get; }
+
+        private IExtensionManager ExtensionManager { get; }
+
         private bool Global { get; set; }
 
         private CacheSubcommand? Subcommand { get; set; }
@@ -51,7 +54,7 @@ namespace WixToolset.Core.ExtensionCache
             }
 
             var success = false;
-            var cacheManager = new ExtensionCacheManager();
+            var cacheManager = new ExtensionCacheManager(this.Messaging, this.ExtensionManager);
 
             switch (this.Subcommand)
             {
diff --git a/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheWarnings.cs b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheWarnings.cs
new file mode 100644
index 00000000..ddcbfdea
--- /dev/null
+++ b/src/wix/WixToolset.Core.ExtensionCache/ExtensionCacheWarnings.cs
@@ -0,0 +1,24 @@
+// 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 WixToolset.Data;
+
+    internal static class ExtensionCacheWarnings
+    {
+        public static Message NugetException(string extensionId, string exceptionMessage)
+        {
+            return Message(new SourceLineNumber(extensionId), Ids.NugetException, "{0}", exceptionMessage);
+        }
+
+        private static Message Message(SourceLineNumber sourceLineNumber, Ids id, string format, params object[] args)
+        {
+            return new Message(sourceLineNumber, MessageLevel.Warning, (int)id, format, args);
+        }
+
+        public enum Ids
+        {
+            NugetException = 6100,
+        } // last available is 6499. 6500 is ExtensionCacheErrors.
+    }
+}
diff --git a/src/wix/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs b/src/wix/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs
index 424fc469..b591661e 100644
--- a/src/wix/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs
+++ b/src/wix/WixToolset.Core.ExtensionCache/WixToolsetCoreServiceProviderExtensions.cs
@@ -27,7 +27,10 @@ namespace WixToolset.Core.ExtensionCache
 
         private static ExtensionCacheManager CreateExtensionCacheManager(IWixToolsetCoreServiceProvider coreProvider, Dictionary<Type, object> singletons)
         {
-            var extensionCacheManager = new ExtensionCacheManager();
+            var messaging = coreProvider.GetService<IMessaging>();
+            var extensionManager = coreProvider.GetService<IExtensionManager>();
+
+            var extensionCacheManager = new ExtensionCacheManager(messaging, extensionManager);
             singletons.Add(typeof(ExtensionCacheManager), extensionCacheManager);
 
             return extensionCacheManager;
diff --git a/src/wix/WixToolset.Core.TestPackage/WixRunner.cs b/src/wix/WixToolset.Core.TestPackage/WixRunner.cs
index ed7c49b8..8ad3af9a 100644
--- a/src/wix/WixToolset.Core.TestPackage/WixRunner.cs
+++ b/src/wix/WixToolset.Core.TestPackage/WixRunner.cs
@@ -7,6 +7,7 @@ namespace WixToolset.Core.TestPackage
     using System.Threading;
     using System.Threading.Tasks;
     using WixToolset.Core.Burn;
+    using WixToolset.Core.ExtensionCache;
     using WixToolset.Core.WindowsInstaller;
     using WixToolset.Data;
     using WixToolset.Extensibility.Services;
@@ -65,7 +66,8 @@ namespace WixToolset.Core.TestPackage
         public static Task<int> Execute(string[] args, IWixToolsetCoreServiceProvider coreProvider, out List<Message> messages, bool warningsAsErrors = true)
         {
             coreProvider.AddWindowsInstallerBackend()
-                        .AddBundleBackend();
+                        .AddBundleBackend()
+                        .AddExtensionCacheManager();
 
             var listener = new TestMessageListener();
 
diff --git a/src/wix/WixToolset.Core.TestPackage/WixToolset.Core.TestPackage.csproj b/src/wix/WixToolset.Core.TestPackage/WixToolset.Core.TestPackage.csproj
index bda54cd1..0c78cc32 100644
--- a/src/wix/WixToolset.Core.TestPackage/WixToolset.Core.TestPackage.csproj
+++ b/src/wix/WixToolset.Core.TestPackage/WixToolset.Core.TestPackage.csproj
@@ -16,6 +16,7 @@
     <ProjectReference Include="..\WixToolset.Core.Native\WixToolset.Core.Native.csproj" PrivateAssets="true" />
     <ProjectReference Include="..\WixToolset.Core\WixToolset.Core.csproj" PrivateAssets="true" />
     <ProjectReference Include="..\WixToolset.Core.Burn\WixToolset.Core.Burn.csproj" PrivateAssets="true" />
+    <ProjectReference Include="..\WixToolset.Core.ExtensionCache\WixToolset.Core.ExtensionCache.csproj" PrivateAssets="true" />
     <ProjectReference Include="..\WixToolset.Core.WindowsInstaller\WixToolset.Core.WindowsInstaller.csproj" PrivateAssets="true" />
   </ItemGroup>
 
diff --git a/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionCacheLocation.cs b/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionCacheLocation.cs
new file mode 100644
index 00000000..b45cd315
--- /dev/null
+++ b/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionCacheLocation.cs
@@ -0,0 +1,19 @@
+// 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.ExtensibilityServices
+{
+    using WixToolset.Extensibility.Data;
+
+    internal class ExtensionCacheLocation : IExtensionCacheLocation
+    {
+        public ExtensionCacheLocation(string path, ExtensionCacheLocationScope scope)
+        {
+            this.Path = path;
+            this.Scope = scope;
+        }
+
+        public string Path { get; }
+
+        public ExtensionCacheLocationScope Scope { get; }
+    }
+}
diff --git a/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionManager.cs b/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionManager.cs
index 2340ed9e..71334841 100644
--- a/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionManager.cs
+++ b/src/wix/WixToolset.Core/ExtensibilityServices/ExtensionManager.cs
@@ -9,6 +9,7 @@ namespace WixToolset.Core.ExtensibilityServices
     using System.Reflection;
     using WixToolset.Data;
     using WixToolset.Extensibility;
+    using WixToolset.Extensibility.Data;
     using WixToolset.Extensibility.Services;
 
     internal class ExtensionManager : IExtensionManager
@@ -16,6 +17,7 @@ namespace WixToolset.Core.ExtensibilityServices
         private const string UserWixFolderName = ".wix4";
         private const string MachineWixFolderName = "WixToolset4";
         private const string ExtensionsFolderName = "extensions";
+        private const string UserEnvironmentName = "WIX_EXTENSIONS";
 
         private readonly List<IExtensionFactory> extensionFactories = new List<IExtensionFactory>();
         private readonly Dictionary<Type, List<object>> loadedExtensionsByType = new Dictionary<Type, List<object>>();
@@ -51,9 +53,14 @@ namespace WixToolset.Core.ExtensibilityServices
                 {
                     if (TryParseExtensionReference(extensionPath, out var extensionId, out var extensionVersion))
                     {
-                        foreach (var cachePath in this.CacheLocations())
+                        foreach (var cacheLocation in this.GetCacheLocations())
                         {
-                            var extensionFolder = Path.Combine(cachePath, extensionId);
+                            var extensionFolder = Path.Combine(cacheLocation.Path, extensionId);
+
+                            if (!Directory.Exists(extensionFolder))
+                            {
+                                continue;
+                            }
 
                             var versionFolder = extensionVersion;
                             if (String.IsNullOrEmpty(versionFolder) && !TryFindLatestVersionInFolder(extensionFolder, out versionFolder))
@@ -94,6 +101,32 @@ namespace WixToolset.Core.ExtensibilityServices
             }
         }
 
+        public IReadOnlyCollection<IExtensionCacheLocation> GetCacheLocations()
+        {
+            var locations = new List<IExtensionCacheLocation>();
+
+            var path = Path.Combine(Environment.CurrentDirectory, UserWixFolderName, ExtensionsFolderName);
+            locations.Add(new ExtensionCacheLocation(path, ExtensionCacheLocationScope.Project));
+
+            path = Environment.GetEnvironmentVariable(UserEnvironmentName) ?? Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
+            path = Path.Combine(path, UserWixFolderName, ExtensionsFolderName);
+            locations.Add(new ExtensionCacheLocation(path, ExtensionCacheLocationScope.User));
+
+            if (Environment.Is64BitOperatingSystem)
+            {
+                path = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonProgramFiles), MachineWixFolderName, ExtensionsFolderName);
+                locations.Add(new ExtensionCacheLocation(path, ExtensionCacheLocationScope.Machine));
+            }
+
+            path = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonProgramFilesX86), MachineWixFolderName, ExtensionsFolderName);
+            locations.Add(new ExtensionCacheLocation(path, ExtensionCacheLocationScope.Machine));
+
+            path = Path.Combine(Path.GetDirectoryName(new Uri(Assembly.GetCallingAssembly().CodeBase).LocalPath), ExtensionsFolderName);
+            locations.Add(new ExtensionCacheLocation(path, ExtensionCacheLocationScope.Machine));
+
+            return locations;
+        }
+
         public IReadOnlyCollection<T> GetServices<T>() where T : class
         {
             if (!this.loadedExtensionsByType.TryGetValue(typeof(T), out var extensions))
@@ -125,43 +158,6 @@ namespace WixToolset.Core.ExtensibilityServices
             return (IExtensionFactory)Activator.CreateInstance(type);
         }
 
-        private IEnumerable<string> CacheLocations()
-        {
-            var path = Path.Combine(Environment.CurrentDirectory, UserWixFolderName, ExtensionsFolderName);
-            if (Directory.Exists(path))
-            {
-                yield return path;
-            }
-
-            path = Environment.GetEnvironmentVariable("WIX_EXTENSIONS") ?? Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
-            path = Path.Combine(path, UserWixFolderName, ExtensionsFolderName);
-            if (Directory.Exists(path))
-            {
-                yield return path;
-            }
-
-            if (Environment.Is64BitOperatingSystem)
-            {
-                path = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonProgramFiles), MachineWixFolderName, ExtensionsFolderName);
-                if (Directory.Exists(path))
-                {
-                    yield return path;
-                }
-            }
-
-            path = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonProgramFilesX86), MachineWixFolderName, ExtensionsFolderName);
-            if (Directory.Exists(path))
-            {
-                yield return path;
-            }
-
-            path = Path.Combine(Path.GetDirectoryName(new Uri(Assembly.GetCallingAssembly().CodeBase).LocalPath), ExtensionsFolderName);
-            if (Directory.Exists(path))
-            {
-                yield return path;
-            }
-        }
-
         private static bool TryParseExtensionReference(string extensionReference, out string extensionId, out string extensionVersion)
         {
             extensionId = extensionReference ?? String.Empty;
diff --git a/src/wix/test/WixToolsetTest.CoreIntegration/ExtensionFixture.cs b/src/wix/test/WixToolsetTest.CoreIntegration/ExtensionFixture.cs
index 225355bf..d814000c 100644
--- a/src/wix/test/WixToolsetTest.CoreIntegration/ExtensionFixture.cs
+++ b/src/wix/test/WixToolsetTest.CoreIntegration/ExtensionFixture.cs
@@ -144,6 +144,52 @@ namespace WixToolsetTest.CoreIntegration
             }
         }
 
+        [Fact]
+        public void CanManipulateExtensionCache()
+        {
+            var currentFolder = Environment.CurrentDirectory;
+
+            try
+            {
+                using (var fs = new DisposableFileSystem())
+                {
+                    var folder = fs.GetFolder(true);
+                    Environment.CurrentDirectory = folder;
+
+                    var result = WixRunner.Execute(new[]
+                    {
+                        "extension", "add", "WixToolset.UI.wixext"
+                    });
+
+                    result.AssertSuccess();
+
+                    var cacheFolder = Path.Combine(folder, ".wix4", "extensions", "WixToolset.UI.wixext");
+                    Assert.True(Directory.Exists(cacheFolder), $"Expected folder '{cacheFolder}' to exist");
+
+                    result = WixRunner.Execute(new[]
+                    {
+                        "extension", "list"
+                    });
+
+                    result.AssertSuccess();
+                    var output = result.Messages.Select(m => m.ToString()).Single();
+                    Assert.StartsWith("WixToolset.UI.wixext 4.", output);
+
+                    result = WixRunner.Execute(new[]
+                    {
+                        "extension", "remove", "WixToolset.UI.wixext"
+                    });
+
+                    result.AssertSuccess();
+                    Assert.False(Directory.Exists(cacheFolder), $"Expected folder '{cacheFolder}' to NOT exist");
+                }
+            }
+            finally
+            {
+                Environment.CurrentDirectory = currentFolder;
+            }
+        }
+
         private static void Build(string[] args)
         {
             var result = WixRunner.Execute(args)