diff --git a/src/Authentication/Authentication.Loader/GraphLoadContextInitializer.cs b/src/Authentication/Authentication.Loader/GraphLoadContextInitializer.cs index 918648b4074..58a30c810c4 100644 --- a/src/Authentication/Authentication.Loader/GraphLoadContextInitializer.cs +++ b/src/Authentication/Authentication.Loader/GraphLoadContextInitializer.cs @@ -22,6 +22,7 @@ public static class GraphLoadContextInitializer private static readonly object s_lock = new object(); private static GraphAssemblyLoadContext s_context; + private static bool s_resolvingHooked; /// /// The isolated context; null until has been called. @@ -29,36 +30,43 @@ public static class GraphLoadContextInitializer public static GraphAssemblyLoadContext Context => s_context; /// - /// Creates the isolated context and registers the default-context redirect. Safe to call multiple times. + /// Creates the isolated context (once) and ensures the default-context redirect is registered. Safe to call + /// multiple times, including after a triggered by a module remove: the isolated context is + /// reused while the hook is re-established so a remove/re-import cycle + /// keeps resolving Microsoft.Graph.Authentication.Core. /// /// Path to the module's Dependencies folder. /// Path to the module's Dependencies/Core folder. public static void Initialize(string dependencyFolder, string psEditionDependencyFolder) { - if (s_context != null) - return; - lock (s_lock) { - if (s_context != null) - return; + if (s_context == null) + { + s_context = new GraphAssemblyLoadContext(dependencyFolder, psEditionDependencyFolder); + } - s_context = new GraphAssemblyLoadContext(dependencyFolder, psEditionDependencyFolder); - AssemblyLoadContext.Default.Resolving += OnDefaultResolving; + if (!s_resolvingHooked) + { + AssemblyLoadContext.Default.Resolving += OnDefaultResolving; + s_resolvingHooked = true; + } } } /// /// Unregisters the default-context redirect. The isolated context itself is not collectible and remains alive - /// for the lifetime of the process, which matches PowerShell's own behaviour for binary modules. + /// for the lifetime of the process, which matches PowerShell's own behaviour for binary modules; a later + /// call re-registers the redirect against the existing context. /// public static void Shutdown() { lock (s_lock) { - if (s_context == null) + if (!s_resolvingHooked) return; AssemblyLoadContext.Default.Resolving -= OnDefaultResolving; + s_resolvingHooked = false; } } diff --git a/src/Authentication/Authentication/ModuleInitializer.cs b/src/Authentication/Authentication/ModuleInitializer.cs index 66b32889721..83e50ce0d7a 100644 --- a/src/Authentication/Authentication/ModuleInitializer.cs +++ b/src/Authentication/Authentication/ModuleInitializer.cs @@ -65,7 +65,16 @@ static ModuleInitializer() public void OnImport() { if (s_isPsCore) + { + // The static constructor only runs on the first load of this assembly (it stays in the default + // AssemblyLoadContext for the process lifetime). On a remove/re-import cycle OnRemove -> Shutdown + // unhooks the default-context redirect, so re-run Initialize here to re-establish it against the + // already-created isolated context. + s_loaderInitializer? + .GetMethod("Initialize", BindingFlags.Public | BindingFlags.Static, null, new[] { typeof(string), typeof(string) }, null)? + .Invoke(null, new object[] { s_dependencyFolder, s_psEditionDependencyFolder }); return; + } AppDomain.CurrentDomain.AssemblyResolve += ResolvingHandler; } diff --git a/src/Authentication/Authentication/Properties/launchSettings.json b/src/Authentication/Authentication/Properties/launchSettings.json index 26a5601995e..124e9418aa2 100644 --- a/src/Authentication/Authentication/Properties/launchSettings.json +++ b/src/Authentication/Authentication/Properties/launchSettings.json @@ -2,7 +2,7 @@ "profiles": { "Graph.Authentication": { "commandName": "Executable", - "executablePath": "C:\\Program Files\\PowerShell\\7\\pwsh.exe", + "executablePath": "pwsh.exe", "commandLineArgs": "-NoExit" } } diff --git a/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 b/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 index 3c4a1828d29..baa41f5fd24 100644 --- a/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 +++ b/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 @@ -72,4 +72,36 @@ Describe "Microsoft.Graph.Authentication module" { $PSModuleInfo.Guid.Guid | Should -Be "883916f2-9184-46ee-b1f8-b6a2fb784cee" } } + + Context "On module re-import" { + # Regression guard for the switch to a private AssemblyLoadContext on PowerShell 7+. + # Importing, removing and re-importing the module must succeed: the second import must not fail + # because the isolated load context (and its already-loaded dependencies) is being initialized again. + It 'Should import, remove and re-import without error' { + { + Import-Module $ModulePath -Force + Remove-Module $ModuleName -Force + Import-Module $ModulePath -Force + } | Should -Not -Throw + } + + It 'Should expose its cmdlets after a re-import' { + Import-Module $ModulePath -Force + Remove-Module $ModuleName -Force + $ReimportedInfo = Import-Module $ModulePath -Force -PassThru + + $ReimportedInfo.ExportedCommands.Keys | Should -Contain "Connect-MgGraph" + $ReimportedInfo.ExportedCommands.Keys | Should -Contain "Invoke-MgGraphRequest" + } + + It 'Should keep the same isolated assemblies loaded after a re-import' { + Import-Module $ModulePath -Force + Remove-Module $ModuleName -Force + Import-Module $ModulePath -Force + + # Invoking a cmdlet forces the isolated dependencies (Azure.Identity / MSAL / Kiota) to bind. + # If the re-import left the load context in a broken state this throws instead of returning $null. + { Get-MgContext } | Should -Not -Throw + } + } } \ No newline at end of file