Shard notification hub

src/Api/Controllers/PushController.cs

@@ -46,15 +46,15 @@ await _pushRegistrationService.CreateOrUpdateRegistrationAsync(model.PushToken,
     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")]
     public async Task PutAddOrganization([FromBody] PushUpdateRequestModel model)
     {
         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 async Task PutDeleteOrganization([FromBody] PushUpdateRequestModel model)
     {
         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/AdminConsole/OrganizationFeatures/OrganizationUsers/RemoveOrganizationUserCommand.cs

@@ -162,12 +162,12 @@ private async Task RepositoryDeleteUserAsync(OrganizationUser orgUser, Guid? del
         }
     }
 
-    private async Task<IEnumerable<KeyValuePair<string, DeviceType>>> GetUserDeviceIdsAsync(Guid userId)
+    private async Task<IEnumerable<string>> GetUserDeviceIdsAsync(Guid userId)
     {
         var devices = await _deviceRepository.GetManyByUserIdAsync(userId);
         return devices
             .Where(d => !string.IsNullOrWhiteSpace(d.PushToken))
-            .Select(d => new KeyValuePair<string, DeviceType>(d.Id.ToString(), d.Type));
+            .Select(d => d.Id.ToString());
     }
 
     private async Task DeleteAndPushUserRegistrationAsync(Guid organizationId, Guid userId)

src/Core/AdminConsole/Services/Implementations/OrganizationService.cs

@@ -1832,12 +1832,12 @@ await _pushRegistrationService.DeleteUserRegistrationOrganizationAsync(devices,
     }
 
 
-    private async Task<IEnumerable<KeyValuePair<string, DeviceType>>> GetUserDeviceIdsAsync(Guid userId)
+    private async Task<IEnumerable<string>> GetUserDeviceIdsAsync(Guid userId)
     {
         var devices = await _deviceRepository.GetManyByUserIdAsync(userId);
         return devices
             .Where(d => !string.IsNullOrWhiteSpace(d.PushToken))
-            .Select(d => new KeyValuePair<string, DeviceType>(d.Id.ToString(), d.Type));
+            .Select(d => d.Id.ToString());
     }
 
     public async Task ReplaceAndUpdateCacheAsync(Organization org, EventType? orgEvent = null)

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

@@ -1,12 +1,9 @@
 using System.ComponentModel.DataAnnotations;
-using Bit.Core.Enums;
 
 namespace Bit.Core.Models.Api;
 
 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,25 @@ public static bool IsInstallationDeviceId(string deviceId)
     {
         return deviceId != null && deviceId.Length == 73 && deviceId[36] == '_';
     }
+    public static bool TryParse(string deviceId, out InstallationDeviceEntity installationDeviceEntity)
+    {
+        installationDeviceEntity = null;
+        var installationId = Guid.Empty;
+        var 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;
+        }
+        installationDeviceEntity = new InstallationDeviceEntity(installationId, deviceIdGuid);
+        return true;
+    }
 }

src/Core/NotificationHub/INotificationHubClientProxy.cs

@@ -0,0 +1,8 @@
+using Microsoft.Azure.NotificationHubs;
+
+namespace Bit.Core.NotificationHub;
+
+public interface INotificationHubProxy
+{
+    Task<(INotificationHubClient Client, NotificationOutcome Outcome)[]> SendTemplateNotificationAsync(IDictionary<string, string> properties, string tagExpression);
+}

src/Core/NotificationHub/INotificationHubPool.cs

@@ -0,0 +1,9 @@
+using Microsoft.Azure.NotificationHubs;
+
+namespace Bit.Core.NotificationHub;
+
+public interface INotificationHubPool
+{
+    NotificationHubClient ClientFor(Guid comb);
+    INotificationHubProxy AllClients { get; }
+}

src/Core/NotificationHub/NotificationHubClientProxy.cs

@@ -0,0 +1,26 @@
+using Microsoft.Azure.NotificationHubs;
+
+namespace Bit.Core.NotificationHub;
+
+public class NotificationHubClientProxy : INotificationHubProxy
+{
+    private readonly IEnumerable<INotificationHubClient> _clients;
+
+    public NotificationHubClientProxy(IEnumerable<INotificationHubClient> clients)
+    {
+        _clients = clients;
+    }
+
+    private async Task<(INotificationHubClient, T)[]> ApplyToAllClientsAsync<T>(Func<INotificationHubClient, Task<T>> action)
+    {
+        var tasks = _clients.Select(async c => (c, await action(c)));
+        return await Task.WhenAll(tasks);
+    }
+
+    // partial proxy of INotificationHubClient implementation
+    // Note: Any other methods that are needed can simply be delegated as done here.
+    public async Task<(INotificationHubClient Client, NotificationOutcome Outcome)[]> SendTemplateNotificationAsync(IDictionary<string, string> properties, string tagExpression)
+    {
+        return await ApplyToAllClientsAsync(async c => await c.SendTemplateNotificationAsync(properties, tagExpression));
+    }
+}

src/Core/NotificationHub/NotificationHubConnection.cs

@@ -0,0 +1,128 @@
+using Bit.Core.Settings;
+using Bit.Core.Utilities;
+using Microsoft.Azure.NotificationHubs;
+
+class NotificationHubConnection
+{
+    public string HubName { get; init; }
+    public string ConnectionString { get; init; }
+    public bool EnableSendTracing { get; init; }
+    private NotificationHubClient _hubClient;
+    /// <summary>
+    /// Gets the NotificationHubClient for this connection.
+    /// 
+    /// If the client is null, it will be initialized.
+    /// 
+    /// <throws>Exception</throws> if the connection is invalid.
+    /// </summary>
+    public NotificationHubClient HubClient
+    {
+        get
+        {
+            if (_hubClient == null)
+            {
+                if (!IsValid)
+                {
+                    throw new Exception("Invalid notification hub settings");
+                }
+                Init();
+            }
+            return _hubClient;
+        }
+        private set
+        {
+            _hubClient = value;
+        }
+    }
+    /// <summary>
+    /// Gets the start date for registration.
+    /// 
+    /// If null, registration is always disabled.
+    /// </summary>
+    public DateTime? RegistrationStartDate { get; init; }
+    /// <summary>
+    /// Gets the end date for registration.
+    /// 
+    /// If null, registration has no end date.
+    /// </summary>
+    public DateTime? RegistrationEndDate { get; init; }
+    /// <summary>
+    /// Gets whether all data needed to generate a connection to Notification Hub is present.
+    /// </summary>
+    public bool IsValid
+    {
+        get
+        {
+            {
+                var invalid = string.IsNullOrWhiteSpace(HubName) || string.IsNullOrWhiteSpace(ConnectionString);
+                return !invalid;
+            }
+        }
+    }
+
+    public string LogString
+    {
+        get
+        {
+            return $"HubName: {HubName}, EnableSendTracing: {EnableSendTracing}, RegistrationStartDate: {RegistrationStartDate}, RegistrationEndDate: {RegistrationEndDate}";
+        }
+    }
+
+    /// <summary>
+    /// Gets whether registration is enabled for the given comb ID.
+    /// This is based off of the generation time encoded in the comb ID.
+    /// </summary>
+    /// <param name="comb"></param>
+    /// <returns></returns>
+    public bool RegistrationEnabled(Guid comb)
+    {
+        var combTime = CoreHelpers.DateFromComb(comb);
+        return RegistrationEnabled(combTime);
+    }
+
+    /// <summary>
+    /// Gets whether registration is enabled for the given time.
+    /// </summary>
+    /// <param name="queryTime">The time to check</param>
+    /// <returns></returns>
+    public bool RegistrationEnabled(DateTime queryTime)
+    {
+        if (queryTime >= RegistrationEndDate || RegistrationStartDate == null)
+        {
+            return false;
+        }
+
+        return RegistrationStartDate < queryTime;
+    }
+
+    private NotificationHubConnection() { }
+
+    /// <summary>
+    /// Creates a new NotificationHubConnection from the given settings.
+    /// </summary>
+    /// <param name="settings"></param>
+    /// <returns></returns>
+    public static NotificationHubConnection From(GlobalSettings.NotificationHubSettings settings)
+    {
+        return new()
+        {
+            HubName = settings.HubName,
+            ConnectionString = settings.ConnectionString,
+            EnableSendTracing = settings.EnableSendTracing,
+            // Comb time is not precise enough for millisecond accuracy
+            RegistrationStartDate = settings.RegistrationStartDate.HasValue ? Truncate(settings.RegistrationStartDate.Value, TimeSpan.FromMilliseconds(10)) : null,
+            RegistrationEndDate = settings.RegistrationEndDate
+        };
+    }
+
+    private NotificationHubConnection Init()
+    {
+        HubClient = NotificationHubClient.CreateClientFromConnectionString(ConnectionString, HubName, EnableSendTracing);
+        return this;
+    }
+
+    private static DateTime Truncate(DateTime dateTime, TimeSpan resolution)
+    {
+        return dateTime.AddTicks(-(dateTime.Ticks % resolution.Ticks));
+    }
+}

src/Core/NotificationHub/NotificationHubPool.cs

@@ -0,0 +1,62 @@
+using Bit.Core.Settings;
+using Bit.Core.Utilities;
+using Microsoft.Azure.NotificationHubs;
+using Microsoft.Extensions.Logging;
+
+namespace Bit.Core.NotificationHub;
+
+public class NotificationHubPool : INotificationHubPool
+{
+    private List<NotificationHubConnection> _connections { get; }
+    private readonly IEnumerable<INotificationHubClient> _clients;
+    private readonly ILogger<NotificationHubPool> _logger;
+    public NotificationHubPool(ILogger<NotificationHubPool> logger, GlobalSettings globalSettings)
+    {
+        _logger = logger;
+        _connections = FilterInvalidHubs(globalSettings.NotificationHubPool.NotificationHubs);
+        _clients = _connections.GroupBy(c => c.ConnectionString).Select(g => g.First().HubClient);
+    }
+
+    private List<NotificationHubConnection> FilterInvalidHubs(IEnumerable<GlobalSettings.NotificationHubSettings> hubs)
+    {
+        List<NotificationHubConnection> result = new();
+        _logger.LogDebug("Filtering {HubCount} notification hubs", hubs.Count());
+        foreach (var hub in hubs)
+        {
+            var connection = NotificationHubConnection.From(hub);
+            if (!connection.IsValid)
+            {
+                _logger.LogWarning("Invalid notification hub settings: {HubName}", hub.HubName ?? "hub name missing");
+                continue;
+            }
+            _logger.LogDebug("Adding notification hub: {ConnectionLogString}", connection.LogString);
+            result.Add(connection);
+        }
+
+        return result;
+    }
+
+
+    /// <summary>
+    /// Gets the NotificationHubClient for the given comb ID.
+    /// </summary>
+    /// <param name="comb"></param>
+    /// <returns></returns>
+    /// <exception cref="InvalidOperationException">Thrown when no notification hub is found for a given comb.</exception>
+    public NotificationHubClient ClientFor(Guid comb)
+    {
+        var possibleConnections = _connections.Where(c => c.RegistrationEnabled(comb)).ToArray();
+        if (possibleConnections.Length == 0)
+        {
+            throw new InvalidOperationException($"No valid notification hubs are available for the given comb ({comb}).\n" +
+                $"The comb's datetime is {CoreHelpers.DateFromComb(comb)}." +
+                $"Hub start and end times are configured as follows:\n" +
+                string.Join("\n", _connections.Select(c => $"Hub {c.HubName} - Start: {c.RegistrationStartDate}, End: {c.RegistrationEndDate}")));
+        }
+        var resolvedConnection = possibleConnections[CoreHelpers.BinForComb(comb, possibleConnections.Length)];
+        _logger.LogTrace("Resolved notification hub for comb {Comb} out of {HubCount} hubs.\n{ConnectionInfo}", comb, possibleConnections.Length, resolvedConnection.LogString);
+        return resolvedConnection.HubClient;
+    }
+
+    public INotificationHubProxy AllClients { get { return new NotificationHubClientProxy(_clients); } }
+}