Extracted the correct graph endpoint for cloud environments

microsoftgraph · Sep 26, 2024 · 4ca11fa · 4ca11fa
1 parent fddf88c
commit 4ca11fa
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/src/Authentication/Authentication/Cmdlets/InvokeMgGraphRequest.cs b/src/Authentication/Authentication/Cmdlets/InvokeMgGraphRequest.cs
@@ -1003,11 +1003,8 @@ private string QualifyFilePath(string path)
         /// </summary>
         private void ResetGraphSessionEnvironment()
         {
-            var currentEnvironment = GraphSession.Instance.Environment;
-            if(currentEnvironment != null && !currentEnvironment.Equals(_originalEnvironment))
-            {
-                GraphSession.Instance.Environment = _originalEnvironment;
-            }
+            _originalEnvironment = GraphSession.Instance.Environment;
+            GraphSession.Instance.Environment = _originalEnvironment;
         }
 
         #region CmdLet LifeCycle
@@ -1046,8 +1043,6 @@ private async Task ProcessRecordAsync()
                             if (ShouldCheckHttpStatus && !isSuccess)
                             {
                                 var httpErrorRecord = await GenerateHttpErrorRecordAsync(httpResponseMessageFormatter, httpRequestMessage);
-                                // A reset of the GraphSession Environment is required to avoid side effects
-                                ResetGraphSessionEnvironment();
                                 ThrowTerminatingError(httpErrorRecord);
                             }
                             await ProcessResponseAsync(httpResponseMessage);

diff --git a/src/Authentication/Authentication/custom/common/Permissions.ps1 b/src/Authentication/Authentication/custom/common/Permissions.ps1
@@ -34,7 +34,20 @@ function Permissions_GetPermissionsData([bool] $online) {
     # Make a REST request to MS Graph to get the permissions data from the Microsoft Graph service principal
     if ( $online -or ! $_permissions.msGraphServicePrincipal -or ! $_permissions.isFromInvokeMgGraphRequest ) {
         try {
-            $restResult = Invoke-MgGraphRequest -method GET -OutputType PSObject $_permissions.msGraphPermissionsRequestUri
+
+            # Get-MgContext is used to get the current context for the request to MS Graph
+            # From the context, we can get the current environment and use it to get the permissions request URI
+            # If the context is not available, then we will use the default permissions request URI
+
+            $context = Get-MgContext
+            $uri = $_permissions.msGraphPermissionsRequestUri
+            if($context){
+                $currentEnv = $context.Environment
+                $allEnv = Get-MgEnvironment
+                $env = $allEnv | Where-Object { $_.Name -eq $currentEnv }
+                $uri = $env.GraphEndpoint + "/v1.0/servicePrincipals?`$filter=appId eq '$Permissions_msGraphApplicationId'"
+            }
+            $restResult = Invoke-MgGraphRequest -method GET -OutputType PSObject $uri
 
             if ( $restResult ) {
                 $_permissions.msGraphServicePrincipal = $restResult | Select-Object -ExpandProperty value