1
0
mirror of https://github.com/bitwarden/server synced 2026-01-09 20:13:24 +00:00

[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>
This commit is contained in:
Kyle Spearrin
2024-04-08 15:39:44 -04:00
committed by GitHub
parent de8b7b14b8
commit 40221f578f
14 changed files with 208 additions and 70 deletions

View File

@@ -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];
}
}