QYZH.InteractiveMagazine/QYZH.InteractiveMagazine.Infrastructure/Middleware/JwtAutoRefreshMiddleware.cs

using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Configuration;
using QYZH.InteractiveMagazine.Infrastructure.Auth;
using QYZH.InteractiveMagazine.Models.Settings;
using System.IdentityModel.Tokens.Jwt;
using System.Security.Claims;

namespace QYZH.InteractiveMagazine.Infrastructure.Middleware;

/// <summary>
/// Refreshes the Redis session TTL for the current JWT when it is still the active token.
/// </summary>
public class JwtAutoRefreshMiddleware
{
    private readonly RequestDelegate _next;
    private readonly IConfiguration _configuration;

    public JwtAutoRefreshMiddleware(RequestDelegate next, IConfiguration configuration)
    {
        _next = next;
        _configuration = configuration;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        var authHeader = context.Request.Headers.Authorization.FirstOrDefault();
        if (!string.IsNullOrEmpty(authHeader) && authHeader.StartsWith("Bearer "))
        {
            var token = authHeader.Substring("Bearer ".Length).Trim();
            await TryRefreshRedisTokenExpiryAsync(token);
        }

        await _next(context);
    }

    private async Task TryRefreshRedisTokenExpiryAsync(string token)
    {
        try
        {
            var tokenHandler = new JwtSecurityTokenHandler();
            if (tokenHandler.ReadToken(token) is not JwtSecurityToken jwtToken)
            {
                return;
            }

            if (jwtToken.ValidTo <= DateTime.UtcNow)
            {
                return;
            }

            var userId = GetClaimValue(jwtToken, ClaimTypes.NameIdentifier, JwtRegisteredClaimNames.NameId);
            if (string.IsNullOrEmpty(userId))
            {
                return;
            }

            var jwtSettings = _configuration.GetSection("JwtSettings").Get<JwtSettings>();
            if (jwtSettings == null || jwtSettings.ExpiryMinutes <= 0)
            {
                return;
            }

            var wxUserId = GetClaimValue(jwtToken, JwtHelper.WxUserIdClaimType);
            await RefreshRedisTokenExpiryAsync(wxUserId, userId, token, jwtSettings.ExpiryMinutes);
        }
        catch
        {
            // Ignore refresh failures. Authentication middleware will validate the request later.
        }
    }

    private static async Task RefreshRedisTokenExpiryAsync(string? wxUserId, string userId, string currentToken, int expiryMinutes)
    {
        var adminTokenKey = JwtHelper.BuildAdminTokenKey(userId);
        var adminToken = await RedisHelper.GetAsync(adminTokenKey);
        if (string.IsNullOrEmpty(wxUserId) && !string.IsNullOrEmpty(adminToken))
        {
            await RedisHelper.SetAsync(adminTokenKey, adminToken, TimeSpan.FromMinutes(expiryMinutes));
            return;
        }

        if (string.IsNullOrEmpty(wxUserId))
        {
            return;
        }

        var wechatTokenKey = JwtHelper.BuildWeChatTokenKey(wxUserId, userId);
        var wechatToken = await RedisHelper.GetAsync(wechatTokenKey);
        if (wechatToken == currentToken)
        {
            await RedisHelper.SetAsync(wechatTokenKey, currentToken, TimeSpan.FromMinutes(expiryMinutes));
        }
    }

    private static string? GetClaimValue(JwtSecurityToken jwtToken, params string[] claimTypes)
    {
        return jwtToken.Claims.FirstOrDefault(c => claimTypes.Contains(c.Type))?.Value;
    }
}