Replace device type sharding with comb + range sharding

src/Api/Controllers/PushController.cs

@@ -46,7 +46,7 @@ public class PushController : Controller
     public async Task PostDelete([FromBody] PushDeviceRequestModel model)
     {
         CheckUsage();
-        await _pushRegistrationService.DeleteRegistrationAsync(Prefix(model.Id), model.Type);
+        await _pushRegistrationService.DeleteRegistrationAsync(Prefix(model.Id));
     }
 
     [HttpPut("add-organization")]
@@ -54,7 +54,7 @@ public class PushController : Controller
     {
         CheckUsage();
         await _pushRegistrationService.AddUserRegistrationOrganizationAsync(
-            model.Devices.Select(d => new KeyValuePair<string, Core.Enums.DeviceType>(Prefix(d.Id), d.Type)),
+            model.Devices.Select(d => Prefix(d.Id)),
             Prefix(model.OrganizationId));
     }
 
@@ -63,7 +63,7 @@ public class PushController : Controller
     {
         CheckUsage();
         await _pushRegistrationService.DeleteUserRegistrationOrganizationAsync(
-            model.Devices.Select(d => new KeyValuePair<string, Core.Enums.DeviceType>(Prefix(d.Id), d.Type)),
+            model.Devices.Select(d => Prefix(d.Id)),
             Prefix(model.OrganizationId));
     }
 

src/Core/Models/Api/Request/PushDeviceRequestModel.cs

@@ -7,6 +7,5 @@ public class PushDeviceRequestModel
 {
     [Required]
     public string Id { get; set; }
-    [Required]
     public DeviceType Type { get; set; }
 }

src/Core/Models/Api/Request/PushUpdateRequestModel.cs

@@ -1,5 +1,4 @@
 using System.ComponentModel.DataAnnotations;
-using Bit.Core.Enums;
 
 namespace Bit.Core.Models.Api;
 
@@ -8,9 +7,9 @@ public class PushUpdateRequestModel
     public PushUpdateRequestModel()
     { }
 
-    public PushUpdateRequestModel(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+    public PushUpdateRequestModel(IEnumerable<string> deviceIds, string organizationId)
     {
-        Devices = devices.Select(d => new PushDeviceRequestModel { Id = d.Key, Type = d.Value });
+        Devices = deviceIds.Select(d => new PushDeviceRequestModel { Id = d });
         OrganizationId = organizationId;
     }
 

src/Core/Models/Data/InstallationDeviceEntity.cs

@@ -37,4 +37,23 @@ public class InstallationDeviceEntity : ITableEntity
     {
         return deviceId != null && deviceId.Length == 73 && deviceId[36] == '_';
     }
+    public static bool TrySplit(string deviceId, out Guid installationId, out Guid deviceIdGuid)
+    {
+        installationId = Guid.Empty;
+        deviceIdGuid = Guid.Empty;
+        if (!IsInstallationDeviceId(deviceId))
+        {
+            return false;
+        }
+        var parts = deviceId.Split("_");
+        if (parts.Length < 2)
+        {
+            return false;
+        }
+        if (!Guid.TryParse(parts[0], out installationId) || !Guid.TryParse(parts[1], out deviceIdGuid))
+        {
+            return false;
+        }
+        return true;
+    }
 }

src/Core/NotificationHub/INotificationHubPool.cs

@@ -5,5 +5,5 @@ namespace Bit.Core.NotificationHub;
 public interface INotificationHubPool
 {
     NotificationHubClient ClientFor(Guid comb);
-
+    INotificationHubProxy AllClients { get; }
 }

src/Core/NotificationHub/NotificationHubPushRegistrationService.cs

@@ -16,7 +16,6 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
     private readonly INotificationHubPool _notificationHubPool;
     private readonly IServiceProvider _serviceProvider;
     private readonly ILogger<NotificationHubPushRegistrationService> _logger;
-    private Dictionary<NotificationHubType, NotificationHubClient> _clients = [];
 
     public NotificationHubPushRegistrationService(
         IInstallationDeviceRepository installationDeviceRepository,
@@ -30,25 +29,6 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         _notificationHubPool = notificationHubPool;
         _serviceProvider = serviceProvider;
         _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,
@@ -121,7 +101,7 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         BuildInstallationTemplate(installation, "badgeMessage", badgeMessageTemplate ?? messageTemplate,
             userId, identifier);
 
-        await GetClient(type).CreateOrUpdateInstallationAsync(installation);
+        await ClientFor(GetComb(deviceId)).CreateOrUpdateInstallationAsync(installation);
         if (InstallationDeviceEntity.IsInstallationDeviceId(deviceId))
         {
             await _installationDeviceRepository.UpsertAsync(new InstallationDeviceEntity(deviceId));
@@ -156,11 +136,11 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         installation.Templates.Add(fullTemplateId, template);
     }
 
-    public async Task DeleteRegistrationAsync(string deviceId, DeviceType deviceType)
+    public async Task DeleteRegistrationAsync(string deviceId)
     {
         try
         {
-            await GetClient(deviceType).DeleteInstallationAsync(deviceId);
+            await ClientFor(GetComb(deviceId)).DeleteInstallationAsync(deviceId);
             if (InstallationDeviceEntity.IsInstallationDeviceId(deviceId))
             {
                 await _installationDeviceRepository.DeleteAsync(new InstallationDeviceEntity(deviceId));
@@ -172,31 +152,31 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         }
     }
 
-    public async Task AddUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+    public async Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
     {
-        await PatchTagsForUserDevicesAsync(devices, UpdateOperationType.Add, $"organizationId:{organizationId}");
-        if (devices.Any() && InstallationDeviceEntity.IsInstallationDeviceId(devices.First().Key))
+        await PatchTagsForUserDevicesAsync(deviceIds, UpdateOperationType.Add, $"organizationId:{organizationId}");
+        if (deviceIds.Any() && InstallationDeviceEntity.IsInstallationDeviceId(deviceIds.First()))
         {
-            var entities = devices.Select(e => new InstallationDeviceEntity(e.Key));
+            var entities = deviceIds.Select(e => new InstallationDeviceEntity(e));
             await _installationDeviceRepository.UpsertManyAsync(entities.ToList());
         }
     }
 
-    public async Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+    public async Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
     {
-        await PatchTagsForUserDevicesAsync(devices, UpdateOperationType.Remove,
+        await PatchTagsForUserDevicesAsync(deviceIds, UpdateOperationType.Remove,
             $"organizationId:{organizationId}");
-        if (devices.Any() && InstallationDeviceEntity.IsInstallationDeviceId(devices.First().Key))
+        if (deviceIds.Any() && InstallationDeviceEntity.IsInstallationDeviceId(deviceIds.First()))
         {
-            var entities = devices.Select(e => new InstallationDeviceEntity(e.Key));
+            var entities = deviceIds.Select(e => new InstallationDeviceEntity(e));
             await _installationDeviceRepository.UpsertManyAsync(entities.ToList());
         }
     }
 
-    private async Task PatchTagsForUserDevicesAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, UpdateOperationType op,
+    private async Task PatchTagsForUserDevicesAsync(IEnumerable<string> deviceIds, UpdateOperationType op,
         string tag)
     {
-        if (!devices.Any())
+        if (!deviceIds.Any())
         {
             return;
         }
@@ -216,11 +196,11 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
             operation.Path += $"/{tag}";
         }
 
-        foreach (var device in devices)
+        foreach (var deviceId in deviceIds)
         {
             try
             {
-                await GetClient(device.Value).PatchInstallationAsync(device.Key, new List<PartialUpdateOperation> { operation });
+                await ClientFor(GetComb(deviceId)).PatchInstallationAsync(deviceId, new List<PartialUpdateOperation> { operation });
             }
             catch (Exception e) when (e.InnerException == null || !e.InnerException.Message.Contains("(404) Not Found"))
             {
@@ -229,53 +209,24 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         }
     }
 
-    private NotificationHubClient GetClient(DeviceType deviceType)
-    {
-        var hubType = NotificationHubType.General;
-        switch (deviceType)
+    private NotificationHubClient ClientFor(Guid deviceId)
     {
-            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;
+        return _notificationHubPool.ClientFor(deviceId);
     }
 
-        if (!_clients.ContainsKey(hubType))
+    private Guid GetComb(string deviceId)
+    {
+        Guid deviceIdGuid;
+        if (InstallationDeviceEntity.TrySplit(deviceId, out _, out deviceIdGuid))
         {
-            _logger.LogWarning("No hub client for '{0}'. Using general hub instead.", hubType);
-            hubType = NotificationHubType.General;
-            if (!_clients.ContainsKey(hubType))
+        }
+        else if (Guid.TryParse(deviceId, out deviceIdGuid))
         {
-                throw new Exception("No general hub client found.");
         }
+        else
+        {
+            throw new Exception($"Invalid device id {deviceId}.");
         }
-        return _clients[hubType];
+        return deviceIdGuid;
     }
 }

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, DeviceType type);
-    Task AddUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId);
-    Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId);
+    Task DeleteRegistrationAsync(string deviceId);
+    Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId);
+    Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId);
 }

src/Core/Services/Implementations/RelayPushRegistrationService.cs

@@ -38,37 +38,36 @@ public class RelayPushRegistrationService : BaseIdentityClientService, IPushRegi
         await SendAsync(HttpMethod.Post, "push/register", requestModel);
     }
 
-    public async Task DeleteRegistrationAsync(string deviceId, DeviceType type)
+    public async Task DeleteRegistrationAsync(string deviceId)
     {
         var requestModel = new PushDeviceRequestModel
         {
             Id = deviceId,
-            Type = type,
         };
         await SendAsync(HttpMethod.Post, "push/delete", requestModel);
     }
 
     public async Task AddUserRegistrationOrganizationAsync(
-        IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+        IEnumerable<string> deviceIds, string organizationId)
     {
-        if (!devices.Any())
+        if (!deviceIds.Any())
         {
             return;
         }
 
-        var requestModel = new PushUpdateRequestModel(devices, organizationId);
+        var requestModel = new PushUpdateRequestModel(deviceIds, organizationId);
         await SendAsync(HttpMethod.Put, "push/add-organization", requestModel);
     }
 
     public async Task DeleteUserRegistrationOrganizationAsync(
-        IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+        IEnumerable<string> deviceIds, string organizationId)
     {
-        if (!devices.Any())
+        if (!deviceIds.Any())
         {
             return;
         }
 
-        var requestModel = new PushUpdateRequestModel(devices, organizationId);
+        var requestModel = new PushUpdateRequestModel(deviceIds, 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<KeyValuePair<string, DeviceType>> devices, string organizationId)
+    public Task AddUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
     {
         return Task.FromResult(0);
     }
@@ -15,12 +15,12 @@ public class NoopPushRegistrationService : IPushRegistrationService
         return Task.FromResult(0);
     }
 
-    public Task DeleteRegistrationAsync(string deviceId, DeviceType deviceType)
+    public Task DeleteRegistrationAsync(string deviceId)
     {
         return Task.FromResult(0);
     }
 
-    public Task DeleteUserRegistrationOrganizationAsync(IEnumerable<KeyValuePair<string, DeviceType>> devices, string organizationId)
+    public Task DeleteUserRegistrationOrganizationAsync(IEnumerable<string> deviceIds, string organizationId)
     {
         return Task.FromResult(0);
     }