1  using System.Collections.Concurrent;
  2  using LiteDB;
  3  using Microsoft.AspNetCore.Identity;
  4  
  5  namespace GUNRPG.Infrastructure.Identity;
  6  
  7  /// <summary>
  8  /// Ensures authenticated users have a stable account identifier for operator ownership isolation.
  9  /// </summary>
 10  public static class AccountIdProvisioning
 11  {
 12      private static readonly ConcurrentDictionary<string, UserLock> UserLocks = new(StringComparer.Ordinal);
 13  
 14      public static async Task<IdentityResult> EnsureAssignedAsync(
 15          UserManager<ApplicationUser> userManager,
 16          ApplicationUser user,
 17          CancellationToken ct = default)
 18      {
 19          if (user.AccountId is { } accountId && accountId != Guid.Empty)
 20              return IdentityResult.Success;
 21  
 22          using (await AcquireUserLockAsync(user.Id, ct))
 23          {
 24              var storedUser = await userManager.FindByIdAsync(user.Id);
 25              if (storedUser is null)
 26              {
 27                  return IdentityResult.Failed(new IdentityError
 28                  {
 29                      Code = "UserNotFound",
 30                      Description = "The user no longer exists in the identity store.",
 31                  });
 32              }
 33  
 34              if (storedUser.AccountId is { } storedAccountId && storedAccountId != Guid.Empty)
 35              {
 36                  user.AccountId = storedAccountId;
 37                  return IdentityResult.Success;
 38              }
 39  
 40              storedUser.AccountId = Guid.NewGuid();
 41              var result = await userManager.UpdateAsync(storedUser);
 42  
 43              if (result.Succeeded && storedUser.AccountId is { } newAccountId && newAccountId != Guid.Empty)
 44              {
 45                  user.AccountId = newAccountId;
 46                  return result;
 47              }
 48  
 49              var reloadedUser = await userManager.FindByIdAsync(user.Id);
 50              if (reloadedUser?.AccountId is { } reloadedAccountId && reloadedAccountId != Guid.Empty)
 51              {
 52                  user.AccountId = reloadedAccountId;
 53                  return IdentityResult.Success;
 54              }
 55  
 56              return result;
 57          }
 58      }
 59  
 60      public static async Task<(IdentityResult Result, Guid? AccountId)> EnsureAssignedAsync(
 61          ILiteCollection<ApplicationUser> users,
 62          string userId,
 63          CancellationToken ct = default)
 64      {
 65          using (await AcquireUserLockAsync(userId, ct))
 66          {
 67              var storedUser = users.FindById(userId);
 68              if (storedUser is null)
 69              {
 70                  return (
 71                      IdentityResult.Failed(new IdentityError
 72                      {
 73                          Code = "UserNotFound",
 74                          Description = "The user no longer exists in the identity store.",
 75                      }),
 76                      null);
 77              }
 78  
 79              if (storedUser.AccountId is { } storedAccountId && storedAccountId != Guid.Empty)
 80                  return (IdentityResult.Success, storedAccountId);
 81  
 82              storedUser.AccountId = Guid.NewGuid();
 83              if (users.Update(storedUser) && storedUser.AccountId is { } newAccountId && newAccountId != Guid.Empty)
 84                  return (IdentityResult.Success, newAccountId);
 85  
 86              var reloadedUser = users.FindById(userId);
 87              if (reloadedUser?.AccountId is { } reloadedAccountId && reloadedAccountId != Guid.Empty)
 88                  return (IdentityResult.Success, reloadedAccountId);
 89  
 90              return (
 91                  IdentityResult.Failed(new IdentityError
 92                  {
 93                      Code = "UserUpdateFailed",
 94                      Description = "Failed to assign an account ID to the user.",
 95                  }),
 96                  null);
 97          }
 98      }
 99  
100      private static async Task<IDisposable> AcquireUserLockAsync(string userId, CancellationToken ct)
101      {
102          while (true)
103          {
104              var userLock = UserLocks.GetOrAdd(userId, static _ => new UserLock());
105              Interlocked.Increment(ref userLock.RefCount);
106  
107              if (UserLocks.TryGetValue(userId, out var currentLock) && ReferenceEquals(userLock, currentLock))
108              {
109                  await userLock.Gate.WaitAsync(ct);
110                  return new Releaser(userId, userLock);
111              }
112  
113              Interlocked.Decrement(ref userLock.RefCount);
114          }
115      }
116  
117      private sealed class UserLock
118      {
119          public SemaphoreSlim Gate { get; } = new(1, 1);
120          public volatile int RefCount;
121      }
122  
123      private readonly struct Releaser(string userId, UserLock userLock) : IDisposable
124      {
125          public void Dispose()
126          {
127              userLock.Gate.Release();
128              if (Interlocked.Decrement(ref userLock.RefCount) == 0)
129                  UserLocks.TryRemove(new KeyValuePair<string, UserLock>(userId, userLock));
130          }
131      }
132  }