Upstream/bitwarden-server
1
0
Fork 0
mirror of https://github.com/bitwarden/server.git synced 2024-11-24 12:35:25 +01:00
Code Issues Projects Releases Wiki Activity

Replace device type sharding with comb + range sharding

Browse Source
This commit is contained in:
Matt Gibson 2024-07-02 10:53:31 -07:00
parent aa8372052e
commit abd67e8ec6
No known key found for this signature in database
GPG Key ID: 963EE038B0581878
9 changed files with 68 additions and 101 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/Api/Controllers/PushController.cs
View File

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

1
src/Core/Models/Api/Request/PushDeviceRequestModel.cs
View File

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

5
src/Core/Models/Api/Request/PushUpdateRequestModel.cs
View File

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

19
src/Core/Models/Data/InstallationDeviceEntity.cs
View File

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

2
src/Core/NotificationHub/INotificationHubPool.cs
View File

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

109
src/Core/NotificationHub/NotificationHubPushRegistrationService.cs
View File

@ -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)
private NotificationHubClient ClientFor(Guid deviceId)
{
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;
}
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))
{
throw new Exception("No general hub client found.");
}
}
return _clients[hubType];
else if (Guid.TryParse(deviceId, out deviceIdGuid))
{
}
else
{
throw new Exception($"Invalid device id {deviceId}.");
}
return deviceIdGuid;
}
}

6
src/Core/Services/IPushRegistrationService.cs
View File

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

15
src/Core/Services/Implementations/RelayPushRegistrationService.cs
View File

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

6
src/Core/Services/NoopImplementations/NoopPushRegistrationService.cs
View File

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