Add JWT multi-scheme support, Cognito package, and Testing package

Trax.Api.slnx

@@ -4,6 +4,8 @@
     <Project Path="src/Trax.Api.Auth/Trax.Api.Auth.csproj" />
     <Project Path="src/Trax.Api.Auth.ApiKey/Trax.Api.Auth.ApiKey.csproj" />
     <Project Path="src/Trax.Api.Auth.Jwt/Trax.Api.Auth.Jwt.csproj" />
+    <Project Path="src/Trax.Api.Auth.Jwt.Cognito/Trax.Api.Auth.Jwt.Cognito.csproj" />
+    <Project Path="src/Trax.Api.Auth.Jwt.Testing/Trax.Api.Auth.Jwt.Testing.csproj" />
     <Project Path="src/Trax.Api.Auth.Oidc/Trax.Api.Auth.Oidc.csproj" />
     <Project Path="src/Trax.Api.GraphQL/Trax.Api.GraphQL.csproj" />
     <Project Path="src/Trax.Api.GraphQL.Audit/Trax.Api.GraphQL.Audit.csproj" />

src/Trax.Api.Auth.Jwt.Cognito/CognitoDefaults.cs

@@ -0,0 +1,52 @@
+namespace Trax.Api.Auth.Jwt.Cognito;
+
+/// <summary>
+/// Constants and well-known claim names emitted by Amazon Cognito user pool
+/// tokens.
+/// </summary>
+/// <remarks>
+/// NO WARRANTY. Trax auth is plumbing, not a security product. You are solely
+/// responsible for securing systems that use it. See SECURITY-DISCLAIMER.md.
+/// </remarks>
+public static class CognitoDefaults
+{
+    /// <summary>
+    /// Discriminator written to <see cref="TraxAuthClaimTypes.PrincipalType"/>
+    /// by <see cref="CognitoJwtPrincipalResolver"/>.
+    /// </summary>
+    public const string PrincipalType = "cognito";
+
+    /// <summary>Cognito ID-token <c>token_use</c> claim value.</summary>
+    public const string TokenUseId = "id";
+
+    /// <summary>Cognito access-token <c>token_use</c> claim value.</summary>
+    public const string TokenUseAccess = "access";
+
+    /// <summary>Claim name carrying the Cognito-internal username.</summary>
+    public const string CognitoUsername = "cognito:username";
+
+    /// <summary>Claim name carrying Cognito group memberships (repeats per group).</summary>
+    public const string CognitoGroups = "cognito:groups";
+
+    /// <summary>Claim name carrying the JSON-encoded federated identities array.</summary>
+    public const string Identities = "identities";
+
+    /// <summary>
+    /// Synthetic claim emitted by <see cref="CognitoJwtPrincipalResolver"/>
+    /// that names the primary federated provider (<c>Google</c>,
+    /// <c>SignInWithApple</c>, etc.) or <c>cognito</c> for native users.
+    /// </summary>
+    public const string IdentityProvider = "identity_provider";
+
+    /// <summary>Claim name carrying <c>access</c> or <c>id</c>.</summary>
+    public const string TokenUse = "token_use";
+
+    /// <summary>Claim name on Cognito access tokens carrying the app client id.</summary>
+    public const string ClientId = "client_id";
+
+    /// <summary>Claim name carrying email-verified state.</summary>
+    public const string EmailVerified = "email_verified";
+
+    /// <summary>Claim name carrying the user's email.</summary>
+    public const string Email = "email";
+}

src/Trax.Api.Auth.Jwt.Cognito/CognitoJwtBuilderExtensions.cs

@@ -0,0 +1,113 @@
+using Microsoft.IdentityModel.JsonWebTokens;
+
+namespace Trax.Api.Auth.Jwt.Cognito;
+
+/// <summary>
+/// Builder helpers that wire <see cref="JwtBuilder"/> to validate tokens
+/// from an Amazon Cognito user pool.
+/// </summary>
+/// <remarks>
+/// NO WARRANTY. Trax auth is plumbing, not a security product. You are solely
+/// responsible for securing systems that use it. See SECURITY-DISCLAIMER.md.
+/// </remarks>
+public static class CognitoJwtBuilderExtensions
+{
+    /// <summary>
+    /// Configures the JWT builder for an Amazon Cognito user pool. Sets the
+    /// authority to <c>https://cognito-idp.{region}.amazonaws.com/{userPoolId}</c>,
+    /// the audience to <paramref name="clientId"/>, and an
+    /// <c>AudienceValidator</c> that accepts both ID tokens (<c>aud</c> claim)
+    /// and access tokens (<c>client_id</c> claim) depending on
+    /// <paramref name="tokenUse"/>. Optionally validates the <c>token_use</c>
+    /// claim itself.
+    /// </summary>
+    /// <param name="builder">The JWT builder.</param>
+    /// <param name="region">AWS region of the user pool, e.g. <c>us-east-1</c>.</param>
+    /// <param name="userPoolId">User pool identifier, e.g. <c>us-east-1_AbCdEfGhI</c>.</param>
+    /// <param name="clientId">App client id registered on the pool.</param>
+    /// <param name="tokenUse">Which token shapes to accept. Defaults to both.</param>
+    public static JwtBuilder UseCognito(
+        this JwtBuilder builder,
+        string region,
+        string userPoolId,
+        string clientId,
+        CognitoTokenUse tokenUse = CognitoTokenUse.IdAndAccess
+    )
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+        ArgumentException.ThrowIfNullOrWhiteSpace(region);
+        ArgumentException.ThrowIfNullOrWhiteSpace(userPoolId);
+        ArgumentException.ThrowIfNullOrWhiteSpace(clientId);
+
+        var authority = $"https://cognito-idp.{region}.amazonaws.com/{userPoolId}";
+        builder.UseAuthority(authority, clientId);
+
+        builder.CustomizeTokenValidation(tvp =>
+        {
+            tvp.AudienceValidator = (audiences, token, _) =>
+            {
+                var acceptId =
+                    tokenUse == CognitoTokenUse.Id || tokenUse == CognitoTokenUse.IdAndAccess;
+                var acceptAccess =
+                    tokenUse == CognitoTokenUse.Access || tokenUse == CognitoTokenUse.IdAndAccess;
+
+                if (acceptId && audiences.Any(a => a == clientId))
+                    return true;
+
+                if (acceptAccess && token is JsonWebToken jwt)
+                {
+                    var claim = jwt.Claims.FirstOrDefault(c => c.Type == CognitoDefaults.ClientId);
+                    if (claim is not null && claim.Value == clientId)
+                        return true;
+                }
+
+                return false;
+            };
+
+            // Enforce token_use to match the configured selection. The
+            // signature/issuer/audience checks above don't otherwise
+            // distinguish id from access tokens.
+            var existingLifetime = tvp.LifetimeValidator;
+            tvp.LifetimeValidator = (notBefore, expires, token, parameters) =>
+            {
+                if (
+                    existingLifetime is not null
+                    && !existingLifetime(notBefore, expires, token, parameters)
+                )
+                    return false;
+
+                if (token is JsonWebToken jwt)
+                {
+                    var use = jwt
+                        .Claims.FirstOrDefault(c => c.Type == CognitoDefaults.TokenUse)
+                        ?.Value;
+
+                    if (
+                        tokenUse == CognitoTokenUse.Id
+                        && use is not null
+                        && use != CognitoDefaults.TokenUseId
+                    )
+                        return false;
+
+                    if (
+                        tokenUse == CognitoTokenUse.Access
+                        && use is not null
+                        && use != CognitoDefaults.TokenUseAccess
+                    )
+                        return false;
+                }
+
+                // Fall back to the default lifetime validator (exp/nbf check).
+                if (expires is null)
+                    return false;
+                var now = DateTime.UtcNow;
+                var skew = parameters.ClockSkew;
+                if (notBefore is { } nbf && nbf > now + skew)
+                    return false;
+                return expires.Value > now - skew;
+            };
+        });
+
+        return builder;
+    }
+}

src/Trax.Api.Auth.Jwt.Cognito/CognitoJwtPrincipalResolver.cs

@@ -0,0 +1,198 @@
+using System.Security.Claims;
+using System.Text.Json;
+
+namespace Trax.Api.Auth.Jwt.Cognito;
+
+/// <summary>
+/// <see cref="ITraxPrincipalResolver{JwtTokenInput}"/> tuned for Amazon
+/// Cognito tokens. Normalizes Cognito-specific claims so downstream code
+/// (authorization, audit, business logic) sees a consistent
+/// <see cref="TraxPrincipal"/> shape regardless of whether the user signed
+/// in with a password or federated through Google/Apple.
+/// </summary>
+/// <remarks>
+/// NO WARRANTY. Trax auth is plumbing, not a security product. You are solely
+/// responsible for securing systems that use it. See SECURITY-DISCLAIMER.md.
+/// <para>Claim handling, beyond the base <see cref="DefaultJwtPrincipalResolver"/>:</para>
+/// <list type="bullet">
+/// <item><c>cognito:groups</c> claims merge into <see cref="TraxPrincipal.Roles"/>.</item>
+/// <item><c>cognito:username</c> participates in display-name fallback.</item>
+/// <item><c>identities</c> JSON array is parsed; the primary provider's
+/// <c>providerName</c> surfaces on the principal as the synthetic
+/// <see cref="CognitoDefaults.IdentityProvider"/> claim.</item>
+/// <item>Native users (no <c>identities</c> claim) get
+/// <see cref="CognitoDefaults.IdentityProvider"/> = <c>cognito</c>.</item>
+/// <item>The <see cref="TraxPrincipal.PrincipalType"/> discriminator is
+/// <see cref="CognitoDefaults.PrincipalType"/>.</item>
+/// </list>
+/// <para>
+/// Apple's second-and-later logins omit <c>email</c>; this resolver tolerates
+/// missing email (the principal is still valid because <c>sub</c> is present).
+/// Hosts that need the email persisted across logins should look it up from
+/// their own database keyed by <c>sub</c>.
+/// </para>
+/// </remarks>
+public sealed class CognitoJwtPrincipalResolver : ITraxPrincipalResolver<JwtTokenInput>
+{
+    private static readonly string[] SubjectClaimTypes = ["sub", ClaimTypes.NameIdentifier];
+
+    private static readonly string[] DisplayNameClaimTypes =
+    [
+        "name",
+        ClaimTypes.Name,
+        "preferred_username",
+        CognitoDefaults.CognitoUsername,
+        CognitoDefaults.Email,
+        ClaimTypes.Email,
+    ];
+
+    private static readonly HashSet<string> ReservedClaimTypes = new(StringComparer.Ordinal)
+    {
+        "sub",
+        ClaimTypes.NameIdentifier,
+        "name",
+        ClaimTypes.Name,
+        "preferred_username",
+        ClaimTypes.Role,
+        "role",
+        "roles",
+        CognitoDefaults.CognitoUsername,
+        CognitoDefaults.CognitoGroups,
+        CognitoDefaults.Identities,
+        TraxAuthClaimTypes.PrincipalId,
+        TraxAuthClaimTypes.PrincipalType,
+    };
+
+    /// <inheritdoc />
+    public ValueTask<TraxPrincipal?> ResolveAsync(JwtTokenInput input, CancellationToken ct)
+    {
+        ArgumentNullException.ThrowIfNull(input);
+
+        var id = FirstClaim(input.Principal, SubjectClaimTypes);
+        if (string.IsNullOrWhiteSpace(id))
+            return new ValueTask<TraxPrincipal?>((TraxPrincipal?)null);
+
+        var displayName = FirstClaim(input.Principal, DisplayNameClaimTypes) ?? id;
+
+        var roles = input
+            .Principal.FindAll(ClaimTypes.Role)
+            .Concat(input.Principal.FindAll("role"))
+            .Concat(input.Principal.FindAll("roles"))
+            .Concat(input.Principal.FindAll(CognitoDefaults.CognitoGroups))
+            .Select(c => c.Value)
+            .Where(v => !string.IsNullOrWhiteSpace(v))
+            .Distinct(StringComparer.Ordinal)
+            .ToArray();
+
+        var claims = new Dictionary<string, string>(StringComparer.Ordinal);
+        foreach (var c in input.Principal.Claims)
+        {
+            if (ReservedClaimTypes.Contains(c.Type))
+                continue;
+            // First write wins to avoid duplicate-key throws while preserving
+            // the resolver's first-occurrence semantics.
+            claims.TryAdd(c.Type, c.Value);
+        }
+
+        var identityProvider = ExtractIdentityProvider(input.Principal);
+        claims[CognitoDefaults.IdentityProvider] = identityProvider;
+
+        var principal = new TraxPrincipal(
+            Id: id,
+            DisplayName: displayName,
+            Roles: roles,
+            Claims: claims.Count > 0 ? claims : null,
+            PrincipalType: CognitoDefaults.PrincipalType
+        );
+        return new ValueTask<TraxPrincipal?>(principal);
+    }
+
+    /// <summary>
+    /// Returns the primary federated provider name from the
+    /// <c>identities</c> claim, or <c>"cognito"</c> for native users.
+    /// Surface defaults that don't break callers: malformed JSON, an empty
+    /// array, or a missing <c>providerName</c> field all resolve to
+    /// <c>"cognito"</c>.
+    /// </summary>
+    internal static string ExtractIdentityProvider(ClaimsPrincipal principal)
+    {
+        var identitiesClaim = principal.FindFirst(CognitoDefaults.Identities);
+        if (identitiesClaim is null)
+            return CognitoDefaults.PrincipalType;
+
+        try
+        {
+            using var doc = JsonDocument.Parse(WrapIfNeeded(identitiesClaim.Value));
+            if (doc.RootElement.ValueKind != JsonValueKind.Array)
+                return CognitoDefaults.PrincipalType;
+
+            // Prefer the entry flagged "primary": true; otherwise pick the
+            // first entry. Cognito normally only emits one identity per
+            // federated user, but the format is an array.
+            JsonElement? chosen = null;
+            foreach (var entry in doc.RootElement.EnumerateArray())
+            {
+                if (entry.ValueKind != JsonValueKind.Object)
+                    continue;
+                if (chosen is null)
+                    chosen = entry;
+                if (entry.TryGetProperty("primary", out var primary) && IsTruthy(primary))
+                {
+                    chosen = entry;
+                    break;
+                }
+            }
+
+            if (
+                chosen is { } picked
+                && picked.TryGetProperty("providerName", out var providerName)
+                && providerName.ValueKind == JsonValueKind.String
+                && providerName.GetString() is { Length: > 0 } name
+            )
+            {
+                return name;
+            }
+        }
+        catch (JsonException)
+        {
+            // Fall through.
+        }
+
+        return CognitoDefaults.PrincipalType;
+    }
+
+    /// <summary>
+    /// Cognito serializes the <c>identities</c> claim as a JSON array when
+    /// emitted by the user pool, but some downstream serializers (mobile
+    /// SDKs, test fixtures) re-emit the same data wrapped in an outer object
+    /// or stringified. Accept both shapes.
+    /// </summary>
+    private static string WrapIfNeeded(string raw)
+    {
+        var trimmed = raw.TrimStart();
+        return trimmed.StartsWith('[') ? raw : "[" + raw + "]";
+    }
+
+    private static bool IsTruthy(JsonElement element) =>
+        element.ValueKind switch
+        {
+            JsonValueKind.True => true,
+            JsonValueKind.String => string.Equals(
+                element.GetString(),
+                "true",
+                StringComparison.OrdinalIgnoreCase
+            ),
+            _ => false,
+        };
+
+    private static string? FirstClaim(ClaimsPrincipal principal, IEnumerable<string> claimTypes)
+    {
+        foreach (var type in claimTypes)
+        {
+            var value = principal.FindFirst(type)?.Value;
+            if (!string.IsNullOrWhiteSpace(value))
+                return value;
+        }
+        return null;
+    }
+}

src/Trax.Api.Auth.Jwt.Cognito/CognitoTokenUse.cs

@@ -0,0 +1,22 @@
+namespace Trax.Api.Auth.Jwt.Cognito;
+
+/// <summary>
+/// Selects which Cognito token shapes <see cref="CognitoJwtBuilderExtensions.UseCognito"/>
+/// configures the validator to accept.
+/// </summary>
+/// <remarks>
+/// Cognito issues both ID tokens (audience in <c>aud</c>) and access tokens
+/// (audience in <c>client_id</c>). The default validator wires acceptance
+/// of either, since most apps mix both depending on the call site.
+/// </remarks>
+public enum CognitoTokenUse
+{
+    /// <summary>Accept Cognito ID tokens (audience claim: <c>aud</c>).</summary>
+    Id,
+
+    /// <summary>Accept Cognito access tokens (audience claim: <c>client_id</c>).</summary>
+    Access,
+
+    /// <summary>Accept both shapes. The default.</summary>
+    IdAndAccess,
+}