Skip to content

Commit 3a55148

Browse files
PrzemyslawKlysPrzemyslawKlys
authored
Isolate Graph authentication assembly loading (#3632)
1 parent 36f2312 commit 3a55148

3 files changed

Lines changed: 216 additions & 10 deletions

File tree

‎src/Authentication/Authentication/Microsoft.Graph.Authentication.psm1‎

Lines changed: 166 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,163 @@
11

2-
# Load the module dll
2+
function Get-GraphAuthenticationLoadContextName {
3+
param(
4+
[Parameter(Mandatory = $true)]
5+
[string] $ModulePath
6+
)
7+
8+
$sha256 = [System.Security.Cryptography.SHA256]::Create()
9+
try {
10+
$pathBytes = [System.Text.Encoding]::UTF8.GetBytes($ModulePath)
11+
$hash = [System.BitConverter]::ToString($sha256.ComputeHash($pathBytes)).Replace('-', '').Substring(0, 16)
12+
"Microsoft.Graph.Authentication.$hash"
13+
}
14+
finally {
15+
$sha256.Dispose()
16+
}
17+
}
18+
19+
function Initialize-GraphAuthenticationAssemblyResolver {
20+
if ('Microsoft.Graph.PowerShell.Authentication.Loader.GraphAuthenticationAssemblyResolver' -as [type]) {
21+
return
22+
}
23+
24+
Add-Type -TypeDefinition @'
25+
using System;
26+
using System.Collections.Concurrent;
27+
using System.IO;
28+
using System.Reflection;
29+
using System.Runtime.Loader;
30+
31+
namespace Microsoft.Graph.PowerShell.Authentication.Loader
32+
{
33+
public static class GraphAuthenticationAssemblyResolver
34+
{
35+
private static readonly ConcurrentDictionary<string, string[]> DependencyFolders = new ConcurrentDictionary<string, string[]>(StringComparer.Ordinal);
36+
private static readonly ConcurrentDictionary<string, bool> RegisteredContexts = new ConcurrentDictionary<string, bool>(StringComparer.Ordinal);
37+
38+
public static void Register(AssemblyLoadContext context, string[] dependencyFolders)
39+
{
40+
if (context == null)
41+
{
42+
throw new ArgumentNullException(nameof(context));
43+
}
44+
45+
string contextName = context.Name ?? string.Empty;
46+
DependencyFolders[contextName] = dependencyFolders ?? Array.Empty<string>();
47+
48+
if (RegisteredContexts.TryAdd(contextName, true))
49+
{
50+
context.Resolving += Resolve;
51+
}
52+
}
53+
54+
private static Assembly Resolve(AssemblyLoadContext context, AssemblyName assemblyName)
55+
{
56+
if (context == null || assemblyName == null)
57+
{
58+
return null;
59+
}
60+
61+
if (!DependencyFolders.TryGetValue(context.Name ?? string.Empty, out string[] dependencyFolders))
62+
{
63+
return null;
64+
}
65+
66+
foreach (string dependencyFolder in dependencyFolders)
67+
{
68+
if (string.IsNullOrWhiteSpace(dependencyFolder) || !Directory.Exists(dependencyFolder))
69+
{
70+
continue;
71+
}
72+
73+
string dependencyPath = Path.Combine(dependencyFolder, assemblyName.Name + ".dll");
74+
if (File.Exists(dependencyPath))
75+
{
76+
return context.LoadFromAssemblyPath(Path.GetFullPath(dependencyPath));
77+
}
78+
}
79+
80+
return null;
81+
}
82+
}
83+
}
84+
'@
85+
}
86+
87+
function Import-GraphAuthenticationAssembly {
88+
param(
89+
[Parameter(Mandatory = $true)]
90+
[string] $ModulePath
91+
)
92+
93+
if ($PSEdition -ne 'Core' -or -not ('System.Runtime.Loader.AssemblyLoadContext' -as [type])) {
94+
return Import-Module -Name $ModulePath -PassThru
95+
}
96+
97+
$loadContextName = Get-GraphAuthenticationLoadContextName -ModulePath $ModulePath
98+
$loadContext = [System.Runtime.Loader.AssemblyLoadContext]::All |
99+
Where-Object { $_.Name -eq $loadContextName } |
100+
Select-Object -First 1
101+
102+
if ($null -eq $loadContext) {
103+
$loadContext = [System.Runtime.Loader.AssemblyLoadContext]::new($loadContextName, $false)
104+
}
105+
106+
$moduleRoot = $PSScriptRoot
107+
$dependencyFolders = @(
108+
(Join-Path $moduleRoot 'Dependencies\Core'),
109+
(Join-Path $moduleRoot 'Dependencies'),
110+
$moduleRoot
111+
)
112+
113+
Initialize-GraphAuthenticationAssemblyResolver
114+
[Microsoft.Graph.PowerShell.Authentication.Loader.GraphAuthenticationAssemblyResolver]::Register($loadContext, [string[]] $dependencyFolders)
115+
116+
$moduleAssembly = $loadContext.Assemblies |
117+
Where-Object { $_.GetName().Name -eq 'Microsoft.Graph.Authentication' } |
118+
Select-Object -First 1
119+
120+
if ($null -eq $moduleAssembly) {
121+
$moduleAssembly = $loadContext.LoadFromAssemblyPath((Resolve-Path -LiteralPath $ModulePath).Path)
122+
}
123+
124+
Import-Module -Assembly $moduleAssembly -PassThru
125+
}
126+
127+
function Test-GraphAuthenticationDoNotExport {
128+
param(
129+
[Parameter(Mandatory = $true)]
130+
[System.Management.Automation.CommandInfo] $Command
131+
)
132+
133+
$implementingType = $Command.ImplementingType
134+
$null -ne $implementingType -and ($implementingType.GetCustomAttributes($false) |
135+
Where-Object { $_.GetType().FullName -eq 'Microsoft.Graph.PowerShell.Authentication.Utilities.Runtime.DoNotExportAttribute' })
136+
}
137+
138+
function New-GraphAuthenticationCmdletAlias {
139+
param(
140+
[Parameter(Mandatory = $true)]
141+
[System.Management.Automation.CmdletInfo] $Command
142+
)
143+
144+
$aliasNames = $Command.ImplementingType.GetCustomAttributes($false) |
145+
Where-Object { $_.GetType().FullName -eq 'System.Management.Automation.AliasAttribute' } |
146+
ForEach-Object { $_.AliasNames } |
147+
Where-Object { -not [string]::IsNullOrWhiteSpace($_) } |
148+
Select-Object -Unique
149+
150+
foreach ($aliasName in $aliasNames) {
151+
New-Alias -Name $aliasName -Value $Command.Name -Force -Scope Script
152+
$aliasName
153+
}
154+
}
155+
156+
# Load the module DLL before exporting cmdlets. On PowerShell Core, loading the
157+
# binary through a custom AssemblyLoadContext keeps its dependencies isolated
158+
# from other Microsoft 365 modules that may already be loaded in the process.
3159
$ModulePath = (Join-Path $PSScriptRoot 'Microsoft.Graph.Authentication.dll')
4-
$null = Import-Module -Name $ModulePath
160+
$ModuleInfo = Import-GraphAuthenticationAssembly -ModulePath $ModulePath
5161

6162
# Export nothing to clear implicit exports.
7163
Export-ModuleMember
@@ -14,4 +170,11 @@ if (Test-Path -Path "$PSScriptRoot\StartupScripts" -ErrorAction Ignore)
14170
}
15171

16172
# Export binary module cmdlets.
17-
Export-ModuleMember -Cmdlet (Get-ModuleCmdlet -ModulePath $ModulePath) -Alias (Get-ModuleCmdlet -ModulePath $ModulePath -AsAlias)
173+
$CmdletsToExport = $ModuleInfo.ExportedCommands.Values |
174+
Where-Object { $_.CommandType -eq 'Cmdlet' -and -not (Test-GraphAuthenticationDoNotExport -Command $_) }
175+
176+
$AliasesToExport = $CmdletsToExport |
177+
ForEach-Object { New-GraphAuthenticationCmdletAlias -Command $_ } |
178+
Select-Object -Unique
179+
180+
Export-ModuleMember -Cmdlet ($CmdletsToExport | Select-Object -ExpandProperty Name -Unique) -Alias $AliasesToExport

‎src/Authentication/Authentication/build-module.ps1‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -186,11 +186,8 @@ Get-ChildItem -Path "$cmdletsSrc/bin/$Configuration/$netStandard/publish/" |
186186
Where-Object { -not $Deps.Contains($_.Name) -and $_.Extension -in $copyExtensions } |
187187
ForEach-Object { Copy-Item -Path $_.FullName -Destination $outDir -Recurse }
188188

189-
# Update module manifest with nested assemblies.
190-
$RequiredAssemblies = @(
191-
'Microsoft.Graph.Authentication.dll',
192-
'Microsoft.Graph.Authentication.Core.dll'
193-
)
194-
Update-ModuleManifest -Path (Join-Path $outDir "$ModulePrefix.$ModuleName.psd1") -NestedModules $RequiredAssemblies
189+
# Keep the script module in charge of loading the binary module. Adding the DLLs
190+
# as NestedModules causes PowerShell to load them before the script can choose
191+
# the AssemblyLoadContext used by Microsoft.Graph.Authentication.psm1.
195192

196193
Write-Host -ForegroundColor Green '-------------Done-------------'

‎src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,5 +75,51 @@ Describe "Microsoft.Graph.Authentication module" {
7575
It 'Should lock GUID' {
7676
$PSModuleInfo.Guid.Guid | Should -Be "883916f2-9184-46ee-b1f8-b6a2fb784cee"
7777
}
78+
79+
It 'Should load the root authentication assembly outside the default AssemblyLoadContext' -Skip:($PSEdition -ne 'Core') {
80+
$assembly = [AppDomain]::CurrentDomain.GetAssemblies() |
81+
Where-Object { $_.GetName().Name -eq $ModuleName } |
82+
Select-Object -First 1
83+
84+
$assembly | Should -Not -BeNullOrEmpty
85+
[System.Runtime.Loader.AssemblyLoadContext]::Default.Assemblies |
86+
Where-Object { $_.GetName().Name -eq $ModuleName } |
87+
Should -BeNullOrEmpty
88+
89+
$loadContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($assembly)
90+
$loadContext.Name | Should -Match '^Microsoft\.Graph\.Authentication\.'
91+
}
92+
93+
It 'Should resolve isolated dependencies from worker threads on PowerShell Core' -Skip:($PSEdition -ne 'Core') {
94+
if (-not ('GraphAuthenticationAssemblyLoadContextTestHelper' -as [type])) {
95+
Add-Type -TypeDefinition @'
96+
using System.Reflection;
97+
using System.Runtime.Loader;
98+
using System.Threading.Tasks;
99+
100+
public static class GraphAuthenticationAssemblyLoadContextTestHelper
101+
{
102+
public static Assembly LoadFromWorker(AssemblyLoadContext context, string assemblyName)
103+
{
104+
return Task.Run(() => context.LoadFromAssemblyName(new AssemblyName(assemblyName))).GetAwaiter().GetResult();
105+
}
106+
}
107+
'@
108+
}
109+
110+
$assembly = [AppDomain]::CurrentDomain.GetAssemblies() |
111+
Where-Object { $_.GetName().Name -eq $ModuleName } |
112+
Select-Object -First 1
113+
114+
$loadContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($assembly)
115+
$dependencyAssembly = [GraphAuthenticationAssemblyLoadContextTestHelper]::LoadFromWorker($loadContext, 'Azure.Core')
116+
$dependencyContext = [System.Runtime.Loader.AssemblyLoadContext]::GetLoadContext($dependencyAssembly)
117+
118+
$dependencyAssembly.GetName().Name | Should -Be 'Azure.Core'
119+
$dependencyContext.Name | Should -Be $loadContext.Name
120+
[System.Runtime.Loader.AssemblyLoadContext]::Default.Assemblies |
121+
Where-Object { $_.GetName().Name -eq 'Azure.Core' } |
122+
Should -BeNullOrEmpty
123+
}
78124
}
79-
}
125+
}

0 commit comments

Comments
 (0)