resolving PR comments

Author: FehintolaObafemi

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

@@ -59,7 +59,7 @@ public class GraphSession : IGraphSession
         /// <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 IGraphRequestProofofPossession GraphRequestProofofPossession { get; set; }
+        public IGraphRequestPopContext GraphRequestPopContext { get; set; }
 
         /// <summary>
         /// Represents a collection of Microsoft Graph PowerShell meta-info.

src/Authentication/Authentication.Core/Interfaces/IGraphRequestProofofPossession.cs renamed to src/Authentication/Authentication.Core/Interfaces/IGraphRequestPopContext.cs

@@ -9,14 +9,14 @@
 
 namespace Microsoft.Graph.PowerShell.Authentication
 {
-    public interface IGraphRequestProofofPossession
+    public interface IGraphRequestPopContext
     {
         Uri Uri { get; set; }
         HttpMethod HttpMethod { get; set; }
         AccessToken AccessToken { get; set; }
         string ProofofPossessionNonce { get; set; }
         PopTokenRequestContext PopTokenContext { get; set; }
         Request Request { get; set; }
-        InteractiveBrowserCredential BrowserCredential { get; set; }
+        InteractiveBrowserCredential PopInteractiveBrowserCredential { get; set; }
     }
 }

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

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

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

@@ -16,6 +16,7 @@
 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;
@@ -88,12 +89,6 @@ private static bool IsWamSupported()
             return GraphSession.Instance.GraphOption.EnableWAMForMSGraph && SharedUtilities.IsWindowsPlatform();
         }
 
-        //Check to see if ATPoP is Supported
-        public static bool IsATPoPSupported()
-        {
-            return GraphSession.Instance.GraphOption.EnableATPoPForMSGraph;
-        }
-
         private static async Task<TokenCredential> GetClientSecretCredentialAsync(IAuthContext authContext)
         {
             if (authContext is null)
@@ -130,27 +125,24 @@ private static async Task<InteractiveBrowserCredential> GetInteractiveBrowserCre
             interactiveOptions.TokenCachePersistenceOptions = GetTokenCachePersistenceOptions(authContext);
 
             var interactiveBrowserCredential = new InteractiveBrowserCredential(interactiveOptions);
-            if (IsATPoPSupported())
+            if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
             {
-                GraphSession.Instance.GraphRequestProofofPossession.PopTokenContext = CreatePopTokenRequestContext(authContext);
-                GraphSession.Instance.GraphRequestProofofPossession.BrowserCredential = interactiveBrowserCredential;
+                GraphSession.Instance.GraphRequestPopContext.PopTokenContext = await CreatePopTokenRequestContext(authContext);
+                GraphSession.Instance.GraphRequestPopContext.PopInteractiveBrowserCredential = interactiveBrowserCredential;
             }
 
             if (!File.Exists(Constants.AuthRecordPath))
             {
                 AuthenticationRecord authRecord;
-                //var interactiveBrowserCredential = new InteractiveBrowserCredential(interactiveOptions);
                 if (IsWamSupported())
                 {
                     // Adding a scenario to account for Access Token Proof of Possession
-                    if (IsATPoPSupported())
+                    if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
                     {
-                        // Logic to implement ATPoP Authentication
                         authRecord = await Task.Run(() =>
                         {
                             // Run the thread in MTA.
-                            //GraphSession.Instance.GraphRequestProofofPossession.AccessToken = interactiveBrowserCredential.GetTokenAsync(GraphSession.Instance.GraphRequestProofofPossession.PopTokenContext, cancellationToken).Result;
-                            return interactiveBrowserCredential.AuthenticateAsync(GraphSession.Instance.GraphRequestProofofPossession.PopTokenContext, cancellationToken);
+                            return interactiveBrowserCredential.AuthenticateAsync(GraphSession.Instance.GraphRequestPopContext.PopTokenContext, cancellationToken);
                         });
                     }
                     else
@@ -478,19 +470,16 @@ public static Task DeleteAuthRecordAsync()
             return Task.CompletedTask;
         }
 
-        public static PopTokenRequestContext CreatePopTokenRequestContext(IAuthContext authContext)
+        private static async Task<PopTokenRequestContext> CreatePopTokenRequestContext(IAuthContext authContext)
         {
             // Creating a httpclient that would handle all pop calls
-            Uri popResourceUri = GraphSession.Instance.GraphRequestProofofPossession.Uri ?? new Uri("https://graph.microsoft.com/beta/organization"); //PPE (https://graph.microsoft-ppe.com) or Canary (https://canary.graph.microsoft.com) or (https://20.190.132.47/beta/me)
+            Uri popResourceUri = GraphSession.Instance.GraphRequestPopContext.Uri ?? new Uri("https://graph.microsoft.com/beta/organization"); //PPE (https://graph.microsoft-ppe.com) or Canary (https://canary.graph.microsoft.com) or (https://20.190.132.47/beta/me)
             HttpClient popHttpClient = new(new HttpClientHandler());
 
-            // Find the WWW-Authenticate header in the response.
-            var popMethod = GraphSession.Instance.GraphRequestProofofPossession.HttpMethod ?? HttpMethod.Get;
-            var popResponse = popHttpClient.SendAsync(new HttpRequestMessage(popMethod, popResourceUri)).Result;
-            var popChallenge = popResponse.Headers.WwwAuthenticate.First(wa => wa.Scheme == "PoP");
-            var nonceStart = popChallenge.Parameter.IndexOf("nonce=\"") + "nonce=\"".Length;
-            var nonceEnd = popChallenge.Parameter.IndexOf('"', nonceStart);
-            GraphSession.Instance.GraphRequestProofofPossession.ProofofPossessionNonce = popChallenge.Parameter.Substring(nonceStart, nonceEnd - nonceStart);
+            // 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));
+            GraphSession.Instance.GraphRequestPopContext.ProofofPossessionNonce = WwwAuthenticateParameters.CreateFromAuthenticationHeaders(popResponse.Headers, "Pop").Nonce;
 
             // Refresh token logic --- start
             var popPipelineOptions = new HttpPipelineOptions(new PopClientOptions()
@@ -499,12 +488,12 @@ public static PopTokenRequestContext CreatePopTokenRequestContext(IAuthContext a
             });
 
             var _popPipeline = HttpPipelineBuilder.Build(popPipelineOptions, new HttpPipelineTransportOptions());
-            GraphSession.Instance.GraphRequestProofofPossession.Request = _popPipeline.CreateRequest();
-            GraphSession.Instance.GraphRequestProofofPossession.Request.Method = ConvertToAzureRequestMethod(popMethod);
-            GraphSession.Instance.GraphRequestProofofPossession.Request.Uri.Reset(popResourceUri);
+            GraphSession.Instance.GraphRequestPopContext.Request = _popPipeline.CreateRequest();
+            GraphSession.Instance.GraphRequestPopContext.Request.Method = ConvertToAzureRequestMethod(popMethod);
+            GraphSession.Instance.GraphRequestPopContext.Request.Uri.Reset(popResourceUri);
 
             // Refresh token logic --- end
-            var popContext = new PopTokenRequestContext(authContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: GraphSession.Instance.GraphRequestProofofPossession.ProofofPossessionNonce, request: GraphSession.Instance.GraphRequestProofofPossession.Request);
+            var popContext = new PopTokenRequestContext(authContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: GraphSession.Instance.GraphRequestPopContext.ProofofPossessionNonce, request: GraphSession.Instance.GraphRequestPopContext.Request);
             return popContext;
         }
         public static RequestMethod ConvertToAzureRequestMethod(HttpMethod httpMethod)

src/Authentication/Authentication/Cmdlets/InvokeMgGraphRequest.cs

@@ -1023,8 +1023,8 @@ private async Task ProcessRecordAsync()
                 try
                 {
                     PrepareSession();
-                    GraphSession.Instance.GraphRequestProofofPossession.Uri = Uri;
-                    GraphSession.Instance.GraphRequestProofofPossession.HttpMethod = GetHttpMethod(Method);
+                    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/Common/GraphSessionInitializer.cs

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

src/Authentication/Authentication/Handlers/AuthenticationHandler.cs

@@ -7,6 +7,7 @@
 using Microsoft.Graph.Authentication;
 using Microsoft.Graph.PowerShell.Authentication.Core.Utilities;
 using Microsoft.Graph.PowerShell.Authentication.Extensions;
+using Microsoft.Identity.Client;
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -23,6 +24,7 @@ internal class AuthenticationHandler : DelegatingHandler
     {
         private const string ClaimsKey = "claims";
         private const string BearerAuthenticationScheme = "Bearer";
+        private const string PopAuthenticationScheme = "Pop";
         private int MaxRetry { get; set; } = 1;
 
         public AzureIdentityAccessTokenProvider AuthenticationProvider { get; set; }
@@ -47,6 +49,13 @@ 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)
+                {
+                    GraphSession.Instance.GraphRequestPopContext.ProofofPossessionNonce = WwwAuthenticateParameters.CreateFromAuthenticationHeaders(response.Headers, PopAuthenticationScheme).Nonce;
+                    GraphSession.Instance.GraphRequestPopContext.PopTokenContext = new PopTokenRequestContext(GraphSession.Instance.AuthContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: GraphSession.Instance.GraphRequestPopContext.ProofofPossessionNonce, request: GraphSession.Instance.GraphRequestPopContext.Request);
+                }
+
                 // Check if response is a 401 & is not a streamed body (is buffered)
                 if (response.StatusCode == HttpStatusCode.Unauthorized && httpRequestMessage.IsBuffered())
                 {
@@ -65,17 +74,17 @@ private async Task AuthenticateRequestAsync(HttpRequestMessage httpRequestMessag
         {
             if (AuthenticationProvider != null)
             {
-                if (AuthenticationHelpers.IsATPoPSupported())
+                if (GraphSession.Instance.GraphOption.EnableATPoPForMSGraph)
                 {
-                    GraphSession.Instance.GraphRequestProofofPossession.Request.Method = AuthenticationHelpers.ConvertToAzureRequestMethod(httpRequestMessage.Method);
-                    GraphSession.Instance.GraphRequestProofofPossession.Request.Uri.Reset(httpRequestMessage.RequestUri);
+                    GraphSession.Instance.GraphRequestPopContext.Request.Method = AuthenticationHelpers.ConvertToAzureRequestMethod(httpRequestMessage.Method);
+                    GraphSession.Instance.GraphRequestPopContext.Request.Uri.Reset(httpRequestMessage.RequestUri);
                     foreach (var header in httpRequestMessage.Headers)
                     {
-                        GraphSession.Instance.GraphRequestProofofPossession.Request.Headers.Add(header.Key, header.Value.First());
+                        GraphSession.Instance.GraphRequestPopContext.Request.Headers.Add(header.Key, header.Value.First());
                     }
 
-                    var accessToken = GraphSession.Instance.GraphRequestProofofPossession.BrowserCredential.GetTokenAsync(GraphSession.Instance.GraphRequestProofofPossession.PopTokenContext, cancellationToken).Result;
-                    httpRequestMessage.Headers.Authorization = new AuthenticationHeaderValue("Pop", accessToken.Token);
+                    var accessToken = await GraphSession.Instance.GraphRequestPopContext.PopInteractiveBrowserCredential.GetTokenAsync(GraphSession.Instance.GraphRequestPopContext.PopTokenContext, cancellationToken).ConfigureAwait(false);
+                    httpRequestMessage.Headers.Authorization = new AuthenticationHeaderValue(PopAuthenticationScheme, accessToken.Token);
                 }
                 else
                 {
@@ -104,15 +113,6 @@ private async Task<HttpResponseMessage> SendRetryAsync(HttpResponseMessage httpR
                 }
                 await DrainAsync(httpResponseMessage).ConfigureAwait(false);
 
-                if (AuthenticationHelpers.IsATPoPSupported())
-                {
-                    var popChallenge = httpResponseMessage.Headers.WwwAuthenticate.First(wa => wa.Scheme == "PoP");
-                    var nonceStart = popChallenge.Parameter.IndexOf("nonce=\"") + "nonce=\"".Length;
-                    var nonceEnd = popChallenge.Parameter.IndexOf('"', nonceStart);
-                    GraphSession.Instance.GraphRequestProofofPossession.ProofofPossessionNonce = popChallenge.Parameter.Substring(nonceStart, nonceEnd - nonceStart);
-                    GraphSession.Instance.GraphRequestProofofPossession.PopTokenContext = new PopTokenRequestContext(GraphSession.Instance.AuthContext.Scopes, isProofOfPossessionEnabled: true, proofOfPossessionNonce: GraphSession.Instance.GraphRequestProofofPossession.ProofofPossessionNonce, request: GraphSession.Instance.GraphRequestProofofPossession.Request);
-                }
-
                 // Authenticate request using auth provider
                 await AuthenticateRequestAsync(newRequest, additionalRequestInfo, cancellationToken).ConfigureAwait(false);
                 httpResponseMessage = await base.SendAsync(newRequest, cancellationToken);

src/Authentication/Authentication/Models/GraphRequestProofofPossession.cs renamed to src/Authentication/Authentication/Models/GraphRequestPopContext.cs

@@ -10,15 +10,15 @@
 
 namespace Microsoft.Graph.PowerShell.Authentication
 {
-    internal class GraphRequestProofofPossession : IGraphRequestProofofPossession
+    internal class GraphRequestPopContext : IGraphRequestPopContext
     {
         public Uri Uri { get; set; }
         public HttpMethod HttpMethod { get; set; }
         public AccessToken AccessToken { get; set; }
         public string ProofofPossessionNonce { get; set; }
         public PopTokenRequestContext PopTokenContext { get; set; }
         public Request Request { get; set; }
-        public InteractiveBrowserCredential BrowserCredential { get; set; }
+        public InteractiveBrowserCredential PopInteractiveBrowserCredential { get; set; }
     }
 
 }