QYZH.InteractiveMagazine/QYZH.InteractiveMagazine.Infrastructure/Extensions/DependencyInjectionExtensions.cs

using Microsoft.AspNetCore.Authentication;
using Microsoft.AspNetCore.Authentication.JwtBearer;
using Microsoft.AspNetCore.Hosting;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Microsoft.IdentityModel.Tokens;
using QYZH.InteractiveMagazine.Infrastructure.Auth;
using QYZH.InteractiveMagazine.Infrastructure.Middleware;
using QYZH.InteractiveMagazine.Models.Settings;
using System.Security.Claims;
using System.Text;
using System.Text.Encodings.Web;

namespace QYZH.InteractiveMagazine.Infrastructure.Extensions;

/// <summary>
/// Infrastructure service registration extensions.
/// </summary>
public static class DependencyInjectionExtensions
{
    private const string NoAuthScheme = "NoAuth";
    private const string DevelopmentAuthScheme = "DevelopmentSmartAuth";

    /// <summary>
    /// Registers infrastructure services.
    /// </summary>
    public static void AddInfrastructureServices(this IServiceCollection services, IConfiguration configuration, IWebHostEnvironment? environment = null)
    {
        AddJwtAuthentication(services, configuration, environment);

        services.AddTransient<GlobalExceptionMiddleware>();
        services.AddTransient<OperationLogMiddleware>();
        services.AddScoped<OperationLogActionFilter>();
    }

    private static void AddJwtAuthentication(IServiceCollection services, IConfiguration configuration, IWebHostEnvironment? environment = null)
    {
        var jwtSettings = configuration.GetSection("JwtSettings").Get<JwtSettings>()!;
        services.AddSingleton(jwtSettings);

        if (environment?.IsDevelopment() == true)
        {
            services.AddAuthentication(options =>
                {
                    options.DefaultScheme = DevelopmentAuthScheme;
                    options.DefaultChallengeScheme = DevelopmentAuthScheme;
                })
                .AddPolicyScheme(DevelopmentAuthScheme, null, options =>
                {
                    options.ForwardDefaultSelector = context =>
                    {
                        var authHeader = context.Request.Headers.Authorization.FirstOrDefault();
                        return !string.IsNullOrWhiteSpace(authHeader) &&
                               authHeader.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase)
                            ? JwtBearerDefaults.AuthenticationScheme
                            : NoAuthScheme;
                    };
                })
                .AddJwtBearer(JwtBearerDefaults.AuthenticationScheme, options => ConfigureJwtBearer(options, jwtSettings))
                .AddScheme<AuthenticationSchemeOptions, NoAuthHandler>(NoAuthScheme, options => { });
            return;
        }

        services.AddAuthentication(JwtBearerDefaults.AuthenticationScheme)
            .AddJwtBearer(options => ConfigureJwtBearer(options, jwtSettings));
    }

    private static void ConfigureJwtBearer(JwtBearerOptions options, JwtSettings jwtSettings)
    {
        options.TokenValidationParameters = new TokenValidationParameters
        {
            ValidateIssuer = true,
            ValidIssuer = jwtSettings.Issuer,
            ValidateAudience = true,
            ValidAudience = jwtSettings.Audience,
            ValidateIssuerSigningKey = true,
            IssuerSigningKey = new SymmetricSecurityKey(Encoding.UTF8.GetBytes(jwtSettings.SecretKey!)),
            ValidateLifetime = true,
            ClockSkew = TimeSpan.Zero
        };

        options.Events = new JwtBearerEvents
        {
            OnTokenValidated = async context =>
            {
                var currentToken = GetBearerToken(context);
                if (string.IsNullOrEmpty(currentToken))
                {
                    context.Fail("Invalid token");
                    return;
                }

                var userId = context.Principal?.FindFirst(ClaimTypes.NameIdentifier)?.Value;
                if (string.IsNullOrEmpty(userId))
                {
                    context.Fail("Invalid token");
                    return;
                }

                var wxUserId = context.Principal?.FindFirst(JwtHelper.WxUserIdClaimType)?.Value;
                if (string.IsNullOrEmpty(wxUserId))
                {
                    var adminToken = await RedisHelper.GetAsync(JwtHelper.BuildAdminTokenKey(userId));
                    if (adminToken == currentToken)
                    {
                        await RefreshRedisTokenExpiryAsync(JwtHelper.BuildAdminTokenKey(userId), currentToken, jwtSettings);
                        return;
                    }

                    context.Fail("Token expired, please login again");
                    return;
                }

                var wechatToken = await RedisHelper.GetAsync(JwtHelper.BuildWeChatTokenKey(wxUserId, userId));
                if (wechatToken == currentToken)
                {
                    await RefreshRedisTokenExpiryAsync(JwtHelper.BuildWeChatTokenKey(wxUserId, userId), currentToken, jwtSettings);
                    return;
                }

                context.Fail("Token expired, please login again");
            }
        };
    }

    private static async Task RefreshRedisTokenExpiryAsync(string tokenKey, string currentToken, JwtSettings jwtSettings)
    {
        if (jwtSettings.ExpiryMinutes > 0)
        {
            await RedisHelper.SetAsync(tokenKey, currentToken, TimeSpan.FromMinutes(jwtSettings.ExpiryMinutes));
        }
    }

    private static string? GetBearerToken(TokenValidatedContext context)
    {
        var authHeader = context.HttpContext.Request.Headers.Authorization.FirstOrDefault();
        if (string.IsNullOrWhiteSpace(authHeader) || !authHeader.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase))
        {
            return null;
        }

        return authHeader.Substring("Bearer ".Length).Trim();
    }
}

/// <summary>
/// Authentication handler used only in development.
/// </summary>
public class NoAuthHandler : AuthenticationHandler<AuthenticationSchemeOptions>
{
    public NoAuthHandler(IOptionsMonitor<AuthenticationSchemeOptions> options, ILoggerFactory logger, UrlEncoder encoder, ISystemClock clock)
        : base(options, logger, encoder, clock)
    {
    }

    protected override Task<AuthenticateResult> HandleAuthenticateAsync()
    {
        var claims = new[]
        {
            new Claim(ClaimTypes.Name, "DevUser"),
            new Claim(ClaimTypes.NameIdentifier, "0")
        };
        var identity = new ClaimsIdentity(claims, Scheme.Name);
        var principal = new ClaimsPrincipal(identity);
        var ticket = new AuthenticationTicket(principal, Scheme.Name);

        return Task.FromResult(AuthenticateResult.Success(ticket));
    }
}