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.Http.Resilience;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Polly;
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;
        }).AddResilienceHandler("authority-discovery", ConfigureResilience);

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

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

        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;
            var timeProvider = provider.GetService<TimeProvider>();
            return new FileTokenCache(cacheDirectory, timeProvider, 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 void ConfigureResilience(ResiliencePipelineBuilder<HttpResponseMessage> builder, ResilienceHandlerContext context)
    {
        context.EnableReloads<StellaOpsAuthClientOptions>();

        var options = context.GetOptions<StellaOpsAuthClientOptions>();
        if (!options.EnableRetries || options.NormalizedRetryDelays.Count == 0)
        {
            return;
        }

        var delays = options.NormalizedRetryDelays;

        builder.AddRetry(new HttpRetryStrategyOptions
        {
            MaxRetryAttempts = delays.Count,
            DelayGenerator = args =>
            {
                var index = args.AttemptNumber < delays.Count ? args.AttemptNumber : delays.Count - 1;
                return ValueTask.FromResult<TimeSpan?>(delays[index]);
            },
            BackoffType = DelayBackoffType.Constant,
            ShouldHandle = static args => ValueTask.FromResult(
                args.Outcome.Exception is not null ||
                args.Outcome.Result?.StatusCode is HttpStatusCode.RequestTimeout
                    or HttpStatusCode.TooManyRequests
                    or >= HttpStatusCode.InternalServerError)
        });
    }

    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);
    }
}