Atomx/Atomx.Admin/Atomx.Admin.Client/Utils/PersistentAuthenticationStateProvider.cs

using Atomx.Common.Constants;
using Atomx.Utils.Extension;
using Blazored.LocalStorage;
using Microsoft.AspNetCore.Components;
using Microsoft.AspNetCore.Components.Authorization;
using System.IdentityModel.Tokens.Jwt;
using System.Security.Claims;

namespace Atomx.Admin.Client.Utils
{
    public class PersistentAuthenticationStateProvider : AuthenticationStateProvider
    {
        readonly ClaimsPrincipal anonymous = new(new ClaimsIdentity());

        // 如果运行在 Server 且在 prerender 时有 Persisted UserInfo，则存储预设的 AuthenticationState
        private Task<AuthenticationState>? _preRenderedAuthState;

        readonly ILocalStorageService _localStorage;

        public PersistentAuthenticationStateProvider(IServiceProvider serviceProvider, ILocalStorageService localStorageService)
        {
            _localStorage = localStorageService;

            // 尝试有条件解析 PersistedComponentState（仅在 Server 交互渲染时可用）
            var state = serviceProvider.GetService<PersistentComponentState>();
            if (state != null)
            {
                if (state.TryTakeFromJson<UserInfo>(nameof(UserInfo), out var userInfo) && userInfo is not null)
                {
                    var claims = new List<Claim>
                    {
                        new(ClaimKeys.Id, userInfo.Id.ToString()),
                        new(ClaimKeys.Name, userInfo.Name),
                        new(ClaimKeys.Email, userInfo.Email),
                        new(ClaimKeys.Mobile, userInfo.MobilePhone),
                        new(ClaimKeys.Role, userInfo.Role),
                    };
                    foreach (var role in userInfo.Permissions ?? Array.Empty<string>())
                    {
                        claims.Add(new Claim(ClaimKeys.Permission, role));
                    }

                    var cp = new ClaimsPrincipal(new ClaimsIdentity(claims, authenticationType: nameof(PersistentAuthenticationStateProvider)));
                    _preRenderedAuthState = Task.FromResult(new AuthenticationState(cp));
                }
            }
        }

        public override async Task<AuthenticationState> GetAuthenticationStateAsync()
        {
            // 如果在 prerender 阶段已从 PersistentComponentState 恢复用户，优先返回该状态（Server prerender）
            if (_preRenderedAuthState != null)
                return await _preRenderedAuthState;

            try
            {
                var jwtToken = await _localStorage.GetItemAsStringAsync(StorageKeys.AccessToken);
                if (string.IsNullOrEmpty(jwtToken))
                    return new AuthenticationState(anonymous);

                var getUserClaims = DecryptToken(jwtToken);
                if (getUserClaims == null || string.IsNullOrEmpty(getUserClaims.Name))
                    return new AuthenticationState(anonymous);

                var claimsPrincipal = SetClaimPrincipal(getUserClaims);
                return new AuthenticationState(claimsPrincipal);
            }
            catch
            {
                return new AuthenticationState(anonymous);
            }
        }

        public ClaimsPrincipal SetClaimPrincipal(UserInfo customUserClaims)
        {
            if (string.IsNullOrEmpty(customUserClaims.Name))
                return new ClaimsPrincipal();
            else
            {
                var claims = new List<Claim>
                {
                    new(ClaimKeys.Id, customUserClaims.Id.ToString()),
                    new(ClaimKeys.Name, customUserClaims.Name),
                    new(ClaimKeys.Email, customUserClaims.Email),
                    new(ClaimKeys.Mobile, customUserClaims.MobilePhone),
                    new(ClaimKeys.Role, customUserClaims.Role.ToString()),
                };
                foreach (var role in customUserClaims.Permissions ?? Array.Empty<string>())
                {
                    claims.Add(new Claim(ClaimKeys.Permission, role));
                }
                return new ClaimsPrincipal(new ClaimsIdentity(claims, authenticationType: nameof(PersistentAuthenticationStateProvider)));
            }
        }

        public void UpdateAuthenticationState(string jwtToken = "")
        {
            var claimsPrincipal = new ClaimsPrincipal();
            if (!string.IsNullOrEmpty(jwtToken))
            {
                var getUserClaims = DecryptToken(jwtToken);
                claimsPrincipal = SetClaimPrincipal(getUserClaims);
            }

            NotifyAuthenticationStateChanged(Task.FromResult(new AuthenticationState(claimsPrincipal)));
        }

        private UserInfo DecryptToken(string jwtToken)
        {
            if (string.IsNullOrEmpty(jwtToken))
                return new UserInfo();
            else
            {
                var handler = new JwtSecurityTokenHandler();
                var token = handler.ReadJwtToken(jwtToken);

                var id = token.Claims.SingleOrDefault(x => x.Type == ClaimKeys.Id)?.Value ?? string.Empty;
                var name = token.Claims.SingleOrDefault(x => x.Type == ClaimKeys.Name)?.Value ?? string.Empty;
                var email = token.Claims.SingleOrDefault(x => x.Type == ClaimKeys.Email)?.Value ?? string.Empty;
                var phone = token.Claims.SingleOrDefault(x => x.Type == ClaimKeys.Mobile)?.Value ?? string.Empty;
                var role = token.Claims.SingleOrDefault(x => x.Type == ClaimKeys.Role)?.Value ?? string.Empty;
                var permissions = token.Claims.Where(x => x.Type == ClaimKeys.Permission).Select(s => s.Value).ToArray();

                return new UserInfo(id.ToLong(), name, email, phone, role, permissions);
            }
        }

        public async Task MarkUserAsLoggedOut()
        {
            await _localStorage.RemoveItemAsync(StorageKeys.AccessToken);
            await _localStorage.RemoveItemAsync(StorageKeys.RefreshToken);

            var authState = Task.FromResult(new AuthenticationState(anonymous));
            NotifyAuthenticationStateChanged(authState);
        }
    }

    public record UserInfo(long Id = 0, string Name = null!, string Email = null!, string MobilePhone = null!, string Role = null!, string[] Permissions = null!);
}