Atomx/Atomx.Admin/Atomx.Admin.Client/Utils/AuthHeaderHandler.cs

using Atomx.Admin.Client.Services;
using Atomx.Common.Models;
using Blazored.LocalStorage;
using Microsoft.AspNetCore.Components;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;

namespace Atomx.Admin.Client.Utils
{
    /// <summary>
    /// WASM 模式下的请求拦截器（DelegatingHandler）
    /// - 在每次请求时将 access token 附带 Authorization header
    /// - 在 401 或响应头包含 "Token-Expired" 时尝试刷新（调用 /api/sign/refresh）
    /// - 防止并发刷新（SemaphoreSlim）
    /// - 注意：刷新请求必须使用一个不包含本 handler 的 HttpClient（避免循环）
    /// </summary>
    public class AuthHeaderHandler : DelegatingHandler
    {
        private readonly ITokenProvider _tokenProvider;
        private readonly NavigationManager _navigationManager;
        private readonly ILogger<AuthHeaderHandler> _logger;
        private readonly ILocalStorageService _localStorage;
        private readonly IHttpClientFactory _httpClientFactory;
        private static readonly SemaphoreSlim _refreshLock = new(1, 1);

        private const string AccessTokenKey = "accessToken";
        private const string RefreshTokenKey = "refreshToken";

        public AuthHeaderHandler(
            ITokenProvider tokenProvider,
            NavigationManager navigationManager,
            ILogger<AuthHeaderHandler> logger,
            ILocalStorageService localStorage,
            IHttpClientFactory httpClientFactory)
        {
            _tokenProvider = tokenProvider;
            _navigationManager = navigationManager;
            _logger = logger;
            _localStorage = localStorage;
            _httpClientFactory = httpClientFactory;
        }

        protected override async Task<HttpResponseMessage> SendAsync(
            HttpRequestMessage request, CancellationToken cancellationToken)
        {
            try
            {
                // 从 ITokenProvider 获取当前 access token（WASM: ClientTokenProvider 从 localStorage 读取）
                var token = await _tokenProvider.GetTokenAsync();
                if (!string.IsNullOrEmpty(token))
                {
                    request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", token);
                }
                else
                {
                    _logger.LogDebug("No token from ITokenProvider for request {Url}", request.RequestUri);
                }

                var response = await base.SendAsync(request, cancellationToken);

                // 当发现未授权或后端标记 token 过期时，尝试刷新
                if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized ||
                    response.Headers.Contains("Token-Expired"))
                {
                    _logger.LogInformation("Unauthorized or Token-Expired detected for {Url}", request.RequestUri);

                    if (OperatingSystem.IsBrowser())
                    {
                        var refreshed = await TryRefreshTokenAsync(cancellationToken);

                        if (refreshed)
                        {
                            var newToken = await _localStorage.GetItemAsync<string>(AccessTokenKey);
                            if (!string.IsNullOrEmpty(newToken))
                            {
                                var clonedRequest = await CloneHttpRequestMessageAsync(request, newToken);
                                return await base.SendAsync(clonedRequest, cancellationToken);
                            }
                        }

                        // 刷新失败或无 token，跳转到登录页（强制刷新页面，清除 SPA 状态）
                        _navigationManager.NavigateTo("/account/login", true);
                    }
                    else
                    {
                        _logger.LogWarning("Unauthorized in server mode for {Url}", request.RequestUri);
                    }
                }

                return response;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error sending HTTP request to {Url}", request.RequestUri);
                throw;
            }
        }

        /// <summary>
        /// 尝试刷新 token
        /// 关键点：
        /// - 使用命名 HttpClient "RefreshClient"（在 WASM Program 中注册，不包含本 handler），避免递归
        /// - 并发刷新时只有一个请求会实际触发刷新，其余等待结果
        /// - 刷新成功后写入 localStorage（accessToken + refreshToken）
        /// </summary>
        private async Task<bool> TryRefreshTokenAsync(CancellationToken cancellationToken)
        {
            await _refreshLock.WaitAsync(cancellationToken);
            try
            {
                var currentAccess = await _localStorage.GetItemAsync<string>(AccessTokenKey);
                var currentRefresh = await _localStorage.GetItemAsync<string>(RefreshTokenKey);

                if (string.IsNullOrEmpty(currentAccess) || string.IsNullOrEmpty(currentRefresh))
                {
                    _logger.LogInformation("No local tokens to refresh");
                    return false;
                }

                // 使用不带本 handler 的 HttpClient 发送刷新请求，避免循环
                var client = _httpClientFactory.CreateClient("RefreshClient");

                var reqModel = new RefreshRequest
                {
                    Token = currentAccess,
                    RefreshToken = currentRefresh
                };
                var reqJson = JsonSerializer.Serialize(reqModel);
                using var req = new HttpRequestMessage(HttpMethod.Post, "api/sign/refresh")
                {
                    Content = new StringContent(reqJson, Encoding.UTF8, "application/json")
                };

                try
                {
                    var resp = await client.SendAsync(req, cancellationToken);
                    if (!resp.IsSuccessStatusCode)
                    {
                        _logger.LogWarning("Refresh request failed with status {Status}", resp.StatusCode);
                        // 如果刷新失败，移除本地 token（防止无限重试）
                        await _localStorage.RemoveItemAsync(AccessTokenKey);
                        await _localStorage.RemoveItemAsync(RefreshTokenKey);
                        return false;
                    }

                    var json = await resp.Content.ReadAsStringAsync(cancellationToken);
                    var authResp = JsonSerializer.Deserialize<AuthResponse>(json, new JsonSerializerOptions
                    {
                        PropertyNameCaseInsensitive = true
                    });

                    if (authResp == null || string.IsNullOrEmpty(authResp.Token) || string.IsNullOrEmpty(authResp.RefreshToken))
                    {
                        _logger.LogWarning("Invalid response from refresh endpoint");
                        await _localStorage.RemoveItemAsync(AccessTokenKey);
                        await _localStorage.RemoveItemAsync(RefreshTokenKey);
                        return false;
                    }

                    // 保存新的 tokens 到 localStorage
                    await _localStorage.SetItemAsync(AccessTokenKey, authResp.Token, cancellationToken);
                    await _localStorage.SetItemAsync(RefreshTokenKey, authResp.RefreshToken, cancellationToken);

                    _logger.LogInformation("Token refreshed successfully");
                    return true;
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Exception while requesting token refresh");
                    return false;
                }
            }
            finally
            {
                _refreshLock.Release();
            }
        }

        /// <summary>
        /// 复制原始请求并用新的 token 替换 Authorization header
        /// </summary>
        private static async Task<HttpRequestMessage> CloneHttpRequestMessageAsync(HttpRequestMessage original, string newToken)
        {
            var clone = new HttpRequestMessage(original.Method, original.RequestUri);

            // 复制内容（如果存在）
            if (original.Content != null)
            {
                var ms = new MemoryStream();
                await original.Content.CopyToAsync(ms).ConfigureAwait(false);
                ms.Position = 0;
                clone.Content = new StreamContent(ms);

                if (original.Content.Headers != null)
                {
                    foreach (var h in original.Content.Headers)
                        clone.Content.Headers.Add(h.Key, h.Value);
                }
            }

            // 复制 headers
            foreach (var header in original.Headers)
                clone.Headers.TryAddWithoutValidation(header.Key, header.Value);

            // 覆盖 Authorization
            clone.Headers.Authorization = new AuthenticationHeaderValue("Bearer", newToken);

            return clone;
        }
    }
}