Gateway Load Balancer Cross tenant issue fix (Azure#15195)

* fix * name fix * put await next to create update Co-authored-by: matyang222 <[email protected]>
lilinvictorms · Jun 10, 2021 · dab7a80 · dab7a80
1 parent f5899ca
commit dab7a80
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 5 deletions.
diff --git a/src/Network/Network/Generated/LoadBalancer/NewAzureRMLoadBalancerCommand.cs b/src/Network/Network/Generated/LoadBalancer/NewAzureRMLoadBalancerCommand.cs
@@ -247,14 +247,37 @@ public override void Execute()
                 }
             }
 
+            List<string> resourceIdsRequiringAuthToken = new List<string>();
+            Dictionary<string, List<string>> auxAuthHeader = null;
+
+            // Get aux token for each gateway lb references
+            foreach (FrontendIPConfiguration frontend in vLoadBalancerModel.FrontendIPConfigurations)
+            {
+                if (frontend.GatewayLoadBalancer != null)
+                {
+                    //Get the aux header for the remote vnet
+                    resourceIdsRequiringAuthToken.Add(frontend.GatewayLoadBalancer.Id);
+                }
+            }
+
+            if (resourceIdsRequiringAuthToken.Count > 0)
+            {
+                var auxHeaderDictionary = GetAuxilaryAuthHeaderFromResourceIds(resourceIdsRequiringAuthToken);
+                if (auxHeaderDictionary != null && auxHeaderDictionary.Count > 0)
+                {
+                    auxAuthHeader = new Dictionary<string, List<string>>(auxHeaderDictionary);
+                }
+            }
+
+
             ConfirmAction(
                 Force.IsPresent,
                 string.Format(Properties.Resources.OverwritingResource, Name),
                 Properties.Resources.CreatingResourceMessage,
                 Name,
             () =>
             {
-                this.NetworkClient.NetworkManagementClient.LoadBalancers.CreateOrUpdate(this.ResourceGroupName, this.Name, vLoadBalancerModel);
+                this.NetworkClient.NetworkManagementClient.LoadBalancers.CreateOrUpdateWithHttpMessagesAsync(this.ResourceGroupName, this.Name, vLoadBalancerModel, auxAuthHeader).GetAwaiter().GetResult();
                 var getLoadBalancer = this.NetworkClient.NetworkManagementClient.LoadBalancers.Get(this.ResourceGroupName, this.Name);
                 var psLoadBalancer = NetworkResourceManagerProfile.Mapper.Map<PSLoadBalancer>(getLoadBalancer);
                 psLoadBalancer.ResourceGroupName = this.ResourceGroupName;

diff --git a/src/Network/Network/Generated/LoadBalancer/SetAzureRMLoadBalancerCommand.cs b/src/Network/Network/Generated/LoadBalancer/SetAzureRMLoadBalancerCommand.cs
@@ -128,8 +128,30 @@ public override void Execute()
             var vLoadBalancerModel = NetworkResourceManagerProfile.Mapper.Map<MNM.LoadBalancer>(this.LoadBalancer);
             vLoadBalancerModel.Tags = TagsConversionHelper.CreateTagDictionary(this.LoadBalancer.Tag, validate: true);
 
+            List<string> resourceIds = new List<string>();
+            Dictionary<string, List<string>> auxAuthHeader = null;
+
+            // Get aux token for each gateway lb references
+            foreach (FrontendIPConfiguration frontend in vLoadBalancerModel.FrontendIPConfigurations)
+            {
+                if (frontend.GatewayLoadBalancer != null)
+                {
+                    //Get the aux header for the remote vnet
+                    resourceIds.Add(frontend.GatewayLoadBalancer.Id);
+                }
+            }
+
+            if (resourceIds.Count > 0)
+            {
+                var auxHeaderDictionary = GetAuxilaryAuthHeaderFromResourceIds(resourceIds);
+                if (auxHeaderDictionary != null && auxHeaderDictionary.Count > 0)
+                {
+                    auxAuthHeader = new Dictionary<string, List<string>>(auxHeaderDictionary);
+                }
+            }
+
             // Execute the PUT LoadBalancer call
-            this.NetworkClient.NetworkManagementClient.LoadBalancers.CreateOrUpdate(this.LoadBalancer.ResourceGroupName, this.LoadBalancer.Name, vLoadBalancerModel);
+            this.NetworkClient.NetworkManagementClient.LoadBalancers.CreateOrUpdateWithHttpMessagesAsync(this.LoadBalancer.ResourceGroupName, this.LoadBalancer.Name, vLoadBalancerModel, auxAuthHeader).GetAwaiter().GetResult();
 
             var getLoadBalancer = this.NetworkClient.NetworkManagementClient.LoadBalancers.Get(this.LoadBalancer.ResourceGroupName, this.LoadBalancer.Name);
             var psLoadBalancer = NetworkResourceManagerProfile.Mapper.Map<PSLoadBalancer>(getLoadBalancer);

diff --git a/src/Network/Network/NetworkInterface/NewAzureNetworkInterfaceCommand.cs b/src/Network/Network/NetworkInterface/NewAzureNetworkInterfaceCommand.cs
@@ -437,13 +437,35 @@ private PSNetworkInterface CreateNetworkInterface()
                 networkInterface.NetworkSecurityGroup.Id = this.NetworkSecurityGroupId;
             }
 
+            List<string> resourceIdsRequiringAuthToken = new List<string>();
+            Dictionary<string, List<string>> auxAuthHeader = null;
+
+            // Get aux token for each gateway lb references
+            foreach (var ipConfiguration in networkInterface.IpConfigurations)
+            {
+                if (ipConfiguration.GatewayLoadBalancer != null)
+                {
+                    //Get the aux header for the remote vnet
+                    resourceIdsRequiringAuthToken.Add(ipConfiguration.GatewayLoadBalancer.Id);
+                }
+            }
+
+            if (resourceIdsRequiringAuthToken.Count > 0)
+            {
+                var auxHeaderDictionary = GetAuxilaryAuthHeaderFromResourceIds(resourceIdsRequiringAuthToken);
+                if (auxHeaderDictionary != null && auxHeaderDictionary.Count > 0)
+                {
+                    auxAuthHeader = new Dictionary<string, List<string>>(auxHeaderDictionary);
+                }
+            }
+
             var networkInterfaceModel = NetworkResourceManagerProfile.Mapper.Map<MNM.NetworkInterface>(networkInterface);
 
 			this.NullifyApplicationSecurityGroupIfAbsent(networkInterfaceModel);
 
 			networkInterfaceModel.Tags = TagsConversionHelper.CreateTagDictionary(this.Tag, validate: true);
 
-            this.NetworkInterfaceClient.CreateOrUpdate(this.ResourceGroupName, this.Name, networkInterfaceModel);
+            this.NetworkInterfaceClient.CreateOrUpdateWithHttpMessagesAsync(this.ResourceGroupName, this.Name, networkInterfaceModel, auxAuthHeader).GetAwaiter().GetResult();
 
             var getNetworkInterface = this.GetNetworkInterface(this.ResourceGroupName, this.Name);
 

diff --git a/src/Network/Network/NetworkInterface/SetAzureNetworkInterfaceCommand.cs b/src/Network/Network/NetworkInterface/SetAzureNetworkInterfaceCommand.cs
@@ -17,6 +17,7 @@
 using Microsoft.Azure.Commands.ResourceManager.Common.Tags;
 using Microsoft.Azure.Management.Network;
 using System;
+using System.Collections.Generic;
 using System.Management.Automation;
 using MNM = Microsoft.Azure.Management.Network.Models;
 
@@ -43,14 +44,32 @@ public override void Execute()
                 throw new ArgumentException(Microsoft.Azure.Commands.Network.Properties.Resources.ResourceNotFound);
             }
 
-            // Verify if PublicIpAddress is empty
+            List<string> resourceIdsRequiringAuthToken = new List<string>();
+            Dictionary<string, List<string>> auxAuthHeader = null;
+
             foreach (var ipconfig in NetworkInterface.IpConfigurations)
             {
+                // Verify if PublicIpAddress is empty
                 if (ipconfig.PublicIpAddress != null &&
                     string.IsNullOrEmpty(ipconfig.PublicIpAddress.Id))
                 {
                     ipconfig.PublicIpAddress = null;
                 }
+
+                if (ipconfig.GatewayLoadBalancer != null)
+                {
+                    //Get the aux header for the remote vnet
+                    resourceIdsRequiringAuthToken.Add(ipconfig.GatewayLoadBalancer.Id);
+                }
+            }
+
+            if (resourceIdsRequiringAuthToken.Count > 0)
+            {
+                var auxHeaderDictionary = GetAuxilaryAuthHeaderFromResourceIds(resourceIdsRequiringAuthToken);
+                if (auxHeaderDictionary != null && auxHeaderDictionary.Count > 0)
+                {
+                    auxAuthHeader = new Dictionary<string, List<string>>(auxHeaderDictionary);
+                }
             }
 
             // Map to the sdk object
@@ -60,7 +79,7 @@ public override void Execute()
 
 			networkInterfaceModel.Tags = TagsConversionHelper.CreateTagDictionary(this.NetworkInterface.Tag, validate: true);
 
-            this.NetworkInterfaceClient.CreateOrUpdate(this.NetworkInterface.ResourceGroupName, this.NetworkInterface.Name, networkInterfaceModel);
+            this.NetworkInterfaceClient.CreateOrUpdateWithHttpMessagesAsync(this.NetworkInterface.ResourceGroupName, this.NetworkInterface.Name, networkInterfaceModel, auxAuthHeader).GetAwaiter().GetResult();
 
             var getNetworkInterface = this.GetNetworkInterface(this.NetworkInterface.ResourceGroupName, this.NetworkInterface.Name);
             WriteObject(getNetworkInterface);