Support Azure Service Bus Geo-Recovery alias for querying queue names (#4938)

* Support ASB Geo-Recovery alias for querying queue names

* Account for other regional clouds

* Tweaks

* Use a dictionary on ArmEnvironment since it is IEquatable

---------

Co-authored-by: Brandon Ording <[email protected]>

Author: DavidBoike

src/Directory.Packages.props

@@ -12,6 +12,7 @@
     <PackageVersion Include="Azure.ResourceManager.ServiceBus" Version="1.1.0" />
     <PackageVersion Include="ByteSize" Version="2.1.2" />
     <PackageVersion Include="Caliburn.Micro" Version="4.0.230" />
+    <PackageVersion Include="DnsClient" Version="1.8.0" />
     <PackageVersion Include="FluentValidation" Version="11.11.0" />
     <PackageVersion Include="Fody" Version="6.9.1" />
     <PackageVersion Include="GitHubActionsTestLogger" Version="2.4.1" />

src/ServiceControl.Transports.ASBS/AzureQuery.cs

@@ -18,6 +18,8 @@ namespace ServiceControl.Transports.ASBS;
 using Azure.ResourceManager.Resources;
 using Azure.ResourceManager.ServiceBus;
 using BrokerThroughput;
+using DnsClient;
+using DnsClient.Protocol;
 using Microsoft.Extensions.Logging;
 
 public class AzureQuery(ILogger<AzureQuery> logger, TimeProvider timeProvider, TransportSettings transportSettings)
@@ -27,6 +29,7 @@ public class AzureQuery(ILogger<AzureQuery> logger, TimeProvider timeProvider, T
     MetricsQueryClient? client;
     ArmClient? armClient;
     string? resourceId;
+    ArmEnvironment armEnvironment;
 
     protected override void InitializeCore(ReadOnlyDictionary<string, string> settings)
     {
@@ -99,11 +102,11 @@ protected override void InitializeCore(ReadOnlyDictionary<string, string> settin
             Diagnostics.AppendLine("Client secret set");
         }
 
-        (ArmEnvironment armEnvironment, MetricsQueryAudience metricsQueryAudience) environment = GetEnvironment();
+        (armEnvironment, var metricsQueryAudience) = GetEnvironment();
 
         if (managementUrl == null)
         {
-            Diagnostics.AppendLine($"Management Url not set, defaulted to \"{environment.armEnvironment.Endpoint}\"");
+            Diagnostics.AppendLine($"Management Url not set, defaulted to \"{armEnvironment.Endpoint}\"");
         }
         else
         {
@@ -126,10 +129,10 @@ protected override void InitializeCore(ReadOnlyDictionary<string, string> settin
             clientCredentials = new ClientSecretCredential(tenantId, clientId, clientSecret);
         }
 
-        client = new MetricsQueryClient(environment.armEnvironment.Endpoint, clientCredentials,
+        client = new MetricsQueryClient(armEnvironment.Endpoint, clientCredentials,
             new MetricsQueryClientOptions
             {
-                Audience = environment.metricsQueryAudience,
+                Audience = metricsQueryAudience,
                 Transport = new HttpClientTransport(
                     new HttpClient(new SocketsHttpHandler
                     {
@@ -139,7 +142,7 @@ protected override void InitializeCore(ReadOnlyDictionary<string, string> settin
         armClient = new ArmClient(clientCredentials, subscriptionId,
             new ArmClientOptions
             {
-                Environment = environment.armEnvironment,
+                Environment = armEnvironment,
                 Transport = new HttpClientTransport(
                     new HttpClient(new SocketsHttpHandler
                     {
@@ -263,14 +266,16 @@ async Task<IReadOnlyList<MetricValue>> GetMetrics(string queueName, DateOnly sta
     public override async IAsyncEnumerable<IBrokerQueue> GetQueueNames(
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
+        var validNamespaces = await GetValidNamespaceNames(cancellationToken);
+
         SubscriptionResource? subscription = await armClient!.GetDefaultSubscriptionAsync(cancellationToken);
         var namespaces =
             subscription.GetServiceBusNamespacesAsync(cancellationToken);
 
         await foreach (var serviceBusNamespaceResource in namespaces.WithCancellation(
                            cancellationToken))
         {
-            if (serviceBusNamespaceResource.Data.Name == serviceBusName)
+            if (validNamespaces.Contains(serviceBusNamespaceResource.Data.Name))
             {
                 resourceId = serviceBusNamespaceResource.Id;
                 await foreach (var queue in serviceBusNamespaceResource.GetServiceBusQueues()
@@ -286,6 +291,43 @@ public override async IAsyncEnumerable<IBrokerQueue> GetQueueNames(
         throw new Exception($"Could not find a ServiceBus named \"{serviceBusName}\"");
     }
 
+    // ArmEnvironment Audience Values: https://github.com/Azure/azure-sdk-for-net/blob/main/sdk/resourcemanager/Azure.ResourceManager/src/ArmEnvironment.cs
+    // Service Bus Domains: https://learn.microsoft.com/en-us/rest/api/servicebus/
+    static readonly Dictionary<ArmEnvironment, string> ServiceBusDomains = new()
+    {
+        { ArmEnvironment.AzurePublicCloud, "servicebus.windows.net" },
+        { ArmEnvironment.AzureGovernment, "servicebus.usgovcloudapi.net" },
+        { ArmEnvironment.AzureGermany, "servicebus.cloudapi.de" },
+        { ArmEnvironment.AzureChina, "servicebus.chinacloudapi.cn" },
+    };
+
+    async Task<HashSet<string>> GetValidNamespaceNames(CancellationToken cancellationToken = default)
+    {
+        var validNamespaces = new HashSet<string>(StringComparer.OrdinalIgnoreCase) { serviceBusName };
+
+        if (!ServiceBusDomains.TryGetValue(armEnvironment, out var serviceBusCloudDomain))
+        {
+            // Worst case: the DNS lookup finds nothing additional to match
+            serviceBusCloudDomain = "servicebus.windows.net";
+        }
+
+        var queryDomain = $"{serviceBusName}.{serviceBusCloudDomain}";
+        var validDomainTail = $".{serviceBusCloudDomain}.";
+
+        var dnsLookup = new LookupClient();
+        var dnsResult = await dnsLookup.QueryAsync(queryDomain, QueryType.CNAME, cancellationToken: cancellationToken);
+        var domain = (dnsResult.Answers.FirstOrDefault() as CNameRecord)?.CanonicalName.Value;
+        if (domain is not null && domain.EndsWith(validDomainTail))
+        {
+            // In some cases, like private networking access, result might be something like `namespacename.private` with a dot in the middle
+            // which is not a big deal because that will not actually match a namespace name in metrics
+            var otherName = domain[..^validDomainTail.Length];
+            validNamespaces.Add(otherName);
+        }
+
+        return validNamespaces;
+    }
+
     public override string SanitizedEndpointNameCleanser(string endpointName) => endpointName.ToLower();
 
     public override KeyDescriptionPair[] Settings =>

src/ServiceControl.Transports.ASBS/ServiceControl.Transports.ASBS.csproj

@@ -14,6 +14,7 @@
     <PackageReference Include="Azure.Identity" />
     <PackageReference Include="Azure.Monitor.Query" />
     <PackageReference Include="Azure.ResourceManager.ServiceBus" />
+    <PackageReference Include="DnsClient" />
     <PackageReference Include="NServiceBus.CustomChecks" />
     <PackageReference Include="NServiceBus.Transport.AzureServiceBus" />
   </ItemGroup>