Academic-Affairs-System/src/Jiaowu.Api/Infrastructure/Auth/AuthSessionService.cs

using System.Security.Cryptography;
using Jiaowu.Api.Domain.Identity;
using Jiaowu.Api.Infrastructure.Persistence;
using Microsoft.AspNetCore.Identity;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Options;

namespace Jiaowu.Api.Infrastructure.Auth;

public interface IAuthSessionService
{
    Task<AuthSessionResult> CreateAsync(
        ApplicationUser user,
        IEnumerable<string> roles,
        AuthenticationClientType clientType,
        CancellationToken cancellationToken = default);

    Task<AuthSessionResult?> RefreshAsync(
        string refreshToken,
        CancellationToken cancellationToken = default);

    Task RevokeAsync(
        string refreshToken,
        CancellationToken cancellationToken = default);
}

public sealed record AuthSessionResult(
    string AccessToken,
    DateTime AccessTokenExpiresAt,
    string RefreshToken,
    DateTime SessionExpiresAt,
    ApplicationUser User,
    IReadOnlyList<string> Roles);

public sealed class AuthSessionService(
    AppDbContext db,
    UserManager<ApplicationUser> userManager,
    ITokenService tokenService,
    IOptions<JwtOptions> options) : IAuthSessionService
{
    private readonly JwtOptions _options = options.Value;

    public async Task<AuthSessionResult> CreateAsync(
        ApplicationUser user,
        IEnumerable<string> roles,
        AuthenticationClientType clientType,
        CancellationToken cancellationToken = default)
    {
        var roleList = roles.ToList();
        var now = DateTime.UtcNow;
        var rawRefreshToken = CreateRefreshToken();
        var session = new RefreshSession
        {
            UserId = user.Id,
            TokenHash = HashToken(rawRefreshToken),
            ClientType = clientType,
            SecurityStamp = user.SecurityStamp ?? string.Empty,
            CreatedAt = now,
            LastRefreshedAt = now,
            ExpiresAt = now.Add(GetIdleTimeout(clientType))
        };

        await RemoveExpiredSessionsAsync(user.Id, now, cancellationToken);
        db.RefreshSessions.Add(session);
        await db.SaveChangesAsync(cancellationToken);

        return BuildResult(user, roleList, rawRefreshToken, session.ExpiresAt);
    }

    public async Task<AuthSessionResult?> RefreshAsync(
        string refreshToken,
        CancellationToken cancellationToken = default)
    {
        var tokenHash = HashToken(refreshToken);
        var now = DateTime.UtcNow;
        var current = await db.RefreshSessions
            .Include(x => x.User)
            .SingleOrDefaultAsync(x => x.TokenHash == tokenHash, cancellationToken);
        var user = current?.User;
        if (current is null || user is null || current.RevokedAt.HasValue ||
            current.ExpiresAt <= now || !user.IsEnabled ||
            await userManager.IsLockedOutAsync(user) ||
            !string.Equals(current.SecurityStamp, user.SecurityStamp ?? string.Empty,
                StringComparison.Ordinal))
        {
            return null;
        }

        var newRawToken = CreateRefreshToken();
        var replacement = new RefreshSession
        {
            UserId = user.Id,
            TokenHash = HashToken(newRawToken),
            ClientType = current.ClientType,
            SecurityStamp = current.SecurityStamp,
            CreatedAt = now,
            LastRefreshedAt = now,
            ExpiresAt = now.Add(GetIdleTimeout(current.ClientType))
        };

        var rotated = await db.ExecuteInRetriableTransactionAsync(
            async transaction =>
            {
                db.ChangeTracker.Clear();
                var updated = await db.RefreshSessions
                    .Where(x => x.Id == current.Id && x.RevokedAt == null && x.ExpiresAt > now)
                    .ExecuteUpdateAsync(setters => setters
                        .SetProperty(x => x.RevokedAt, now)
                        .SetProperty(x => x.ReplacedBySessionId, replacement.Id),
                        cancellationToken);
                if (updated != 1)
                {
                    await transaction.RollbackAsync(cancellationToken);
                    return false;
                }

                db.RefreshSessions.Add(replacement);
                await db.SaveChangesAsync(cancellationToken);
                await transaction.CommitAsync(cancellationToken);
                return true;
            },
            cancellationToken);
        if (!rotated) return null;

        var roles = await userManager.GetRolesAsync(user);
        return BuildResult(user, roles.ToList(), newRawToken, replacement.ExpiresAt);
    }

    public async Task RevokeAsync(
        string refreshToken,
        CancellationToken cancellationToken = default)
    {
        var tokenHash = HashToken(refreshToken);
        var now = DateTime.UtcNow;
        await db.RefreshSessions
            .Where(x => x.TokenHash == tokenHash && x.RevokedAt == null)
            .ExecuteUpdateAsync(
                setters => setters.SetProperty(x => x.RevokedAt, now),
                cancellationToken);
    }

    private AuthSessionResult BuildResult(
        ApplicationUser user,
        IReadOnlyList<string> roles,
        string refreshToken,
        DateTime sessionExpiresAt)
    {
        var accessToken = tokenService.Create(user, roles);
        return new AuthSessionResult(
            accessToken.Token,
            accessToken.ExpiresAt,
            refreshToken,
            sessionExpiresAt,
            user,
            roles);
    }

    private TimeSpan GetIdleTimeout(AuthenticationClientType clientType) =>
        TimeSpan.FromMinutes(clientType == AuthenticationClientType.App
            ? _options.AppIdleMinutes
            : _options.WebIdleMinutes);

    private async Task RemoveExpiredSessionsAsync(
        Guid userId,
        DateTime now,
        CancellationToken cancellationToken)
    {
        var retentionCutoff = now.AddDays(-7);
        await db.RefreshSessions
            .Where(x => x.UserId == userId &&
                (x.ExpiresAt < now || x.RevokedAt < retentionCutoff))
            .ExecuteDeleteAsync(cancellationToken);
    }

    private static string CreateRefreshToken() =>
        Convert.ToBase64String(RandomNumberGenerator.GetBytes(48));

    private static string HashToken(string token) =>
        Convert.ToHexString(SHA256.HashData(System.Text.Encoding.UTF8.GetBytes(token)));
}