From faa068f48fd959824e4ef26e526c233e47756a95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20K=C5=82ys?= Date: Wed, 3 Jun 2026 09:57:44 +0200 Subject: [PATCH] Isolate Graph authentication assembly loading --- .../Microsoft.Graph.Authentication.psm1 | 169 +++++++++++++++++- .../Authentication/build-module.ps1 | 9 +- .../Microsoft.Graph.Authentication.Tests.ps1 | 48 ++++- 3 files changed, 216 insertions(+), 10 deletions(-) diff --git a/src/Authentication/Authentication/Microsoft.Graph.Authentication.psm1 b/src/Authentication/Authentication/Microsoft.Graph.Authentication.psm1 index 792636e9907..77d328013c0 100644 --- a/src/Authentication/Authentication/Microsoft.Graph.Authentication.psm1 +++ b/src/Authentication/Authentication/Microsoft.Graph.Authentication.psm1 @@ -1,7 +1,163 @@ -# Load the module dll +function Get-GraphAuthenticationLoadContextName { + param( + [Parameter(Mandatory = $true)] + [string] $ModulePath + ) + + $sha256 = [System.Security.Cryptography.SHA256]::Create() + try { + $pathBytes = [System.Text.Encoding]::UTF8.GetBytes($ModulePath) + $hash = [System.BitConverter]::ToString($sha256.ComputeHash($pathBytes)).Replace('-', '').Substring(0, 16) + "Microsoft.Graph.Authentication.$hash" + } + finally { + $sha256.Dispose() + } +} + +function Initialize-GraphAuthenticationAssemblyResolver { + if ('Microsoft.Graph.PowerShell.Authentication.Loader.GraphAuthenticationAssemblyResolver' -as [type]) { + return + } + + Add-Type -TypeDefinition @' +using System; +using System.Collections.Concurrent; +using System.IO; +using System.Reflection; +using System.Runtime.Loader; + +namespace Microsoft.Graph.PowerShell.Authentication.Loader +{ + public static class GraphAuthenticationAssemblyResolver + { + private static readonly ConcurrentDictionary DependencyFolders = new ConcurrentDictionary(StringComparer.Ordinal); + private static readonly ConcurrentDictionary RegisteredContexts = new ConcurrentDictionary(StringComparer.Ordinal); + + public static void Register(AssemblyLoadContext context, string[] dependencyFolders) + { + if (context == null) + { + throw new ArgumentNullException(nameof(context)); + } + + string contextName = context.Name ?? string.Empty; + DependencyFolders[contextName] = dependencyFolders ?? Array.Empty(); + + if (RegisteredContexts.TryAdd(contextName, true)) + { + context.Resolving += Resolve; + } + } + + private static Assembly Resolve(AssemblyLoadContext context, AssemblyName assemblyName) + { + if (context == null || assemblyName == null) + { + return null; + } + + if (!DependencyFolders.TryGetValue(context.Name ?? string.Empty, out string[] dependencyFolders)) + { + return null; + } + + foreach (string dependencyFolder in dependencyFolders) + { + if (string.IsNullOrWhiteSpace(dependencyFolder) || !Directory.Exists(dependencyFolder)) + { + continue; + } + + string dependencyPath = Path.Combine(dependencyFolder, assemblyName.Name + ".dll"); + if (File.Exists(dependencyPath)) + { + return context.LoadFromAssemblyPath(Path.GetFullPath(dependencyPath)); + } + } + + return null; + } + } +} +'@ +} + +function Import-GraphAuthenticationAssembly { + param( + [Parameter(Mandatory = $true)] + [string] $ModulePath + ) + + if ($PSEdition -ne 'Core' -or -not ('System.Runtime.Loader.AssemblyLoadContext' -as [type])) { + return Import-Module -Name $ModulePath -PassThru + } + + $loadContextName = Get-GraphAuthenticationLoadContextName -ModulePath $ModulePath + $loadContext = [System.Runtime.Loader.AssemblyLoadContext]::All | + Where-Object { $_.Name -eq $loadContextName } | + Select-Object -First 1 + + if ($null -eq $loadContext) { + $loadContext = [System.Runtime.Loader.AssemblyLoadContext]::new($loadContextName, $false) + } + + $moduleRoot = $PSScriptRoot + $dependencyFolders = @( + (Join-Path $moduleRoot 'Dependencies\Core'), + (Join-Path $moduleRoot 'Dependencies'), + $moduleRoot + ) + + Initialize-GraphAuthenticationAssemblyResolver + [Microsoft.Graph.PowerShell.Authentication.Loader.GraphAuthenticationAssemblyResolver]::Register($loadContext, [string[]] $dependencyFolders) + + $moduleAssembly = $loadContext.Assemblies | + Where-Object { $_.GetName().Name -eq 'Microsoft.Graph.Authentication' } | + Select-Object -First 1 + + if ($null -eq $moduleAssembly) { + $moduleAssembly = $loadContext.LoadFromAssemblyPath((Resolve-Path -LiteralPath $ModulePath).Path) + } + + Import-Module -Assembly $moduleAssembly -PassThru +} + +function Test-GraphAuthenticationDoNotExport { + param( + [Parameter(Mandatory = $true)] + [System.Management.Automation.CommandInfo] $Command + ) + + $implementingType = $Command.ImplementingType + $null -ne $implementingType -and ($implementingType.GetCustomAttributes($false) | + Where-Object { $_.GetType().FullName -eq 'Microsoft.Graph.PowerShell.Authentication.Utilities.Runtime.DoNotExportAttribute' }) +} + +function New-GraphAuthenticationCmdletAlias { + param( + [Parameter(Mandatory = $true)] + [System.Management.Automation.CmdletInfo] $Command + ) + + $aliasNames = $Command.ImplementingType.GetCustomAttributes($false) | + Where-Object { $_.GetType().FullName -eq 'System.Management.Automation.AliasAttribute' } | + ForEach-Object { $_.AliasNames } | + Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | + Select-Object -Unique + + foreach ($aliasName in $aliasNames) { + New-Alias -Name $aliasName -Value $Command.Name -Force -Scope Script + $aliasName + } +} + +# Load the module DLL before exporting cmdlets. On PowerShell Core, loading the +# binary through a custom AssemblyLoadContext keeps its dependencies isolated +# from other Microsoft 365 modules that may already be loaded in the process. $ModulePath = (Join-Path $PSScriptRoot 'Microsoft.Graph.Authentication.dll') -$null = Import-Module -Name $ModulePath +$ModuleInfo = Import-GraphAuthenticationAssembly -ModulePath $ModulePath # Export nothing to clear implicit exports. Export-ModuleMember @@ -14,4 +170,11 @@ if (Test-Path -Path "$PSScriptRoot\StartupScripts" -ErrorAction Ignore) } # Export binary module cmdlets. -Export-ModuleMember -Cmdlet (Get-ModuleCmdlet -ModulePath $ModulePath) -Alias (Get-ModuleCmdlet -ModulePath $ModulePath -AsAlias) +$CmdletsToExport = $ModuleInfo.ExportedCommands.Values | + Where-Object { $_.CommandType -eq 'Cmdlet' -and -not (Test-GraphAuthenticationDoNotExport -Command $_) } + +$AliasesToExport = $CmdletsToExport | + ForEach-Object { New-GraphAuthenticationCmdletAlias -Command $_ } | + Select-Object -Unique + +Export-ModuleMember -Cmdlet ($CmdletsToExport | Select-Object -ExpandProperty Name -Unique) -Alias $AliasesToExport diff --git a/src/Authentication/Authentication/build-module.ps1 b/src/Authentication/Authentication/build-module.ps1 index 44d580f4009..10d74473a8b 100644 --- a/src/Authentication/Authentication/build-module.ps1 +++ b/src/Authentication/Authentication/build-module.ps1 @@ -186,11 +186,8 @@ Get-ChildItem -Path "$cmdletsSrc/bin/$Configuration/$netStandard/publish/" | Where-Object { -not $Deps.Contains($_.Name) -and $_.Extension -in $copyExtensions } | ForEach-Object { Copy-Item -Path $_.FullName -Destination $outDir -Recurse } -# Update module manifest with nested assemblies. -$RequiredAssemblies = @( - 'Microsoft.Graph.Authentication.dll', - 'Microsoft.Graph.Authentication.Core.dll' -) -Update-ModuleManifest -Path (Join-Path $outDir "$ModulePrefix.$ModuleName.psd1") -NestedModules $RequiredAssemblies +# Keep the script module in charge of loading the binary module. Adding the DLLs +# as NestedModules causes PowerShell to load them before the script can choose +# the AssemblyLoadContext used by Microsoft.Graph.Authentication.psm1. Write-Host -ForegroundColor Green '-------------Done-------------' diff --git a/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 b/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 index e53ce9051bb..87420fe05b1 100644 --- a/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 +++ b/src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1 @@ -75,5 +75,51 @@ Describe "Microsoft.Graph.Authentication module" { It 'Should lock GUID' { $PSModuleInfo.Guid.Guid | Should -Be "883916f2-9184-46ee-b1f8-b6a2fb784cee" } + + It 'Should load the root authentication assembly outside the default AssemblyLoadContext' -Skip:($PSEdition -ne 'Core') { + $assembly = [AppDomain]::CurrentDomain.GetAssemblies() | + Where-Object { $_.GetName().Name -eq $ModuleName } | + Select-Object -First 1 + + $assembly | Should -Not -BeNullOrEmpty + [System.Runtime.Loader.AssemblyLoadContext]::Default.Assemblies | + Where-Object { $_.GetName().Name -eq $ModuleName } | + Should -BeNullOrEmpty + + $loadContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($assembly) + $loadContext.Name | Should -Match '^Microsoft\.Graph\.Authentication\.' + } + + It 'Should resolve isolated dependencies from worker threads on PowerShell Core' -Skip:($PSEdition -ne 'Core') { + if (-not ('GraphAuthenticationAssemblyLoadContextTestHelper' -as [type])) { + Add-Type -TypeDefinition @' +using System.Reflection; +using System.Runtime.Loader; +using System.Threading.Tasks; + +public static class GraphAuthenticationAssemblyLoadContextTestHelper +{ + public static Assembly LoadFromWorker(AssemblyLoadContext context, string assemblyName) + { + return Task.Run(() => context.LoadFromAssemblyName(new AssemblyName(assemblyName))).GetAwaiter().GetResult(); + } +} +'@ + } + + $assembly = [AppDomain]::CurrentDomain.GetAssemblies() | + Where-Object { $_.GetName().Name -eq $ModuleName } | + Select-Object -First 1 + + $loadContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($assembly) + $dependencyAssembly = [GraphAuthenticationAssemblyLoadContextTestHelper]::LoadFromWorker($loadContext, 'Azure.Core') + $dependencyContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($dependencyAssembly) + + $dependencyAssembly.GetName().Name | Should -Be 'Azure.Core' + $dependencyContext.Name | Should -Be $loadContext.Name + [System.Runtime.Loader.AssemblyLoadContext]::Default.Assemblies | + Where-Object { $_.GetName().Name -eq 'Azure.Core' } | + Should -BeNullOrEmpty + } } -} \ No newline at end of file +}