git.stella-ops.org/src/Authority/StellaOps.Authority/StellaOps.Auth.Client/ServiceCollectionExtensions.cs

using System;
using System.Net;
using System.Net.Http;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.DependencyInjection.Extensions;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Polly;
using Polly.Extensions.Http;
using StellaOps.AirGap.Policy;

namespace StellaOps.Auth.Client;

/// <summary>
/// DI helpers for the StellaOps auth client.
/// </summary>
public static class ServiceCollectionExtensions
{
    /// <summary>
    /// Registers the StellaOps auth client with the provided configuration.
    /// </summary>
    public static IServiceCollection AddStellaOpsAuthClient(this IServiceCollection services, Action<StellaOpsAuthClientOptions> configure)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentNullException.ThrowIfNull(configure);

        services.AddOptions<StellaOpsAuthClientOptions>()
            .Configure(configure)
            .PostConfigure(static options => options.Validate());

        services.TryAddSingleton<IStellaOpsTokenCache, InMemoryTokenCache>();

        services.AddHttpClient<StellaOpsDiscoveryCache>((provider, client) =>
        {
            var options = provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>().CurrentValue;
            EnsureEgressAllowed(provider, options, "authority-discovery");
            client.Timeout = options.HttpTimeout;
        }).AddPolicyHandler(static (provider, _) => CreateRetryPolicy(provider));

        services.AddHttpClient<StellaOpsJwksCache>((provider, client) =>
        {
            var options = provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>().CurrentValue;
            EnsureEgressAllowed(provider, options, "authority-jwks");
            client.Timeout = options.HttpTimeout;
        }).AddPolicyHandler(static (provider, _) => CreateRetryPolicy(provider));

        services.AddHttpClient<IStellaOpsTokenClient, StellaOpsTokenClient>((provider, client) =>
        {
            var options = provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>().CurrentValue;
            EnsureEgressAllowed(provider, options, "authority-token");
            client.Timeout = options.HttpTimeout;
        }).AddPolicyHandler(static (provider, _) => CreateRetryPolicy(provider));

        return services;
    }

    /// <summary>
    /// Registers a file-backed token cache implementation.
    /// </summary>
    public static IServiceCollection AddStellaOpsFileTokenCache(this IServiceCollection services, string cacheDirectory)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentException.ThrowIfNullOrWhiteSpace(cacheDirectory);

        services.Replace(ServiceDescriptor.Singleton<IStellaOpsTokenCache>(provider =>
        {
            var logger = provider.GetService<Microsoft.Extensions.Logging.ILogger<FileTokenCache>>();
            var options = provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>().CurrentValue;
            return new FileTokenCache(cacheDirectory, TimeProvider.System, options.ExpirationSkew, logger);
        }));

        return services;
    }

    /// <summary>
    /// Adds authentication and tenancy header handling for an <see cref="HttpClient"/> registered via <see cref="IHttpClientBuilder"/>.
    /// </summary>
    public static IHttpClientBuilder AddStellaOpsApiAuthentication(this IHttpClientBuilder builder, Action<StellaOpsApiAuthenticationOptions> configure)
    {
        ArgumentNullException.ThrowIfNull(builder);
        ArgumentNullException.ThrowIfNull(configure);

        builder.Services.AddOptions<StellaOpsApiAuthenticationOptions>(builder.Name)
            .Configure(configure)
            .PostConfigure(static options => options.Validate());

        builder.AddHttpMessageHandler(provider => new StellaOpsBearerTokenHandler(
            builder.Name,
            provider.GetRequiredService<IOptionsMonitor<StellaOpsApiAuthenticationOptions>>(),
            provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>(),
            provider.GetRequiredService<IStellaOpsTokenClient>(),
            provider.GetService<TimeProvider>(),
            provider.GetService<ILogger<StellaOpsBearerTokenHandler>>()));

        return builder;
    }

    private static IAsyncPolicy<HttpResponseMessage> CreateRetryPolicy(IServiceProvider provider)
    {
        var options = provider.GetRequiredService<IOptionsMonitor<StellaOpsAuthClientOptions>>().CurrentValue;
        var delays = options.NormalizedRetryDelays;
        if (delays.Count == 0)
        {
            return Policy.NoOpAsync<HttpResponseMessage>();
        }

        var logger = provider.GetService<ILoggerFactory>()?.CreateLogger("StellaOps.Auth.Client.HttpRetry");

        return HttpPolicyExtensions
            .HandleTransientHttpError()
            .OrResult(static response => response.StatusCode == HttpStatusCode.TooManyRequests)
            .WaitAndRetryAsync(
                delays.Count,
                attempt => delays[attempt - 1],
                (outcome, delay, attempt, _) =>
                {
                    if (logger is null)
                    {
                        return;
                    }

                    if (outcome.Exception is not null)
                    {
                        logger.LogWarning(
                            outcome.Exception,
                            "Retrying Authority HTTP call ({Attempt}/{TotalAttempts}) after exception; waiting {Delay}.",
                            attempt,
                            delays.Count,
                            delay);
                    }
                    else
                    {
                        logger.LogWarning(
                            "Retrying Authority HTTP call ({Attempt}/{TotalAttempts}) due to status {StatusCode}; waiting {Delay}.",
                            attempt,
                            delays.Count,
                            outcome.Result!.StatusCode,
                            delay);
                    }
                });
    }

    private static void EnsureEgressAllowed(
        IServiceProvider provider,
        StellaOpsAuthClientOptions options,
        string intent)
    {
        ArgumentNullException.ThrowIfNull(provider);
        ArgumentNullException.ThrowIfNull(options);
        ArgumentException.ThrowIfNullOrWhiteSpace(intent);

        if (options.AuthorityUri is null)
        {
            return;
        }

        var policy = provider.GetService<IEgressPolicy>();
        if (policy is null)
        {
            return;
        }

        var request = new EgressRequest("StellaOpsAuthClient", options.AuthorityUri, intent);
        policy.EnsureAllowed(request);
    }
}