ClaudeDo/src/ClaudeDo.Worker/Online/ZitadelAuthProvider.cs

using System.Net.Http.Json;
using System.Runtime.Versioning;
using System.Text.Json;
using System.Text.Json.Serialization;
using ClaudeDo.Worker.Online.Interfaces;
using Microsoft.Extensions.Logging;

namespace ClaudeDo.Worker.Online;

[SupportedOSPlatform("windows")]
public sealed class ZitadelAuthProvider : IOnlineAuthProvider
{
    private readonly IHttpClientFactory _httpClientFactory;
    private readonly OnlineTokenStore _tokenStore;
    private readonly OnlineInboxConfig _config;
    private readonly ILogger<ZitadelAuthProvider> _logger;

    private readonly SemaphoreSlim _lock = new(1, 1);

    // Cached access token state.
    private string? _cachedAccessToken;
    private DateTimeOffset _cacheExpiry;
    // The refresh token that minted the cached access token. When the stored refresh token
    // changes (sign-out, or signing in as a different user), the cache is no longer valid.
    private string? _refreshTokenUsed;

    // Cached token endpoint URL (discovered once).
    private string? _tokenEndpoint;

    public ZitadelAuthProvider(
        IHttpClientFactory httpClientFactory,
        OnlineTokenStore tokenStore,
        OnlineInboxConfig config,
        ILogger<ZitadelAuthProvider> logger)
    {
        _httpClientFactory = httpClientFactory;
        _tokenStore = tokenStore;
        _config = config;
        _logger = logger;
    }

    public Task<string?> GetAccessTokenAsync(CancellationToken ct = default)
        => GetAccessTokenAsync(false, ct);

    public async Task<string?> GetAccessTokenAsync(bool forceRefresh, CancellationToken ct = default)
    {
        var refreshToken = _tokenStore.Read();

        // Fast path: cached token is valid, not forced, and was minted from the still-current
        // refresh token (i.e. the signed-in user hasn't changed).
        if (IsCacheUsable(forceRefresh, refreshToken))
            return _cachedAccessToken;

        await _lock.WaitAsync(ct);
        try
        {
            // Re-read + re-check inside the lock (double-checked locking).
            refreshToken = _tokenStore.Read();
            if (IsCacheUsable(forceRefresh, refreshToken))
                return _cachedAccessToken;

            // Drop any stale access token so a fresh one is minted for the current user.
            _cachedAccessToken = null;
            _cacheExpiry = default;

            if (refreshToken is null)
            {
                _refreshTokenUsed = null;
                _logger.LogDebug("No refresh token stored; skipping token refresh.");
                return null;
            }

            return await RefreshAsync(refreshToken, ct);
        }
        finally
        {
            _lock.Release();
        }
    }

    private bool IsCacheUsable(bool forceRefresh, string? storedRefreshToken) =>
        !forceRefresh
        && _cachedAccessToken is not null
        && DateTimeOffset.UtcNow < _cacheExpiry
        && storedRefreshToken == _refreshTokenUsed;

    private async Task<string?> RefreshAsync(string refreshToken, CancellationToken ct)
    {
        var tokenEndpoint = await GetTokenEndpointAsync(ct);
        if (tokenEndpoint is null)
            return null;

        using var http = _httpClientFactory.CreateClient(nameof(ZitadelAuthProvider));

        var form = new Dictionary<string, string>
        {
            ["grant_type"]     = "refresh_token",
            ["refresh_token"]  = refreshToken,
            ["client_id"]      = _config.Zitadel.ClientId,
            ["scope"]          = _config.Zitadel.Scopes,
        };

        HttpResponseMessage response;
        try
        {
            response = await http.PostAsync(tokenEndpoint, new FormUrlEncodedContent(form), ct);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Token refresh request failed.");
            return null;
        }

        if (!response.IsSuccessStatusCode)
        {
            var body = await response.Content.ReadAsStringAsync(ct);
            if ((int)response.StatusCode == 400 && body.Contains("invalid_grant"))
            {
                _logger.LogWarning("Refresh token rejected (invalid_grant). Will retry once a new token is stored.");
            }
            else
            {
                _logger.LogWarning("Token refresh returned {Status}: {Body}", (int)response.StatusCode, body);
            }
            return null;
        }

        var tokenResponse = await response.Content.ReadFromJsonAsync<TokenResponse>(ct);
        if (tokenResponse?.AccessToken is null)
        {
            _logger.LogWarning("Token refresh response missing access_token.");
            return null;
        }

        // If Zitadel rotated the refresh token, persist the new one.
        var persistedRefreshToken = refreshToken;
        if (tokenResponse.RefreshToken is not null && tokenResponse.RefreshToken != refreshToken)
        {
            _logger.LogDebug("Refresh token rotated; persisting new token.");
            _tokenStore.Save(tokenResponse.RefreshToken);
            persistedRefreshToken = tokenResponse.RefreshToken;
        }

        // Cache the access token (subtract 60 s safety margin; minimum 0 to avoid far-future expiry on zero).
        // Remember which refresh token it was minted from so the cache invalidates on a user switch.
        _cachedAccessToken = tokenResponse.AccessToken;
        _cacheExpiry = DateTimeOffset.UtcNow.AddSeconds(tokenResponse.ExpiresIn - 60);
        _refreshTokenUsed = persistedRefreshToken;

        return _cachedAccessToken;
    }

    private async Task<string?> GetTokenEndpointAsync(CancellationToken ct)
    {
        if (_tokenEndpoint is not null)
            return _tokenEndpoint;

        var discoveryUrl = _config.Zitadel.Authority.TrimEnd('/') + "/.well-known/openid-configuration";

        using var http = _httpClientFactory.CreateClient(nameof(ZitadelAuthProvider));
        try
        {
            var doc = await http.GetFromJsonAsync<OidcDiscovery>(discoveryUrl, ct);
            _tokenEndpoint = doc?.TokenEndpoint;
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to discover OIDC configuration from {Url}.", discoveryUrl);
            return null;
        }

        if (_tokenEndpoint is null)
            _logger.LogWarning("OIDC discovery at {Url} did not return a token_endpoint.", discoveryUrl);

        return _tokenEndpoint;
    }

    private sealed class OidcDiscovery
    {
        [JsonPropertyName("token_endpoint")]
        public string? TokenEndpoint { get; init; }
    }

    private sealed class TokenResponse
    {
        [JsonPropertyName("access_token")]
        public string? AccessToken { get; init; }

        [JsonPropertyName("expires_in")]
        public int ExpiresIn { get; init; }

        [JsonPropertyName("refresh_token")]
        public string? RefreshToken { get; init; }
    }
}