Flawless-Version-Control/Flawless.Server/Controllers/AuthenticationController.cs

using System.Security.Claims;
using Flawless.Communication.Request.Auth;
using Flawless.Communication.Response;
using Flawless.Communication.Shared;
using Flawless.Server.Models;
using Flawless.Server.Services;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Mvc;
using Microsoft.EntityFrameworkCore;

namespace Flawless.Server.Controllers;

[ApiController, Route("api/auth")]
public class AuthenticationController(
    UserManager<AppUser> userManager,
    SignInManager<AppUser> signInManager,
    TokenGenerationService tokenService,
    AppDbContext dbContext,
    IConfiguration config,
    ILogger<AuthenticationController> logger)
    : ControllerBase
{
    
    [HttpPost("register")]
    public async Task<IActionResult> PublicRegisterAsync(RegisterRequest request)
    {
        if (!config.GetValue("User:PublicRegister", true)) 
            return BadRequest(new FailedResponse("Not opened for public register."));
        
        
        var user = new AppUser
        {
            UserName = request.Username, 
            Email = request.Email,
            EmailConfirmed = true,
            CreatedOn = DateTime.UtcNow,
            SecurityStamp = Guid.NewGuid().ToString()
        };
        
        var result = await userManager.CreateAsync(user, request.Password);
        if (result.Succeeded) 
        {
            logger.LogInformation("User '{0}' created (PUBLIC REGISTER)", user.UserName);
            return Ok();
        }
        
        logger.LogInformation("User '{0}' NOT created (PUBLIC REGISTER) : {1}", user.UserName, result.Errors);
        return BadRequest(new FailedResponse(result.Errors));
    }

    [HttpPost("login")]
    public async Task<IActionResult> LoginAsync(LoginRequest r)
    {
        var user = await userManager.FindByNameAsync(r.Username);
        if (user == null) return BadRequest(new FailedResponse("Invalid username or password.")); 
        
        var result = await signInManager.CheckPasswordSignInAsync(user, r.Password, false);
        if (result.Succeeded)
        {
            var refreshToken = tokenService.GenerateRefreshToken();
            var claims = await GetClaimsAsync(user, refreshToken);
            var jwtToken = tokenService.GenerateToken(claims);
            var refKey = new AppUserRefreshKey
            {
                UserId = user.Id,
                RefreshToken = refreshToken,
                ExpireIn = DateTime.UtcNow.AddDays(tokenService.RefreshTokenLifeTime),
            };

            dbContext.RefreshTokens.Add(refKey);
            await dbContext.SaveChangesAsync();
            await userManager.AddLoginAsync(user, new UserLoginInfo("login-interface", "", null));
            
            return Ok(new TokenInfo { Token = jwtToken });
        }

        if (result.IsLockedOut)
        {
            return BadRequest(new FailedResponse("Account is locked out."));
        }
        
        return BadRequest(new FailedResponse("Invalid username or password.")); 
    }
    
    [HttpPost("refresh")]
    public async Task<IActionResult> RefreshAsync(TokenInfo r)
    {
        var now = DateTime.UtcNow;
        var set = dbContext.RefreshTokens;
        var principal = tokenService.GetPrincipalFromExpiredToken(r.Token);
        var user = await signInManager.ValidateSecurityStampAsync(principal);
        if (user == null) return BadRequest(new FailedResponse("Token is ban. Please login again."));

        try
        {
            // Remove timeout guys
            await set.Where(k => k.ExpireIn < now).ExecuteDeleteAsync();
        
            // Find valid expired refresh token
            var refreshToken = principal.FindFirst(FlawlessClaimsType.RefreshToken)?.Value;
            var tk = await set.FirstOrDefaultAsync(k => k.RefreshToken == refreshToken && k.UserId.ToString() == user.Id.ToString());
            if (tk == null) return BadRequest(new FailedResponse("Token is ban. Please login again."));

            // Renew keys
            refreshToken = tokenService.GenerateRefreshToken();
            var claims = await GetClaimsAsync(user, refreshToken);
            var newJwtToken = tokenService.GenerateToken(claims);
        
            // Reassign a new key.
            set.Remove(tk);
            set.Add(new AppUserRefreshKey
            {
                UserId = user.Id,
                RefreshToken = refreshToken,
                ExpireIn = DateTime.UtcNow.AddDays(tokenService.RefreshTokenLifeTime),
            });
    
            return Ok(new TokenInfo { Token = newJwtToken });
        }
        finally
        {
            await dbContext.SaveChangesAsync();
        }
        
    }

    private async ValueTask<IList<Claim>> GetClaimsAsync(AppUser user, string refreshToken)
    {
        var c = await userManager.GetClaimsAsync(user);
        c.Add(new (FlawlessClaimsType.SecurityStamp, user.SecurityStamp ?? string.Empty));
        c.Add(new (ClaimTypes.NameIdentifier, user.Id.ToString()));
        c.Add(new (ClaimTypes.Name, user.UserName!));
        c.Add(new (FlawlessClaimsType.RefreshToken, refreshToken));

        return c;
    }
}