perf(auth): allocation-free token parsing; single partition-key build per RPC

2026-08-15 12:23:25 -04:00
parent 3ff073d1ea
commit 88d38bb900
5 changed files with 221 additions and 25 deletions
@@ -213,7 +213,10 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
    // DashboardApiKeyManagementService.ValidateKeyId each restrict a key id to
    // char.IsAsciiLetterOrDigit || '.' || '-'. Key ids are never library-generated, so no path can
    // mint one containing '_'.
-    private static string? TryParseKeyId(string? authorizationHeader)
+    //
+    // Internal rather than private so the parse rules can be pinned directly by test: the guard's
+    // correctness depends on this returning the full key id.
+    internal static string? TryParseKeyId(string? authorizationHeader)
    {
        if (string.IsNullOrEmpty(authorizationHeader))
        {
@@ -226,15 +229,29 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
            ? header[bearer.Length..].Trim()
            : header;

-        string[] parts = token.ToString().Split('_');
-        if (parts.Length < 3
-            || !string.Equals(parts[0], TokenPrefix, StringComparison.Ordinal)
-            || parts[1].Length == 0)
+        // Scanned rather than split, for the same reason as the interceptor's copy: Split would
+        // allocate a token copy, an array and a string per segment on every cache miss to produce
+        // one key id. Two IndexOf scans allocate only that key id.
+        //
+        // The '_' checked immediately after the prefix is what makes "mxgw" the whole first segment
+        // (so "mxgwabc_..." is still rejected), and the second separator must exist because the
+        // split form required three segments — a token with no secret delimiter is not a key token.
+        if (!token.StartsWith(TokenPrefix, StringComparison.Ordinal)
+            || token.Length <= TokenPrefix.Length
+            || token[TokenPrefix.Length] != '_')
        {
            return null;
        }

-        return parts[1];
+        ReadOnlySpan<char> afterPrefix = token[(TokenPrefix.Length + 1)..];
+        int separator = afterPrefix.IndexOf('_');
+        if (separator <= 0)
+        {
+            // -1 is a token with no second separator; 0 is an empty key id.
+            return null;
+        }
+
+        return new string(afterPrefix[..separator]);
    }

    private void IndexCacheKey(string keyId, string cacheKey)
@@ -127,16 +127,31 @@ public sealed class ApiKeyFailureLimiter
    /// <summary>Decides whether an authentication attempt may reach the verifier.</summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    /// <returns>The admission decision for this attempt.</returns>
-    public ApiKeyThrottleDecision Check(ApiKeyThrottlePartition partition)
+    public ApiKeyThrottleDecision Check(ApiKeyThrottlePartition partition) => Check(partition, out _);
+
+    /// <summary>
+    /// Decides whether an authentication attempt may reach the verifier, handing back the storage
+    /// key it resolved so a caller that goes on to <see cref="Reset(ApiKeyThrottlePartition, PartitionResolution)"/>
+    /// the same request does not resolve — and rebuild the composite key string — a second time.
+    /// </summary>
+    /// <param name="partition">The throttle partition derived from the request.</param>
+    /// <param name="resolution">
+    /// The resolved storage key, or <see cref="PartitionResolution.Unresolved"/> when the limiter is
+    /// disabled and never resolved one.
+    /// </param>
+    /// <returns>The admission decision for this attempt.</returns>
+    internal ApiKeyThrottleDecision Check(ApiKeyThrottlePartition partition, out PartitionResolution resolution)
    {
        string peer = RequirePeer(partition);
        if (_limit <= 0)
        {
+            resolution = PartitionResolution.Unresolved;
            return ApiKeyThrottleDecision.Allowed;
        }

        long now = _clock.GetUtcNow().UtcTicks;
        (string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: false);
+        resolution = new PartitionResolution(partitionKey, effectiveKeyId);

        WindowState? peerState = _partitions.TryGetValue(partitionKey, out WindowState? tracked) ? tracked : null;
        WindowState? aggregateState = null;
@@ -221,10 +236,25 @@ public sealed class ApiKeyFailureLimiter

    /// <summary>Clears both limiter layers for the partition after a successful verification.</summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
-    public void Reset(ApiKeyThrottlePartition partition)
+    public void Reset(ApiKeyThrottlePartition partition) => Reset(partition, PartitionResolution.Unresolved);
+
+    /// <summary>
+    /// Clears both limiter layers for the partition after a successful verification, reusing the
+    /// storage key <see cref="Check(ApiKeyThrottlePartition, out PartitionResolution)"/> already
+    /// resolved for this request.
+    /// </summary>
+    /// <param name="partition">The throttle partition derived from the request.</param>
+    /// <param name="resolution">
+    /// The resolution handed back by <c>Check</c>; <see cref="PartitionResolution.Unresolved"/>
+    /// resolves here instead. Reusing the check-time resolution is deliberate: it is the partition
+    /// this request was admitted against, so the reset clears exactly what the check consulted.
+    /// </param>
+    internal void Reset(ApiKeyThrottlePartition partition, PartitionResolution resolution)
    {
        string peer = RequirePeer(partition);
-        (string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: false);
+        (string partitionKey, string? effectiveKeyId) = resolution.IsResolved
+            ? (resolution.PartitionKey!, resolution.EffectiveKeyId)
+            : ResolvePartitionKey(peer, partition.KeyId, mint: false);

        // Clear only a partition this caller actually owns. When its key id was squeezed into the
        // address's shared fallback bucket by the per-peer cap, that bucket also holds failures
@@ -542,6 +572,25 @@ public sealed class ApiKeyFailureLimiter
        public long ProbeVersion;
    }

+    /// <summary>
+    /// A partition's resolved storage key, carried from the check to the reset of the same request.
+    /// The composite key is a fresh string per build, so resolving once per RPC rather than once per
+    /// call keeps the successful auth path (check, then reset) to a single allocation.
+    /// </summary>
+    /// <param name="PartitionKey">The storage key, or <see langword="null"/> when unresolved.</param>
+    /// <param name="EffectiveKeyId">
+    /// The key id that actually earned a partition, or <see langword="null"/> when the token carried
+    /// none or the per-peer cap collapsed it onto the transport-peer fallback.
+    /// </param>
+    internal readonly record struct PartitionResolution(string? PartitionKey, string? EffectiveKeyId)
+    {
+        /// <summary>Gets the sentinel for "not resolved yet"; the receiving call resolves it itself.</summary>
+        internal static PartitionResolution Unresolved => default;
+
+        /// <summary>Gets a value indicating whether this carries a resolved storage key.</summary>
+        internal bool IsResolved => PartitionKey is not null;
+    }
+
    /// <summary>A probe slot reservation: what to restore, and the stamp proving it is still ours.</summary>
    /// <param name="PreviousProbeAtTicks">The slot value replaced when the claim was made.</param>
    /// <param name="Version">The <see cref="WindowState.ProbeVersion"/> stamped by this claim.</param>
@@ -77,8 +77,13 @@ public sealed class GatewayGrpcAuthorizationInterceptor(
        // 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);
+        ApiKeyThrottleDecision decision = failureLimiter.Check(
+            throttlePartition,
+            out ApiKeyFailureLimiter.PartitionResolution throttleResolution);
        if (decision is ApiKeyThrottleDecision.ThrottledByPeer or ApiKeyThrottleDecision.ThrottledByAggregate)
        {
            metrics.RecordAuthThrottled(
@@ -107,7 +112,7 @@ public sealed class GatewayGrpcAuthorizationInterceptor(
        // 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);
+        failureLimiter.Reset(throttlePartition, throttleResolution);

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

@@ -137,7 +142,11 @@ public sealed class GatewayGrpcAuthorizationInterceptor(
    // 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 (SEC-32). Anything that fails the check
    // falls back to the sender's transport-peer partition.
-    private static string? TryResolveKeyId(string? authorizationHeader)
+    //
+    // Internal rather than private so the parse rules can be pinned directly by test; the shape
+    // check is a security boundary (SEC-32) and is worth asserting without routing every case
+    // through a full RPC.
+    internal static string? TryResolveKeyId(string? authorizationHeader)
    {
        if (string.IsNullOrWhiteSpace(authorizationHeader))
        {
@@ -150,22 +159,36 @@ public sealed class GatewayGrpcAuthorizationInterceptor(
            ? header[bearer.Length..].Trim()
            : header;

-        string[] parts = token.ToString().Split('_');
-        if (parts.Length < 3)
+        // 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;
        }

-        if (!string.Equals(parts[0], AuthStoreServiceCollectionExtensions.TokenPrefix, StringComparison.Ordinal))
-        {
-            return null;
-        }
-
-        if (parts[1].Length == 0 || parts[1].Length > MaxKeyIdLength || parts[2].Length == 0)
-        {
-            return null;
-        }
-
-        return parts[1];
+        return new string(afterPrefix[..separator]);
    }
 }
@@ -196,6 +196,55 @@ public sealed class CachingApiKeyVerifierTests
        Assert.Equal(2, inner.MarkUsedCount);
    }

+    /// <summary>
+    /// Pins the header parse that arms the revoke-vs-in-flight generation guard. It runs on every
+    /// cache miss and is written to allocate only the returned key id; these cases hold it to the
+    /// rules the original <c>Split('_')</c> form applied. Note this parse is deliberately laxer than
+    /// the interceptor's partition parse — it applies no key-id length cap and does not require a
+    /// non-empty third segment — because a wrong <em>id</em> here would disarm the guard, whereas
+    /// an over-long one merely fails to match any generation.
+    /// </summary>
+    /// <param name="authorizationHeader">The presented header value.</param>
+    /// <param name="expected">The key id the parse must yield, or <see langword="null"/>.</param>
+    [Theory]
+    [InlineData(null, null)]
+    [InlineData("", null)]
+    [InlineData("   ", null)]
+    [InlineData("Bearer", null)]
+    [InlineData("Bearer ", null)]
+    [InlineData("Bearer mxgw_operator01_super-secret", "operator01")]
+    [InlineData("bearer mxgw_operator01_super-secret", "operator01")]
+    [InlineData("  Bearer   mxgw_operator01_super-secret  ", "operator01")]
+    [InlineData("mxgw_operator01_super-secret", "operator01")]
+    [InlineData("Bearer mxgw_abc_sec_ret", "abc")]
+    [InlineData("Bearer mxgw_a_b_c", "a")]
+
+    // Laxer than the interceptor: an empty or absent third segment still yields the key id.
+    [InlineData("Bearer mxgw_abc_", "abc")]
+    [InlineData("Bearer mxgw_abc__secret", "abc")]
+    [InlineData("Bearer mxgw_abc", null)]
+    [InlineData("Bearer mxgwabcsecret", null)]
+    [InlineData("Bearer mxgw__secret", null)]
+    [InlineData("Bearer _mxgw_abc_secret", null)]
+    [InlineData("Bearer MXGW_abc_secret", null)]
+    [InlineData("Bearer xmxgw_abc_secret", null)]
+    [InlineData("Bearer mxgw", null)]
+    [InlineData("Bearer mxgw_", null)]
+    [InlineData("Bearer ___", null)]
+    public void TryParseKeyId_MatchesTokenShapeRules(string? authorizationHeader, string? expected)
+    {
+        Assert.Equal(expected, CachingApiKeyVerifier.TryParseKeyId(authorizationHeader));
+    }
+
+    /// <summary>The guard parse applies no key-id length cap: an over-long id is still returned whole.</summary>
+    [Fact]
+    public void TryParseKeyId_LongKeyId_ReturnedWhole()
+    {
+        string keyId = new('a', 65);
+
+        Assert.Equal(keyId, CachingApiKeyVerifier.TryParseKeyId($"Bearer mxgw_{keyId}_secret"));
+    }
+
    private static MemoryCache NewCache() => new(new MemoryCacheOptions());

    private static ApiKeyVerification Success(string keyId) => new(
@@ -659,6 +659,64 @@ public sealed class GatewayGrpcAuthorizationInterceptorTests
        }
    }

+    /// <summary>
+    /// Pins the token shape the limiter partition is minted from. The parser is on every
+    /// authenticated request, so it is written to allocate only the returned key id; these cases
+    /// hold it to the rules the original <c>Split('_')</c> form applied — literal <c>mxgw</c> first
+    /// segment, non-empty second and third segments, and a bounded key id — including the ones a
+    /// hand-rolled scanner is most likely to drift on (a key id read from the second segment even
+    /// when the secret itself contains separators).
+    /// </summary>
+    /// <param name="authorizationHeader">The presented header value.</param>
+    /// <param name="expected">The key id the parse must yield, or <see langword="null"/>.</param>
+    [Theory]
+    [InlineData(null, null)]
+    [InlineData("", null)]
+    [InlineData("   ", null)]
+    [InlineData("Bearer", null)]
+    [InlineData("Bearer ", null)]
+    [InlineData("Bearer mxgw_abc_secret", "abc")]
+    [InlineData("bearer mxgw_abc_secret", "abc")]
+    [InlineData("  Bearer   mxgw_abc_secret  ", "abc")]
+    [InlineData("mxgw_abc_secret", "abc")]
+
+    // The secret may carry separators of its own; the key id is still the second segment.
+    [InlineData("Bearer mxgw_abc_sec_ret", "abc")]
+    [InlineData("Bearer mxgw_a_b_c", "a")]
+    [InlineData("Bearer mxgw_abc_secret_", "abc")]
+
+    // Shape failures: no separators, too few segments, empty segments, wrong prefix.
+    [InlineData("Bearer mxgwabcsecret", null)]
+    [InlineData("Bearer mxgw_abc", null)]
+    [InlineData("Bearer mxgw_abc_", null)]
+    [InlineData("Bearer mxgw_abc__secret", null)]
+    [InlineData("Bearer mxgw__secret", null)]
+    [InlineData("Bearer _mxgw_abc_secret", null)]
+    [InlineData("Bearer _abc_secret", null)]
+    [InlineData("Bearer MXGW_abc_secret", null)]
+    [InlineData("Bearer xmxgw_abc_secret", null)]
+    [InlineData("Bearer mxgw", null)]
+    [InlineData("Bearer mxgw_", null)]
+    [InlineData("Bearer ___", null)]
+    public void TryResolveKeyId_MatchesTokenShapeRules(string? authorizationHeader, string? expected)
+    {
+        Assert.Equal(expected, GatewayGrpcAuthorizationInterceptor.TryResolveKeyId(authorizationHeader));
+    }
+
+    /// <summary>
+    /// The key-id length cap is what stops an invented id of arbitrary length from becoming a
+    /// limiter partition, so the boundary is pinned on both sides.
+    /// </summary>
+    [Fact]
+    public void TryResolveKeyId_HonoursKeyIdLengthCap()
+    {
+        string atCap = new('a', 64);
+        string overCap = new('a', 65);
+
+        Assert.Equal(atCap, GatewayGrpcAuthorizationInterceptor.TryResolveKeyId($"Bearer mxgw_{atCap}_secret"));
+        Assert.Null(GatewayGrpcAuthorizationInterceptor.TryResolveKeyId($"Bearer mxgw_{overCap}_secret"));
+    }
+
    private static MxAccessGatewayService CreateService(
        ISessionManager sessionManager,
        IGatewayRequestIdentityAccessor identityAccessor)