Enable hosts to provide custom assembly resolution

src/Compilers/Core/CodeAnalysisTest/AnalyzerAssemblyLoaderTests.cs

@@ -54,7 +54,7 @@ public enum AnalyzerTestKind
     ///
     /// Limitation 1: .NET Framework probing path.
     ///
-    /// The .NET Framework assembly loader will only call AppDomain.AssemblyResolve when it cannot satifisfy a load
+    /// The .NET Framework assembly loader will only call AppDomain.AssemblyResolve when it cannot satisfy a load
     /// request. One of the places the assembly loader will always consider when looking for dependencies of A.dll
     /// is the directory that A.dll was loading from (it's added to the probing path). That means if B.dll is in the
     /// same directory then the runtime will silently load it without a way for us to intervene.
@@ -95,25 +95,27 @@ public AnalyzerAssemblyLoaderTests(ITestOutputHelper testOutputHelper, AssemblyL
 
 #if NETCOREAPP
 
-        private void Run(AnalyzerTestKind kind, Action<AnalyzerAssemblyLoader, AssemblyLoadTestFixture> testAction, [CallerMemberName] string? memberName = null) =>
+        private void Run(AnalyzerTestKind kind, Action<AnalyzerAssemblyLoader, AssemblyLoadTestFixture> testAction, IAnalyzerAssemblyResolver[]? externalResolvers = null, [CallerMemberName] string? memberName = null) =>
             Run(
                 kind,
                 static (_, _) => { },
                 testAction,
+                externalResolvers,
                 memberName);
 
         private void Run(
             AnalyzerTestKind kind,
             Action<AssemblyLoadContext, AssemblyLoadTestFixture> prepLoadContextAction,
             Action<AnalyzerAssemblyLoader, AssemblyLoadTestFixture> testAction,
+            IAnalyzerAssemblyResolver[]? externalResolvers = null,
             [CallerMemberName] string? memberName = null)
         {
             var alc = new AssemblyLoadContext($"Test {memberName}", isCollectible: true);
             try
             {
                 prepLoadContextAction(alc, TestFixture);
                 var util = new InvokeUtil();
-                util.Exec(TestOutputHelper, alc, TestFixture, kind, testAction.Method.DeclaringType!.FullName!, testAction.Method.Name);
+                util.Exec(TestOutputHelper, alc, TestFixture, kind, testAction.Method.DeclaringType!.FullName!, testAction.Method.Name, externalResolvers ?? []);
             }
             finally
             {
@@ -126,6 +128,7 @@ public AnalyzerAssemblyLoaderTests(ITestOutputHelper testOutputHelper, AssemblyL
         private void Run(
             AnalyzerTestKind kind,
             Action<AnalyzerAssemblyLoader, AssemblyLoadTestFixture> testAction,
+            IAnalyzerAssemblyResolver[]? externalResolvers = null,
             [CallerMemberName] string? memberName = null)
         {
             AppDomain? appDomain = null;
@@ -135,7 +138,7 @@ public AnalyzerAssemblyLoaderTests(ITestOutputHelper testOutputHelper, AssemblyL
                 var testOutputHelper = new AppDomainTestOutputHelper(TestOutputHelper);
                 var type = typeof(InvokeUtil);
                 var util = (InvokeUtil)appDomain.CreateInstanceAndUnwrap(type.Assembly.FullName, type.FullName);
-                util.Exec(testOutputHelper, TestFixture, kind, testAction.Method.DeclaringType.FullName, testAction.Method.Name);
+                util.Exec(testOutputHelper, TestFixture, kind, testAction.Method.DeclaringType.FullName, testAction.Method.Name, externalResolvers ?? []);
             }
             finally
             {
@@ -1421,5 +1424,79 @@ public void AssemblyLoadingInNonDefaultContext_AnalyzerReferencesSystemCollectio
                 });
         }
 #endif
+
+        [Theory]
+        [CombinatorialData]
+        public void ExternalResolver_CanIntercept_ReturningNull(AnalyzerTestKind kind)
+        {
+            var resolver = new TestAnalyzerAssemblyResolver(n => null);
+            Run(kind, (AnalyzerAssemblyLoader loader, AssemblyLoadTestFixture testFixture) =>
+            {
+                loader.AddDependencyLocation(testFixture.Delta1);
+                Assembly delta = loader.LoadFromPath(testFixture.Delta1);
+                Assert.NotNull(delta);
+                VerifyDependencyAssemblies(loader, testFixture.Delta1);
+
+            }, externalResolvers: [resolver]);
+            Assert.Collection(resolver.CalledFor, (a => Assert.Equal("Delta", a.Name)));
+        }
+
+        [Theory]
+        [CombinatorialData]
+        public void ExternalResolver_CanIntercept_ReturningAssembly(AnalyzerTestKind kind)
+        {
+            var resolver = new TestAnalyzerAssemblyResolver(n => GetType().Assembly);
+            Run(kind, (AnalyzerAssemblyLoader loader, AssemblyLoadTestFixture testFixture) =>
+            {
+                // net core assembly loader checks that the resolved assembly name is the same as the requested one
+                // so we use the assembly the tests are contained in as its already be loaded
+                var thisAssembly = typeof(AnalyzerAssemblyLoaderTests).Assembly;
+                loader.AddDependencyLocation(thisAssembly.Location);
+                Assembly loaded = loader.LoadFromPath(thisAssembly.Location);
+                Assert.Equal(thisAssembly, loaded);
+
+            }, externalResolvers: [resolver]);
+            Assert.Collection(resolver.CalledFor, (a => Assert.Equal(GetType().Assembly.GetName().Name, a.Name)));
+        }
+
+        [Theory]
+        [CombinatorialData]
+        public void ExternalResolver_CanIntercept_ReturningAssembly_Or_Null(AnalyzerTestKind kind)
+        {
+            var thisAssemblyName = GetType().Assembly.GetName();
+            var resolver = new TestAnalyzerAssemblyResolver(n => n == thisAssemblyName ? GetType().Assembly : null);
+            Run(kind, (AnalyzerAssemblyLoader loader, AssemblyLoadTestFixture testFixture) =>
+            {
+                var thisAssembly = typeof(AnalyzerAssemblyLoaderTests).Assembly;
+
+                loader.AddDependencyLocation(testFixture.Alpha);
+                Assembly alpha = loader.LoadFromPath(testFixture.Alpha);
+                Assert.NotNull(alpha);
+
+                loader.AddDependencyLocation(thisAssembly.Location);
+                Assembly loaded = loader.LoadFromPath(thisAssembly.Location);
+                Assert.Equal(thisAssembly, loaded);
+
+                loader.AddDependencyLocation(testFixture.Delta1);
+                Assembly delta = loader.LoadFromPath(testFixture.Delta1);
+                Assert.NotNull(delta);
+
+            }, externalResolvers: [resolver]);
+            Assert.Collection(resolver.CalledFor, (a => Assert.Equal("Alpha", a.Name)), a => Assert.Equal(thisAssemblyName.Name, a.Name), a => Assert.Equal("Delta", a.Name));
+        }
+
+        [Serializable]
+        class TestAnalyzerAssemblyResolver(Func<AssemblyName, Assembly?> func) : MarshalByRefObject, IAnalyzerAssemblyResolver
+        {
+            private readonly Func<AssemblyName, Assembly?> _func = func;
+
+            public List<AssemblyName> CalledFor { get; } = [];
+
+            public Assembly? ResolveAssembly(AssemblyName assemblyName)
+            {
+                CalledFor.Add(assemblyName);
+                return _func(assemblyName);
+            }
+        }
     }
 }

src/Compilers/Core/CodeAnalysisTest/InvokeUtil.cs

@@ -35,7 +35,7 @@ namespace Microsoft.CodeAnalysis.UnitTests
 
     public sealed class InvokeUtil
     {
-        public void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadContext compilerContext, AssemblyLoadTestFixture fixture, AnalyzerTestKind kind, string typeName, string methodName)
+        internal void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadContext compilerContext, AssemblyLoadTestFixture fixture, AnalyzerTestKind kind, string typeName, string methodName, IAnalyzerAssemblyResolver[] externalResolvers)
         {
             // Ensure that the test did not load any of the test fixture assemblies into 
             // the default load context. That should never happen. Assemblies should either 
@@ -48,9 +48,9 @@ public void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadContext compile
             using var tempRoot = new TempRoot();
             AnalyzerAssemblyLoader loader = kind switch
             {
-                AnalyzerTestKind.LoadDirect => new DefaultAnalyzerAssemblyLoader(compilerContext, AnalyzerLoadOption.LoadFromDisk),
-                AnalyzerTestKind.LoadStream => new DefaultAnalyzerAssemblyLoader(compilerContext, AnalyzerLoadOption.LoadFromStream),
-                AnalyzerTestKind.ShadowLoad => new ShadowCopyAnalyzerAssemblyLoader(compilerContext, tempRoot.CreateDirectory().Path),
+                AnalyzerTestKind.LoadDirect => new DefaultAnalyzerAssemblyLoader(compilerContext, AnalyzerLoadOption.LoadFromDisk, externalResolvers.ToImmutableArray()),
+                AnalyzerTestKind.LoadStream => new DefaultAnalyzerAssemblyLoader(compilerContext, AnalyzerLoadOption.LoadFromStream, externalResolvers.ToImmutableArray()),
+                AnalyzerTestKind.ShadowLoad => new ShadowCopyAnalyzerAssemblyLoader(compilerContext, tempRoot.CreateDirectory().Path, externalResolvers.ToImmutableArray()),
                 _ => throw ExceptionUtilities.Unreachable()
             };
 
@@ -93,13 +93,13 @@ public void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadContext compile
 
     public sealed class InvokeUtil : MarshalByRefObject
     {
-        public void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadTestFixture fixture, AnalyzerTestKind kind, string typeName, string methodName)
+        internal void Exec(ITestOutputHelper testOutputHelper, AssemblyLoadTestFixture fixture, AnalyzerTestKind kind, string typeName, string methodName, IAnalyzerAssemblyResolver[] externalResolvers)
         {
             using var tempRoot = new TempRoot();
             AnalyzerAssemblyLoader loader = kind switch
             {
-                AnalyzerTestKind.LoadDirect => new DefaultAnalyzerAssemblyLoader(),
-                AnalyzerTestKind.ShadowLoad => new ShadowCopyAnalyzerAssemblyLoader(tempRoot.CreateDirectory().Path),
+                AnalyzerTestKind.LoadDirect => new DefaultAnalyzerAssemblyLoader(externalResolvers.ToImmutableArray()),
+                AnalyzerTestKind.ShadowLoad => new ShadowCopyAnalyzerAssemblyLoader(tempRoot.CreateDirectory().Path, externalResolvers.ToImmutableArray()),
                 _ => throw ExceptionUtilities.Unreachable()
             };
 

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerAssemblyLoader.Core.cs

@@ -43,15 +43,16 @@ internal partial class AnalyzerAssemblyLoader
         internal AssemblyLoadContext CompilerLoadContext => _compilerLoadContext;
         internal AnalyzerLoadOption AnalyzerLoadOption => _loadOption;
 
-        internal AnalyzerAssemblyLoader()
-            : this(null, AnalyzerLoadOption.LoadFromDisk)
+        internal AnalyzerAssemblyLoader(ImmutableArray<IAnalyzerAssemblyResolver> externalResolvers)
+            : this(null, AnalyzerLoadOption.LoadFromDisk, externalResolvers)
         {
         }
 
-        internal AnalyzerAssemblyLoader(AssemblyLoadContext? compilerLoadContext, AnalyzerLoadOption loadOption)
+        internal AnalyzerAssemblyLoader(AssemblyLoadContext? compilerLoadContext, AnalyzerLoadOption loadOption, ImmutableArray<IAnalyzerAssemblyResolver> externalResolvers)
         {
             _loadOption = loadOption;
             _compilerLoadContext = compilerLoadContext ?? AssemblyLoadContext.GetLoadContext(typeof(AnalyzerAssemblyLoader).GetTypeInfo().Assembly)!;
+            _externalResolvers = [.. externalResolvers, new CompilerAnalyzerAssemblyResolver(_compilerLoadContext)];
         }
 
         public bool IsHostAssembly(Assembly assembly)
@@ -69,7 +70,7 @@ public bool IsHostAssembly(Assembly assembly)
             {
                 if (!_loadContextByDirectory.TryGetValue(fullDirectoryPath, out loadContext))
                 {
-                    loadContext = new DirectoryLoadContext(fullDirectoryPath, this, _compilerLoadContext);
+                    loadContext = new DirectoryLoadContext(fullDirectoryPath, this);
                     _loadContextByDirectory[fullDirectoryPath] = loadContext;
                 }
             }
@@ -107,33 +108,23 @@ internal sealed class DirectoryLoadContext : AssemblyLoadContext
         {
             internal string Directory { get; }
             private readonly AnalyzerAssemblyLoader _loader;
-            private readonly AssemblyLoadContext _compilerLoadContext;
 
-            public DirectoryLoadContext(string directory, AnalyzerAssemblyLoader loader, AssemblyLoadContext compilerLoadContext)
+            public DirectoryLoadContext(string directory, AnalyzerAssemblyLoader loader)
                 : base(isCollectible: true)
             {
                 Directory = directory;
                 _loader = loader;
-                _compilerLoadContext = compilerLoadContext;
             }
 
             protected override Assembly? Load(AssemblyName assemblyName)
             {
-                var simpleName = assemblyName.Name!;
-                try
-                {
-                    if (_compilerLoadContext.LoadFromAssemblyName(assemblyName) is { } compilerAssembly)
-                    {
-                        return compilerAssembly;
-                    }
-                }
-                catch
+                if (_loader.ResolveAssemblyExternally(assemblyName) is { } externallyResolvedAssembly)
                 {
-                    // Expected to happen when the assembly cannot be resolved in the compiler / host
-                    // AssemblyLoadContext.
+                    return externallyResolvedAssembly;
                 }
 
                 // Prefer registered dependencies in the same directory first.
+                var simpleName = assemblyName.Name!;
                 var assemblyPath = Path.Combine(Directory, simpleName + ".dll");
                 if (_loader.IsAnalyzerDependencyPath(assemblyPath))
                 {
@@ -147,7 +138,7 @@ public DirectoryLoadContext(string directory, AnalyzerAssemblyLoader loader, Ass
                 // Note: when loading from disk the .NET runtime has a fallback step that will handle
                 // satellite assembly loading if the call to Load(satelliteAssemblyName) fails. This
                 // loader has a mode where it loads from Stream though and the runtime will not handle
-                // that automatically. Rather than bifurate our loading behavior between Disk and
+                // that automatically. Rather than bifurcate our loading behavior between Disk and
                 // Stream both modes just handle satellite loading directly
                 if (assemblyName.CultureInfo is not null && simpleName.EndsWith(".resources", StringComparison.Ordinal))
                 {
@@ -201,6 +192,13 @@ protected override IntPtr LoadUnmanagedDll(string unmanagedDllName)
                 return IntPtr.Zero;
             }
         }
+
+        internal sealed class CompilerAnalyzerAssemblyResolver(AssemblyLoadContext compilerContext) : IAnalyzerAssemblyResolver
+        {
+            private readonly AssemblyLoadContext _compilerAlc = compilerContext;
+
+            public Assembly? ResolveAssembly(AssemblyName assemblyName) => _compilerAlc.LoadFromAssemblyName(assemblyName);
+        }
     }
 }
 

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerAssemblyLoader.Desktop.cs

@@ -5,6 +5,7 @@
 #if !NETCOREAPP
 
 using System;
+using System.Collections.Immutable;
 using System.Globalization;
 using System.IO;
 using System.Reflection;
@@ -28,8 +29,9 @@ internal partial class AnalyzerAssemblyLoader
     {
         private bool _hookedAssemblyResolve;
 
-        internal AnalyzerAssemblyLoader()
+        internal AnalyzerAssemblyLoader(ImmutableArray<IAnalyzerAssemblyResolver> externalResolvers)
         {
+            _externalResolvers = externalResolvers;
         }
 
         public bool IsHostAssembly(Assembly assembly)
@@ -56,6 +58,11 @@ public bool IsHostAssembly(Assembly assembly)
 
         private partial Assembly? Load(AssemblyName assemblyName, string assemblyOriginalPath)
         {
+            if (ResolveAssemblyExternally(assemblyName) is { } externallyResolvedAssembly)
+            {
+                return externallyResolvedAssembly;
+            }
+
             EnsureResolvedHooked();
 
             return AppDomain.CurrentDomain.Load(assemblyName);

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerAssemblyLoader.cs

@@ -61,6 +61,14 @@ internal abstract partial class AnalyzerAssemblyLoader : IAnalyzerAssemblyLoader
         /// </remarks>
         private readonly Dictionary<string, ImmutableHashSet<string>> _knownAssemblyPathsBySimpleName = new(StringComparer.OrdinalIgnoreCase);
 
+        /// <summary>
+        /// A collection of <see cref="IAnalyzerAssemblyResolver"/>s that can be used to override the assembly resolution process.
+        /// </summary>
+        /// <remarks>
+        /// When multiple resolvers are present they are consulted in-order, with the first resolver to return a non-null
+        /// <see cref="Assembly"/> winning.</remarks>
+        private readonly ImmutableArray<IAnalyzerAssemblyResolver> _externalResolvers;
+
         /// <summary>
         /// The implementation needs to load an <see cref="Assembly"/> with the specified <see cref="AssemblyName"/>. The
         /// <paramref name="assemblyOriginalPath"/> parameter is the original path. It may be different than
@@ -330,5 +338,33 @@ internal string GetRealAnalyzerLoadPath(string originalFullPath)
                     .ToArray();
             }
         }
+
+        /// <summary>
+        /// Iterates the <see cref="_externalResolvers"/> if any, to see if any of them can resolve
+        /// the given <see cref="AssemblyName"/> to an <see cref="Assembly"/>.
+        /// </summary>
+        /// <param name="assemblyName">The name of the assembly to resolve</param>
+        /// <returns>An <see langword="assembly"/> if one of the resolvers is successful, or <see langword="null"/></returns>
+        internal Assembly? ResolveAssemblyExternally(AssemblyName assemblyName)
+        {
+            if (!_externalResolvers.IsDefaultOrEmpty)
+            {
+                foreach (var resolver in _externalResolvers)
+                {
+                    try
+                    {
+                        if (resolver.ResolveAssembly(assemblyName) is { } resolvedAssembly)
+                        {
+                            return resolvedAssembly;
+                        }
+                    }
+                    catch
+                    {
+                        // Ignore if the external resolver throws
+                    }
+                }
+            }
+            return null;
+        }
     }
 }