Add AssemblyLoadContext.SetAssemblyLocationOverride

src/coreclr/System.Private.CoreLib/src/System/Reflection/RuntimeAssembly.cs

@@ -519,7 +519,7 @@ public override string Location
                 RuntimeAssembly runtimeAssembly = this;
                 GetLocation(new QCallAssembly(ref runtimeAssembly), new StringHandleOnStack(ref location));
 
-                return location!;
+                return AssemblyLoadContext.ResolveAssemblyLocation(this, location!);
             }
         }
 

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Reflection/Runtime/General/ThunkedApis.cs

@@ -63,7 +63,7 @@ public override string Location
         {
             get
             {
-                return string.Empty;
+                return System.Runtime.Loader.AssemblyLoadContext.ResolveAssemblyLocation(this, string.Empty);
             }
         }
 

src/libraries/System.Private.CoreLib/src/Resources/Strings.resx

@@ -4052,6 +4052,9 @@
   <data name="ResourceManager_ReflectionNotAllowed" xml:space="preserve">
     <value>Use of ResourceManager for custom types is disabled. Set the MSBuild Property CustomResourceTypesSupport to true in order to enable it.</value>
   </data>
+  <data name="InvalidOperation_AssemblyLocationOverrideAlreadySet" xml:space="preserve">
+    <value>The assembly location override has already been set.</value>
+  </data>
   <data name="InvalidOperation_AssemblyNotEditable" xml:space="preserve">
     <value>The assembly can not be edited or changed.</value>
   </data>

src/libraries/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.cs

@@ -309,6 +309,45 @@ public static AssemblyName GetAssemblyName(string assemblyPath)
             return AssemblyName.GetAssemblyName(assemblyPath);
         }
 
+        // Callback that, when set, can override the value returned by Assembly.Location.
+        private static Func<Assembly, string, string>? s_assemblyLocationOverride;
+
+        /// <summary>
+        /// Sets a process-wide callback that overrides the value returned by <see cref="Assembly.Location"/>.
+        /// </summary>
+        /// <remarks>
+        /// The callback can only be set once for the lifetime of the process. The callback should not
+        /// call <see cref="Assembly.Location"/> on the provided assembly to avoid recursion.
+        /// </remarks>
+        /// <param name="locationOverride">
+        /// A callback that receives an <see cref="Assembly"/> and the location the runtime computed for it, and
+        /// returns the value that <see cref="Assembly.Location"/> should report.
+        /// </param>
+        /// <exception cref="ArgumentNullException"><paramref name="locationOverride"/> is <see langword="null"/>.</exception>
+        /// <exception cref="InvalidOperationException">The location override has already been set.</exception>
+        public static void SetAssemblyLocationOverride(Func<Assembly, string, string> locationOverride)
+        {
+            ArgumentNullException.ThrowIfNull(locationOverride);
+
+            if (Interlocked.CompareExchange(ref s_assemblyLocationOverride, locationOverride, null) is not null)
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_AssemblyLocationOverrideAlreadySet);
+            }
+        }
+
+        // Applies the location override callback (if any) to the location the runtime computed for the assembly.
+        // Called from each runtime's RuntimeAssembly.Location implementation.
+        internal static string ResolveAssemblyLocation(Assembly assembly, string originalLocation)
+        {
+            Func<Assembly, string, string>? locationOverride = s_assemblyLocationOverride;
+            if (locationOverride is null)
+            {
+                return originalLocation;
+            }
+
+            return locationOverride(assembly, originalLocation);
+        }
+
         // Custom AssemblyLoadContext implementations can override this
         // method to perform custom processing and use one of the protected
         // helpers above to load the assembly.

src/libraries/System.Runtime.Loader/ref/System.Runtime.Loader.cs

@@ -87,6 +87,7 @@ public event System.Action<System.Runtime.Loader.AssemblyLoadContext>? Unloading
         public System.Reflection.Assembly LoadFromStream(System.IO.Stream assembly, System.IO.Stream? assemblySymbols) { throw null; }
         protected virtual System.IntPtr LoadUnmanagedDll(string unmanagedDllName) { throw null; }
         protected System.IntPtr LoadUnmanagedDllFromPath(string unmanagedDllPath) { throw null; }
+        public static void SetAssemblyLocationOverride(System.Func<System.Reflection.Assembly, string, string> locationOverride) { }
         public void SetProfileOptimizationRoot(string directoryPath) { }
         public void StartProfileOptimization(string? profile) { }
         public override string ToString() { throw null; }

src/libraries/System.Runtime.Loader/tests/AssemblyLoadContextTest.cs

@@ -418,6 +418,63 @@ public static void LoadFromAssemblyPath_DefaultAlc_NonTpaAssembly_MvidMismatch()
             }).Dispose();
         }
 
+        [Fact]
+        public static void SetAssemblyLocationOverride_NullArgument_Throws()
+        {
+            // The null check happens before any global state is mutated, so this is safe to run in-process.
+            AssertExtensions.Throws<ArgumentNullException>("locationOverride", () => AssemblyLoadContext.SetAssemblyLocationOverride(null));
+        }
+
+        [ConditionalFact(typeof(AssemblyLoadContextTest), nameof(IsRemoteExecutorSupportedAndAssemblyLoadingSupported))]
+        public static void SetAssemblyLocationOverride_OverridesLocationForStreamLoadedAssembly()
+        {
+            // The override is process-wide and set-once, so it must run in its own process.
+            RemoteExecutor.Invoke(static () =>
+            {
+                // Override the location only for assemblies the runtime reports no location for (e.g. loaded
+                // from a stream / memory). Embed the assembly name so we also verify the Assembly argument.
+                AssemblyLoadContext.SetAssemblyLocationOverride(
+                    static (assembly, location) => string.IsNullOrEmpty(location) ? $"/overridden/{assembly.GetName().Name}.dll" : location);
+
+                string asmPath = ExtractEmbeddedAssembly("System.Runtime.Loader.Tests.AssemblyVersion1");
+                try
+                {
+                    // Loading from a stream means the runtime has no location, so the override kicks in.
+                    var streamAlc = new AssemblyLoadContext("LocationOverride_Stream", isCollectible: true);
+                    Assembly streamLoaded;
+                    using (FileStream fs = File.OpenRead(asmPath))
+                    {
+                        streamLoaded = streamAlc.LoadFromStream(fs);
+                    }
+                    Assert.Equal($"/overridden/{streamLoaded.GetName().Name}.dll", streamLoaded.Location);
+
+                    // The same assembly loaded from a path (into a separate context) has a real location,
+                    // so the callback leaves it untouched.
+                    var pathAlc = new AssemblyLoadContext("LocationOverride_Path", isCollectible: true);
+                    Assembly fileLoaded = pathAlc.LoadFromAssemblyPath(asmPath);
+                    Assert.False(string.IsNullOrEmpty(fileLoaded.Location));
+                    Assert.DoesNotContain("/overridden/", fileLoaded.Location);
+                }
+                finally
+                {
+                    try { File.Delete(asmPath); } catch { }
+                }
+            }).Dispose();
+        }
+
+        [ConditionalFact(typeof(AssemblyLoadContextTest), nameof(IsRemoteExecutorSupportedAndAssemblyLoadingSupported))]
+        public static void SetAssemblyLocationOverride_CalledTwice_Throws()
+        {
+            RemoteExecutor.Invoke(static () =>
+            {
+                AssemblyLoadContext.SetAssemblyLocationOverride(static (assembly, location) => location);
+                Assert.Throws<InvalidOperationException>(
+                    () => AssemblyLoadContext.SetAssemblyLocationOverride(static (assembly, location) => location));
+            }).Dispose();
+        }
+
+        private static bool IsRemoteExecutorSupportedAndAssemblyLoadingSupported => RemoteExecutor.IsSupported && PlatformDetection.IsAssemblyLoadingSupported;
+
         private static bool IsRemoteExecutorSupportedAndCoreCLR => RemoteExecutor.IsSupported && PlatformDetection.IsAssemblyLoadingSupported && PlatformDetection.IsCoreCLR;
 
         private static string ExtractEmbeddedAssembly(string name)

src/mono/System.Private.CoreLib/src/System/Reflection/RuntimeAssembly.cs

@@ -131,7 +131,7 @@ public override Module ManifestModule
 
         public override string ImageRuntimeVersion => GetInfo(AssemblyInfoKind.ImageRuntimeVersion)!;
 
-        public override string Location => GetInfo(AssemblyInfoKind.Location)!;
+        public override string Location => AssemblyLoadContext.ResolveAssemblyLocation(this, GetInfo(AssemblyInfoKind.Location)!);
 
         // TODO: consider a dedicated icall instead
         public override bool IsCollectible => AssemblyLoadContext.GetLoadContext((Assembly)this)!.IsCollectible;