OpenWarehouse/OpenWarehouse.Common.Lib/AddAuthServices.cs

using System.Security.Cryptography;
using System.Text;
using Microsoft.AspNetCore.Authentication.JwtBearer;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.IdentityModel.Tokens;
using OpenWarehouse.auth.Common.ApiInterfeiceLib;
using OpenWarehouse.auth.lib.Model.Auth;
using OpenWarehouse.auth.lib.Model.WellKnow;

namespace OpenWarehouse.Common.Lib;

public static class AddAuthServices
{
    public static async Task<T> AddAuthService<T>(this T builder, string authUrl, ILogger logger)
        where T : IHostApplicationBuilder
    {
        var jwks = await builder.GetKey(authUrl, logger);
        var conf = await builder.GetConfig(authUrl, logger);

        var jwk = jwks.Keys[0];

        var rsaParameters = new RSAParameters
        {
            Modulus = Base64UrlEncoder.DecodeBytes(jwk.N),
            Exponent = Base64UrlEncoder.DecodeBytes(jwk.E)
        };

        builder.Services.AddSingleton(_ =>
        {
            var rsa = RSA.Create();
            rsa.ImportParameters(rsaParameters);
            return rsa;
        });

        builder.Services.AddSingleton<RsaSecurityKey>(provider =>
        {
            var rsa = provider.GetRequiredService<RSA>();
            return new RsaSecurityKey(rsa) { KeyId = jwk.Kid };
        });

        builder.Services.AddAuthentication(options =>
            {
                options.DefaultScheme = JwtBearerDefaults.AuthenticationScheme;
                options.DefaultAuthenticateScheme = JwtBearerDefaults.AuthenticationScheme;
                options.DefaultChallengeScheme = JwtBearerDefaults.AuthenticationScheme;
            })
            .AddJwtBearer(options =>
            {
                var rsaSecurityKey = builder.Services.BuildServiceProvider().GetRequiredService<RsaSecurityKey>();

                options.TokenValidationParameters = new TokenValidationParameters
                {
                    ValidateIssuerSigningKey = true,
                    IssuerSigningKey = rsaSecurityKey,
                    TokenDecryptionKey =
                        new SymmetricSecurityKey(Encoding.UTF8.GetBytes(builder.Configuration["Jwt:AesKey"] ?? "")),
                    ValidateIssuer = true,
                    ValidIssuer = conf.Issuer,
                    ValidateAudience = true,
                    ValidAudience = builder.Configuration["Jwt:Audience"],
                    ValidateLifetime = true,
                    SaveSigninToken = true
                };
                options.Events = new JwtBearerEvents
                {
                    OnTokenValidated = async context =>
                    {
                        var client = new OpenWarehouseAuthClient(authUrl, new HttpClient());
                        context.Request.Headers.TryGetValue("Authorization", out var token);
                        var result = await client.TokenIsLockAsync(new TokenAskIfIsLockModel { Token = token });
                        if (result)
                        {
                            context.Fail("This token has been revoked.");
                        }
                    },
                    OnAuthenticationFailed = context =>
                    {
                        logger.LogError("Authentication failed: {Exception}", context.Exception);
                        return Task.CompletedTask;
                    },
                };
            });

        return builder;
    }


    private static async Task<OpenIdConfiguration> GetConfig<T>(this T builder, string url, ILogger logger)
        where T : IHostApplicationBuilder
    {
        try
        {
            logger.LogInformation("Requesting OICD from: {Url}", url);

            using var httpClient = new HttpClient();
            var client = new OpenWarehouseAuthClient(url, httpClient);

            var result = await client.OpenidConfigurationAsync();

            logger.LogInformation("Successfully retrieved OICD from: {Url}", url);
            return result;
        }
        catch (Exception ex)
        {
            logger.LogError(ex, "Failed to retrieve OICD from: {Url}", url);
            throw;
        }
    }

    private static async Task<JwksResponse> GetKey<T>(this T builder, string url, ILogger logger)
        where T : IHostApplicationBuilder
    {
        try
        {
            logger.LogInformation("Requesting JWKS from: {Url}", url);

            using var httpClient = new HttpClient();
            var client = new OpenWarehouseAuthClient(url, httpClient);

            var result = await client.JwksAsync();

            logger.LogInformation("Successfully retrieved JWKS from: {Url}", url);
            return result;
        }
        catch (Exception ex)
        {
            logger.LogError(ex, "Failed to retrieve JWKS from: {Url}", url);
            throw;
        }
    }
}