mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Security/Authorization/GatewayGrpcAuthorizationInterceptor.cs

using Grpc.Core;
using Grpc.Core.Interceptors;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Metrics;
using ZB.MOM.WW.MxGateway.Server.Security.Authentication;

// The handler pushes the gateway's constraint-bearing identity; alias away the shared library's
// ApiKeyIdentity so the unqualified name resolves to the gateway type.
using ApiKeyIdentity = ZB.MOM.WW.MxGateway.Server.Security.Authentication.ApiKeyIdentity;

namespace ZB.MOM.WW.MxGateway.Server.Security.Authorization;

public sealed class GatewayGrpcAuthorizationInterceptor(
    IApiKeyVerifier apiKeyVerifier,
    GatewayGrpcScopeResolver scopeResolver,
    IGatewayRequestIdentityAccessor identityAccessor,
    IOptions<GatewayOptions> options,
    ApiKeyFailureLimiter failureLimiter,
    GatewayMetrics metrics) : Interceptor
{
    // Generated key ids are far shorter; the cap only exists to stop an invented "key id" of
    // arbitrary length from becoming a limiter partition.
    private const int MaxKeyIdLength = 64;

    /// <inheritdoc />
    public override async Task<TResponse> UnaryServerHandler<TRequest, TResponse>(
        TRequest request,
        ServerCallContext context,
        UnaryServerMethod<TRequest, TResponse> continuation)
    {
        ApiKeyIdentity? identity = await AuthenticateAndAuthorizeAsync(request, context).ConfigureAwait(false);
        IDisposable? identityScope = identity is null ? null : identityAccessor.Push(identity);
        using (identityScope)
        {
            return await continuation(request, context).ConfigureAwait(false);
        }
    }

    /// <inheritdoc />
    public override async Task ServerStreamingServerHandler<TRequest, TResponse>(
        TRequest request,
        IServerStreamWriter<TResponse> responseStream,
        ServerCallContext context,
        ServerStreamingServerMethod<TRequest, TResponse> continuation)
    {
        ApiKeyIdentity? identity = await AuthenticateAndAuthorizeAsync(request, context).ConfigureAwait(false);
        IDisposable? identityScope = identity is null ? null : identityAccessor.Push(identity);
        using (identityScope)
        {
            await continuation(request, responseStream, context).ConfigureAwait(false);
        }
    }

    /// <summary>Authenticates the API key and authorizes the RPC call by required scope.</summary>
    /// <typeparam name="TRequest">Request message type.</typeparam>
    /// <param name="request">Request payload.</param>
    /// <param name="context">RPC server call context.</param>
    /// <returns>Authenticated API key identity, or null if authentication is disabled.</returns>
    private async Task<ApiKeyIdentity?> AuthenticateAndAuthorizeAsync<TRequest>(
        TRequest request,
        ServerCallContext context)
        where TRequest : class
    {
        if (options.Value.Authentication.Mode == AuthenticationMode.Disabled)
        {
            return null;
        }

        string? authorizationHeader = context.RequestHeaders.GetValue("authorization");

        // Short-circuit a partition that has already failed too many times inside the sliding window
        // BEFORE the verification store read, so online guessing cannot spend a SQLite read per
        // attempt. The partition is the composite (transport peer, key id) — never the key id alone,
        // which is public and would let any peer deny a key it does not own — plus a cross-peer
        // aggregate for the key id. An over-limit state still admits one probe per interval, so the
        // holder of the correct secret always reaches the verifier and resets the state.
        // ResourceExhausted signals throttling without revealing whether any secret was valid.
        //
        // The check hands back the storage key it resolved so the reset below reuses it instead of
        // rebuilding the composite (peer, key id) string a second time on every successful RPC.
        ApiKeyThrottlePartition throttlePartition = ResolveThrottlePartition(authorizationHeader, context);
        ApiKeyThrottleDecision decision = failureLimiter.Check(
            throttlePartition,
            out ApiKeyFailureLimiter.PartitionResolution throttleResolution);
        if (decision is ApiKeyThrottleDecision.ThrottledByPeer or ApiKeyThrottleDecision.ThrottledByAggregate)
        {
            metrics.RecordAuthThrottled(
                decision == ApiKeyThrottleDecision.ThrottledByPeer ? "peer" : "aggregate");
            throw new RpcException(new Status(
                StatusCode.ResourceExhausted,
                "Too many authentication attempts. Try again later."));
        }

        // The shared verifier owns parse + pepper + lookup + revocation + constant-time compare,
        // returning a discriminated failure rather than throwing. Every authentication failure maps
        // to Unauthenticated with an opaque message; the client never learns which stage failed.
        ApiKeyVerification verification = await apiKeyVerifier
            .VerifyAsync(authorizationHeader ?? string.Empty, context.CancellationToken)
            .ConfigureAwait(false);

        if (!verification.Succeeded || verification.Identity is null)
        {
            failureLimiter.RecordFailure(throttlePartition);
            throw new RpcException(new Status(
                StatusCode.Unauthenticated,
                "Missing or invalid API key."));
        }

        // Successful authentication clears both limiter layers, so a legitimate client that
        // fat-fingered a few attempts is not penalised once it recovers — and, because the check
        // above admits a probe rather than blocking absolutely, this reset stays reachable while the
        // key is under an active spray.
        failureLimiter.Reset(throttlePartition, throttleResolution);

        ApiKeyIdentity identity = GatewayApiKeyIdentityMapper.ToGatewayIdentity(verification.Identity);

        string requiredScope = scopeResolver.ResolveRequiredScope(request);
        if (!identity.Scopes.Contains(requiredScope))
        {
            throw new RpcException(new Status(
                StatusCode.PermissionDenied,
                $"API key is missing required scope '{requiredScope}'."));
        }

        return identity;
    }

    // Resolves the failure-limiter partition: always the transport peer, plus the presented key id
    // when — and only when — the token is validly shaped. Both halves are unauthenticated input, so
    // the peer half is what keeps a throttle bound to the address that earned it.
    private static ApiKeyThrottlePartition ResolveThrottlePartition(
        string? authorizationHeader,
        ServerCallContext context)
    {
        return new ApiKeyThrottlePartition(context.Peer, TryResolveKeyId(authorizationHeader));
    }

    // Token shape mxgw_<keyId>_<secret>: the key id is the second underscore-delimited segment (the
    // secret may itself contain underscores, which does not affect the key id). The shape is checked
    // before a key-id partition is minted so a spray of invented tokens cannot mint one tracked
    // partition each and flush the limiter's bounded map. Anything that fails the check
    // falls back to the sender's transport-peer partition.
    //
    // Internal rather than private so the parse rules can be pinned directly by test; the shape
    // check is a security boundary and is worth asserting without routing every case
    // through a full RPC.
    /// <summary>
    /// Extracts the API key id from a bearer authorization header.
    /// </summary>
    /// <param name="authorizationHeader">The raw <c>authorization</c> header value, or <see langword="null"/>.</param>
    /// <returns>The key id when the header has the <c>mxgw_&lt;keyId&gt;_&lt;secret&gt;</c> shape; otherwise <see langword="null"/>.</returns>
    internal static string? TryResolveKeyId(string? authorizationHeader)
    {
        if (string.IsNullOrWhiteSpace(authorizationHeader))
        {
            return null;
        }

        ReadOnlySpan<char> header = authorizationHeader.AsSpan().Trim();
        const string bearer = "Bearer ";
        ReadOnlySpan<char> token = header.StartsWith(bearer, StringComparison.OrdinalIgnoreCase)
            ? header[bearer.Length..].Trim()
            : header;

        // Scanned rather than split: this runs on every authenticated RPC, and Split would copy the
        // token out of the header and allocate an array plus a string per segment to reach a key id
        // that is then usually a dictionary-lookup miss. Two IndexOf scans reach the same answer and
        // allocate only the key id itself.
        const string prefix = AuthStoreServiceCollectionExtensions.TokenPrefix;
        if (!token.StartsWith(prefix, StringComparison.Ordinal)
            || token.Length <= prefix.Length
            || token[prefix.Length] != '_')
        {
            // Guards the whole first segment, not just its start: the '_' immediately after the
            // prefix is what makes "mxgw" the entire segment, so "mxgwabc_..." is still rejected.
            return null;
        }

        ReadOnlySpan<char> afterPrefix = token[(prefix.Length + 1)..];
        int separator = afterPrefix.IndexOf('_');
        if (separator <= 0 || separator > MaxKeyIdLength)
        {
            // -1 is a token with no second separator (too few segments); 0 is an empty key id.
            return null;
        }

        // The third segment must be non-empty, which the split form expressed as parts[2].Length: it
        // ends at the NEXT separator, so a secret beginning with '_' fails the same way it always did.
        ReadOnlySpan<char> afterKeyId = afterPrefix[(separator + 1)..];
        if (afterKeyId.IsEmpty || afterKeyId[0] == '_')
        {
            return null;
        }

        return new string(afterPrefix[..separator]);
    }
}