Clean up rotated OAuth tokens

src/Exceptionless.Core/Jobs/CleanupDataJob.cs

@@ -21,12 +21,15 @@ namespace Exceptionless.Core.Jobs;
 [Job(Description = "Deletes soft deleted data and enforces data retention.", IsContinuous = false)]
 public class CleanupDataJob : JobWithLockBase, IHealthCheck
 {
+    private static readonly TimeSpan OAuthTokenCleanupSafetyWindow = TimeSpan.FromDays(1);
+
     private readonly IOrganizationRepository _organizationRepository;
     private readonly OrganizationService _organizationService;
     private readonly IProjectRepository _projectRepository;
     private readonly IStackRepository _stackRepository;
     private readonly IEventRepository _eventRepository;
     private readonly ITokenRepository _tokenRepository;
+    private readonly IOAuthTokenRepository _oauthTokenRepository;
     private readonly IWebHookRepository _webHookRepository;
     private readonly BillingManager _billingManager;
     private readonly UsageService _usageService;
@@ -43,6 +46,7 @@ public CleanupDataJob(
         IStackRepository stackRepository,
         IEventRepository eventRepository,
         ITokenRepository tokenRepository,
+        IOAuthTokenRepository oauthTokenRepository,
         IWebHookRepository webHookRepository,
         ILockProvider lockProvider,
         ICacheClient cacheClient,
@@ -61,6 +65,7 @@ ILoggerFactory loggerFactory
         _stackRepository = stackRepository;
         _eventRepository = eventRepository;
         _tokenRepository = tokenRepository;
+        _oauthTokenRepository = oauthTokenRepository;
         _webHookRepository = webHookRepository;
         _billingManager = billingManager;
         _usageService = usageService;
@@ -80,6 +85,7 @@ protected override async Task<JobResult> RunInternalAsync(JobContext context)
         _lastRun = _timeProvider.GetUtcNow().UtcDateTime;
 
         await MarkTokensSuspended(context);
+        await CleanupOAuthTokensAsync(context);
         await CleanupSoftDeletedOrganizationsAsync(context);
         await CleanupSoftDeletedProjectsAsync(context);
         await CleanupSoftDeletedStacksAsync(context);
@@ -107,6 +113,13 @@ private async Task MarkTokensSuspended(JobContext context)
         } while (!context.CancellationToken.IsCancellationRequested && await suspendedOrganizations.NextPageAsync());
     }
 
+    private async Task CleanupOAuthTokensAsync(JobContext context)
+    {
+        var utcCutoff = _timeProvider.GetUtcNow().UtcDateTime.Subtract(OAuthTokenCleanupSafetyWindow);
+        long removed = await _oauthTokenRepository.RemoveExpiredDisabledAsync(utcCutoff, context.CancellationToken);
+        _logger.LogInformation("Removed {OAuthTokenCount} expired disabled OAuth token(s)", removed);
+    }
+
     private async Task CleanupSoftDeletedOrganizationsAsync(JobContext context)
     {
         var organizationResults = await _organizationRepository.GetAllAsync(o => o.SoftDeleteMode(SoftDeleteQueryMode.DeletedOnly).SearchAfterPaging().PageLimit(5));

src/Exceptionless.Core/Repositories/Interfaces/IOAuthTokenRepository.cs

@@ -9,7 +9,10 @@ public interface IOAuthTokenRepository : ISearchableRepository<OAuthToken>
     Task<FindResults<OAuthToken>> GetByAccessTokenHashAsync(string accessTokenHash, CommandOptionsDescriptor<OAuthToken>? options = null);
     Task<FindResults<OAuthToken>> GetByRefreshTokenHashAsync(string refreshTokenHash, CommandOptionsDescriptor<OAuthToken>? options = null);
     Task<FindResults<OAuthToken>> GetByGrantIdAsync(string grantId, CommandOptionsDescriptor<OAuthToken>? options = null);
+    Task<FindResults<OAuthToken>> GetByGrantIdForUpdateAsync(string grantId, CommandOptionsDescriptor<OAuthToken>? options = null);
     Task<FindResults<OAuthToken>> GetByUserIdAsync(string userId, CommandOptionsDescriptor<OAuthToken>? options = null);
     Task<FindResults<OAuthToken>> GetByUserIdAndClientIdAsync(string userId, string clientId, CommandOptionsDescriptor<OAuthToken>? options = null);
+    Task<FindResults<OAuthToken>> GetByUserIdAndClientIdForUpdateAsync(string userId, string clientId, CommandOptionsDescriptor<OAuthToken>? options = null);
+    Task<long> RemoveExpiredDisabledAsync(DateTime utcCutoff, CancellationToken cancellationToken = default);
     Task<long> RemoveAllByUserIdAsync(string userId, CommandOptionsDescriptor<OAuthToken>? options = null);
 }

src/Exceptionless.Core/Repositories/OAuthTokenRepository.cs

@@ -8,6 +8,8 @@ namespace Exceptionless.Core.Repositories;
 
 public class OAuthTokenRepository : RepositoryBase<OAuthToken>, IOAuthTokenRepository
 {
+    private const int CleanupBatchSize = 500;
+
     public OAuthTokenRepository(ExceptionlessElasticConfiguration configuration, MiniValidationValidator validator, AppOptions options)
         : base(configuration.OAuthTokens, validator, options)
     {
@@ -31,6 +33,11 @@ public Task<FindResults<OAuthToken>> GetByGrantIdAsync(string grantId, CommandOp
         return FindAsync(q => q.FieldEquals(t => t.GrantId, grantId).SortDescending(t => t.UpdatedUtc), options);
     }
 
+    public Task<FindResults<OAuthToken>> GetByGrantIdForUpdateAsync(string grantId, CommandOptionsDescriptor<OAuthToken>? options = null)
+    {
+        return FindAsync(q => q.FieldEquals(t => t.GrantId, grantId).SortAscending(t => t.Id), options);
+    }
+
     public Task<FindResults<OAuthToken>> GetByUserIdAsync(string userId, CommandOptionsDescriptor<OAuthToken>? options = null)
     {
         return FindAsync(q => q.FieldEquals(t => t.UserId, userId).SortDescending(t => t.UpdatedUtc), options);
@@ -41,6 +48,47 @@ public Task<FindResults<OAuthToken>> GetByUserIdAndClientIdAsync(string userId,
         return FindAsync(q => q.FieldEquals(t => t.UserId, userId).FieldEquals(t => t.ClientId, clientId).SortDescending(t => t.UpdatedUtc), options);
     }
 
+    public Task<FindResults<OAuthToken>> GetByUserIdAndClientIdForUpdateAsync(string userId, string clientId, CommandOptionsDescriptor<OAuthToken>? options = null)
+    {
+        return FindAsync(q => q.FieldEquals(t => t.UserId, userId).FieldEquals(t => t.ClientId, clientId).SortAscending(t => t.Id), options);
+    }
+
+    public async Task<long> RemoveExpiredDisabledAsync(DateTime utcCutoff, CancellationToken cancellationToken = default)
+    {
+        var expiredSpentTokenResults = await FindAsync(q => q
+            .FieldEquals(t => t.IsDisabled, true)
+            .DateRange(null, utcCutoff, (OAuthToken t) => t.RefreshExpiresUtc)
+            .SortAscending(t => t.Id), o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(CleanupBatchSize));
+
+        long removed = await RemoveMatchingAsync(expiredSpentTokenResults, t => !String.IsNullOrEmpty(t.RefreshTokenHash));
+
+        var clearedRefreshTokenResults = await FindAsync(q => q
+            .FieldEquals(t => t.IsDisabled, true)
+            .SortAscending(t => t.Id), o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(CleanupBatchSize));
+
+        removed += await RemoveMatchingAsync(clearedRefreshTokenResults, t => String.IsNullOrEmpty(t.RefreshTokenHash) && t.UpdatedUtc < utcCutoff);
+        return removed;
+
+        async Task<long> RemoveMatchingAsync(FindResults<OAuthToken> results, Func<OAuthToken, bool> shouldRemove)
+        {
+            long removedCount = 0;
+            do
+            {
+                var tokensToRemove = results.Documents
+                    .Where(shouldRemove)
+                    .ToArray();
+
+                if (tokensToRemove.Length > 0)
+                {
+                    await RemoveAsync(tokensToRemove, o => o.ImmediateConsistency());
+                    removedCount += tokensToRemove.Length;
+                }
+            } while (!cancellationToken.IsCancellationRequested && await results.NextPageAsync());
+
+            return removedCount;
+        }
+    }
+
     public Task<long> RemoveAllByUserIdAsync(string userId, CommandOptionsDescriptor<OAuthToken>? options = null)
     {
         return RemoveAllAsync(q => q.FieldEquals(t => t.UserId, userId), options);

src/Exceptionless.Core/Services/OAuthService.cs

@@ -517,10 +517,19 @@ private async Task RevokeOAuthGrantFamilyAsync(OAuthToken token)
             return;
         }
 
-        var results = await oauthTokenRepository.GetByGrantIdAsync(token.GrantId, o => o.ImmediateConsistency().PageLimit(OAuthGrantFamilyPageLimit));
-        IEnumerable<OAuthToken> familyTokens = results.Documents.Count > 0 ? results.Documents : [token];
-        foreach (var familyToken in familyTokens.Where(t => String.Equals(t.GrantId, token.GrantId, StringComparison.Ordinal)))
-            await DisableTokenAsync(familyToken);
+        bool disabledAny = false;
+        var results = await oauthTokenRepository.GetByGrantIdForUpdateAsync(token.GrantId, o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(OAuthGrantFamilyPageLimit));
+        do
+        {
+            foreach (var familyToken in results.Documents.Where(t => String.Equals(t.GrantId, token.GrantId, StringComparison.Ordinal)))
+            {
+                disabledAny = true;
+                await DisableTokenAsync(familyToken);
+            }
+        } while (await results.NextPageAsync());
+
+        if (!disabledAny)
+            await DisableTokenAsync(token);
     }
 
     private static RefreshScopeValidationResult ValidateRefreshScopes(OAuthToken token, OAuthClientOptions client)

src/Exceptionless.Web/Controllers/UserController.cs

@@ -75,9 +75,14 @@ public async Task<ActionResult<ViewCurrentUser>> GetCurrentUserAsync()
     [HttpGet("me/oauth-grants")]
     public async Task<ActionResult<IReadOnlyCollection<ViewOAuthGrant>>> GetOAuthGrantsAsync()
     {
-        var results = await _oauthTokenRepository.GetByUserIdAsync(CurrentUser.Id, o => o.PageLimit(MAXIMUM_SKIP));
-        var tokens = results.Documents.Where(IsActiveOAuthGrantToken).ToArray();
-        if (tokens.Length == 0)
+        var tokens = new List<OAuthToken>();
+        var results = await _oauthTokenRepository.GetByUserIdAsync(CurrentUser.Id, o => o.SearchAfterPaging().PageLimit(MAXIMUM_SKIP));
+        do
+        {
+            tokens.AddRange(results.Documents.Where(IsActiveOAuthGrantToken));
+        } while (!HttpContext.RequestAborted.IsCancellationRequested && await results.NextPageAsync());
+
+        if (tokens.Count == 0)
             return Ok(Array.Empty<ViewOAuthGrant>());
 
         var applicationsByClientId = new Dictionary<string, OAuthApplication?>(StringComparer.Ordinal);
@@ -102,24 +107,34 @@ public async Task<ActionResult<IReadOnlyCollection<ViewOAuthGrant>>> GetOAuthGra
     [HttpDelete("me/oauth-grants/{id:minlength(1)}")]
     public async Task<IActionResult> RevokeOAuthGrantAsync(string id)
     {
-        var grantResults = await _oauthTokenRepository.GetByGrantIdAsync(id, o => o.ImmediateConsistency().PageLimit(MAXIMUM_SKIP));
-        var token = grantResults.Documents.FirstOrDefault(t =>
-            String.Equals(t.UserId, CurrentUser.Id, StringComparison.Ordinal)
-            && !String.IsNullOrWhiteSpace(t.ClientId));
+        OAuthToken? token = null;
+        var grantResults = await _oauthTokenRepository.GetByGrantIdForUpdateAsync(id, o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(MAXIMUM_SKIP));
+        do
+        {
+            token = grantResults.Documents.FirstOrDefault(t =>
+                String.Equals(t.UserId, CurrentUser.Id, StringComparison.Ordinal)
+                && !String.IsNullOrWhiteSpace(t.ClientId));
+            if (token is not null)
+                break;
+        } while (!HttpContext.RequestAborted.IsCancellationRequested && await grantResults.NextPageAsync());
+
         if (token is null)
             return NotFound();
 
         string clientId = token.ClientId;
 
-        var results = await _oauthTokenRepository.GetByUserIdAndClientIdAsync(CurrentUser.Id, clientId, o => o.ImmediateConsistency().PageLimit(MAXIMUM_SKIP));
+        var results = await _oauthTokenRepository.GetByUserIdAndClientIdForUpdateAsync(CurrentUser.Id, clientId, o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(MAXIMUM_SKIP));
         var utcNow = _timeProvider.GetUtcNow().UtcDateTime;
-        foreach (var oauthToken in results.Documents)
+        do
         {
-            oauthToken.IsDisabled = true;
-            oauthToken.RefreshTokenHash = null;
-            oauthToken.UpdatedUtc = utcNow;
-            await _oauthTokenRepository.SaveAsync(oauthToken, o => o.ImmediateConsistency());
-        }
+            foreach (var oauthToken in results.Documents)
+            {
+                oauthToken.IsDisabled = true;
+                oauthToken.RefreshTokenHash = null;
+                oauthToken.UpdatedUtc = utcNow;
+                await _oauthTokenRepository.SaveAsync(oauthToken, o => o.ImmediateConsistency());
+            }
+        } while (!HttpContext.RequestAborted.IsCancellationRequested && await results.NextPageAsync());
 
         return NoContent();
     }

tests/Exceptionless.Tests/Controllers/OAuthControllerTests.cs

@@ -1052,6 +1052,56 @@ public async Task TokenAsync_RefreshTokenWithRemovedOfflineAccess_RevokesGrantFa
         await AssertRefreshFailsAndRevokesGrantFamilyAsync(token);
     }
 
+    [Fact]
+    public async Task TokenAsync_RefreshTokenWithLargeGrantFamily_RevokesAllTokens()
+    {
+        var token = await IssueTokenAsync(
+            resource: RestApiResource,
+            scope: $"{AuthorizationRoles.ProjectsRead} {AuthorizationRoles.OfflineAccess}");
+        var storedToken = await GetStoredOAuthTokenAsync(token.AccessToken);
+        Assert.NotNull(storedToken);
+        string grantId = storedToken.GrantId;
+        var utcNow = TimeProvider.GetUtcNow().UtcDateTime;
+        var familyTokens = Enumerable.Range(0, 1005)
+            .Select(i => new OAuthToken
+            {
+                Id = ObjectId.GenerateNewId().ToString(),
+                UserId = storedToken.UserId,
+                ClientId = storedToken.ClientId,
+                GrantId = grantId,
+                Resource = storedToken.Resource,
+                AccessTokenHash = OAuthService.CreateTokenHash(StringExtensions.GetRandomString(OAuthService.OAuthTokenLength)),
+                RefreshTokenHash = OAuthService.CreateTokenHash(StringExtensions.GetRandomString(OAuthService.OAuthTokenLength)),
+                Scopes = [AuthorizationRoles.ProjectsRead, AuthorizationRoles.OfflineAccess],
+                OrganizationIds = [TestConstants.OrganizationId],
+                ExpiresUtc = utcNow.AddHours(1),
+                RefreshExpiresUtc = utcNow.AddDays(30),
+                IsDisabled = false,
+                CreatedBy = storedToken.UserId,
+                CreatedUtc = utcNow,
+                UpdatedUtc = utcNow
+            })
+            .ToArray();
+        await _oauthTokenRepository.AddAsync(familyTokens, o => o.ImmediateConsistency());
+        await SetStoredOAuthApplicationScopesAsync(ClientId, AuthorizationRoles.ProjectsRead);
+
+        await AssertRefreshFailsAndRevokesGrantFamilyAsync(token);
+
+        var results = await _oauthTokenRepository.GetByGrantIdForUpdateAsync(grantId, o => o.ImmediateConsistency().SearchAfterPaging().PageLimit(1000));
+        int tokenCount = 0;
+        do
+        {
+            foreach (var familyToken in results.Documents)
+            {
+                tokenCount++;
+                Assert.True(familyToken.IsDisabled);
+                Assert.Null(familyToken.RefreshTokenHash);
+            }
+        } while (await results.NextPageAsync());
+
+        Assert.Equal(familyTokens.Length + 1, tokenCount);
+    }
+
     [Fact]
     public async Task TokenAsync_RefreshTokenWithRemovedRequiredResourceScope_RevokesGrantFamily()
     {