[PM-6339] Shard notification hub clients across multiple accounts (#3812)

* WIP registration updates

* fix deviceHubs

* addHub inline in ctor

* adjust setttings for hub reg

* send to all clients

* fix multiservice push

* use notification hub type

* feedback

---------

Co-authored-by: Matt Bishop <mbishop@bitwarden.com>

src/Core/Services/IPushRegistrationService.cs

@@ -6,7 +6,7 @@ public interface IPushRegistrationService
 {
     Task CreateOrUpdateRegistrationAsync(string pushToken, string deviceId, string userId,
         string identifier, DeviceType type);
-    Task DeleteRegistrationAsync(string deviceId);
-    Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId);
-    Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId);
+    Task DeleteRegistrationAsync(string deviceId, DeviceType type);
+    Task AddUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId);
+    Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId);
 }

src/Core/Services/Implementations/DeviceService.cs

@@ -38,13 +38,13 @@ public class DeviceService : IDeviceService
     public async Task ClearTokenAsync(Device device)
     {
         await _deviceRepository.ClearPushTokenAsync(device.Id);
-        await _pushRegistrationService.DeleteRegistrationAsync(device.Id.ToString());
+        await _pushRegistrationService.DeleteRegistrationAsync(device.Id.ToString(), device.Type);
     }
 
     public async Task DeleteAsync(Device device)
     {
         await _deviceRepository.DeleteAsync(device);
-        await _pushRegistrationService.DeleteRegistrationAsync(device.Id.ToString());
+        await _pushRegistrationService.DeleteRegistrationAsync(device.Id.ToString(), device.Type);
     }
 
     public async Task UpdateDevicesTrustAsync(string currentDeviceIdentifier,

src/Core/Services/Implementations/MultiServicePushNotificationService.cs

@@ -43,7 +43,8 @@ public class MultiServicePushNotificationService : IPushNotificationService
         }
         else
         {
-            if (CoreHelpers.SettingHasValue(globalSettings.NotificationHub.ConnectionString))
+            var generalHub = globalSettings.NotificationHubs?.FirstOrDefault(h => h.HubType == NotificationHubType.General);
+            if (CoreHelpers.SettingHasValue(generalHub?.ConnectionString))
             {
                 _services.Add(new NotificationHubPushNotificationService(installationDeviceRepository,
                     globalSettings, httpContextAccessor, hubLogger));

src/Core/Services/Implementations/NotificationHubPushNotificationService.cs

@@ -20,8 +20,9 @@ public class NotificationHubPushNotificationService : IPushNotificationService
     private readonly IInstallationDeviceRepository _installationDeviceRepository;
     private readonly GlobalSettings _globalSettings;
     private readonly IHttpContextAccessor _httpContextAccessor;
-    private NotificationHubClient _client = null;
-    private ILogger _logger;
+    private readonly List<NotificationHubClient> _clients = [];
+    private readonly bool _enableTracing = false;
+    private readonly ILogger _logger;
 
     public NotificationHubPushNotificationService(
         IInstallationDeviceRepository installationDeviceRepository,
@@ -32,10 +33,18 @@ public class NotificationHubPushNotificationService : IPushNotificationService
         _installationDeviceRepository = installationDeviceRepository;
         _globalSettings = globalSettings;
         _httpContextAccessor = httpContextAccessor;
-        _client = NotificationHubClient.CreateClientFromConnectionString(
-            _globalSettings.NotificationHub.ConnectionString,
-            _globalSettings.NotificationHub.HubName,
-            _globalSettings.NotificationHub.EnableSendTracing);
+
+        foreach (var hub in globalSettings.NotificationHubs)
+        {
+            var client = NotificationHubClient.CreateClientFromConnectionString(
+                hub.ConnectionString,
+                hub.HubName,
+                hub.EnableSendTracing);
+            _clients.Add(client);
+
+            _enableTracing = _enableTracing || hub.EnableSendTracing;
+        }
+
         _logger = logger;
     }
 
@@ -255,16 +264,31 @@ public class NotificationHubPushNotificationService : IPushNotificationService
 
     private async Task SendPayloadAsync(string tag, PushType type, object payload)
     {
-        var outcome = await _client.SendTemplateNotificationAsync(
-            new Dictionary<string, string>
-            {
-                { "type",  ((byte)type).ToString() },
-                { "payload", JsonSerializer.Serialize(payload) }
-            }, tag);
-        if (_globalSettings.NotificationHub.EnableSendTracing)
+        var tasks = new List<Task<NotificationOutcome>>();
+        foreach (var client in _clients)
         {
-            _logger.LogInformation("Azure Notification Hub Tracking ID: {id} | {type} push notification with {success} successes and {failure} failures with a payload of {@payload} and result of {@results}",
-                outcome.TrackingId, type, outcome.Success, outcome.Failure, payload, outcome.Results);
+            var task = client.SendTemplateNotificationAsync(
+                new Dictionary<string, string>
+                {
+                    { "type",  ((byte)type).ToString() },
+                    { "payload", JsonSerializer.Serialize(payload) }
+                }, tag);
+            tasks.Add(task);
+        }
+
+        await Task.WhenAll(tasks);
+
+        if (_enableTracing)
+        {
+            for (var i = 0; i < tasks.Count; i++)
+            {
+                if (_clients[i].EnableTestSend)
+                {
+                    var outcome = await tasks[i];
+                    _logger.LogInformation("Azure Notification Hub Tracking ID: {id} | {type} push notification with {success} successes and {failure} failures with a payload of {@payload} and result of {@results}",
+                        outcome.TrackingId, type, outcome.Success, outcome.Failure, payload, outcome.Results);
+                }
+            }
         }
     }
 

src/Core/Services/Implementations/NotificationHubPushRegistrationService.cs

@@ -3,6 +3,7 @@ using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
 using Bit.Core.Settings;
 using Microsoft.Azure.NotificationHubs;
+using Microsoft.Extensions.Logging;
 
 namespace Bit.Core.Services;
 
@@ -10,18 +11,36 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
 {
     private readonly IInstallationDeviceRepository _installationDeviceRepository;
     private readonly GlobalSettings _globalSettings;
-
-    private NotificationHubClient _client = null;
+    private readonly ILogger<NotificationHubPushRegistrationService> _logger;
+    private Dictionary<NotificationHubType, NotificationHubClient> _clients = [];
 
     public NotificationHubPushRegistrationService(
         IInstallationDeviceRepository installationDeviceRepository,
-        GlobalSettings globalSettings)
+        GlobalSettings globalSettings,
+        ILogger<NotificationHubPushRegistrationService> logger)
     {
         _installationDeviceRepository = installationDeviceRepository;
         _globalSettings = globalSettings;
-        _client = NotificationHubClient.CreateClientFromConnectionString(
-            _globalSettings.NotificationHub.ConnectionString,
-            _globalSettings.NotificationHub.HubName);
+        _logger = logger;
+
+        // Is this dirty to do in the ctor?
+        void addHub(NotificationHubType type)
+        {
+            var hubRegistration = globalSettings.NotificationHubs.FirstOrDefault(
+                h => h.HubType == type && h.EnableRegistration);
+            if (hubRegistration != null)
+            {
+                var client = NotificationHubClient.CreateClientFromConnectionString(
+                    hubRegistration.ConnectionString,
+                    hubRegistration.HubName,
+                    hubRegistration.EnableSendTracing);
+                _clients.Add(type, client);
+            }
+        }
+
+        addHub(NotificationHubType.General);
+        addHub(NotificationHubType.iOS);
+        addHub(NotificationHubType.Android);
     }
 
     public async Task CreateOrUpdateRegistrationAsync(string pushToken, string deviceId, string userId,
@@ -84,7 +103,7 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         BuildInstallationTemplate(installation, "badgeMessage", badgeMessageTemplate ?? messageTemplate,
             userId, identifier);
 
-        await _client.CreateOrUpdateInstallationAsync(installation);
+        await GetClient(type).CreateOrUpdateInstallationAsync(installation);
         if (InstallationDeviceEntity.IsInstallationDeviceId(deviceId))
         {
             await _installationDeviceRepository.UpsertAsync(new InstallationDeviceEntity(deviceId));
@@ -119,11 +138,11 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         installation.Templates.Add(fullTemplateId, template);
     }
 
-    public async Task DeleteRegistrationAsync(string deviceId)
+    public async Task DeleteRegistrationAsync(string deviceId, DeviceType deviceType)
     {
         try
         {
-            await _client.DeleteInstallationAsync(deviceId);
+            await GetClient(deviceType).DeleteInstallationAsync(deviceId);
             if (InstallationDeviceEntity.IsInstallationDeviceId(deviceId))
             {
                 await _installationDeviceRepository.DeleteAsync(new InstallationDeviceEntity(deviceId));
@@ -135,31 +154,31 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         }
     }
 
-    public async Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public async Task AddUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
-        await PatchTagsForUserDevicesAsync(deviceIds, UpdateOperationType.Add, $"organizationId:{organizationId}");
-        if (deviceIds.Any() && InstallationDeviceEntity.IsInstallationDeviceId(deviceIds.First()))
+        await PatchTagsForUserDevicesAsync(devices, UpdateOperationType.Add, $"organizationId:{organizationId}");
+        if (devices.Any() && InstallationDeviceEntity.IsInstallationDeviceId(devices.First().Key))
         {
-            var entities = deviceIds.Select(e => new InstallationDeviceEntity(e));
+            var entities = devices.Select(e => new InstallationDeviceEntity(e.Key));
             await _installationDeviceRepository.UpsertManyAsync(entities.ToList());
         }
     }
 
-    public async Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public async Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
-        await PatchTagsForUserDevicesAsync(deviceIds, UpdateOperationType.Remove,
+        await PatchTagsForUserDevicesAsync(devices, UpdateOperationType.Remove,
             $"organizationId:{organizationId}");
-        if (deviceIds.Any() && InstallationDeviceEntity.IsInstallationDeviceId(deviceIds.First()))
+        if (devices.Any() && InstallationDeviceEntity.IsInstallationDeviceId(devices.First().Key))
         {
-            var entities = deviceIds.Select(e => new InstallationDeviceEntity(e));
+            var entities = devices.Select(e => new InstallationDeviceEntity(e.Key));
             await _installationDeviceRepository.UpsertManyAsync(entities.ToList());
         }
     }
 
-    private async Task PatchTagsForUserDevicesAsync(IEnumerable<string> deviceIds, UpdateOperationType op,
+    private async Task PatchTagsForUserDevicesAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, UpdateOperationType op,
         string tag)
     {
-        if (!deviceIds.Any())
+        if (!devices.Any())
         {
             return;
         }
@@ -179,11 +198,11 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
             operation.Path += $"/{tag}";
         }
 
-        foreach (var id in deviceIds)
+        foreach (var device in devices)
         {
             try
             {
-                await _client.PatchInstallationAsync(id, new List<PartialUpdateOperation> { operation });
+                await GetClient(device.Value).PatchInstallationAsync(device.Key, new List<PartialUpdateOperation> { operation });
             }
             catch (Exception e) when (e.InnerException == null || !e.InnerException.Message.Contains("(404) Not Found"))
             {
@@ -191,4 +210,54 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
             }
         }
     }
+
+    private NotificationHubClient GetClient(DeviceType deviceType)
+    {
+        var hubType = NotificationHubType.General;
+        switch (deviceType)
+        {
+            case DeviceType.Android:
+                hubType = NotificationHubType.Android;
+                break;
+            case DeviceType.iOS:
+                hubType = NotificationHubType.iOS;
+                break;
+            case DeviceType.ChromeExtension:
+            case DeviceType.FirefoxExtension:
+            case DeviceType.OperaExtension:
+            case DeviceType.EdgeExtension:
+            case DeviceType.VivaldiExtension:
+            case DeviceType.SafariExtension:
+                hubType = NotificationHubType.GeneralBrowserExtension;
+                break;
+            case DeviceType.WindowsDesktop:
+            case DeviceType.MacOsDesktop:
+            case DeviceType.LinuxDesktop:
+                hubType = NotificationHubType.GeneralDesktop;
+                break;
+            case DeviceType.ChromeBrowser:
+            case DeviceType.FirefoxBrowser:
+            case DeviceType.OperaBrowser:
+            case DeviceType.EdgeBrowser:
+            case DeviceType.IEBrowser:
+            case DeviceType.UnknownBrowser:
+            case DeviceType.SafariBrowser:
+            case DeviceType.VivaldiBrowser:
+                hubType = NotificationHubType.GeneralWeb;
+                break;
+            default:
+                break;
+        }
+
+        if (!_clients.ContainsKey(hubType))
+        {
+            _logger.LogWarning("No hub client for '{0}'. Using general hub instead.", hubType);
+            hubType = NotificationHubType.General;
+            if (!_clients.ContainsKey(hubType))
+            {
+                throw new Exception("No general hub client found.");
+            }
+        }
+        return _clients[hubType];
+    }
 }

src/Core/Services/Implementations/RelayPushRegistrationService.cs

@@ -38,30 +38,37 @@ public class RelayPushRegistrationService : BaseIdentityClientService, IPushRegi
         await SendAsync(HttpMethod.Post, "push/register", requestModel);
     }
 
-    public async Task DeleteRegistrationAsync(string deviceId)
+    public async Task DeleteRegistrationAsync(string deviceId, DeviceType type)
     {
-        await SendAsync(HttpMethod.Delete, string.Concat("push/", deviceId));
+        var requestModel = new PushDeviceRequestModel
+        {
+            Id = deviceId,
+            Type = type,
+        };
+        await SendAsync(HttpMethod.Post, "push/delete", requestModel);
     }
 
-    public async Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public async Task AddUserRegistrationOrganizationAsync(
+        IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
-        if (!deviceIds.Any())
+        if (!devices.Any())
         {
             return;
         }
 
-        var requestModel = new PushUpdateRequestModel(deviceIds, organizationId);
+        var requestModel = new PushUpdateRequestModel(devices, organizationId);
         await SendAsync(HttpMethod.Put, "push/add-organization", requestModel);
     }
 
-    public async Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public async Task DeleteUserRegistrationOrganizationAsync(
+        IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
-        if (!deviceIds.Any())
+        if (!devices.Any())
         {
             return;
         }
 
-        var requestModel = new PushUpdateRequestModel(deviceIds, organizationId);
+        var requestModel = new PushUpdateRequestModel(devices, organizationId);
         await SendAsync(HttpMethod.Put, "push/delete-organization", requestModel);
     }
 }

src/Core/Services/NoopImplementations/NoopPushRegistrationService.cs

@@ -4,7 +4,7 @@ namespace Bit.Core.Services;
 
 public class NoopPushRegistrationService : IPushRegistrationService
 {
-    public Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public Task AddUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
         return Task.FromResult(0);
     }
@@ -15,12 +15,12 @@ public class NoopPushRegistrationService : IPushRegistrationService
         return Task.FromResult(0);
     }
 
-    public Task DeleteRegistrationAsync(string deviceId)
+    public Task DeleteRegistrationAsync(string deviceId, DeviceType deviceType)
     {
         return Task.FromResult(0);
     }
 
-    public Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
+    public Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
     {
         return Task.FromResult(0);
     }