From 2aea69215c681691281f4e7114f07c65fe4d6327 Mon Sep 17 00:00:00 2001
From: Mikayla Hutchinson <m.j.hutchinson@gmail.com>
Date: Wed, 10 Jan 2024 18:34:25 -0500
Subject: [PATCH] Fix AssemblyLoadContext assembly loading behavior

The behavior should now be equivalent to that of the AppDomain codepath:
assemblies are loaded from referenced assembly paths only after
failing to resolve them from loaded assemblies in the host context
and after the runtime fails to resolve them from the default context.
---
 .../TemplateAssemblyLoadContext.cs            | 53 +++++++++++++++++--
 1 file changed, 48 insertions(+), 5 deletions(-)

diff --git a/Mono.TextTemplating/Mono.TextTemplating/TemplateAssemblyLoadContext.cs b/Mono.TextTemplating/Mono.TextTemplating/TemplateAssemblyLoadContext.cs
index 7453e0c..fb48b73 100644
--- a/Mono.TextTemplating/Mono.TextTemplating/TemplateAssemblyLoadContext.cs
+++ b/Mono.TextTemplating/Mono.TextTemplating/TemplateAssemblyLoadContext.cs
@@ -6,6 +6,7 @@
 using System;
 using System.CodeDom.Compiler;
 using System.IO;
+using System.Linq;
 using System.Reflection;
 using System.Runtime.Loader;
 
@@ -15,6 +16,8 @@ namespace Mono.TextTemplating
 {
 	class TemplateAssemblyLoadContext : AssemblyLoadContext
 	{
+		readonly AssemblyLoadContext hostAssemblyLoadContext;
+
 		readonly string[] templateAssemblyFiles;
 		readonly ITextTemplatingEngineHost host;
 		readonly Assembly hostAssembly;
@@ -28,19 +31,43 @@ public TemplateAssemblyLoadContext (string[] templateAssemblyFiles, ITextTemplat
 			: base (isCollectible: true)
 #endif
 		{
-			this.templateAssemblyFiles = templateAssemblyFiles;
-			this.host = host;
 			hostAssembly = host.GetType ().Assembly;
 
 			codeDomAsmName = typeof (CompilerErrorCollection).Assembly.GetName ();
 			textTemplatingAsmName = typeof (TemplateGenerator).Assembly.GetName ();
 			hostAsmName = hostAssembly.GetName ();
+
+			// the parent load context is the context that loaded the host assembly
+			this.hostAssemblyLoadContext = AssemblyLoadContext.GetLoadContext (hostAssembly);
+
+			this.templateAssemblyFiles = templateAssemblyFiles;
+			this.host = host;
+
+			this.Resolving += ResolveAssembly;
 		}
 
+		// Load order is as follows:
+		//
+		// First, the Load(AssemblyName) override is called. Our impl of this ensures that the CodeDom and TextTemplating
+		// and other host assemblies are loaded from the host AssemblyLoadContext, so that we can interchange types.
+		//
+		// For assemblies that are not handled by Load(AssemblyName), the runtime next attempts to resolve them
+		// from AssemblyLoadContext.Default, which may load assemblies into AssemblyLoadContext.Default via
+		// assembly probing. This is where runtime assemblies wil be loaded.
+		//
+		// Finally, if the runtime fails to resolve the assembly, the Resolving event is raised, which we handle
+		// to resolve assemblies explicitly referenced by the template. The priority of this event is equivalent
+		// to AppDomain.AssemblyResolve, so using this matches the behavior of the AppDomain codepath.
+		//
+		// See https://learn.microsoft.com/en-us/dotnet/core/dependency-loading/loading-managed#algorithm
+
 		protected override Assembly Load (AssemblyName assemblyName)
 		{
-			// CodeDom and TextTemplating MUST come from the same context as the host as we need to be able to reflect and cast
-			// this is an issue with MSBuild which loads tasks in their own load contexts
+			// The CodeDom and TextTemplating assemblies in the template load context MUST be the same as the
+			// ones used by the host as we need to be able to reflect and cast and interchange types.
+			// We cannot rely on falling back to resolving them from AssemblyLoadContext.Default, as this fails
+			// in cases such as running the templating engine in an MSBuild task, as MSBuild loads tasks in
+			// their own load contexts.
 			if (assemblyName.Name == codeDomAsmName.Name) {
 				return typeof (CompilerErrorCollection).Assembly;
 			}
@@ -48,11 +75,27 @@ protected override Assembly Load (AssemblyName assemblyName)
 				return typeof (TemplateGenerator).Assembly;
 			}
 
-			// the host may be a custom host, and must also be in the same context
+			// The host may be a custom host, and must also be in the same context, so that host-specific
+			// templates can access the host instance.
 			if (assemblyName.Name == hostAsmName.Name) {
 				return hostAssembly;
 			}
 
+			// Resolve any more assemblies from the parent context that we can. There may be a custom host,
+			// and it may expose types from other assemblies that may also need to be interchangeable.
+			// Technically this loops makes the explicit checks above redundant but it's better
+			// to be absolutely clear about what we're doing and why.
+			var fromParent = hostAssemblyLoadContext.Assemblies.FirstOrDefault (a => a.GetName ().Name == assemblyName.Name);
+			if (fromParent is not null) {
+				return fromParent;
+			}
+
+			// let the runtime resolve from AssemblyLoadContext.Default
+			return null;
+		}
+
+		Assembly ResolveAssembly (AssemblyLoadContext context, AssemblyName assemblyName)
+		{
 			for (int i = 0; i < templateAssemblyFiles.Length; i++) {
 				var asmFile = templateAssemblyFiles[i];
 				if (asmFile is null) {