mirror of
https://github.com/okieselbach/AutopilotMonitor.git
synced 2026-10-24 04:32:32 +00:00
perf(backend): cache Graph app token + stream-parse V2 ingest
Ingest hot path + Graph auth latency. IngestTelemetry is the #1 backend endpoint (~35.6k/day); its p99 tail was dominated not by parsing but by an uncached Graph token: GraphTokenService minted a fresh client-credentials token (500-1900ms POST) on every device-validation cache miss and threw it away. Cache it per tenant. GraphTokenService: - Per-tenant token cache keyed on the validated tenantId; TTL from the token's own expires_in minus a 5min skew margin, clamped to [60s, 55min] (expires_in was previously ignored). Only successful tokens are cached. - Per-tenant SemaphoreSlim stampede gate (bounded wait) so an enrollment wave does not fan a burst of identical token POSTs at AAD. - InvalidateTenant() drops the token + consent-status entries so the consent/Graph-permission fresh-read contract still holds. - Per-tenant invalidation generation closes the write-after-invalidate race: a POST in flight when InvalidateTenant runs tags its entry with the pre-invalidation generation, which readers reject. GraphFeatureDetector: - InvalidateTenant() now also invalidates the underlying token cache, and applies the same generation guard to its own parsed-roles cache. - Cross-layer order is load-bearing: invalidate the token layer FIRST, then bump the detector generation, so a concurrent acquire cannot pair a new detector generation with a stale token read. IngestTelemetryFunction: - Deserialize the request body directly from the buffered MemoryStream via JsonTextReader instead of materializing a full UTF-16 string + ToArray() copy (peak ~4N -> ~1N). Cap still applies before parsing; same parse semantics. Pure allocation hygiene on the hottest endpoint. Tests: token caching, cross-tenant isolation, failure-not-cached, stampede coalescing (barrier-gated), in-flight write-after-invalidate, cross-layer invalidation, and cross-layer ordering -- each verified red without its fix. Parse hygiene: PayloadJson stays a raw string, malformed -> 400. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
+218
@@ -313,6 +313,145 @@ public class GraphFeatureDetectorTests
|
||||
Assert.Equal(2, tokenService.CallCount);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task InvalidateTenant_drops_underlying_token_cache_so_post_grant_refresh_sees_new_roles()
|
||||
{
|
||||
// Regression for the exact consent-propagation bug: a token acquired BEFORE the grant
|
||||
// carries an empty roles claim, and GraphTokenService caches that raw JWT (~55 min). After
|
||||
// the admin grants the role and hits Refresh (→ InvalidateTenant), the detector must mint a
|
||||
// FRESH token and report the new role — NOT re-parse the stale cached JWT. This exercises
|
||||
// the cross-layer invalidation wiring with a REAL GraphTokenService (not the stub), so a
|
||||
// future change that forgets to clear the token cache fails here.
|
||||
const string tenantId = "66666666-6666-6666-6666-666666666666";
|
||||
|
||||
var handler = new SwappableTokenHandler(
|
||||
BuildJwt(Array.Empty<Claim>(), expires: DateTime.UtcNow.AddHours(1))); // pre-grant: no roles
|
||||
var cache = new MemoryCache(new MemoryCacheOptions());
|
||||
var tokenService = new GraphTokenService(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
new SingleHandlerFactory(handler),
|
||||
cache,
|
||||
new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build());
|
||||
var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true });
|
||||
var detector = new GraphFeatureDetector(
|
||||
tokenService, cache, NullLogger<GraphFeatureDetector>.Instance, telemetry);
|
||||
|
||||
// Pre-grant read: empty roles, and the raw JWT is now cached in GraphTokenService.
|
||||
var before = await detector.GetSnapshotAsync(tenantId);
|
||||
Assert.Empty(before.GrantedRoles);
|
||||
Assert.Equal(1, handler.Requests);
|
||||
|
||||
// Admin grants the role; the wire now returns a JWT WITH it.
|
||||
handler.Body = BuildJwt(
|
||||
new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) },
|
||||
expires: DateTime.UtcNow.AddHours(1));
|
||||
|
||||
// The fresh-read trigger. Without the cross-layer wiring the stale empty-roles JWT would
|
||||
// still be served and this would keep reporting "not granted".
|
||||
detector.InvalidateTenant(tenantId);
|
||||
|
||||
var after = await detector.GetSnapshotAsync(tenantId);
|
||||
Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, after.GrantedRoles);
|
||||
Assert.Equal(2, handler.Requests); // fresh token minted, not served from cache
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task InvalidateTenant_during_in_flight_acquire_does_not_cache_stale_roles()
|
||||
{
|
||||
// Write-after-invalidate race on the DETECTOR's own cache: an acquire that started before
|
||||
// InvalidateTenant must not leave a stale parsed-roles ctx behind. The token acquire is
|
||||
// parked in flight; the admin grants + Refreshes (InvalidateTenant) while it is parked; on
|
||||
// release the detector writes a (stale-generation-tagged) ctx. A fresh read MUST reject it
|
||||
// and re-acquire, seeing the newly granted role.
|
||||
const string tenantId = "77777777-7777-7777-7777-777777777777";
|
||||
var tokenService = new GatedStubGraphTokenService(
|
||||
firstToken: BuildJwt(Array.Empty<Claim>(), expires: DateTime.UtcNow.AddHours(1)), // pre-grant: no roles
|
||||
laterToken: BuildJwt(new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) },
|
||||
expires: DateTime.UtcNow.AddHours(1)));
|
||||
var cache = new MemoryCache(new MemoryCacheOptions());
|
||||
var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true });
|
||||
var detector = new GraphFeatureDetector(
|
||||
tokenService, cache, NullLogger<GraphFeatureDetector>.Instance, telemetry);
|
||||
|
||||
// D1 starts; its token acquire parks in flight (genAtStart captured at method entry).
|
||||
var inflight = detector.GetSnapshotAsync(tenantId);
|
||||
await WaitUntilAsync(() => tokenService.CallCount >= 1, TimeSpan.FromSeconds(2));
|
||||
|
||||
// Admin grants + Refresh WHILE D1 is mid-acquire.
|
||||
detector.InvalidateTenant(tenantId);
|
||||
|
||||
// D1 completes → writes a CachedVerdict tagged with the pre-invalidation generation.
|
||||
tokenService.Release();
|
||||
var before = await inflight;
|
||||
Assert.Empty(before.GrantedRoles); // D1 itself saw the pre-grant token
|
||||
|
||||
// Fresh read must reject the stale-tagged entry (gen mismatch) and re-acquire → new role.
|
||||
var after = await detector.GetSnapshotAsync(tenantId);
|
||||
Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, after.GrantedRoles);
|
||||
Assert.Equal(2, tokenService.CallCount); // a real re-acquire happened (no stale cache hit)
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task InvalidateTenant_token_layer_first_prevents_stale_detector_write_in_overlap()
|
||||
{
|
||||
// Cross-layer ordering race: InvalidateTenant must invalidate the token layer BEFORE bumping
|
||||
// its own generation. We pause INSIDE the token-layer invalidation and slip a concurrent
|
||||
// GetSnapshotAsync into the gap. With the correct (token-first) order, that overlapping read
|
||||
// cannot leave an accepted stale ctx behind; with the buggy (gen-first) order it would, and
|
||||
// the post-refresh read below would still report the old (no-roles) verdict.
|
||||
const string tenantId = "88888888-8888-8888-8888-888888888888";
|
||||
var handler = new SwappableTokenHandler(
|
||||
BuildJwt(Array.Empty<Claim>(), expires: DateTime.UtcNow.AddHours(1))); // pre-grant: no roles
|
||||
var cache = new MemoryCache(new MemoryCacheOptions());
|
||||
|
||||
var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
using var release = new System.Threading.ManualResetEventSlim(false);
|
||||
var tokenService = new BarrieredInvalidateTokenService(
|
||||
new SingleHandlerFactory(handler), cache,
|
||||
onInvalidate: () => { entered.TrySetResult(); release.Wait(TimeSpan.FromSeconds(5)); });
|
||||
var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true });
|
||||
var detector = new GraphFeatureDetector(
|
||||
tokenService, cache, NullLogger<GraphFeatureDetector>.Instance, telemetry);
|
||||
|
||||
// Prime: cache the old (no-roles) token + parsed ctx on both layers.
|
||||
var primed = await detector.GetSnapshotAsync(tenantId);
|
||||
Assert.Empty(primed.GrantedRoles);
|
||||
|
||||
// Admin grants the role on the wire.
|
||||
handler.Body = BuildJwt(
|
||||
new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) },
|
||||
expires: DateTime.UtcNow.AddHours(1));
|
||||
|
||||
// Admin Refresh → InvalidateTenant, blocked inside the token-layer invalidation step.
|
||||
var invalidateTask = Task.Run(() => detector.InvalidateTenant(tenantId));
|
||||
await entered.Task;
|
||||
|
||||
// A concurrent read slips into the gap, then we let InvalidateTenant finish.
|
||||
var overlap = await detector.GetSnapshotAsync(tenantId);
|
||||
release.Set();
|
||||
await invalidateTask;
|
||||
|
||||
// Post-refresh read MUST reflect the granted role — no stale entry survived the overlap.
|
||||
var afterRefresh = await detector.GetSnapshotAsync(tenantId);
|
||||
Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, afterRefresh.GrantedRoles);
|
||||
_ = overlap; // the overlapping read itself may legitimately observe pre-refresh state
|
||||
}
|
||||
|
||||
private static async Task WaitUntilAsync(Func<bool> condition, TimeSpan timeout)
|
||||
{
|
||||
var deadline = DateTime.UtcNow + timeout;
|
||||
while (!condition() && DateTime.UtcNow < deadline)
|
||||
{
|
||||
await Task.Delay(10);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Test seam: token service stub + detector factory ──────────────────────────────
|
||||
|
||||
/// <summary>
|
||||
@@ -349,6 +488,85 @@ public class GraphFeatureDetectorTests
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Token-service stub whose FIRST acquire parks until <see cref="Release"/> is called (modelling
|
||||
/// a token fetch in flight while an InvalidateTenant runs), returning <c>firstToken</c>; every
|
||||
/// later acquire returns <c>laterToken</c> immediately. Drives the detector's in-flight
|
||||
/// write-after-invalidate test. CallCount is incremented BEFORE parking and read thread-safely.
|
||||
/// </summary>
|
||||
private sealed class GatedStubGraphTokenService : GraphTokenService
|
||||
{
|
||||
private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private readonly string _firstToken;
|
||||
private readonly string _laterToken;
|
||||
private int _calls;
|
||||
|
||||
public GatedStubGraphTokenService(string firstToken, string laterToken)
|
||||
: base(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
new NoopHttpClientFactory(),
|
||||
new MemoryCache(new MemoryCacheOptions()),
|
||||
new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client-id",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build())
|
||||
{
|
||||
_firstToken = firstToken;
|
||||
_laterToken = laterToken;
|
||||
}
|
||||
|
||||
public int CallCount => System.Threading.Volatile.Read(ref _calls);
|
||||
public void Release() => _release.TrySetResult();
|
||||
|
||||
public override async Task<GraphTokenResult> GetAccessTokenAsync(string tenantId, CancellationToken ct = default)
|
||||
{
|
||||
var n = System.Threading.Interlocked.Increment(ref _calls);
|
||||
if (n == 1)
|
||||
{
|
||||
await _release.Task.ConfigureAwait(false);
|
||||
return GraphTokenResult.Success(_firstToken);
|
||||
}
|
||||
return GraphTokenResult.Success(_laterToken);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Real GraphTokenService (real token cache + generation + HTTP path) whose InvalidateTenant is
|
||||
/// wrapped with a test-controlled barrier, run BEFORE delegating to the base. Lets a test pause
|
||||
/// inside the token-layer invalidation and interleave a concurrent acquire, to prove the
|
||||
/// cross-layer ordering in GraphFeatureDetector.InvalidateTenant.
|
||||
/// </summary>
|
||||
private sealed class BarrieredInvalidateTokenService : GraphTokenService
|
||||
{
|
||||
private readonly Action _onInvalidate;
|
||||
|
||||
public BarrieredInvalidateTokenService(
|
||||
System.Net.Http.IHttpClientFactory factory, IMemoryCache cache, Action onInvalidate)
|
||||
: base(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
factory,
|
||||
cache,
|
||||
new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client-id",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build())
|
||||
{
|
||||
_onInvalidate = onInvalidate;
|
||||
}
|
||||
|
||||
public override void InvalidateTenant(string tenantId)
|
||||
{
|
||||
_onInvalidate();
|
||||
base.InvalidateTenant(tenantId);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Placeholder factory — the detector never causes the base GraphTokenService HTTP path to run.</summary>
|
||||
private sealed class NoopHttpClientFactory : System.Net.Http.IHttpClientFactory
|
||||
{
|
||||
|
||||
+91
@@ -47,3 +47,94 @@ internal sealed class StubHttpClientFactory : System.Net.Http.IHttpClientFactory
|
||||
public StubHttpClientFactory(StubHttpMessageHandler handler) => _handler = handler;
|
||||
public HttpClient CreateClient(string name) => new(_handler, disposeHandler: false);
|
||||
}
|
||||
|
||||
/// <summary>Factory that wires every client to one arbitrary <see cref="HttpMessageHandler"/>.</summary>
|
||||
internal sealed class SingleHandlerFactory : System.Net.Http.IHttpClientFactory
|
||||
{
|
||||
private readonly HttpMessageHandler _handler;
|
||||
public SingleHandlerFactory(HttpMessageHandler handler) => _handler = handler;
|
||||
public HttpClient CreateClient(string name) => new(_handler, disposeHandler: false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Returns a token response whose body can be swapped between calls — models AAD issuing a token
|
||||
/// with a different <c>roles</c> claim after an admin grants a permission. Counts requests so a
|
||||
/// test can assert a fresh acquire actually hit the wire (vs. being served from cache).
|
||||
/// </summary>
|
||||
internal sealed class SwappableTokenHandler : HttpMessageHandler
|
||||
{
|
||||
private int _requests;
|
||||
public string Body { get; set; }
|
||||
public int Requests => System.Threading.Volatile.Read(ref _requests);
|
||||
|
||||
public SwappableTokenHandler(string body) => Body = body;
|
||||
|
||||
protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
|
||||
{
|
||||
System.Threading.Interlocked.Increment(ref _requests);
|
||||
return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
|
||||
{
|
||||
Content = new StringContent($"{{\"access_token\":\"{Body}\",\"expires_in\":3599}}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Parks the FIRST request until <see cref="Release"/> is called (modelling a token POST that is
|
||||
/// in flight while an InvalidateTenant runs), then answers it with <c>firstBody</c>; every later
|
||||
/// request answers immediately with <c>laterBody</c> (the post-grant re-mint). Lets a test prove a
|
||||
/// stale in-flight write is rejected and a fresh read re-mints.
|
||||
/// </summary>
|
||||
internal sealed class InflightRaceHandler : HttpMessageHandler
|
||||
{
|
||||
private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private readonly string _firstBody;
|
||||
private readonly string _laterBody;
|
||||
private int _entered;
|
||||
|
||||
public InflightRaceHandler(string firstBody, string laterBody)
|
||||
{
|
||||
_firstBody = firstBody;
|
||||
_laterBody = laterBody;
|
||||
}
|
||||
|
||||
public int Entered => System.Threading.Volatile.Read(ref _entered);
|
||||
public void Release() => _release.TrySetResult();
|
||||
|
||||
protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
|
||||
{
|
||||
var n = System.Threading.Interlocked.Increment(ref _entered);
|
||||
if (n == 1)
|
||||
{
|
||||
await _release.Task.ConfigureAwait(false);
|
||||
}
|
||||
return new HttpResponseMessage(HttpStatusCode.OK)
|
||||
{
|
||||
Content = new StringContent(n == 1 ? _firstBody : _laterBody),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Parks every request until <see cref="Release"/> is called, counting how many reached the wire.
|
||||
/// Lets a stampede test prove that N simultaneous cache-misses collapse to a SINGLE in-flight POST
|
||||
/// — which a synchronous stub cannot show, because the first call can populate the cache before the
|
||||
/// others even start.
|
||||
/// </summary>
|
||||
internal sealed class GatedHttpMessageHandler : HttpMessageHandler
|
||||
{
|
||||
private readonly string _body;
|
||||
private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private int _entered;
|
||||
|
||||
public GatedHttpMessageHandler(string body) => _body = body;
|
||||
public int Entered => System.Threading.Volatile.Read(ref _entered);
|
||||
public void Release() => _release.TrySetResult();
|
||||
|
||||
protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
|
||||
{
|
||||
System.Threading.Interlocked.Increment(ref _entered);
|
||||
await _release.Task.ConfigureAwait(false);
|
||||
return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(_body) };
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Linq;
|
||||
using System.Net;
|
||||
using System.Threading.Tasks;
|
||||
using AutopilotMonitor.Functions.Security;
|
||||
using AutopilotMonitor.Functions.Tests.GraphResolution;
|
||||
using Microsoft.Extensions.Caching.Memory;
|
||||
using Microsoft.Extensions.Configuration;
|
||||
using Microsoft.Extensions.Logging.Abstractions;
|
||||
|
||||
namespace AutopilotMonitor.Functions.Tests;
|
||||
|
||||
/// <summary>
|
||||
/// Tests for the per-tenant app-token cache on <see cref="GraphTokenService"/>. The security
|
||||
/// contract: a successful client-credentials token is reused until shortly before expiry, keyed
|
||||
/// strictly per tenant (no cross-tenant reuse), and failures are never cached. The
|
||||
/// <see cref="StubHttpMessageHandler"/> records every outbound request so we can assert exactly how
|
||||
/// many token POSTs hit AAD.
|
||||
/// </summary>
|
||||
public class GraphTokenServiceTests
|
||||
{
|
||||
private const string TenantA = "11111111-1111-1111-1111-111111111111";
|
||||
private const string TenantB = "22222222-2222-2222-2222-222222222222";
|
||||
|
||||
private static (GraphTokenService svc, StubHttpMessageHandler handler) Create(StubHttpMessageHandler handler)
|
||||
{
|
||||
var config = new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build();
|
||||
|
||||
var svc = new GraphTokenService(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
new StubHttpClientFactory(handler),
|
||||
new MemoryCache(new MemoryCacheOptions()),
|
||||
config);
|
||||
|
||||
return (svc, handler);
|
||||
}
|
||||
|
||||
private static int TokenPosts(StubHttpMessageHandler handler, string tenantId)
|
||||
=> handler.Requests.Count(u => u.Contains($"/{tenantId}/oauth2/v2.0/token"));
|
||||
|
||||
[Fact]
|
||||
public async Task GetAccessTokenAsync_caches_token_across_calls()
|
||||
{
|
||||
var handler = new StubHttpMessageHandler()
|
||||
.When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}");
|
||||
var (svc, _) = Create(handler);
|
||||
|
||||
var first = await svc.GetAccessTokenAsync(TenantA);
|
||||
var second = await svc.GetAccessTokenAsync(TenantA);
|
||||
|
||||
Assert.Equal("tok-A", first.AccessToken);
|
||||
Assert.Equal("tok-A", second.AccessToken);
|
||||
// Second call is served from cache — only one POST ever reached AAD.
|
||||
Assert.Equal(1, TokenPosts(handler, TenantA));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task GetAccessTokenAsync_isolates_tokens_per_tenant()
|
||||
{
|
||||
// Cross-tenant isolation: tenant A's token must never be served for tenant B, and each
|
||||
// tenant caches independently.
|
||||
var handler = new StubHttpMessageHandler()
|
||||
.When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}")
|
||||
.When($"/{TenantB}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-B\",\"expires_in\":3599}");
|
||||
var (svc, _) = Create(handler);
|
||||
|
||||
var a1 = await svc.GetAccessTokenAsync(TenantA);
|
||||
var a2 = await svc.GetAccessTokenAsync(TenantA);
|
||||
var b1 = await svc.GetAccessTokenAsync(TenantB);
|
||||
|
||||
Assert.Equal("tok-A", a1.AccessToken);
|
||||
Assert.Equal("tok-A", a2.AccessToken);
|
||||
Assert.Equal("tok-B", b1.AccessToken);
|
||||
Assert.Equal(1, TokenPosts(handler, TenantA));
|
||||
Assert.Equal(1, TokenPosts(handler, TenantB));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task GetAccessTokenAsync_does_not_cache_failure()
|
||||
{
|
||||
// A non-retryable 400 (not a consent-propagation error) returns immediately and must NOT be
|
||||
// cached — the next request re-acquires so a transient outage doesn't pin a failure.
|
||||
var handler = new StubHttpMessageHandler()
|
||||
.When($"/{TenantA}/oauth2", HttpStatusCode.BadRequest, "{\"error\":\"invalid_request\"}");
|
||||
var (svc, _) = Create(handler);
|
||||
|
||||
var first = await svc.GetAccessTokenAsync(TenantA);
|
||||
var second = await svc.GetAccessTokenAsync(TenantA);
|
||||
|
||||
Assert.Null(first.AccessToken);
|
||||
Assert.Null(second.AccessToken);
|
||||
Assert.Equal(2, TokenPosts(handler, TenantA));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task GetAccessTokenAsync_caches_token_without_expires_in()
|
||||
{
|
||||
// AAD always returns expires_in, but the fallback lifetime path must still cache so a quirky
|
||||
// response doesn't silently disable caching.
|
||||
var handler = new StubHttpMessageHandler()
|
||||
.When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\"}");
|
||||
var (svc, _) = Create(handler);
|
||||
|
||||
await svc.GetAccessTokenAsync(TenantA);
|
||||
await svc.GetAccessTokenAsync(TenantA);
|
||||
|
||||
Assert.Equal(1, TokenPosts(handler, TenantA));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task InvalidateTenant_forces_fresh_token_on_next_call()
|
||||
{
|
||||
// The fresh-read contract: after an admin grants consent/roles, the cached token still
|
||||
// carries the OLD roles claim. InvalidateTenant (called by GraphFeatureDetector.InvalidateTenant
|
||||
// on the admin Refresh / access-check path) must drop it so the next acquire mints fresh.
|
||||
var handler = new StubHttpMessageHandler()
|
||||
.When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}");
|
||||
var (svc, _) = Create(handler);
|
||||
|
||||
await svc.GetAccessTokenAsync(TenantA);
|
||||
svc.InvalidateTenant(TenantA);
|
||||
await svc.GetAccessTokenAsync(TenantA);
|
||||
|
||||
Assert.Equal(2, TokenPosts(handler, TenantA));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task InvalidateTenant_during_in_flight_acquire_is_not_overwritten_by_stale_token()
|
||||
{
|
||||
// Write-after-invalidate race: a token POST that STARTED before InvalidateTenant must not
|
||||
// repopulate the cache with its now-stale token. The first POST is parked in flight; the
|
||||
// admin grants + Refreshes (InvalidateTenant) while it is parked; on release it writes a
|
||||
// (stale-generation-tagged) token. A subsequent read MUST reject it and re-mint fresh.
|
||||
var handler = new InflightRaceHandler(
|
||||
firstBody: "{\"access_token\":\"stale\",\"expires_in\":3599}",
|
||||
laterBody: "{\"access_token\":\"fresh\",\"expires_in\":3599}");
|
||||
var config = new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build();
|
||||
var svc = new GraphTokenService(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
new SingleHandlerFactory(handler),
|
||||
new MemoryCache(new MemoryCacheOptions()),
|
||||
config);
|
||||
|
||||
// R1 starts; its POST parks in flight (holding the gate, before writing to the cache).
|
||||
var inflight = svc.GetAccessTokenAsync(TenantA);
|
||||
await WaitUntilAsync(() => handler.Entered >= 1, TimeSpan.FromSeconds(2));
|
||||
|
||||
// Admin grants + Refresh WHILE R1's POST is in flight.
|
||||
svc.InvalidateTenant(TenantA);
|
||||
|
||||
// R1 completes → writes "stale", tagged with the pre-invalidation generation.
|
||||
handler.Release();
|
||||
var r1 = await inflight;
|
||||
Assert.Equal("stale", r1.AccessToken); // R1's own caller still gets the token it acquired
|
||||
|
||||
// Fresh read must NOT serve the stale entry — the generation mismatch makes it a miss → re-mint.
|
||||
var r2 = await svc.GetAccessTokenAsync(TenantA);
|
||||
Assert.Equal("fresh", r2.AccessToken);
|
||||
Assert.Equal(2, handler.Entered); // a real second POST happened (no stale cache hit)
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task GetAccessTokenAsync_coalesces_concurrent_misses()
|
||||
{
|
||||
// Stampede guard, proven properly: a barrier-gated handler parks the first request on the
|
||||
// wire so the other nine genuinely pile up at the miss BEFORE any response/cache-fill. With
|
||||
// the per-tenant semaphore only ONE request ever reaches the wire; without it all ten would.
|
||||
// (A synchronous stub could pass even without the gate, because the first call may fill the
|
||||
// cache before the rest start — hence this stronger construction.)
|
||||
var handler = new GatedHttpMessageHandler("{\"access_token\":\"tok-A\",\"expires_in\":3599}");
|
||||
var config = new ConfigurationBuilder()
|
||||
.AddInMemoryCollection(new Dictionary<string, string?>
|
||||
{
|
||||
["EntraId:ClientId"] = "test-client",
|
||||
["EntraId:ClientSecret"] = "test-secret",
|
||||
})
|
||||
.Build();
|
||||
var svc = new GraphTokenService(
|
||||
NullLogger<GraphTokenService>.Instance,
|
||||
new SingleHandlerFactory(handler),
|
||||
new MemoryCache(new MemoryCacheOptions()),
|
||||
config);
|
||||
|
||||
var tasks = Enumerable.Range(0, 10).Select(_ => svc.GetAccessTokenAsync(TenantA)).ToArray();
|
||||
|
||||
// Wait for the first request to reach the wire, then give the other nine ample time to block
|
||||
// on the semaphore. If coalescing works, exactly one request is at the handler.
|
||||
await WaitUntilAsync(() => handler.Entered >= 1, TimeSpan.FromSeconds(2));
|
||||
await Task.Delay(300);
|
||||
Assert.Equal(1, handler.Entered);
|
||||
|
||||
handler.Release();
|
||||
var results = await Task.WhenAll(tasks);
|
||||
|
||||
Assert.All(results, r => Assert.Equal("tok-A", r.AccessToken));
|
||||
Assert.Equal(1, handler.Entered);
|
||||
}
|
||||
|
||||
private static async Task WaitUntilAsync(Func<bool> condition, TimeSpan timeout)
|
||||
{
|
||||
var deadline = DateTime.UtcNow + timeout;
|
||||
while (!condition() && DateTime.UtcNow < deadline)
|
||||
{
|
||||
await Task.Delay(10);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -48,53 +48,88 @@ public class IngestTelemetryFunctionTests
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_accepts_body_under_cap()
|
||||
{
|
||||
var payload = Encoding.UTF8.GetBytes("[{\"kind\":\"Event\"}]");
|
||||
var payload = Encoding.UTF8.GetBytes("[{\"Kind\":\"Event\"}]");
|
||||
using var stream = new MemoryStream(payload);
|
||||
|
||||
var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024);
|
||||
var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024);
|
||||
|
||||
Assert.False(exceeded);
|
||||
Assert.Equal("[{\"kind\":\"Event\"}]", body);
|
||||
Assert.NotNull(items);
|
||||
var item = Assert.Single(items!);
|
||||
Assert.Equal("Event", item.Kind);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_preserves_PayloadJson_as_raw_string()
|
||||
{
|
||||
// Hot-path contract: the inner payload stays an opaque escaped string on the envelope —
|
||||
// the batch parse must NOT re-parse/re-serialise it (that's TelemetryPayloadParser's job,
|
||||
// exactly once, downstream). Guards against an accidental round-trip that would mangle it.
|
||||
var payload = Encoding.UTF8.GetBytes(
|
||||
"[{\"Kind\":\"Event\",\"PayloadJson\":\"{\\\"EventType\\\":\\\"x\\\",\\\"n\\\":1}\"}]");
|
||||
using var stream = new MemoryStream(payload);
|
||||
|
||||
var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024);
|
||||
|
||||
Assert.False(exceeded);
|
||||
var item = Assert.Single(items!);
|
||||
Assert.Equal("{\"EventType\":\"x\",\"n\":1}", item.PayloadJson);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_accepts_body_exactly_at_cap()
|
||||
{
|
||||
// Strict greater-than semantics (matches legacy NdjsonParser): equal-to is OK.
|
||||
var payload = new byte[100];
|
||||
// Strict greater-than semantics (matches legacy NdjsonParser): equal-to is OK. Pad a valid
|
||||
// batch with trailing whitespace (ignored by the parser) to land exactly on the cap.
|
||||
var json = "[{\"Kind\":\"Signal\",\"PayloadJson\":\"{}\"}]".PadRight(100);
|
||||
var payload = Encoding.UTF8.GetBytes(json);
|
||||
Assert.Equal(100, payload.Length); // ASCII → 1 byte/char
|
||||
using var stream = new MemoryStream(payload);
|
||||
|
||||
var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100);
|
||||
var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100);
|
||||
|
||||
Assert.False(exceeded);
|
||||
Assert.Equal(100, Encoding.UTF8.GetByteCount(body));
|
||||
Assert.Single(items!);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_rejects_body_over_cap_without_draining_full_stream()
|
||||
{
|
||||
// 1 MB payload with a 100-byte cap → short-circuits after a single buffer read.
|
||||
// 1 MB payload with a 100-byte cap → short-circuits after a single buffer read, before
|
||||
// the deserialiser is ever invoked.
|
||||
var payload = new byte[1_000_000];
|
||||
using var stream = new MemoryStream(payload);
|
||||
|
||||
var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100);
|
||||
var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100);
|
||||
|
||||
Assert.True(exceeded);
|
||||
Assert.Equal(string.Empty, body);
|
||||
Assert.Null(items);
|
||||
// Stream position advanced past the cap but not to EOF — the helper bails early to
|
||||
// bound memory on hostile senders.
|
||||
Assert.True(stream.Position < stream.Length);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_empty_stream_returns_empty_body()
|
||||
public async Task ReadBodyWithSizeCapAsync_empty_stream_returns_null_items()
|
||||
{
|
||||
using var stream = new MemoryStream(Array.Empty<byte>());
|
||||
|
||||
var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024);
|
||||
var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024);
|
||||
|
||||
Assert.False(exceeded);
|
||||
Assert.Equal(string.Empty, body);
|
||||
Assert.Null(items);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadBodyWithSizeCapAsync_throws_on_malformed_json()
|
||||
{
|
||||
// Malformed JSON surfaces as JsonException so the HTTP trigger maps it to 400 — the cap
|
||||
// path must not swallow it as an empty/successful parse.
|
||||
var payload = Encoding.UTF8.GetBytes("[{\"Kind\":\"Event\"");
|
||||
using var stream = new MemoryStream(payload);
|
||||
|
||||
await Assert.ThrowsAnyAsync<Newtonsoft.Json.JsonException>(
|
||||
() => IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024));
|
||||
}
|
||||
|
||||
// ============================================================ Batch PartitionKey uniformity
|
||||
|
||||
+36
-19
@@ -136,7 +136,18 @@ namespace AutopilotMonitor.Functions.Functions.Ingest
|
||||
var tenantConfig = await _configService.GetConfigurationAsync(tenantIdHeader);
|
||||
var maxPayloadBytes = (tenantConfig?.MaxNdjsonPayloadSizeMB ?? 5) * 1024 * 1024;
|
||||
|
||||
var (exceeded, body) = await ReadBodyWithSizeCapAsync(req.Body, maxPayloadBytes);
|
||||
bool exceeded;
|
||||
List<TelemetryItemDto>? items;
|
||||
try
|
||||
{
|
||||
(exceeded, items) = await ReadBodyWithSizeCapAsync(req.Body, maxPayloadBytes);
|
||||
}
|
||||
catch (JsonException ex)
|
||||
{
|
||||
_logger.LogWarning(ex, "IngestTelemetry: malformed JSON body");
|
||||
return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "Malformed JSON body"));
|
||||
}
|
||||
|
||||
if (exceeded)
|
||||
{
|
||||
_logger.LogWarning(
|
||||
@@ -147,17 +158,6 @@ namespace AutopilotMonitor.Functions.Functions.Ingest
|
||||
$"Payload exceeds {maxPayloadBytes / (1024 * 1024)} MB cap"));
|
||||
}
|
||||
|
||||
List<TelemetryItemDto>? items;
|
||||
try
|
||||
{
|
||||
items = JsonConvert.DeserializeObject<List<TelemetryItemDto>>(body);
|
||||
}
|
||||
catch (JsonException ex)
|
||||
{
|
||||
_logger.LogWarning(ex, "IngestTelemetry: malformed JSON body");
|
||||
return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "Malformed JSON body"));
|
||||
}
|
||||
|
||||
if (items == null || items.Count == 0)
|
||||
{
|
||||
return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "No telemetry items provided"));
|
||||
@@ -502,12 +502,24 @@ namespace AutopilotMonitor.Functions.Functions.Ingest
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reads a stream into a string with a strict byte cap. Returns <c>(true, "")</c> as
|
||||
/// soon as the cap would be exceeded — the stream isn't fully drained, which bounds
|
||||
/// memory use even for a malicious sender. Matches the legacy NDJSON parser's guard
|
||||
/// semantics (strict greater-than, so a payload equal to the cap is accepted).
|
||||
/// Buffers the request body under a strict byte cap, then deserialises the telemetry batch
|
||||
/// directly off the buffered bytes. Returns <c>(true, null)</c> as soon as the cap would be
|
||||
/// exceeded — the stream isn't fully drained, which bounds memory use even for a malicious
|
||||
/// sender (strict greater-than, so a payload equal to the cap is accepted; matches the
|
||||
/// legacy NDJSON parser's guard semantics).
|
||||
/// <para>
|
||||
/// Hot-path note: this deserialises through a <see cref="JsonTextReader"/> over the buffered
|
||||
/// <see cref="MemoryStream"/> rather than materialising a full UTF-16 <c>string</c> + a
|
||||
/// <c>ToArray()</c> copy first. That removes ~3× the body size in transient allocations per
|
||||
/// request on the highest-volume backend endpoint — the reader decodes UTF-8 incrementally
|
||||
/// in 8 KB chunks. The cap is still applied <i>before</i> parsing, so an over-cap body
|
||||
/// never reaches the deserialiser.
|
||||
/// </para>
|
||||
/// Throws <see cref="JsonException"/> on malformed JSON (caller maps to 400). The
|
||||
/// deserialiser uses the same default settings as <c>JsonConvert.DeserializeObject</c>, so
|
||||
/// the parse semantics are unchanged — only the intermediate allocations are gone.
|
||||
/// </summary>
|
||||
internal static async Task<(bool exceeded, string body)> ReadBodyWithSizeCapAsync(
|
||||
internal static async Task<(bool exceeded, List<TelemetryItemDto>? items)> ReadBodyWithSizeCapAsync(
|
||||
Stream source, int maxBytes)
|
||||
{
|
||||
using var buffered = new MemoryStream();
|
||||
@@ -517,10 +529,15 @@ namespace AutopilotMonitor.Functions.Functions.Ingest
|
||||
while ((read = await source.ReadAsync(buffer, 0, buffer.Length)) > 0)
|
||||
{
|
||||
total += read;
|
||||
if (total > maxBytes) return (true, string.Empty);
|
||||
if (total > maxBytes) return (true, null);
|
||||
await buffered.WriteAsync(buffer, 0, read);
|
||||
}
|
||||
return (false, System.Text.Encoding.UTF8.GetString(buffered.ToArray()));
|
||||
|
||||
buffered.Position = 0;
|
||||
using var textReader = new StreamReader(buffered, System.Text.Encoding.UTF8);
|
||||
using var jsonReader = new JsonTextReader(textReader);
|
||||
var items = JsonSerializer.CreateDefault().Deserialize<List<TelemetryItemDto>>(jsonReader);
|
||||
return (false, items);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
using System;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Collections.Generic;
|
||||
using System.Net.Http;
|
||||
using System.Threading;
|
||||
@@ -19,6 +20,32 @@ namespace AutopilotMonitor.Functions.Security
|
||||
{
|
||||
private static readonly TimeSpan ConsentStatusTtl = TimeSpan.FromMinutes(2);
|
||||
|
||||
// App-token cache tuning. TTL = expires_in - skew margin, clamped so a missing/pathological
|
||||
// expires_in can neither thrash the cache nor pin a stale token. Max stays under the typical
|
||||
// ~60-90 min AAD client-credentials lifetime; the 5-min margin covers clock skew between
|
||||
// this worker and the Graph resource server plus the slowest downstream call in a request.
|
||||
private static readonly TimeSpan TokenSkewMargin = TimeSpan.FromMinutes(5);
|
||||
private static readonly TimeSpan MinTokenTtl = TimeSpan.FromSeconds(60);
|
||||
private static readonly TimeSpan MaxTokenTtl = TimeSpan.FromMinutes(55);
|
||||
private static readonly TimeSpan DefaultTokenLifetime = TimeSpan.FromMinutes(35); // when AAD omits expires_in
|
||||
private static readonly TimeSpan TokenGateWait = TimeSpan.FromSeconds(10);
|
||||
|
||||
// One coalescing gate per tenant (see GetAccessTokenAsync). A SemaphoreSlim is tiny and
|
||||
// tenant cardinality is bounded + validated upstream, so these are never evicted.
|
||||
private readonly ConcurrentDictionary<string, SemaphoreSlim> _tokenLocks = new();
|
||||
|
||||
// Per-tenant invalidation generation. Closes the write-after-invalidate race: a token POST
|
||||
// that STARTED before InvalidateTenant captures the old generation and tags its cache entry
|
||||
// with it; InvalidateTenant bumps the generation, so that late stale write is rejected by
|
||||
// every reader (entry.Gen != current) even though it physically lands after the Remove.
|
||||
// Lock-free — the comparison happens at read time against the latest generation.
|
||||
private readonly ConcurrentDictionary<string, long> _tokenCacheGen = new();
|
||||
|
||||
private long CurrentGen(string tenantId) => _tokenCacheGen.TryGetValue(tenantId, out var g) ? g : 0;
|
||||
|
||||
/// <summary>Cached app token tagged with the invalidation generation in effect when its acquire started.</summary>
|
||||
private sealed record CachedAppToken(string Token, long Gen);
|
||||
|
||||
private readonly ILogger<GraphTokenService> _logger;
|
||||
private readonly IHttpClientFactory _httpClientFactory;
|
||||
private readonly IMemoryCache _cache;
|
||||
@@ -38,7 +65,7 @@ namespace AutopilotMonitor.Functions.Security
|
||||
|
||||
public async Task<GraphConsentStatusResult> GetConsentStatusAsync(string tenantId, CancellationToken ct = default)
|
||||
{
|
||||
var cacheKey = $"graph-consent-status:{tenantId}";
|
||||
var cacheKey = ConsentCacheKey(tenantId);
|
||||
if (_cache.TryGetValue(cacheKey, out GraphConsentStatusResult? cached) && cached != null)
|
||||
{
|
||||
return cached;
|
||||
@@ -76,6 +103,14 @@ namespace AutopilotMonitor.Functions.Security
|
||||
|
||||
public virtual async Task<GraphTokenResult> GetAccessTokenAsync(string tenantId, CancellationToken ct = default)
|
||||
{
|
||||
// Serve a still-valid cached app token for this tenant. The cache key is keyed strictly
|
||||
// on the (already authenticated upstream) tenantId — the same value that drives the token
|
||||
// URL below — so a token minted for one tenant can never be served to another.
|
||||
if (TryGetCachedToken(tenantId, out var cachedToken))
|
||||
{
|
||||
return GraphTokenResult.Success(cachedToken);
|
||||
}
|
||||
|
||||
var clientId = _configuration["EntraId:ClientId"];
|
||||
var clientSecret = _configuration["EntraId:ClientSecret"];
|
||||
|
||||
@@ -86,6 +121,54 @@ namespace AutopilotMonitor.Functions.Security
|
||||
return GraphTokenResult.PermanentFailure();
|
||||
}
|
||||
|
||||
// Stampede guard: coalesce concurrent cache-misses for the same tenant so an enrollment
|
||||
// wave doesn't fan a burst of identical token POSTs at AAD (which would trip its own
|
||||
// throttling). One acquirer fetches + populates the cache; the rest hit it after the
|
||||
// gate. The wait is bounded so a pathological slow acquire (e.g. a no-consent tenant on
|
||||
// the long retry chain) degrades to today's parallel-fetch behaviour instead of pinning
|
||||
// callers indefinitely.
|
||||
var gate = _tokenLocks.GetOrAdd(tenantId, _ => new SemaphoreSlim(1, 1));
|
||||
bool gotGate;
|
||||
try
|
||||
{
|
||||
gotGate = await gate.WaitAsync(TokenGateWait, ct);
|
||||
}
|
||||
catch (OperationCanceledException) when (ct.IsCancellationRequested)
|
||||
{
|
||||
// Caller budget exhausted while waiting — transient by definition (never cache).
|
||||
return GraphTokenResult.TransientFailure();
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
// Double-check: another thread may have populated the cache while we waited. The
|
||||
// gen-aware read also rejects a stale entry written by an acquire that overlapped an
|
||||
// InvalidateTenant, so we fall through to a fresh mint instead of serving it.
|
||||
if (gotGate && TryGetCachedToken(tenantId, out var tokenAfterWait))
|
||||
{
|
||||
return GraphTokenResult.Success(tokenAfterWait);
|
||||
}
|
||||
|
||||
return await AcquireTokenViaClientCredentialsAsync(tenantId, clientId, clientSecret, ct);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (gotGate) gate.Release();
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Performs the client-credentials POST (with the consent-propagation retry chain) and
|
||||
/// caches a successful token until shortly before its expiry. Failures are never cached, so
|
||||
/// a tenant that loses consent re-acquires (and fails) on the next request after expiry.
|
||||
/// </summary>
|
||||
private async Task<GraphTokenResult> AcquireTokenViaClientCredentialsAsync(
|
||||
string tenantId, string clientId, string clientSecret, CancellationToken ct)
|
||||
{
|
||||
// Capture the invalidation generation BEFORE the POST. If an InvalidateTenant lands
|
||||
// while this POST is in flight, the generation changes and the token we cache below is
|
||||
// tagged with the now-stale value — so readers reject it and re-mint fresh.
|
||||
var genAtStart = CurrentGen(tenantId);
|
||||
var tokenUrl = $"https://login.microsoftonline.com/{tenantId}/oauth2/v2.0/token";
|
||||
|
||||
// Retry with backoff to handle Azure AD consent propagation delays.
|
||||
@@ -124,6 +207,12 @@ namespace AutopilotMonitor.Functions.Security
|
||||
{
|
||||
var tokenJson = JsonConvert.DeserializeObject<JObject>(responseBody);
|
||||
var accessToken = tokenJson?["access_token"]?.ToString();
|
||||
if (!string.IsNullOrEmpty(accessToken))
|
||||
{
|
||||
// Cache the freshly minted token (never log its value). TTL derives from
|
||||
// the token's own expires_in, minus a skew margin, clamped to a safe band.
|
||||
CacheToken(tenantId, accessToken!, genAtStart, (long?)tokenJson?["expires_in"]);
|
||||
}
|
||||
return GraphTokenResult.Success(accessToken);
|
||||
}
|
||||
|
||||
@@ -208,6 +297,64 @@ namespace AutopilotMonitor.Functions.Security
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static string TokenCacheKey(string tenantId) => $"graph-token:{tenantId}";
|
||||
private static string ConsentCacheKey(string tenantId) => $"graph-consent-status:{tenantId}";
|
||||
|
||||
/// <summary>
|
||||
/// Drops this tenant's cached app token (and consent-status verdict) so the next acquire
|
||||
/// mints a fresh JWT from AAD. Required for the consent / Graph-permission fresh-read
|
||||
/// contract: after an admin grants consent or assigns app roles, the OLD token still carries
|
||||
/// the OLD <c>roles</c> claim until it expires — serving it would make a post-grant "Refresh"
|
||||
/// keep reporting "not granted" for up to the cache TTL. The single fresh-read entry point
|
||||
/// <c>GraphFeatureDetector.InvalidateTenant</c> calls this so BOTH cache layers are cleared
|
||||
/// together (the detector's parsed-roles cache AND this raw-token cache).
|
||||
/// </summary>
|
||||
public virtual void InvalidateTenant(string tenantId)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(tenantId)) return;
|
||||
// Bump the generation FIRST so any token POST already in flight (which captured the old
|
||||
// generation) writes an entry that readers reject as stale — this closes the
|
||||
// write-after-invalidate race, not just the already-written entry. Then drop the current
|
||||
// entries to free memory and clear the coarse consent-status verdict.
|
||||
_tokenCacheGen.AddOrUpdate(tenantId, 1, (_, v) => v + 1);
|
||||
_cache.Remove(TokenCacheKey(tenantId));
|
||||
_cache.Remove(ConsentCacheKey(tenantId));
|
||||
}
|
||||
|
||||
private bool TryGetCachedToken(string tenantId, out string token)
|
||||
{
|
||||
// Accept the cached token only if its generation still matches the tenant's current one.
|
||||
// A mismatch means an InvalidateTenant ran since this entry's acquire started, so the
|
||||
// token may carry a stale roles claim — treat it as a miss and re-mint.
|
||||
if (_cache.TryGetValue(TokenCacheKey(tenantId), out CachedAppToken? cached)
|
||||
&& cached != null
|
||||
&& !string.IsNullOrEmpty(cached.Token)
|
||||
&& cached.Gen == CurrentGen(tenantId))
|
||||
{
|
||||
token = cached.Token;
|
||||
return true;
|
||||
}
|
||||
token = string.Empty;
|
||||
return false;
|
||||
}
|
||||
|
||||
private void CacheToken(string tenantId, string accessToken, long genAtStart, long? expiresInSeconds)
|
||||
{
|
||||
// expires_in is reported by AAD over TLS, so it is trustworthy — but still clamped so a
|
||||
// missing or pathological value cannot pin a stale token (cap) or cause thrash (floor).
|
||||
var lifetime = expiresInSeconds is > 0
|
||||
? TimeSpan.FromSeconds(expiresInSeconds.Value)
|
||||
: DefaultTokenLifetime;
|
||||
var ttl = lifetime - TokenSkewMargin;
|
||||
if (ttl < MinTokenTtl) ttl = MinTokenTtl;
|
||||
if (ttl > MaxTokenTtl) ttl = MaxTokenTtl;
|
||||
|
||||
_cache.Set(TokenCacheKey(tenantId), new CachedAppToken(accessToken, genAtStart), new MemoryCacheEntryOptions
|
||||
{
|
||||
AbsoluteExpirationRelativeToNow = ttl,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
||||
+48
-9
@@ -1,4 +1,5 @@
|
||||
using System;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Collections.Generic;
|
||||
using System.Globalization;
|
||||
using System.IdentityModel.Tokens.Jwt;
|
||||
@@ -43,6 +44,21 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector
|
||||
public static readonly NegativeAcquireMarker Instance = new();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Cache entry tagged with the per-tenant invalidation generation in effect when its acquire
|
||||
/// started. <see cref="Payload"/> is either a <see cref="GraphTenantTokenContext"/> (success)
|
||||
/// or <see cref="NegativeAcquireMarker"/> (definitive no-consent). Readers accept it only while
|
||||
/// <see cref="Gen"/> still matches the tenant's current generation — a stale write that lands
|
||||
/// after an <see cref="InvalidateTenant"/> carries an older generation and is rejected.
|
||||
/// </summary>
|
||||
private sealed record CachedVerdict(object Payload, long Gen);
|
||||
|
||||
// Per-tenant invalidation generation (see CachedVerdict + InvalidateTenant). Mirrors the
|
||||
// GraphTokenService fix one layer down so the fresh-read contract holds on this cache too.
|
||||
private readonly ConcurrentDictionary<string, long> _invalidationGen = new();
|
||||
|
||||
private long CurrentGen(string tenantId) => _invalidationGen.TryGetValue(tenantId, out var g) ? g : 0;
|
||||
|
||||
private readonly GraphTokenService _tokenService;
|
||||
private readonly IMemoryCache _cache;
|
||||
private readonly ILogger<GraphFeatureDetector> _logger;
|
||||
@@ -111,15 +127,23 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(tenantId)) return (null, false);
|
||||
|
||||
// Capture the invalidation generation BEFORE the acquire. If InvalidateTenant runs while we
|
||||
// fetch/parse the token below, the generation changes and the entry we write is tagged with
|
||||
// the now-stale value — so readers reject it and re-acquire. Mirrors the GraphTokenService
|
||||
// fix one layer down, closing the write-after-invalidate race on this cache too.
|
||||
var genAtStart = CurrentGen(tenantId);
|
||||
|
||||
var cacheKey = BuildCacheKey(tenantId);
|
||||
if (_cache.TryGetValue(cacheKey, out object? cached) && cached != null)
|
||||
if (_cache.TryGetValue(cacheKey, out CachedVerdict? cached)
|
||||
&& cached != null
|
||||
&& cached.Gen == CurrentGen(tenantId))
|
||||
{
|
||||
if (cached is NegativeAcquireMarker)
|
||||
if (cached.Payload is NegativeAcquireMarker)
|
||||
{
|
||||
// Definitive "no consent / no SP in this tenant" verdict from a recent attempt.
|
||||
return (null, false);
|
||||
}
|
||||
if (cached is GraphTenantTokenContext cachedCtx
|
||||
if (cached.Payload is GraphTenantTokenContext cachedCtx
|
||||
&& cachedCtx.ExpiresAt > _time.GetUtcNow() + ExpirySafetyMargin)
|
||||
{
|
||||
return (cachedCtx, false);
|
||||
@@ -145,7 +169,7 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector
|
||||
// revoked, or the app credentials are misconfigured. Cache the verdict briefly so
|
||||
// repeated cross-tenant clicks don't keep replaying the consent-propagation
|
||||
// retry chain for the same dead-end tenant.
|
||||
CacheNegativeVerdict(cacheKey);
|
||||
CacheNegativeVerdict(cacheKey, genAtStart);
|
||||
return (null, false);
|
||||
}
|
||||
|
||||
@@ -165,7 +189,10 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector
|
||||
};
|
||||
|
||||
var ttl = ComputeTtl(expiresAt);
|
||||
_cache.Set(cacheKey, ctx, new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = ttl });
|
||||
// Tag with the generation captured before the acquire (see top of method). A concurrent
|
||||
// InvalidateTenant bumps the generation, so this write — if it lands after the invalidation —
|
||||
// is rejected by readers via the gen check, closing the write-after-invalidate race.
|
||||
_cache.Set(cacheKey, new CachedVerdict(ctx, genAtStart), new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = ttl });
|
||||
|
||||
_logger.LogDebug(
|
||||
"GraphFeatureDetector: cached token for tenant {TenantId} with {RoleCount} roles, TTL={Ttl}",
|
||||
@@ -201,15 +228,27 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector
|
||||
public void InvalidateTenant(string tenantId)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(tenantId)) return;
|
||||
// Single cache key carries both the success token AND the negative marker, so one
|
||||
// Remove call clears whichever verdict is currently in flight.
|
||||
// ORDER IS LOAD-BEARING. Invalidate the UNDERLYING token layer FIRST, then bump OUR
|
||||
// generation, then drop our cache. Token-first is what closes the cross-layer overlap race:
|
||||
// otherwise a GetSnapshotAsync starting in the gap could capture our NEW generation yet
|
||||
// still read the OLD (not-yet-invalidated) GraphTokenService cache, and then write a
|
||||
// stale-roles ctx tagged with the new generation — which readers would ACCEPT, re-opening
|
||||
// the stale "not granted" window until our TTL. Because an acquire captures its generation
|
||||
// at method entry (before its token read), invalidating the token layer first guarantees
|
||||
// any acquire able to read a stale token captured its generation before this bump, so its
|
||||
// write loses the gen check. (This also drops the raw-token cache so our own next read can't
|
||||
// re-parse a stale JWT — the original cross-layer fix.)
|
||||
_tokenService.InvalidateTenant(tenantId);
|
||||
_invalidationGen.AddOrUpdate(tenantId, 1, (_, v) => v + 1);
|
||||
// Single cache key carries both the success token AND the negative marker, so one Remove
|
||||
// clears whichever verdict is currently in flight.
|
||||
_cache.Remove(BuildCacheKey(tenantId));
|
||||
_logger.LogInformation("GraphFeatureDetector: cache invalidated for tenant {TenantId}", tenantId);
|
||||
}
|
||||
|
||||
private void CacheNegativeVerdict(string cacheKey)
|
||||
private void CacheNegativeVerdict(string cacheKey, long gen)
|
||||
{
|
||||
_cache.Set(cacheKey, NegativeAcquireMarker.Instance, new MemoryCacheEntryOptions
|
||||
_cache.Set(cacheKey, new CachedVerdict(NegativeAcquireMarker.Instance, gen), new MemoryCacheEntryOptions
|
||||
{
|
||||
AbsoluteExpirationRelativeToNow = NegativeCacheTtl,
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user