ColaFlow/colaflow-api/src/Modules/Identity/ColaFlow.Modules.Identity.Infrastructure/Services/RefreshTokenService.cs

using ColaFlow.Modules.Identity.Application.Services;
using ColaFlow.Modules.Identity.Domain.Aggregates.Tenants;
using ColaFlow.Modules.Identity.Domain.Aggregates.Users;
using ColaFlow.Modules.Identity.Domain.Repositories;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Logging;
using System.Security.Cryptography;
using System.Text;

namespace ColaFlow.Modules.Identity.Infrastructure.Services;

public class RefreshTokenService : IRefreshTokenService
{
    private readonly IRefreshTokenRepository _refreshTokenRepository;
    private readonly IUserRepository _userRepository;
    private readonly ITenantRepository _tenantRepository;
    private readonly IUserTenantRoleRepository _userTenantRoleRepository;
    private readonly IJwtService _jwtService;
    private readonly IConfiguration _configuration;
    private readonly ILogger<RefreshTokenService> _logger;

    public RefreshTokenService(
        IRefreshTokenRepository refreshTokenRepository,
        IUserRepository userRepository,
        ITenantRepository tenantRepository,
        IUserTenantRoleRepository userTenantRoleRepository,
        IJwtService jwtService,
        IConfiguration configuration,
        ILogger<RefreshTokenService> logger)
    {
        _refreshTokenRepository = refreshTokenRepository;
        _userRepository = userRepository;
        _tenantRepository = tenantRepository;
        _userTenantRoleRepository = userTenantRoleRepository;
        _jwtService = jwtService;
        _configuration = configuration;
        _logger = logger;
    }

    public async Task<string> GenerateRefreshTokenAsync(
        User user,
        string? ipAddress = null,
        string? userAgent = null,
        CancellationToken cancellationToken = default)
    {
        // Generate cryptographically secure random token (64 bytes)
        var randomBytes = new byte[64];
        using var rng = RandomNumberGenerator.Create();
        rng.GetBytes(randomBytes);
        var token = Convert.ToBase64String(randomBytes);

        // Hash token before storage (SHA-256)
        var tokenHash = ComputeSha256Hash(token);

        // Get expiration from configuration (default 7 days)
        var expirationDays = _configuration.GetValue<int>("Jwt:RefreshTokenExpirationDays", 7);
        var expiresAt = DateTime.UtcNow.AddDays(expirationDays);

        // Create refresh token entity
        var refreshToken = RefreshToken.Create(
            tokenHash: tokenHash,
            userId: UserId.Create(user.Id),
            tenantId: user.TenantId.Value,
            expiresAt: expiresAt,
            ipAddress: ipAddress,
            userAgent: userAgent,
            deviceInfo: userAgent // Use user agent as device info for now
        );

        // Save to database
        await _refreshTokenRepository.AddAsync(refreshToken, cancellationToken);

        _logger.LogInformation(
            "Generated refresh token for user {UserId}, expires at {ExpiresAt}",
            user.Id, expiresAt);

        // Return plain text token (only time we return it)
        return token;
    }

    public async Task<(string accessToken, string refreshToken)> RefreshTokenAsync(
        string refreshToken,
        string? ipAddress = null,
        string? userAgent = null,
        CancellationToken cancellationToken = default)
    {
        // Hash the provided token to look it up
        var tokenHash = ComputeSha256Hash(refreshToken);

        // Find existing token
        var existingToken = await _refreshTokenRepository.GetByTokenHashAsync(tokenHash, cancellationToken);

        if (existingToken == null)
        {
            _logger.LogWarning("Refresh token not found: {TokenHash}", tokenHash[..10] + "...");
            throw new UnauthorizedAccessException("Invalid refresh token");
        }

        // Check if token is active (not expired and not revoked)
        if (!existingToken.IsActive())
        {
            _logger.LogWarning(
                "Attempted to use invalid refresh token for user {UserId}. Expired: {IsExpired}, Revoked: {IsRevoked}",
                existingToken.UserId.Value, existingToken.IsExpired(), existingToken.IsRevoked());

            // SECURITY: Token reuse detection - revoke all user tokens
            if (existingToken.IsRevoked())
            {
                _logger.LogWarning(
                    "SECURITY ALERT: Revoked token reused for user {UserId}. Revoking all tokens.",
                    existingToken.UserId.Value);
                await RevokeAllUserTokensAsync(existingToken.UserId.Value, cancellationToken);
            }

            throw new UnauthorizedAccessException("Invalid or expired refresh token");
        }

        // Get user and tenant
        var user = await _userRepository.GetByIdAsync(existingToken.UserId, cancellationToken);
        if (user == null || user.Status != UserStatus.Active)
        {
            _logger.LogWarning("User not found or inactive: {UserId}", existingToken.UserId.Value);
            throw new UnauthorizedAccessException("User not found or inactive");
        }

        var tenant = await _tenantRepository.GetByIdAsync(TenantId.Create(existingToken.TenantId), cancellationToken);
        if (tenant == null || tenant.Status != TenantStatus.Active)
        {
            _logger.LogWarning("Tenant not found or inactive: {TenantId}", existingToken.TenantId);
            throw new UnauthorizedAccessException("Tenant not found or inactive");
        }

        // Get user's tenant role
        var userTenantRole = await _userTenantRoleRepository.GetByUserAndTenantAsync(
            user.Id,
            tenant.Id,
            cancellationToken);

        if (userTenantRole == null)
        {
            _logger.LogWarning("User {UserId} has no role assigned for tenant {TenantId}", user.Id, tenant.Id);
            throw new UnauthorizedAccessException("User role not found");
        }

        // Generate new access token with role
        var newAccessToken = _jwtService.GenerateToken(user, tenant, userTenantRole.Role);

        // Generate new refresh token (token rotation)
        var newRefreshToken = await GenerateRefreshTokenAsync(user, ipAddress, userAgent, cancellationToken);

        // Mark old token as replaced
        var newTokenHash = ComputeSha256Hash(newRefreshToken);
        existingToken.MarkAsReplaced(newTokenHash);
        await _refreshTokenRepository.UpdateAsync(existingToken, cancellationToken);

        _logger.LogInformation(
            "Rotated refresh token for user {UserId}",
            user.Id);

        return (newAccessToken, newRefreshToken);
    }

    public async Task RevokeTokenAsync(
        string refreshToken,
        string? ipAddress = null,
        CancellationToken cancellationToken = default)
    {
        var tokenHash = ComputeSha256Hash(refreshToken);
        var token = await _refreshTokenRepository.GetByTokenHashAsync(tokenHash, cancellationToken);

        if (token == null)
        {
            _logger.LogWarning("Attempted to revoke non-existent token");
            return; // Silent failure for security
        }

        if (token.IsRevoked())
        {
            _logger.LogWarning("Token already revoked: {TokenId}", token.Id);
            return;
        }

        var reason = ipAddress != null
            ? $"User logout from {ipAddress}"
            : "User logout";

        token.Revoke(reason);
        await _refreshTokenRepository.UpdateAsync(token, cancellationToken);

        _logger.LogInformation(
            "Revoked refresh token {TokenId} for user {UserId}",
            token.Id, token.UserId.Value);
    }

    public async Task RevokeAllUserTokensAsync(
        Guid userId,
        CancellationToken cancellationToken = default)
    {
        await _refreshTokenRepository.RevokeAllUserTokensAsync(
            userId,
            "User requested logout from all devices",
            cancellationToken);

        _logger.LogInformation(
            "Revoked all refresh tokens for user {UserId}",
            userId);
    }

    private static string ComputeSha256Hash(string input)
    {
        using var sha256 = SHA256.Create();
        var bytes = Encoding.UTF8.GetBytes(input);
        var hash = sha256.ComputeHash(bytes);
        return Convert.ToBase64String(hash);
    }
}