PowerShell · daxian-dbw · Apr 4, 2024 · Mar 21, 2024
diff --git a/src/System.Management.Automation/CoreCLR/CorePsAssemblyLoadContext.cs b/src/System.Management.Automation/CoreCLR/CorePsAssemblyLoadContext.cs
@@ -36,16 +36,19 @@ internal sealed partial class PowerShellAssemblyLoadContext
         /// <summary>
         /// Initialize a singleton of PowerShellAssemblyLoadContext.
         /// </summary>
-        internal static PowerShellAssemblyLoadContext InitializeSingleton(string basePaths)
+        internal static PowerShellAssemblyLoadContext InitializeSingleton(string basePaths, bool throwOnReentry)
         {
             lock (s_syncObj)
             {
-                if (Instance != null)
+                if (Instance is null)
+                {
+                    Instance = new PowerShellAssemblyLoadContext(basePaths);
+                }
+                else if (throwOnReentry)
                 {
                     throw new InvalidOperationException(SingletonAlreadyInitialized);
                 }
 
-                Instance = new PowerShellAssemblyLoadContext(basePaths);
                 return Instance;
             }
         }
@@ -581,7 +584,8 @@ public static void SetPowerShellAssemblyLoadContext([MarshalAs(UnmanagedType.LPW
         {
             ArgumentException.ThrowIfNullOrEmpty(basePaths);
 
-            PowerShellAssemblyLoadContext.InitializeSingleton(basePaths);
+            // Disallow calling this method from native code for more than once.
+            PowerShellAssemblyLoadContext.InitializeSingleton(basePaths, throwOnReentry: true);
         }
     }
 

diff --git a/src/System.Management.Automation/engine/hostifaces/ConnectionBase.cs b/src/System.Management.Automation/engine/hostifaces/ConnectionBase.cs
@@ -25,6 +25,28 @@ internal abstract class RunspaceBase : Runspace
     {
         #region constructors
 
+        /// <summary>
+        /// Initialize powershell AssemblyLoadContext and register the 'Resolving' event, if it's not done already.
+        /// If powershell is hosted by a native host such as DSC, then PS ALC may be initialized via 'SetPowerShellAssemblyLoadContext' before loading S.M.A.
+        /// </summary>
+        /// <remarks>
+        /// We do this both here and during the initialization of the 'ClrFacade' type.
+        /// This is because we want to make sure the assembly/library resolvers are:
+        ///  1. registered before any script/cmdlet can run.
+        ///  2. registered before 'ClrFacade' gets used for assembly related operations.
+        ///
+        /// The 'ClrFacade' type may be used without a Runspace created, for example, by calling type conversion methods in the 'LanguagePrimitive' type.
+        /// And at the mean time, script or cmdlet may run without the 'ClrFacade' type initialized.
+        /// That's why we attempt to create the singleton of 'PowerShellAssemblyLoadContext' at both places.
+        /// </remarks>
+        static RunspaceBase()
+        {
+            if (PowerShellAssemblyLoadContext.Instance is null)
+            {
+                PowerShellAssemblyLoadContext.InitializeSingleton(string.Empty, throwOnReentry: false);
+            }
+        }
+
         /// <summary>
         /// Construct an instance of an Runspace using a custom
         /// implementation of PSHost.

diff --git a/src/System.Management.Automation/utils/ClrFacade.cs b/src/System.Management.Automation/utils/ClrFacade.cs
@@ -22,13 +22,23 @@ internal static class ClrFacade
     {
         /// <summary>
         /// Initialize powershell AssemblyLoadContext and register the 'Resolving' event, if it's not done already.
-        /// If powershell is hosted by a native host such as DSC, then PS ALC might be initialized via 'SetPowerShellAssemblyLoadContext' before loading S.M.A.
+        /// If powershell is hosted by a native host such as DSC, then PS ALC may be initialized via 'SetPowerShellAssemblyLoadContext' before loading S.M.A.
         /// </summary>
+        /// <remarks>
+        /// We do this both here and during the initialization of the 'RunspaceBase' type.
+        /// This is because we want to make sure the assembly/library resolvers are:
+        ///  1. registered before any script/cmdlet can run.
+        ///  2. registered before 'ClrFacade' gets used for assembly related operations.
+        ///
+        /// The 'ClrFacade' type may be used without a Runspace created, for example, by calling type conversion methods in the 'LanguagePrimitive' type.
+        /// And at the mean time, script or cmdlet may run without the 'ClrFacade' type initialized.
+        /// That's why we attempt to create the singleton of 'PowerShellAssemblyLoadContext' at both places.
+        /// </remarks>
         static ClrFacade()
         {
-            if (PowerShellAssemblyLoadContext.Instance == null)
+            if (PowerShellAssemblyLoadContext.Instance is null)
             {
-                PowerShellAssemblyLoadContext.InitializeSingleton(string.Empty);
+                PowerShellAssemblyLoadContext.InitializeSingleton(string.Empty, throwOnReentry: false);
             }
         }
 

diff --git a/test/powershell/engine/Basic/RegisterAssemblyResolverEarly.ps1 b/test/powershell/engine/Basic/RegisterAssemblyResolverEarly.ps1
@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+Describe "Assembly resolvers should be registered early at startup" -Tags "CI" {
+
+    ## The PKI module requires loading an assembly from GAC, and thus depends on the PowerShell assembly resolver to work.
+    It "Can load the PKI module with 'pwsh -ExecutionPolicy bypass -NoProfile -c `"Import-Module PKI`"'" -Skip:(!$IsWindows) {
+        ## Use 'Bypass' execution policy so that it doesn't trigger 'AuthorizationManager' which would trigger 'ClrFacade' initialization.
+        ## We want to make sure even if 'ClrFacade' is not hit during startup, the resolvers are still registered early enough.
+        $out = pwsh -ExecutionPolicy bypass -NoProfile -c "Import-Module PKI; Get-Module | % name"
+        $out | Should -BeExactly "PKI"
+    }
+}