git.stella-ops.org/src/Policy/StellaOps.Policy.Gateway/Services/PolicyEngineTokenProvider.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StellaOps.Auth.Client;
using StellaOps.Policy.Gateway.Options;
using System;
using System.Collections.Generic;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;

namespace StellaOps.Policy.Gateway.Services;

internal sealed class PolicyEngineTokenProvider
{
    private readonly IStellaOpsTokenClient tokenClient;
    private readonly IOptionsMonitor<PolicyGatewayOptions> optionsMonitor;
    private readonly PolicyGatewayDpopProofGenerator dpopGenerator;
    private readonly TimeProvider timeProvider;
    private readonly ILogger<PolicyEngineTokenProvider> logger;
    private readonly SemaphoreSlim mutex = new(1, 1);
    private CachedToken? cachedToken;

    public PolicyEngineTokenProvider(
        IStellaOpsTokenClient tokenClient,
        IOptionsMonitor<PolicyGatewayOptions> optionsMonitor,
        PolicyGatewayDpopProofGenerator dpopGenerator,
        TimeProvider timeProvider,
        ILogger<PolicyEngineTokenProvider> logger)
    {
        this.tokenClient = tokenClient ?? throw new ArgumentNullException(nameof(tokenClient));
        this.optionsMonitor = optionsMonitor ?? throw new ArgumentNullException(nameof(optionsMonitor));
        this.dpopGenerator = dpopGenerator ?? throw new ArgumentNullException(nameof(dpopGenerator));
        this.timeProvider = timeProvider ?? TimeProvider.System;
        this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public bool IsEnabled => optionsMonitor.CurrentValue.PolicyEngine.ClientCredentials.Enabled;

    public async ValueTask<PolicyGatewayAuthorization?> GetAuthorizationAsync(HttpMethod method, Uri targetUri, CancellationToken cancellationToken)
    {
        if (!IsEnabled)
        {
            return null;
        }

        var tokenResult = await GetTokenAsync(cancellationToken).ConfigureAwait(false);
        if (tokenResult is null)
        {
            return null;
        }

        var token = tokenResult.Value;
        string? proof = null;
        if (dpopGenerator.Enabled)
        {
            proof = dpopGenerator.CreateProof(method, targetUri, token.AccessToken);
        }

        var scheme = string.Equals(token.TokenType, "dpop", StringComparison.OrdinalIgnoreCase)
            ? "DPoP"
            : token.TokenType;

        var authorization = $"{scheme} {token.AccessToken}";
        return new PolicyGatewayAuthorization(authorization, proof, "service");
    }

    private async ValueTask<CachedToken?> GetTokenAsync(CancellationToken cancellationToken)
    {
        var options = optionsMonitor.CurrentValue.PolicyEngine;
        if (!options.ClientCredentials.Enabled)
        {
            return null;
        }

        var now = timeProvider.GetUtcNow();
        if (cachedToken is { } existing && existing.ExpiresAt > now + TimeSpan.FromSeconds(30))
        {
            return existing;
        }

        await mutex.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (cachedToken is { } cached && cached.ExpiresAt > now + TimeSpan.FromSeconds(30))
            {
                return cached;
            }

            var scopeString = BuildScopeClaim(options);
            try
            {
                var result = await tokenClient.RequestClientCredentialsTokenAsync(scopeString, null, cancellationToken).ConfigureAwait(false);
                var expiresAt = result.ExpiresAtUtc;
                cachedToken = new CachedToken(result.AccessToken, string.IsNullOrWhiteSpace(result.TokenType) ? "Bearer" : result.TokenType, expiresAt);
                logger.LogInformation("Issued Policy Engine client credentials token; expires at {ExpiresAt:o}.", expiresAt);
                return cachedToken;
            }
            catch (Exception ex)
            {
                logger.LogWarning(
                    ex,
                    "Unable to issue Policy Engine client credentials token for scopes '{Scopes}'.",
                    scopeString);
                return null;
            }
        }
        finally
        {
            mutex.Release();
        }
    }

    private string BuildScopeClaim(PolicyGatewayPolicyEngineOptions options)
    {
        var scopeSet = new SortedSet<string>(StringComparer.Ordinal)
        {
            $"aud:{options.Audience.Trim().ToLowerInvariant()}"
        };

        foreach (var scope in options.ClientCredentials.Scopes)
        {
            if (string.IsNullOrWhiteSpace(scope))
            {
                continue;
            }

            scopeSet.Add(scope.Trim());
        }

        return string.Join(' ', scopeSet);
    }

    private readonly record struct CachedToken(string AccessToken, string TokenType, DateTimeOffset ExpiresAt);
}