Handle DNS Double Async Headers

Author: deathly809

src/StackAdmin/Dns/Commands.Dns/Commands.Dns.csproj

@@ -53,6 +53,7 @@
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="Models\DnsClient.cs" />
+    <Compile Include="Models\DoubleFetchHandler.cs" />
     <Compile Include="Models\DnsBaseCmdlet.cs" />
     <Compile Include="Models\DnsRecordSet.cs" />
     <Compile Include="Models\DnsZone.cs" />
@@ -108,4 +109,4 @@
     </ItemGroup>
     <Copy SourceFiles="@(MarkdownFiles)" DestinationFolder="$(OutputPath)\help\" ContinueOnError="false" />
   </Target>
-</Project>
+</Project>

src/StackAdmin/Dns/Commands.Dns/Models/DnsClient.cs

@@ -23,6 +23,7 @@
 using System.Linq;
 using ProjectResources = Microsoft.Azure.Commands.Dns.Properties.Resources;
 using Sdk = Microsoft.Azure.Management.Dns.Models;
+using System.Net.Http;
 
 namespace Microsoft.Azure.Commands.Dns.Models
 {
@@ -51,8 +52,21 @@ public class DnsClient
         };
 
         public DnsClient(IAzureContext context)
-            : this(AzureSession.Instance.ClientFactory.CreateArmClient<DnsManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager))
         {
+             // Factories
+            var authFactory = AzureSession.Instance.AuthenticationFactory;
+            var clientFactory = AzureSession.Instance.ClientFactory;
+
+            var endpoint = AzureEnvironment.Endpoint.ResourceManager;
+
+            // Get parameters
+            var handler = new DelegatingHandler[] { new DoubleFetchHandler() };
+            var creds = authFactory.GetServiceClientCredentials(context, endpoint);
+            var baseUri = context.Environment.GetEndpointAsUri(endpoint);
+
+            // Construct client
+            this.DnsManagementClient = clientFactory.CreateCustomArmClient<DnsManagementClient>(baseUri, creds, handler);
+            this.DnsManagementClient.SubscriptionId = context.Subscription.Id.ToString();
         }
 
         public DnsClient(IDnsManagementClient managementClient)
@@ -127,9 +141,9 @@ public List<DnsZone> ListDnsZonesInResourceGroup(string resourceGroupName)
                 }
                 else
                 {
-                    getResponse = this.DnsManagementClient.Zones.ListInResourceGroup(resourceGroupName);    
+                    getResponse = this.DnsManagementClient.Zones.ListInResourceGroup(resourceGroupName);
                 }
-                
+
                 results.AddRange(getResponse.Select(ToDnsZone));
             } while (getResponse != null && getResponse.NextPageLink != null);
 
@@ -204,7 +218,7 @@ private RecordSet ConstructRecordSetPropeties(string recordSetName, RecordType r
             }
             else
             {
-                FillEmptyRecordsForType( properties, recordType);                
+                FillEmptyRecordsForType( properties, recordType);
             }
             return properties;
         }
@@ -354,7 +368,7 @@ public List<DnsRecordSet> ListRecordSets(string zoneName, string resourceGroupNa
                 }
 
                 results.AddRange(listResponse.Select(recordSet => GetPowerShellRecordSet(zoneName, resourceGroupName, recordSet)));
-                
+
             } while (listResponse != null && listResponse.NextPageLink != null);
 
             return results;

src/StackAdmin/Dns/Commands.Dns/Models/DoubleFetchHandler.cs

@@ -0,0 +1,29 @@
+using System;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Azure.Commands.Dns.Models
+{
+    public class DoubleFetchHandler : DelegatingHandler, ICloneable
+    {
+        public object Clone() {
+            return new DoubleFetchHandler();
+        }
+
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) {
+
+            return base.SendAsync(request, cancellationToken).ContinueWith<HttpResponseMessage>(response =>
+            {
+                var result = response.Result;
+                    if (result.Headers.Contains("Location") && result.Headers.Contains("Azure-AsyncOperation"))
+                    {
+                        result.Headers.Remove("Location");
+                    }
+
+                return result;
+            });
+        }
+
+    }
+}