Isolate Microsoft.Graph.Authentication assembly loading

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<string, string[]> DependencyFolders = new ConcurrentDictionary<string, string[]>(StringComparer.Ordinal);
+        private static readonly ConcurrentDictionary<string, bool> RegisteredContexts = new ConcurrentDictionary<string, bool>(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<string>();
+
+            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

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-------------'

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
+        }
     }
-}
+}