From 75dbcec41a4dd5ca57c0420eb41fdacf2cf21a11 Mon Sep 17 00:00:00 2001 From: Peter Ombwa Date: Wed, 30 Mar 2022 15:23:11 -0700 Subject: [PATCH] Reuse session HTTP client. --- .../Common/GraphSession.cs | 3 + .../Helpers/HttpHelpersTests.cs | 108 ++++++++++++++++++ .../Authentication/Cmdlets/ConnectMgGraph.cs | 1 + .../Cmdlets/DisconnectMgGraph.cs | 1 + .../Authentication/Helpers/HttpHelpers.cs | 6 +- 5 files changed, 118 insertions(+), 1 deletion(-) diff --git a/src/Authentication/Authentication.Core/Common/GraphSession.cs b/src/Authentication/Authentication.Core/Common/GraphSession.cs index f25c9de2854..1d8fab28866 100644 --- a/src/Authentication/Authentication.Core/Common/GraphSession.cs +++ b/src/Authentication/Authentication.Core/Common/GraphSession.cs @@ -9,6 +9,7 @@ namespace Microsoft.Graph.PowerShell.Authentication using System; using System.Collections; + using System.Net.Http; using System.Security; using System.Threading; @@ -68,6 +69,8 @@ public byte[] MSALToken /// public Hashtable[] MgCommandMetadata { get; set; } + public HttpClient GraphHttpClient { get; set; } + /// /// Gets an instance of . /// diff --git a/src/Authentication/Authentication.Test/Helpers/HttpHelpersTests.cs b/src/Authentication/Authentication.Test/Helpers/HttpHelpersTests.cs index c42ee246ff3..eec17d36e24 100644 --- a/src/Authentication/Authentication.Test/Helpers/HttpHelpersTests.cs +++ b/src/Authentication/Authentication.Test/Helpers/HttpHelpersTests.cs @@ -5,6 +5,7 @@ namespace Microsoft.Graph.Authentication.Test.Helpers { using System; using System.Net.Http; + using Microsoft.Graph.Authentication.Core; using Microsoft.Graph.PowerShell.Authentication; using Microsoft.Graph.PowerShell.Authentication.Helpers; using Xunit; @@ -71,5 +72,112 @@ public void GetGraphHttpClientShouldReturnHttpClientWithCustomerProvidedTimeout( // reset static instance. GraphSession.Reset(); } + + [Fact] + public void GetGraphHttpClientShouldReturnInSessionHttpClientWhenSessionHasAClient() + { + GraphSession.Initialize(() => new GraphSession()); + GraphSession.Instance.AuthContext = new AuthContext + { + AuthType = AuthenticationType.UserProvidedAccessToken, + ContextScope = ContextScope.Process, + }; + GraphSession.Instance.GraphHttpClient = new HttpClient + { + BaseAddress = new Uri("https://test.contoso.com/v1.0/") + }; + + HttpClient httpClient = HttpHelpers.GetGraphHttpClient(); + + Assert.NotNull(httpClient); + Assert.NotNull(GraphSession.Instance.GraphHttpClient); + Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress); + + // reset static instance. + GraphSession.Reset(); + } + + [Fact] + public void GetGraphHttpClientShouldReturnNewHttpClientWhenSessionHasNoClient() + { + GraphSession.Initialize(() => new GraphSession()); + GraphSession.Instance.AuthContext = new AuthContext + { + AuthType = AuthenticationType.UserProvidedAccessToken, + ContextScope = ContextScope.Process, + }; + GraphSession.Instance.GraphHttpClient = null; + + HttpClient httpClient = HttpHelpers.GetGraphHttpClient(); + + Assert.NotNull(httpClient); + Assert.NotNull(GraphSession.Instance.GraphHttpClient); + Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress); + + // reset static instance. + GraphSession.Reset(); + } + + [Fact] + public void GetGraphHttpClientShouldReturnNewHttpClientWhenSessionIsNew() + { + GraphSession.Initialize(() => new GraphSession()); + GraphSession.Instance.AuthContext = new AuthContext + { + AuthType = AuthenticationType.UserProvidedAccessToken, + ContextScope = ContextScope.Process, + }; + + HttpClient httpClient = HttpHelpers.GetGraphHttpClient(); + + Assert.NotNull(httpClient); + Assert.NotNull(GraphSession.Instance.GraphHttpClient); + Assert.Equal(httpClient.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress); + + // reset static instance. + GraphSession.Reset(); + } + + [Fact] + public void GetGraphHttpClientShouldReturnNewHttpClientSignOutThenSignIn() + { + GraphSession.Initialize(() => new GraphSession()); + GraphSession.Instance.AuthContext = new AuthContext + { + AuthType = AuthenticationType.UserProvidedAccessToken, + ContextScope = ContextScope.Process, + }; + + var dummyClient = new HttpClient + { + BaseAddress = new Uri("https://test.contoso.com/v1.0/") + }; + GraphSession.Instance.GraphHttpClient = dummyClient; + HttpClient httpClientAttempt1 = HttpHelpers.GetGraphHttpClient(); + + // Mock sign out. + Authenticator.LogOut(GraphSession.Instance.AuthContext); + GraphSession.Instance.AuthContext = null; + GraphSession.Instance.GraphHttpClient = null; + + // Mock sign in. + GraphSession.Initialize(() => new GraphSession()); + GraphSession.Instance.AuthContext = new AuthContext + { + AuthType = AuthenticationType.UserProvidedAccessToken, + ContextScope = ContextScope.Process, + }; + HttpClient httpClientAttempt2 = HttpHelpers.GetGraphHttpClient(); + + Assert.NotNull(httpClientAttempt1); + Assert.NotNull(httpClientAttempt2); + Assert.NotNull(GraphSession.Instance.GraphHttpClient); + Assert.NotEqual(httpClientAttempt2.BaseAddress, httpClientAttempt1.BaseAddress); + Assert.Equal(httpClientAttempt1.BaseAddress, dummyClient.BaseAddress); + Assert.Equal(httpClientAttempt2.BaseAddress, GraphSession.Instance.GraphHttpClient.BaseAddress); + + // reset static instance. + GraphSession.Reset(); + } } } diff --git a/src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs b/src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs index 587421fe440..cb9b9c9c977 100644 --- a/src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs +++ b/src/Authentication/Authentication/Cmdlets/ConnectMgGraph.cs @@ -189,6 +189,7 @@ private async Task ProcessRecordAsync() if (MyInvocation.BoundParameters.ContainsKey(nameof(ClientTimeout))) { authContext.ClientTimeout = TimeSpan.FromSeconds(ClientTimeout); } // Set selected environment to the session object. GraphSession.Instance.Environment = environment; + GraphSession.Instance.GraphHttpClient = null; switch (ParameterSetName) { case Constants.UserParameterSet: diff --git a/src/Authentication/Authentication/Cmdlets/DisconnectMgGraph.cs b/src/Authentication/Authentication/Cmdlets/DisconnectMgGraph.cs index a059cc595aa..a7c88cfc1cd 100644 --- a/src/Authentication/Authentication/Cmdlets/DisconnectMgGraph.cs +++ b/src/Authentication/Authentication/Cmdlets/DisconnectMgGraph.cs @@ -34,6 +34,7 @@ protected override void ProcessRecord() Authenticator.LogOut(authContext); GraphSession.Instance.AuthContext = null; + GraphSession.Instance.GraphHttpClient = null; } protected override void StopProcessing() diff --git a/src/Authentication/Authentication/Helpers/HttpHelpers.cs b/src/Authentication/Authentication/Helpers/HttpHelpers.cs index 7a94e1981bc..a5a447c4598 100644 --- a/src/Authentication/Authentication/Helpers/HttpHelpers.cs +++ b/src/Authentication/Authentication/Helpers/HttpHelpers.cs @@ -60,6 +60,8 @@ public static HttpClient GetGraphHttpClient(InvocationInfo invocationInfo, IAuth /// public static HttpClient GetGraphHttpClient(IAuthContext authContext = null) { + if (GraphSession.Instance?.GraphHttpClient != null) + return GraphSession.Instance.GraphHttpClient; authContext = authContext ?? GraphSession.Instance.AuthContext; if (authContext is null) { @@ -67,7 +69,9 @@ public static HttpClient GetGraphHttpClient(IAuthContext authContext = null) } IAuthenticationProvider authProvider = AuthenticationHelpers.GetAuthProvider(authContext); - return GetGraphHttpClient(authProvider, authContext.ClientTimeout); + var newHttpClient = GetGraphHttpClient(authProvider, authContext.ClientTimeout); + GraphSession.Instance.GraphHttpClient = newHttpClient; + return newHttpClient; } ///