Nexus.Reader/src/NexusReader.Infrastructure/Identity/TokenLimitHandler.cs

using System.Security.Claims;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Identity;
using NexusReader.Domain.Entities;
using NexusReader.Data.Persistence;
using Microsoft.EntityFrameworkCore;

namespace NexusReader.Infrastructure.Identity;

/// <summary>
/// Handler that validates if the user has available AI tokens.
/// </summary>
public class TokenLimitHandler : AuthorizationHandler<TokenLimitRequirement>
{
    private readonly IDbContextFactory<AppDbContext> _dbContextFactory;
    private readonly UserManager<NexusUser> _userManager;

    public TokenLimitHandler(IDbContextFactory<AppDbContext> dbContextFactory, UserManager<NexusUser> userManager)
    {
        _dbContextFactory = dbContextFactory;
        _userManager = userManager;
    }

    protected override async Task HandleRequirementAsync(
        AuthorizationHandlerContext context,
        TokenLimitRequirement requirement)
    {
        var userId = context.User.FindFirstValue(ClaimTypes.NameIdentifier);
        if (userId == null)
        {
            return;
        }

        using var db = _dbContextFactory.CreateDbContext();
        var user = await db.Users
            .Include(u => u.SubscriptionPlan)
            .FirstOrDefaultAsync(u => u.Id == userId);

        if (user == null)
        {
            return;
        }

        // Check if user has available tokens or unlimited plan
        if (user.SubscriptionPlan?.IsUnlimitedTokens == true || user.AITokensUsed < user.AITokenLimit)
        {
            context.Succeed(requirement);
        }
    }
}