// Copyright Epic Games, Inc. All Rights Reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Security.Claims;
using System.Threading;
using System.Threading.Tasks;
using EpicGames.Horde.Jobs;
using EpicGames.Horde.Jobs.Bisect;
using EpicGames.Horde.Users;
using HordeServer.Server;
using HordeServer.Utilities;
using MongoDB.Bson;
using MongoDB.Bson.Serialization;
using MongoDB.Bson.Serialization.Attributes;
using MongoDB.Driver;
namespace HordeServer.Users
{
///
/// Manages user documents
///
class UserCollectionV1 : IUserCollection
{
class UserDocument : IUser, IUserClaims, IUserSettings
{
public UserId Id { get; set; }
public ClaimDocument PrimaryClaim { get; set; } = null!;
public List Claims { get; set; } = new List();
[BsonDefaultValue(false), BsonIgnoreIfDefault]
public bool EnableExperimentalFeatures { get; set; }
[BsonDefaultValue(false), BsonIgnoreIfDefault]
public bool AlwaysTagPreflightCL { get; set; }
public BsonValue DashboardSettings { get; set; } = BsonNull.Value;
public List PinnedJobIds { get; set; } = new List();
public List PinnedBisectTaskIds { get; set; } = new List();
string IUser.Name => Claims.FirstOrDefault(x => String.Equals(x.Type, "name", StringComparison.Ordinal))?.Value ?? PrimaryClaim.Value;
string IUser.Login => Claims.FirstOrDefault(x => String.Equals(x.Type, ClaimTypes.Name, StringComparison.Ordinal))?.Value ?? PrimaryClaim.Value;
string? IUser.Email => Claims.FirstOrDefault(x => String.Equals(x.Type, ClaimTypes.Email, StringComparison.Ordinal))?.Value;
UserId IUserClaims.UserId => Id;
IReadOnlyList IUserClaims.Claims => Claims;
UserId IUserSettings.UserId => Id;
IReadOnlyList IUserSettings.PinnedJobIds => PinnedJobIds;
IReadOnlyList IUserSettings.PinnedBisectTaskIds => PinnedBisectTaskIds;
IReadOnlyList? IUserSettings.JobTemplateSettings => null;
}
class ClaimDocument : IUserClaim
{
public string Type { get; set; }
public string Value { get; set; }
private ClaimDocument()
{
Type = null!;
Value = null!;
}
public ClaimDocument(string type, string value)
{
Type = type;
Value = value;
}
public ClaimDocument(IUserClaim other)
{
Type = other.Type;
Value = other.Value;
}
}
readonly IMongoCollection _users;
///
/// Static constructor
///
static UserCollectionV1()
{
BsonSerializer.RegisterDiscriminatorConvention(typeof(UserDocument), NullDiscriminatorConvention.Instance);
BsonSerializer.RegisterDiscriminatorConvention(typeof(ClaimDocument), NullDiscriminatorConvention.Instance);
}
///
/// Constructor
///
///
public UserCollectionV1(IMongoService mongoService)
{
_users = mongoService.GetCollection("Users", keys => keys.Ascending(x => x.PrimaryClaim), unique: true);
}
///
public async Task GetUserAsync(UserId id, CancellationToken cancellationToken)
{
return await _users.Find(x => x.Id == id).FirstOrDefaultAsync(cancellationToken);
}
///
public async ValueTask GetCachedUserAsync(UserId? id, CancellationToken cancellationToken)
{
if (id == null)
{
return null;
}
else
{
return await GetUserAsync(id.Value, cancellationToken);
}
}
///
public async Task> FindUsersAsync(IEnumerable? ids, string? nameRegex, int? index, int? count, CancellationToken cancellationToken)
{
FilterDefinition filter = Builders.Filter.In(x => x.Id, ids);
return await _users.Find(filter).Range(index, count).ToListAsync(cancellationToken);
}
///
public async Task FindUserByLoginAsync(string login, CancellationToken cancellationToken)
{
ClaimDocument primaryClaim = new ClaimDocument(HordeClaimTypes.User, login);
return await _users.Find(x => x.PrimaryClaim == primaryClaim).FirstOrDefaultAsync(cancellationToken);
}
///
public Task FindUserByEmailAsync(string email, CancellationToken cancellationToken)
{
return Task.FromResult(null);
}
///
public async Task FindOrAddUserByLoginAsync(string login, string? name, string? email, CancellationToken cancellationToken)
{
ClaimDocument newPrimaryClaim = new ClaimDocument(ClaimTypes.Name, login);
UpdateDefinition update = Builders.Update.SetOnInsert(x => x.Id, new UserId(BinaryIdUtils.CreateNew()));
return await _users.FindOneAndUpdateAsync(x => x.PrimaryClaim == newPrimaryClaim, update, new FindOneAndUpdateOptions { IsUpsert = true, ReturnDocument = ReturnDocument.After }, cancellationToken);
}
///
public async Task GetClaimsAsync(UserId userId, CancellationToken cancellationToken)
{
return await _users.Find(x => x.Id == userId).FirstOrDefaultAsync(cancellationToken) ?? new UserDocument { Id = userId };
}
///
public async Task UpdateClaimsAsync(UserId userId, IEnumerable claims, CancellationToken cancellationToken)
{
List newClaims = claims.Select(x => new ClaimDocument(x)).ToList();
await _users.FindOneAndUpdateAsync(x => x.Id == userId, Builders.Update.Set(x => x.Claims, newClaims), cancellationToken: cancellationToken);
}
///
public async Task GetSettingsAsync(UserId userId, CancellationToken cancellationToken)
{
return await _users.Find(x => x.Id == userId).FirstOrDefaultAsync(cancellationToken) ?? new UserDocument { Id = userId };
}
///
public async Task UpdateSettingsAsync(UserId userId, bool? enableExperimentalFeatures, bool? alwaysTagPreflightCL = null, BsonValue? dashboardSettings = null, IEnumerable? addPinnedJobIds = null, IEnumerable? removePinnedJobIds = null, UpdateUserJobTemplateOptions? templateOptions = null, IEnumerable? addBisectTaskIds = null, IEnumerable? removeBisectTaskIds = null, CancellationToken cancellationToken = default)
{
if (addPinnedJobIds != null)
{
foreach (JobId pinnedJobId in addPinnedJobIds)
{
FilterDefinition filter = Builders.Filter.Eq(x => x.Id, userId) & Builders.Filter.AnyNin(x => x.PinnedJobIds, new[] { pinnedJobId });
UpdateDefinition update = Builders.Update.PushEach(x => x.PinnedJobIds, new[] { pinnedJobId }, 50);
await _users.UpdateOneAsync(filter, update, null, cancellationToken);
}
}
List> updates = new List>();
if (enableExperimentalFeatures != null)
{
updates.Add(Builders.Update.SetOrUnsetNull(x => x.EnableExperimentalFeatures, enableExperimentalFeatures));
}
if (alwaysTagPreflightCL != null)
{
updates.Add(Builders.Update.SetOrUnsetNull(x => x.AlwaysTagPreflightCL, alwaysTagPreflightCL));
}
if (dashboardSettings != null)
{
updates.Add(Builders.Update.Set(x => x.DashboardSettings, dashboardSettings));
}
if (removePinnedJobIds != null && removePinnedJobIds.Any())
{
updates.Add(Builders.Update.PullAll(x => x.PinnedJobIds, removePinnedJobIds));
}
if (updates.Count > 0)
{
await _users.UpdateOneAsync(x => x.Id == userId, Builders.Update.Combine(updates), (UpdateOptions?)null, cancellationToken);
}
}
///
/// Enumerate all the documents in this collection
///
///
public async IAsyncEnumerable<(IUser, IUserClaims, IUserSettings)> EnumerateDocumentsAsync([EnumeratorCancellation] CancellationToken cancellationToken)
{
using (IAsyncCursor cursor = await _users.Find(FilterDefinition.Empty).ToCursorAsync(cancellationToken))
{
while (await cursor.MoveNextAsync(cancellationToken))
{
foreach (UserDocument document in cursor.Current)
{
if (document.Claims.Count > 0)
{
yield return (document, document, document);
}
}
}
}
}
}
}