Rebasing to keep up to date

Author: FehintolaObafemi

src/Authentication/Authentication.Core/Common/GraphSession.cs

@@ -56,11 +56,6 @@ public class GraphSession : IGraphSession
         /// </summary>
         public IGraphOption GraphOption { get; set; }
 
-        /// <summary>
-        /// Temporarily stores the user's Graph request details such as Method and Uri. Essential as part of the Proof of Possession efforts.
-        /// </summary>
-        public IGraphRequestPopContext GraphRequestPopContext { get; set; }
-
         /// <summary>
         /// Represents a collection of Microsoft Graph PowerShell meta-info.
         /// </summary>

src/Authentication/Authentication.Core/Interfaces/IGraphRequestPopContext.cs

@@ -12,6 +12,5 @@ public interface IGraphSession
         IDataStore DataStore { get; set; }
         IRequestContext RequestContext { get; set; }
         IGraphOption GraphOption { get; set; }
-        IGraphRequestPopContext GraphRequestPopContext { get; set; }
     }
 }

src/Authentication/Authentication.Core/Interfaces/IGraphSession.cs

@@ -1,10 +1,10 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <Import Project="$(MSBuildThisFileDirectory)..\..\..\Repo.props" />
   <PropertyGroup>
     <LangVersion>9.0</LangVersion>
     <TargetFrameworks>netstandard2.0;net6.0;net472</TargetFrameworks>
     <RootNamespace>Microsoft.Graph.PowerShell.Authentication.Core</RootNamespace>
-    <Version>2.31.0</Version>
+    <Version>2.32.0</Version>
     <!-- Suppress .NET Target Framework Moniker (TFM) Support Build Warnings -->
     <SuppressTfmSupportBuildWarnings>true</SuppressTfmSupportBuildWarnings>
   </PropertyGroup>
@@ -15,7 +15,9 @@
   <ItemGroup>
     <PackageReference Include="Azure.Identity" Version="1.13.2" />
     <PackageReference Include="Azure.Identity.Broker" Version="1.2.0" />
-    <PackageReference Include="Microsoft.Graph.Core" Version="3.2.4" />
+    <PackageReference Include="Microsoft.Graph.Core" Version="3.2.2" />
+    <PackageReference Include="Microsoft.Identity.Client" Version="4.67.2" />
+    <PackageReference Include="Microsoft.Identity.Client.Broker" Version="4.67.2" />
     <PackageReference Include="Newtonsoft.Json" Version="13.0.3" />
     <PackageReference Include="System.Text.Json" Version="8.0.5" />
   </ItemGroup>

src/Authentication/Authentication.Core/Microsoft.Graph.Authentication.Core.csproj

@@ -15,8 +15,6 @@
 using System.Globalization;
 using System.IO;
 using System.Linq;
-using System.Net.Http;
-using System.Net.Http.Headers;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
@@ -126,35 +124,16 @@ private static async Task<InteractiveBrowserCredential> GetInteractiveBrowserCre
             interactiveOptions.TokenCachePersistenceOptions = GetTokenCachePersistenceOptions(authContext);
 
             var interactiveBrowserCredential = new InteractiveBrowserCredential(interactiveOptions);
-            var popTokenRequestContext = new PopTokenRequestContext();
-            if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
-            {
-                popTokenRequestContext = await CreatePopTokenRequestContext(authContext);
-                GraphSession.Instance.GraphRequestPopContext.PopInteractiveBrowserCredential = interactiveBrowserCredential;
-            }
-
             if (!File.Exists(Constants.AuthRecordPath))
             {
                 AuthenticationRecord authRecord;
                 if (IsWamSupported())
                 {
-                    // Adding a scenario to account for Access Token Proof of Possession
-                    if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
-                    {
-                        authRecord = await Task.Run(() =>
-                        {
-                            // Run the thread in MTA.
-                            return interactiveBrowserCredential.AuthenticateAsync(popTokenRequestContext, cancellationToken);
-                        });
-                    }
-                    else
+                    authRecord = await Task.Run(() =>
                     {
-                        authRecord = await Task.Run(() =>
-                        {
-                            // Run the thread in MTA.
-                            return interactiveBrowserCredential.Authenticate(new TokenRequestContext(authContext.Scopes), cancellationToken);
-                        });
-                    }
+                        // Run the thread in MTA.
+                        return interactiveBrowserCredential.AuthenticateAsync(new TokenRequestContext(authContext.Scopes), cancellationToken);
+                    });
                 }
                 else
                 {
@@ -471,34 +450,5 @@ public static Task DeleteAuthRecordAsync()
                 File.Delete(Constants.AuthRecordPath);
             return Task.CompletedTask;
         }
-
-        private static async Task<PopTokenRequestContext> CreatePopTokenRequestContext(IAuthContext authContext)
-        {
-            // Creating a httpclient that would handle all pop calls
-            Uri popResourceUri = GraphSession.Instance.GraphRequestPopContext.Uri ?? new Uri("https://graph.microsoft.com/beta/organization");
-            HttpClient popHttpClient = new(new HttpClientHandler());
-
-            // Find the nonce in the WWW-Authenticate header in the response.
-            var popMethod = GraphSession.Instance.GraphRequestPopContext.HttpMethod ?? HttpMethod.Get;
-            var popResponse = await popHttpClient.SendAsync(new HttpRequestMessage(popMethod, popResourceUri));
-            
-            // Refresh token logic --- start
-            var popPipelineOptions = new HttpPipelineOptions(new PopClientOptions()
-            {
-
-            });
-
-            GraphSession.Instance.GraphRequestPopContext.PopPipeline = HttpPipelineBuilder.Build(popPipelineOptions, new HttpPipelineTransportOptions());
-            var popRequest = GraphSession.Instance.GraphRequestPopContext.PopPipeline.CreateRequest();
-            popRequest.Method = RequestMethod.Parse(popMethod.Method.ToUpper());
-            popRequest.Uri.Reset(popResourceUri);
-
-            // Refresh token logic --- end
-            var popContext = new PopTokenRequestContext(authContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: WwwAuthenticateParameters.CreateFromAuthenticationHeaders(popResponse.Headers, "Pop").Nonce, request: popRequest);
-            return popContext;
-        }
-    }
-    internal class PopClientOptions : ClientOptions
-    {
     }
 }

src/Authentication/Authentication.Core/Utilities/AuthenticationHelpers.cs

@@ -1026,8 +1026,6 @@ private async Task ProcessRecordAsync()
                 try
                 {
                     PrepareSession();
-                    GraphSession.Instance.GraphRequestPopContext.Uri = Uri;
-                    GraphSession.Instance.GraphRequestPopContext.HttpMethod = GetHttpMethod(Method);
                     var client = HttpHelpers.GetGraphHttpClient();
                     ValidateRequestUri();
                     using (var httpRequestMessage = GetRequest(client, Uri))

src/Authentication/Authentication/Cmdlets/InvokeMgGraphRequest.cs

@@ -47,8 +47,7 @@ internal static GraphSession CreateInstance(IDataStore dataStore = null)
             {
                 DataStore = dataStore ?? new DiskDataStore(),
                 RequestContext = new RequestContext(),
-                GraphOption = graphOptions ?? new GraphOption(),
-                GraphRequestPopContext = new GraphRequestPopContext()
+                GraphOption = graphOptions ?? new GraphOption()
             };
         }
         /// <summary>

src/Authentication/Authentication/Common/GraphSessionInitializer.cs

@@ -27,8 +27,8 @@ internal class AuthenticationHandler : DelegatingHandler
         private const string BearerAuthenticationScheme = "Bearer";
         private const string PopAuthenticationScheme = "Pop";
         private int MaxRetry { get; set; } = 1;
-        private PopTokenRequestContext popTokenRequestContext;
-        private Request popRequest = GraphSession.Instance.GraphRequestPopContext.PopPipeline.CreateRequest();
+        private TokenRequestContext popTokenRequestContext;
+        private string cachedNonce;
 
         public AzureIdentityAccessTokenProvider AuthenticationProvider { get; set; }
 
@@ -52,10 +52,22 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
 
                 HttpResponseMessage response = await base.SendAsync(httpRequestMessage, cancellationToken).ConfigureAwait(false);
 
-                // Continuous nonce extraction on each request
-                if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
+                // Extract nonce from API responses for future PoP requests
+                if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph && IsApiRequest(httpRequestMessage.RequestUri))
                 {
-                    popTokenRequestContext = new PopTokenRequestContext(GraphSession.Instance.AuthContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: WwwAuthenticateParameters.CreateFromAuthenticationHeaders(response.Headers, PopAuthenticationScheme).Nonce, request: popRequest);
+                    try
+                    {
+                        var wwwAuthParams = WwwAuthenticateParameters.CreateFromAuthenticationHeaders(response.Headers, PopAuthenticationScheme);
+                        if (wwwAuthParams?.Nonce != null && !string.IsNullOrEmpty(wwwAuthParams.Nonce))
+                        {
+                            cachedNonce = wwwAuthParams.Nonce;
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        System.Diagnostics.Debug.WriteLine($"AuthenticationHandler: Failed to extract PoP nonce: {ex.Message}");
+                        // Don't throw - nonce extraction failure shouldn't break the response
+                    }
                 }
 
                 // Check if response is a 401 & is not a streamed body (is buffered)
@@ -76,17 +88,48 @@ private async Task AuthenticateRequestAsync(HttpRequestMessage httpRequestMessag
         {
             if (AuthenticationProvider != null)
             {
+                // Determine if this is an API request that should use PoP (when enabled)
+                // vs an authentication request that should always use Bearer
+                bool isApiRequest = IsApiRequest(httpRequestMessage.RequestUri);
+                bool shouldUsePoP = GraphSession.Instance.GraphOption.EnableATPoPForMSGraph && isApiRequest;
+                
+                // Debug logging for flow routing
                 if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
                 {
-                    popRequest.Method = RequestMethod.Parse(httpRequestMessage.Method.Method.ToUpper());
-                    popRequest.Uri.Reset(httpRequestMessage.RequestUri);
-                    foreach (var header in httpRequestMessage.Headers)
+                    var requestType = isApiRequest ? "API" : "Auth";
+                    var tokenType = shouldUsePoP ? "PoP" : "Bearer";
+                    System.Diagnostics.Debug.WriteLine($"AuthenticationHandler: {requestType} request to {httpRequestMessage.RequestUri?.Host} using {tokenType} token");
+                }
+                
+                if (shouldUsePoP)
+                {
+                    // API Request with PoP enabled - use PoP tokens ONLY
+                    try
+                    {
+                        // Create proper TokenRequestContext for PoP
+                        // Note: cachedNonce may be null for initial requests - this is expected
+                        popTokenRequestContext = new TokenRequestContext(
+                            scopes: GraphSession.Instance.AuthContext.Scopes,
+                            parentRequestId: null,
+                            claims: additionalAuthenticationContext?.ContainsKey(ClaimsKey) == true ? additionalAuthenticationContext[ClaimsKey]?.ToString() : null,
+                            tenantId: null,
+                            isCaeEnabled: false,
+                            isProofOfPossessionEnabled: true,
+                            proofOfPossessionNonce: cachedNonce // May be null for initial requests
+                        );
+                        
+                        // Get TokenCredential from existing AuthenticationProvider
+                        var tokenCredential = await AuthenticationHelpers.GetTokenCredentialAsync(
+                            GraphSession.Instance.AuthContext, cancellationToken).ConfigureAwait(false);
+                        
+                        var accessToken = await tokenCredential.GetTokenAsync(popTokenRequestContext, cancellationToken).ConfigureAwait(false);
+                        httpRequestMessage.Headers.Authorization = new AuthenticationHeaderValue(PopAuthenticationScheme, accessToken.Token);
+                    }
+                    catch (Exception ex) when (!(ex is OperationCanceledException))
                     {
-                        popRequest.Headers.Add(header.Key, header.Value.First());
+                        // Re-throw with context for PoP-specific failures
+                        throw new AuthenticationException($"Failed to acquire PoP token for {httpRequestMessage.RequestUri}: {ex.Message}", ex);
                     }
-                    
-                    var accessToken = await GraphSession.Instance.GraphRequestPopContext.PopInteractiveBrowserCredential.GetTokenAsync(popTokenRequestContext, cancellationToken).ConfigureAwait(false);
-                    httpRequestMessage.Headers.Authorization = new AuthenticationHeaderValue(PopAuthenticationScheme, accessToken.Token);
                 }
                 else
                 {
@@ -156,5 +199,52 @@ private static async Task DrainAsync(HttpResponseMessage response)
             }
             response.Dispose();
         }
+
+        /// <summary>
+        /// Determines if the request is an API request that should use PoP when enabled,
+        /// vs an authentication/token request that should always use Bearer tokens.
+        /// This method implements the core routing logic for AT-PoP:
+        /// - Graph API endpoints → PoP tokens (when enabled) 
+        /// - Authentication endpoints → Bearer tokens (always)
+        /// - Unknown endpoints → Bearer tokens (safe default)
+        /// </summary>
+        /// <param name="requestUri">The request URI to evaluate</param>
+        /// <returns>True if this is an API request, false if it's an authentication request</returns>
+        private static bool IsApiRequest(Uri requestUri)
+        {
+            if (requestUri == null) return false;
+            
+            var host = requestUri.Host?.ToLowerInvariant();
+            var path = requestUri.AbsolutePath?.ToLowerInvariant();
+            
+            // Microsoft Graph API endpoints that should use PoP
+            if (host?.Contains("graph.microsoft.com") == true ||
+                host?.Contains("graph.microsoft.us") == true ||
+                host?.Contains("microsoftgraph.chinacloudapi.cn") == true ||
+                host?.Contains("graph.microsoft.de") == true)
+            {
+                // Exclude authentication/token endpoints - these should always use Bearer
+                if (path?.Contains("/oauth2/") == true ||
+                    path?.Contains("/token") == true ||
+                    path?.Contains("/authorize") == true ||
+                    path?.Contains("/devicecode") == true)
+                {
+                    return false; // Authentication request
+                }
+                return true; // API request
+            }
+            
+            // Azure AD/authentication endpoints - never use PoP
+            if (host?.Contains("login.microsoftonline.com") == true ||
+                host?.Contains("login.microsoft.com") == true ||
+                host?.Contains("login.chinacloudapi.cn") == true ||
+                host?.Contains("login.microsoftonline.de") == true ||
+                host?.Contains("login.microsoftonline.us") == true)
+            {
+                return false; // Authentication request
+            }
+            
+            return false; // Default to authentication request for unknown endpoints
+        }
     }
 }

src/Authentication/Authentication/Handlers/AuthenticationHandler.cs

@@ -50,14 +50,12 @@
     <file src="artifacts\Dependencies\Core\runtimes\win-arm64\native\msalruntime_arm64.dll" target="Dependencies" />
     <!-- Core-->
     <file src="artifacts\Dependencies\Core\Azure.Core.dll" target="Dependencies\Core" />
-    <file src="artifacts\Dependencies\Core\Azure.Core.Experimental.dll" target="Dependencies\Core" />
     <file src="artifacts\Dependencies\Core\Microsoft.Graph.Core.dll" target="Dependencies\Core" />
     <file src="artifacts\Dependencies\Core\Microsoft.Identity.Client.dll" target="Dependencies\Core" />
     <file src="artifacts\Dependencies\Core\Microsoft.Identity.Client.Extensions.Msal.dll" target="Dependencies\Core" />
     <file src="artifacts\Dependencies\Core\Newtonsoft.Json.dll" target="Dependencies\Core" />
     <!-- Desktop -->
     <file src="artifacts\Dependencies\Desktop\Azure.Core.dll" target="Dependencies\Desktop" />
-    <file src="artifacts\Dependencies\Desktop\Azure.Core.Experimental.dll" target="Dependencies\Desktop" />
     <file src="artifacts\Dependencies\Desktop\Microsoft.Graph.Core.dll" target="Dependencies\Desktop" />
     <file src="artifacts\Dependencies\Desktop\Microsoft.Identity.Client.dll" target="Dependencies\Desktop" />
     <file src="artifacts\Dependencies\Desktop\Microsoft.Identity.Client.Extensions.Msal.dll" target="Dependencies\Desktop" />